diff --git a/codecov.yml b/codecov.yml index 4d6cf5961..70d65c692 100644 --- a/codecov.yml +++ b/codecov.yml @@ -1,3 +1,10 @@ +codecov: + notify: + after_n_builds: 7 + +comment: + after_n_builds: 7 + flags: unittests: carryforward: false diff --git a/dascore/__init__.py b/dascore/__init__.py index 6a47de753..dd7f24f73 100644 --- a/dascore/__init__.py +++ b/dascore/__init__.py @@ -8,7 +8,7 @@ from dascore.core.patch import Patch from dascore.core.attrs import PatchAttrs from dascore.core.summary import PatchSummary -from dascore.core.spool import BaseSpool, spool +from dascore.core.spool import BaseSpool, Spool, spool from dascore.core.coordmanager import get_coord_manager, CoordManager from dascore.core.coords import get_coord from dascore.config import DascoreConfig, get_config, reset_config, set_config diff --git a/dascore/clients/__init__.py b/dascore/clients/__init__.py deleted file mode 100644 index 2c6634da4..000000000 --- a/dascore/clients/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -""" -DAS Core module for accessing remote resources. -""" -from __future__ import annotations diff --git a/dascore/clients/dirspool.py b/dascore/clients/dirspool.py deleted file mode 100644 index 4397ddb57..000000000 --- a/dascore/clients/dirspool.py +++ /dev/null @@ -1,176 +0,0 @@ -""" -A spool for working with file systems. - -The spool uses a simple hdf5 index for keeping track of files. -""" - -from __future__ import annotations - -import copy -from pathlib import Path - -import pandas as pd -from rich.text import Text -from typing_extensions import Self - -import dascore as dc -from dascore.compat import UPath -from dascore.constants import PROGRESS_LEVELS -from dascore.core.spool import BaseSpool, DataFrameSpool, MemorySpool -from dascore.exceptions import MissingPatchError -from dascore.io.indexer import AbstractIndexer, DirectoryIndexer -from dascore.utils.docs import compose_docstring -from dascore.utils.pd import adjust_segments - - -class DirectorySpool(DataFrameSpool): - """ - A spool for interacting with DAS files on disk. - - FileSpool creates and index of all files then allows for simple querying - and bulk processing of the files. - - Parameters - ---------- - base_path - The path to the directory to index. - index_path - The path to the index file containing the contents of the directory. - By default it will be created in the top-level of the data directory. - preferred_format - A string to specify the format of the data. Specifying this parameter - will save time in indexing. - select_kwargs - Dict of keyword arguments to restrict output contents. - """ - - _drop_columns = ("file_format", "file_version", "path", "source_patch_id") - - def __init__( - self, - base_path: str | Path | UPath | Self | AbstractIndexer = ".", - *, - index_path: Path | None = None, - preferred_format: str | None = None, - select_kwargs: dict | None = None, - merge_kwargs: dict | None = None, - ): - super().__init__(select_kwargs=select_kwargs, merge_kwargs=merge_kwargs) - # Init file spool from another file spool - if isinstance(base_path, self.__class__): - self.__dict__.update(copy.deepcopy(base_path.__dict__)) - return - # Init file spool from indexer - elif isinstance(base_path, AbstractIndexer): - self.indexer = base_path - elif isinstance(base_path, Path | str | UPath): - self.indexer = DirectoryIndexer(base_path, index_path=index_path) - assert hasattr(self, "indexer"), "indexer not set." - self._preferred_format = preferred_format - - def __rich__(self): - """Augment rich string directory spool stuff.""" - base = super().__rich__() - path = self.indexer.path - kwargs = self._select_kwargs - out = base + Text(f"\n Path: {path}") - out += Text(f"\n Select kwargs: {kwargs}") if kwargs else Text("") - return out - - def _get_df(self): - """Get the dataframe of current contents.""" - out = adjust_segments( - self._source_df, ignore_bad_kwargs=True, **self._select_kwargs - ) - return out - - def _get_instruction_df(self): - """Return instruction df on how to get from source_df to df.""" - _, _, instruction = self._get_dummy_dataframes(self._df) - return instruction - - def _get_source_df(self): - """Return a dataframe of sources in spool.""" - return self.indexer(**self._select_kwargs).reset_index(drop=True) - - @property - def spool_path(self): - """Return the path in which the spool contents are found.""" - return self.indexer.path - - @compose_docstring(doc=BaseSpool.get_contents.__doc__) - def get_contents(self) -> pd.DataFrame: - """{doc}.""" - return self._df - - @compose_docstring(doc=BaseSpool.update.__doc__) - def update(self, progress: PROGRESS_LEVELS = "standard") -> Self: - """{doc}.""" - out = self.__class__( - base_path=self.indexer.update(progress=progress), - preferred_format=self._preferred_format, - select_kwargs=self._select_kwargs, - ) - return out - - def _df_to_dict_list(self, df): - """ - Convert the dataframe to a list of dicts for iteration. - - This is significantly faster than iterating rows. - """ - df = df.copy(deep=False).replace("", None) - # note: need to add extra / here since we no longer store it in db. - df["path"] = (str(self.spool_path) + "/") + df["path"] - return super()._df_to_dict_list(df) - - def _load_patch(self, kwargs) -> Self: - """Given a row from the managed dataframe, return a patch.""" - final_kwargs = dict(kwargs) - final_kwargs.update(self._select_kwargs) - patches = self._read_patches(final_kwargs) - if patches is None: # fast path doesn't apply, use generic read. - return self._read_and_resolve_patch(final_kwargs) - if not patches: - # Iteration skips these with a warning, see #583. - msg = ( - f"No patch in {final_kwargs.get('path')} matches the " - f"requested range; it may have been trimmed to nothing." - ) - raise MissingPatchError(msg) - return patches[0] - - def _read_patches(self, kwargs) -> list[dc.Patch] | None: - """ - Read patches directly through the file's FiberIO. - - This skips the format detection of dc.read and the spool indexing - machinery applied to its output, which add up when loading many - files. Returns None when the fast path can't be safely used. - """ - fmt, version = kwargs.get("file_format"), kwargs.get("file_version") - if not fmt or not version: - # Without a concrete version get_fiberio would return the newest - # reader; let dc.read detect the file's actual version instead. - return None - fiber_io = dc.io.FiberIO.manager.get_fiberio(format=fmt, version=version) - # Only apply select kwargs when the source patch is trimmed by the - # instruction df or the spool itself; otherwise the whole file is - # wanted and selection is wasted work. - if kwargs.get("_modified") or self._select_kwargs: - select = { - k: v - for k, v in kwargs.items() - if k not in self._drop_columns and not k.startswith("_") - } - else: - select = {} - spool = fiber_io.read(kwargs["path"], **select) - if not isinstance(spool, MemorySpool): - return None - patches = list(spool) - # A multi-patch file is ambiguous: the row refers to one specific - # patch. Let the generic path resolve source_patch_id. - if len(patches) > 1: - return None - return patches diff --git a/dascore/clients/filespool.py b/dascore/clients/filespool.py deleted file mode 100644 index 6bfbe9322..000000000 --- a/dascore/clients/filespool.py +++ /dev/null @@ -1,84 +0,0 @@ -"""A spool for working with a single file.""" - -from __future__ import annotations - -import copy -from pathlib import Path - -from rich.text import Text -from typing_extensions import Self - -import dascore as dc -from dascore.compat import UPath -from dascore.constants import PROGRESS_LEVELS, SpoolType -from dascore.core.spool import BaseSpool, DataFrameSpool -from dascore.exceptions import MissingPatchError -from dascore.io.core import FiberIO -from dascore.utils.docs import compose_docstring - - -class FileSpool(DataFrameSpool): - """ - A spool for a single file. - - Parameters - ---------- - path - The path to the file. - file_format - The format name, optional. - file_version - The version string of the format, optional. - - Notes - ----- - Some file formats support storing multiple patches, this is most useful - for those formats, but should work on all dascore supported formats. - """ - - _drop_columns = ("source_patch_id",) - - def __init__( - self, - path: str | Path | UPath, - file_format: str | None = None, - file_version: str | None = None, - ): - super().__init__() - # Init file spool from another file spool - if isinstance(path, self.__class__): - self.__dict__.update(copy.deepcopy(path.__dict__)) - return - # Support UPaths, but keep standard Path support here because it is faster - self._path = path if isinstance(path, UPath) else Path(path) - if not self._path.exists() or self._path.is_dir(): - msg = f"{path} does not exist or is a directory" - raise FileNotFoundError(msg) - - _format, _version = dc.get_format(path, file_format, file_version) - source_df = dc.scan_to_df(path, file_format=_format, file_version=_version) - dfs = self._get_dummy_dataframes(source_df) - self._df, self._source_df, self._instruction_df = dfs - self._file_format = _format - self._file_version = _version - - def __rich__(self): - """Augment rich string with path.""" - base = super().__rich__() - out = base + Text(f" Path: {self._path}") - return out - - def _load_patch(self, kwargs) -> Self: - """Given a row from the managed dataframe, return a patch.""" - return self._read_and_resolve_patch(dict(kwargs)) - - @compose_docstring(doc=BaseSpool.update.__doc__) - def update(self: SpoolType, progress: PROGRESS_LEVELS = "standard") -> Self: - """ - {doc}. - """ - formatter = FiberIO.manager.get_fiberio( - format=self._file_format, version=self._file_version - ) - getattr(formatter, "index", lambda x: None)(self._path) - return self diff --git a/dascore/config.py b/dascore/config.py index 8e0979279..84e7f1f86 100644 --- a/dascore/config.py +++ b/dascore/config.py @@ -51,6 +51,34 @@ class DascoreConfig(BaseModel): default="standard", description="Controls whether DASCore appends processing history to patches.", ) + sampling_group_tolerance: float = Field( + default=0.05, + gt=0, + description=( + "Relative sampling-interval difference above which patches are " + "never combined during chunk/merge operations. E.g. the default " + "0.05 keeps patches whose steps differ by more than 5% in " + "separate groups." + ), + ) + groupby_attrs: tuple[str, ...] = Field( + default=( + "network", + "station", + "data_type", + "data_category", + "tag", + "instrument_id", + "acquisition_id", + ), + description=( + "Attributes which partition patches into separate groups for " + "chunk/merge operations. Patches whose values differ on any of " + "these are never combined (no error); the per-call `group` " + "argument overrides this default. Names missing from a spool " + "are ignored." + ), + ) # Local cache and index locations. downloader_cache_dir: Path = Field( @@ -66,20 +94,6 @@ class DascoreConfig(BaseModel): description="Time buffer applied when querying cached directory indexes.", ) - # HDF index writing. - hdf_index_complib: str = Field( - default="blosc:lz4", - description="Compression library used when writing DASCore HDF index files.", - ) - hdf_index_complevel: int = Field( - default=5, - description="Compression level used when writing DASCore HDF index files.", - ) - hdf_index_max_retries: int = Field( - default=3, - description="Maximum number of retries for concurrent HDF index access.", - ) - # Progress display. progress_basic_refresh_per_second: float = Field( default=0.25, diff --git a/dascore/core/coords.py b/dascore/core/coords.py index 5c015e1cd..a4bdcc705 100644 --- a/dascore/core/coords.py +++ b/dascore/core/coords.py @@ -1600,8 +1600,8 @@ def select( return self.empty(), out if np.all(out): return self, slice(None, None) - # Convert boolean to int indexes because these are supported for - # indexing pytables arrays but booleans are not. + # Convert boolean to int indexes; some consumers (eg lazy file + # readers) index with these where booleans are not supported. if len(self.shape) == 1: out = np.arange(len(out))[out] return self.new(values=values[out]), out diff --git a/dascore/core/patch.py b/dascore/core/patch.py index 6d128e599..27ae737bc 100644 --- a/dascore/core/patch.py +++ b/dascore/core/patch.py @@ -5,6 +5,7 @@ from collections.abc import Mapping, Sequence from functools import cached_property from typing import Final +from uuid import uuid4 import numpy as np from rich.text import Text @@ -99,6 +100,9 @@ def __init__( self._coords = coords self._attrs = attrs self._data = array(self.coords.validate_data(data)) + # Lineage identity: minted eagerly so copies made at any point + # (deepcopy/pickle carry __dict__) share it deterministically. + self._instance_id = uuid4().hex def __eq__(self, other): """Compare one Patch.""" diff --git a/dascore/core/spool.py b/dascore/core/spool.py index 4de9cc822..90f5665b3 100644 --- a/dascore/core/spool.py +++ b/dascore/core/spool.py @@ -4,7 +4,7 @@ import abc import warnings -from collections.abc import Callable, Generator, Mapping, Sequence +from collections.abc import Callable, Generator, Sequence from functools import singledispatch from pathlib import Path from typing import ClassVar, Literal, TypeVar @@ -27,89 +27,27 @@ timeable_types, ) from dascore.exceptions import ( - CoordMergeError, InvalidSpoolError, MissingPatchError, ParameterError, ) -from dascore.utils.attrs import combine_patch_attrs -from dascore.utils.chunk import ChunkManager from dascore.utils.display import get_dascore_text, get_nice_text from dascore.utils.docs import compose_docstring -from dascore.utils.mapping import FrozenDict from dascore.utils.misc import ( - CacheDescriptor, _spool_map, - broadcast_for_index, deep_equality_check, ) from dascore.utils.namespace import NamespaceOwner from dascore.utils.patch import ( - _force_patch_merge, - _get_merge_dim, - _get_merged_coord, - _spool_up, concatenate_patches, get_patch_names, - patches_to_df, stack_patches, ) from dascore.utils.paths import coerce_to_upath, requires_local_directory -from dascore.utils.pd import ( - _column_or_value, - _convert_min_max_in_kwargs, - adjust_segments, - filter_df, - get_column_names_from_dim, - get_dim_names_from_columns, - split_df_query, -) T = TypeVar("T") -def _get_varying_dim(df) -> str | None: - """ - Get the single dimension whose range varies across rows of df. - - Returns None when no dimension varies, several do, or the dataframe - doesn't carry range columns for the varying dimension; those cases - need the fully materialized merge to sort out. - """ - dims = get_dim_names_from_columns(df) - varying = [] - for dim in dims: - mins, maxs = df.get(f"{dim}_min"), df.get(f"{dim}_max") - if mins.nunique(dropna=False) > 1 or maxs.nunique(dropna=False) > 1: - varying.append(dim) - return varying[0] if len(varying) == 1 else None - - -def _estimate_merge_samples(df, dim) -> int | None: - """ - Estimate the total number of samples along dim of the merged rows. - - Returns None if the estimate cannot be made (eg unknown steps), in - which case streaming the merge isn't possible. - """ - if dim is None: - return None - cols = [f"{dim}_min", f"{dim}_max", f"{dim}_step"] - if not set(cols).issubset(df.columns): - return None - mins, maxs, steps = (df[x] for x in cols) - if mins.isnull().any() or maxs.isnull().any() or steps.isnull().any(): - return None - ratios = (maxs - mins) / steps - # Degenerate steps (eg 0) make the sample counts meaningless. - if not np.isfinite(ratios.astype(np.float64)).all(): - return None - counts = np.round(ratios).astype(np.int64) + 1 - if (counts < 0).any(): - return None - return int(counts.sum()) - - class BaseSpool(NamespaceOwner, abc.ABC): """Spool Abstract Base Class (ABC) for defining Spool interface.""" @@ -165,6 +103,46 @@ def __eq__(self, other) -> bool: other_dict = getattr(other, "__dict__", {}) return deep_equality_check(my_dict, other_dict) + def __add__(self, other) -> BaseSpool: + """ + Combine two spools into one containing the patches of both. + + The result is a lazy spool over the union of both spools' + metadata: file-backed patches stay unloaded, in-memory patches + are shared (not copied), and the same source appearing in both + spools keeps a single entry. Selections on the inputs carry over + by row membership. + + Examples + -------- + >>> import dascore as dc + >>> sp1 = dc.get_example_spool("random_das") + >>> sp2 = dc.get_example_spool("diverse_das") + >>> combined = sp1 + sp2 + >>> assert len(combined) == len(sp1) + len(sp2) + """ + if not isinstance(other, BaseSpool): + return NotImplemented + from dascore.io.index.catalog import PatchCatalog + + members = [self._as_catalog_member(), other._as_catalog_member()] + union = PatchCatalog.union(members) + new = Spool() + new._catalog = union + return new + + def _as_catalog_member(self): + """ + Return (catalog, patch_ids) describing this spool for a union. + + `patch_ids` limits membership to the spool's current rows; None + means the whole catalog (or the catalog view itself carries the + selection). The base implementation materializes the patches. + """ + from dascore.io.index.catalog import PatchCatalog + + return PatchCatalog.from_patches(list(self)), None + @abc.abstractmethod @compose_docstring(conflict_desc=attr_conflict_description) def chunk( @@ -174,6 +152,8 @@ def chunk( snap_coords: bool = True, tolerance: float = 1.5, conflict: Literal["drop", "raise", "keep_first"] = "raise", + group: str | Sequence[str] | None = None, + missing_dim: Literal["raise", "drop"] = "raise", **kwargs, ) -> Self: """ @@ -188,13 +168,24 @@ def chunk( If True, keep the segments which are smaller than chunk size. This often occurs because of data gaps or at end of chunks. snap_coords - If True, snap the coords on joined patches such that the spacing - remains constant. + If True (default), simplify the coordinates of joined patches to + an evenly sampled range when doing so moves no coordinate value + by more than `tolerance` samples. Merges whose gaps exceed that + keep an exact segmented coordinate instead. tolerance - The maximum number of samples a block of data can be spaced (gap) and - still be considered contiguous. + The maximum number of samples a block of data can be spaced (gap) + and still be considered contiguous. conflict {conflict_desc} + group + Attributes which partition patches into separate outputs (their + values differing is never an error). Defaults to the config + option `groupby_attrs`; unlike the default, explicitly passed + names must exist on at least one patch. Dimensions and + coordinate identities always partition implicitly. + missing_dim + What to do when patches lack the chunked dimension: "raise" + (default) or "drop" (exclude them from the output). kwargs kwargs are used to specify the dimension along which to chunk, eg: `time=10` chunks along the time axis in 10 second increments. @@ -214,6 +205,13 @@ def chunk( ----- [`Spool.concatenate`](`dascore.BaseSpool.concatenate`) performs a similar operation but disregards the coordinate values. + + To inspect what a chunk call will do before running it — which + output patches it produces and which slice of which source patch + feeds each one — use + [`Spool.chunk_plan`](`dascore.core.spool.Spool.chunk_plan`), + which takes the same arguments and returns the plan without + touching any data. """ @abc.abstractmethod @@ -222,10 +220,26 @@ def select(self, **kwargs) -> Self: Sub-select parts of the spool. Can be used to specify dimension ranges, or unix-style matches - on string attributes. + on string attributes. Bare keyword names resolve against + attributes first, then coordinates; unknown names raise. Parameters ---------- + _attrs + Attribute selections: a dict of ``name -> selector`` (the + general form — required when a name cannot be a Python + keyword) or a name/collection of names tagging bare kwargs + as attributes (disambiguates names shared with coordinates). + _coords + Coordinate selections; same forms as ``_attrs``, validating + names as coordinates only. + samples + If True, selections are coordinate-only and given in sample + indices; they never exclude patches, but are applied to each + patch as it loads. + relative + If True, range bounds are relative to the spool's coordinate + envelope: positive from the start, negative from the end. **kwargs Specifies query. Can be of the form {dim_name=(start, stop)} or {attr_name=query}. @@ -403,288 +417,329 @@ def viz(self): raise AttributeError(msg) -class DataFrameSpool(BaseSpool): - """An abstract class for spools whose contents are managed by a dataframe.""" - - # A dataframe which represents contents as they will be output - _df: pd.DataFrame = CacheDescriptor("_cache", "_get_df") - # A dataframe which shows patches in the source - _source_df: pd.DataFrame = CacheDescriptor("_cache", "_get_source_df") - # A dataframe of instructions for going from source_df to df - _instruction_df: pd.DataFrame = CacheDescriptor("_cache", "_get_instruction_df") - # kwargs for filtering contents - _select_kwargs: Mapping | None = FrozenDict() - # kwargs for merging patches - _merge_kwargs: Mapping | None = FrozenDict() - # attributes which effect merge groups for internal patches - _group_columns = ("network", "station", "dims", "data_type", "tag") - _drop_columns = ("patch",) +class Spool(BaseSpool): + """ + The concrete spool: a view over a `PatchCatalog`. - def _get_df(self): - """Function to get the current df.""" + Constructed from in-memory patches directly (or via + [`dascore.spool`](`dascore.spool`)), from a directory of files with + [`Spool.from_directory`](`dascore.core.spool.Spool.from_directory`), + or from a single file with + [`Spool.from_file`](`dascore.core.spool.Spool.from_file`). - def _get_source_df(self): - """Function to get the current df.""" + Parameters + ---------- + data + A patch, sequence of patches, or another spool whose (in-memory) + patches this spool should hold; None creates an empty spool. + + Notes + ----- + The catalog is the spool's entire state: live patches sit in its + resolver registry, file-backed patches in its index tables, and + restructured views (chunk/concat) are derived in-memory catalogs + whose rows are the plan outputs. Selection, ordering, and windowing + are lazy specs composed on the catalog; one engine serves every + construction path. + """ - def _get_instruction_df(self): - """Function to get the current df.""" + # synthetic catalog identity columns must not join patch kwargs + # comparisons or chunk merge-compatibility checks + _drop_columns = ("patch", "path", "file_format", "file_version", "source_patch_id") + # The catalog backing this spool. + _catalog = None + # single-file provenance (set by from_file; drives update()) + _file_path = None + _file_format = None + _file_version = None def __init__( - self, select_kwargs: dict | None = None, merge_kwargs: dict | None = None + self, + data: PatchType | Sequence[PatchType] | BaseSpool | None = None, ): - self._cache = {} - self._select_kwargs = {} if select_kwargs is None else select_kwargs - self._merge_kwargs = {} if merge_kwargs is None else merge_kwargs - - def _select_from_array(self, array) -> Self: - """Create new spool with contents changed from array input.""" - if np.issubdtype(array.dtype, np.bool_): # boolean select - df = self._df[array] - elif np.issubdtype(array.dtype, np.integer): - df = self._df.iloc[array] - else: - msg = "Only bool or int dtypes are supported for spool array selection." - raise ValueError(msg) - source = self._source_df - inst = self._instruction_df - select_kwargs, merge_kwargs = self._select_kwargs, self._merge_kwargs - new = self.new_from_df( - df, - source_df=source, - instruction_df=inst, - select_kwargs=select_kwargs, - merge_kwargs=merge_kwargs, - ) - return new + from dascore.io.index.catalog import PatchCatalog - def __getitem__(self, item) -> PatchType | BaseSpool: - if isinstance(item, slice): # a slice was used, return a sub-spool - new_df = self._df.iloc[item] - inst, source = self._instruction_df, self._source_df - new_inst = inst[inst["current_index"].isin(new_df.index)] - new_source = source.loc[new_inst.index] - out = self.new_from_df( - df=new_df, - instruction_df=new_inst, - source_df=new_source, + if isinstance(data, Spool): + # copy-construction: share the catalog and provenance + self.__dict__.update(data.__dict__) + return + if data is None: + patches = () + elif isinstance(data, dc.Patch): + patches = (data,) + elif isinstance(data, BaseSpool): + # e.g. wrapping dc.read output; the patches are in memory + patches = tuple(data) + elif isinstance(data, Sequence) and all(isinstance(x, dc.Patch) for x in data): + patches = data + else: + msg = ( + "Spool accepts a Patch, a sequence of patches, or a " + f"spool; got {type(data)}." ) - elif is_array(item): # An array was passed use np type selection. - return self._select_from_array(np.asarray(item)) - else: # a single index was used, should return a single patch - out = self._unbox_patch(self._get_patches_from_index(item)) - return out + raise InvalidSpoolError(msg) + self._catalog = PatchCatalog.from_patches(patches) + + # --- presented relation -------------------------------------------- + + @property + def _df(self) -> pd.DataFrame: + """The realized flat relation (cached by the catalog).""" + return self._catalog.to_df() + + @compose_docstring(doc=BaseSpool.get_contents.__doc__) + def get_contents(self) -> pd.DataFrame: + """{doc}.""" + return self._df def __len__(self): - return len(self._df) + # counting pushes to SQL (or the cold live registry); the flat + # relation is never realized just for a length + return len(self._catalog) + + def __getitem__(self, item) -> PatchType | BaseSpool: + if isinstance(item, slice): + # a lazy id-membership window (D2); never realizes the flat + # relation, and keeps split()/map() parts cheap + return self._new_from_catalog(self._catalog.window(item)) + if is_array(item): + array = np.asarray(item) + if not ( + np.issubdtype(array.dtype, np.bool_) + or np.issubdtype(array.dtype, np.integer) + ): + msg = ( + "Only bool or int dtypes are supported for spool " + "array selection." + ) + raise ValueError(msg) + return self._new_from_catalog(self._catalog.restrict(array)) + try: + return self._catalog.get_patch(int(item)) + except MissingPatchError: + # MissingPatchError subclasses IndexError for backwards + # compatibility; it must never masquerade as out-of-bounds + raise + except IndexError: + msg = f"index of [{item}] is out of bounds for spool." + raise IndexError(msg) from None def __iter__(self): - for ind in range(len(self._df)): + for ind in range(len(self._catalog)): try: - yield self._unbox_patch(self._get_patches_from_index(ind)) + yield self._catalog.get_patch(ind) except MissingPatchError as e: # The patch couldn't be produced, usually because a # coordinate mismatch trimmed it to nothing (see #583). msg = f"Skipping patch at index {ind} (see #583): {e}" warnings.warn(msg, UserWarning, stacklevel=2) - def _unbox_patch(self, patch_list): - """Unbox a single patch from a patch list, check len.""" - assert len(patch_list) == 1 - return patch_list[0] - - def _get_patches_from_index(self, df_ind): - """Given an index (from current df), return the corresponding patch.""" - source = self._source_df - instruction = self._instruction_df - # handle negative index. - df_ind = df_ind if df_ind >= 0 else len(self._df) + df_ind - try: - inds = self._df.index[df_ind] - except IndexError: - msg = f"index of [{df_ind}] is out of bounds for spool." - raise IndexError(msg) from None - # Group positional instruction rows by current index (and cache) to - # avoid a full instruction df scan for each requested patch. - indices = self._cache.get("_instruction_indices") - if indices is None: - indices = instruction.groupby("current_index").indices - self._cache["_instruction_indices"] = indices - positions = indices.get(inds) - assert positions is not None and len(positions), "no instructions found" - df1 = instruction.iloc[positions] - joined = df1.join(source.drop(columns=df1.columns, errors="ignore")) - # Occasionally, duplicates can creep into the source_df, - # but it costs a bit to check for duplicates, so only check and drop - # duplicates on large joined dataframes where performance might be - # affected. - if len(joined) > 10: - cols = set(joined.columns) - set(self._drop_columns) - joined = joined.drop_duplicates(subset=list(cols), keep="first") - return self._patch_from_instruction_df(joined) - - def _patch_from_instruction_df(self, joined): - """Get the patches joined columns of instruction df.""" - df_dict_list = self._df_to_dict_list(joined) - expected_len = len(joined["current_index"].unique()) - if len(df_dict_list) > expected_len: - # Several sources merge into one patch. When the output size can - # be determined from the instructions, stream the sources into a - # pre-allocated array so they don't all need to be in memory with - # the merged output at once. - merge_dim = _get_varying_dim(joined) - samples = _estimate_merge_samples(joined, merge_dim) - if samples is not None: - patch = self._merge_patches_streaming( - joined, df_dict_list, merge_dim, samples - ) - return [patch] - out = [] - for patch_kwargs in df_dict_list: - patch = self._load_trimmed_patch(patch_kwargs, joined) - # The index doesn't carry all the dimensional info, so get what - # merging needs from the patch coords (cheaper than attr dumps). - info = patch.coords._get_dim_summary() - info["patch"] = patch - out.append(info) - if len(out) > expected_len: - out = _force_patch_merge(out, merge_kwargs=self._merge_kwargs) - return [x["patch"] for x in out] - - def _load_trimmed_patch(self, patch_kwargs, joined) -> dc.Patch: - """Load a single patch and trim it to its instruction range.""" - # convert kwargs to format understood by parser/patch.select - kwargs = _convert_min_max_in_kwargs(patch_kwargs, joined) - patch = self._load_patch(kwargs) - # If the limits of the source patch were not modified, we can just - # use the select kwargs. This is important for missing coordinates - # (NaN values) to not get trimmed out. - if kwargs.get("_modified"): - select_kwargs = { - i: v - for i, v in kwargs.items() - if i in patch.dims or i in patch.coords.coord_map - } - else: - select_kwargs = self._select_kwargs - if select_kwargs: - patch = patch.select(**select_kwargs) - return patch + # --- selection and presentation specs ------------------------------- - def _merge_patches_streaming(self, joined, df_dict_list, merge_dim, samples): - """ - Merge the patches described by the instructions along merge_dim. + @compose_docstring(doc=BaseSpool.select.__doc__) + def select( + self, + *, + _attrs: dict | None = None, + _coords: dict | None = None, + samples: bool = False, + relative: bool = False, + **kwargs, + ) -> Self: + """{doc}.""" + catalog = self._catalog.select( + _attrs=_attrs, + _coords=_coords, + samples=samples, + relative=relative, + **kwargs, + ) + return self._new_from_catalog(catalog) + + @compose_docstring(doc=BaseSpool.sort.__doc__) + def sort(self, attribute) -> Self: + """{doc}.""" + # a lazy ORDER BY spec (D2): no copy, no realization; the + # ordinal contract supplies the deterministic tiebreak + return self._new_from_catalog(self._catalog.order_by(attribute)) + + @compose_docstring(doc=BaseSpool.split.__doc__) + def split( + self, + size: int | None = None, + count: int | None = None, + ) -> Generator[Self, None, None]: + """{doc}.""" + if not ((count is not None) ^ (size is not None)): + msg = "Spool.split requires either spool_count or spool_size." + raise ParameterError(msg) + start = 0 + step = int(np.ceil(len(self) / count if count else size)) + while start < len(self): + yield self[start : start + step] + start += step - Each patch is copied into a pre-allocated output array as it is - loaded, then released; this avoids holding all source patches and - the merged output in memory at the same time, as concatenating - would. + def _new_from_catalog(self, catalog) -> Self: + """Create a spool view over a (possibly derived) catalog.""" + new = self.__class__(self) + new._catalog = catalog + return new + + def _as_catalog_member(self): """ - buffer, offset, axis, dims = None, 0, None, None - coords, attrs, summaries = [], [], [] - for patch_kwargs in df_dict_list: - patch = self._load_trimmed_patch(patch_kwargs, joined) - if dims is None: - dims = patch.dims - axis = patch.get_axis(merge_dim) - elif patch.dims != dims: - patch = patch.transpose(*dims) - data = patch.data - if buffer is None: - shape = list(data.shape) - shape[axis] = samples - buffer = np.empty(shape, dtype=data.dtype) - # Mixed dtypes upcast, mirroring np.concatenate behavior. - dtype = np.result_type(buffer.dtype, data.dtype) - if dtype != buffer.dtype: - buffer = buffer.astype(dtype) - end = offset + data.shape[axis] - if end > buffer.shape[axis]: - # The estimate came up short (eg from slightly uneven - # sampling); grow the buffer to fit. - shape = list(buffer.shape) - shape[axis] = end - new_buffer = np.empty(shape, dtype=buffer.dtype) - head = broadcast_for_index(buffer.ndim, axis, slice(0, offset)) - new_buffer[head] = buffer[head] - buffer = new_buffer - try: - index = broadcast_for_index(buffer.ndim, axis, slice(offset, end)) - buffer[index] = data - except ValueError as e: - msg = ( - f"Cannot merge patches; their shapes are incompatible " - f"along the dimensions not being merged ({merge_dim})." - ) - raise CoordMergeError(msg) from e - offset = end - coords.append(patch.coords) - attrs.append(patch.attrs) - summaries.append(patch.coords._get_dim_summary()) - if offset != buffer.shape[axis]: # over-estimated; trim excess. - buffer = buffer[broadcast_for_index(buffer.ndim, axis, slice(0, offset))] - # Ensure the loaded patches only vary along the expected dimension, - # the same requirement _force_patch_merge enforces. - summary_df = pd.DataFrame(summaries) - found_dim = _get_merge_dim(summary_df) - if found_dim != merge_dim: - msg = ( - f"Cannot merge patches; expected them to vary along " - f"{merge_dim} but found {found_dim}." + Return (catalog, patch_ids) describing this spool for a union. + + Row membership (attr predicates, windows, id arrays) survives a + table union as-is, but residual trims and order specs live + Python-side and would silently vanish; a spool carrying those + first bakes them into a derived catalog (tables only — no patch + data is loaded). A catalog default order (directory time + presentation) bakes only when the source-record transfer would + actually present rows differently — an interleaved multi-patch + file — so ordinary archives keep record-grain transfer and its + same-source deduplication. + """ + catalog = self._catalog + if catalog._residuals or catalog._order is not None: + return self._materialize_lossy(), None + if catalog._default_order is not None and not self._transfer_keeps_order(): + return self._materialize_lossy(), None + return catalog, None + + def _transfer_keeps_order(self) -> bool: + """True when ordinal-grain transfer matches the presented order.""" + catalog = self._catalog + presented = catalog._ordered_ids() + by_ordinal = tuple( + catalog.backend.query_ids( + list(catalog._queries) or None, + order_by=None, + patch_ids=catalog._ids, ) - raise CoordMergeError(msg) - conf = self._merge_kwargs.get("conflicts", None) - drop_conflicting = conf in {"drop", "keep_first"} - new_coord = _get_merged_coord(summary_df, merge_dim, coords, drop_conflicting) - new_attrs = combine_patch_attrs(attrs, **self._merge_kwargs) - return dc.Patch(data=buffer, coords=new_coord, attrs=new_attrs, dims=list(dims)) - - def _get_dummy_dataframes(self, current): + ) + return tuple(presented) == by_ordinal + + def _materialize_lossy(self): """ - Return dummy current, source, and instruction dataframes. + Bake residual trims and presentation order into a derived catalog. - Dummy because the source and current df are the same, so the - instruction df is a straight mapping between the two. + An identity plan over the view's presented rows: one output per + row (in presentation order, so ordinals record the order spec), + with trimmed envelopes as the output envelopes and the trims + themselves re-applied at load through the plan resolver. """ - source = current.copy(deep=False) # shallow to not copy patches - dims = get_dim_names_from_columns(source) - cols2keep = get_column_names_from_dim(dims) - instruction = ( - current.copy(deep=False)[cols2keep] - .assign( - source_index=source.index, - # This tracks the current spool row after spool operations. - # It is not the source patch identity within a file. - current_index=source.index, - _modified=lambda x: _column_or_value(x, "_modified", False), - ) - .set_index("source_index") - .sort_values("current_index") + from dascore.io.index.planned import derived_catalog + from dascore.utils.chunk_plan import ( + _SOURCE_COLUMNS, + ChunkPlan, + samples_adjusted_envelopes, + ) + + rows = self._df.reset_index(drop=True) + working = samples_adjusted_envelopes(rows, self._catalog._residuals) + working = working.reset_index(drop=True) + ids = np.arange(len(working), dtype=np.int64) + # outputs are not file rows: source bookkeeping stays on the + # members (where loading needs it), never on the derived rows + outputs = working.drop( + columns=["_patch_id", *_SOURCE_COLUMNS], errors="ignore" + ).assign(output_id=ids) + members = pd.DataFrame( + { + "output_id": ids, + "_patch_id": working.get("_patch_id", pd.Series(dtype=object)).values, + "_modified": False, + } ) - return current, source, instruction + plan = ChunkPlan(outputs, members, "", None, {}) + return derived_catalog( + source_rows=working, + plan=plan, + parent=self._catalog, + merge_kwargs={}, + mode="identity", + origin_path=self.spool_path, + ) + + # --- restructuring (materializing) operations ----------------------- - def _df_to_dict_list(self, df): + def _plan_frames(self, dim: str | None = None) -> tuple[pd.DataFrame, pd.DataFrame]: """ - Convert the dataframe to a list of dicts for iteration. + Return (source_rows, working) frames for planning along ``dim``. + + Re-planning the *same* dimension collapses (never nests): a + derived catalog re-plans from its members — the trimmed source + rows — restricted to the outputs the current view presents. + Planning a *different* dimension must keep the already-assembled + boundaries, so it plans over the current output rows themselves + (loaded through the plan resolver). Patch-local samples + residuals adjust the working envelopes so plans reflect the + loading truth. + """ + from dascore.io.index.planned import PlanResolver, collapse_working_df + from dascore.utils.chunk_plan import ( + _ensure_patch_id, + samples_adjusted_envelopes, + ) - This is significantly faster than iterating rows. + resolver = self._catalog.resolver + same_dim = isinstance(resolver, PlanResolver) and resolver.dim == dim + base = collapse_working_df(self._catalog) if same_dim else None + if base is None: + base = self._catalog.to_df().reset_index(drop=True) + base = _ensure_patch_id(base) + working = base.drop(columns=list(self._drop_columns), errors="ignore") + working = samples_adjusted_envelopes(working, self._catalog._residuals) + base = base[base["_patch_id"].isin(working["_patch_id"])] + return base.reset_index(drop=True), working.reset_index(drop=True) + + def chunk_plan( + self, + overlap: numeric_types | timeable_types | None = None, + keep_partial: bool = False, + snap_coords: bool = True, + tolerance: float = 1.5, + conflict: Literal["drop", "raise", "keep_first"] = "raise", + group: str | Sequence[str] | None = None, + missing_dim: Literal["raise", "drop"] = "raise", + **kwargs, + ): """ - return df.to_dict("records") + Return the plan `chunk` would execute, without touching any data. - @abc.abstractmethod - def _load_patch(self, kwargs) -> dc.Patch: - """Given a row from the managed dataframe, return a patch.""" + The returned [`ChunkPlan`](`dascore.utils.chunk_plan.ChunkPlan`) is a + read-only diagnostic: its `outputs` table describes each patch the + chunked spool would contain (envelopes, step, carried attributes), + its `members` table shows exactly which slice of which source patch + feeds each output, and `params` records every resolved parameter + (including the group attributes and sampling tolerance in effect). + Accepts the same arguments as + [`chunk`](`dascore.BaseSpool.chunk`). - def _read_and_resolve_patch(self, final_kwargs) -> dc.Patch: - """Read patches for one instruction row and resolve to one patch.""" - from dascore.io.core import _select_patch_from_spool + Examples + -------- + >>> import dascore as dc + >>> spool = dc.get_example_spool("random_das") + >>> plan = spool.chunk_plan(time=3) + >>> assert len(plan.outputs) == len(spool.chunk(time=3)) + >>> # See which sources contribute to the first output patch. + >>> members = plan.members + >>> first = members[members["output_id"] == 0] + """ + from dascore.utils.chunk_plan import build_chunk_plan - source_patch_id = final_kwargs.get("source_patch_id", "") - spool = dc.read(**final_kwargs) - # Some readers consume source_patch_id internally and return the one - # matching patch without preserving that reload metadata on the patch. - if source_patch_id and len(spool) == 1: - return spool[0] - return _select_patch_from_spool(spool, source_patch_id=source_patch_id) + _, working = self._plan_frames(next(iter(kwargs), None)) + return build_chunk_plan( + working, + overlap=overlap, + keep_partial=keep_partial, + snap_coords=snap_coords, + tolerance=tolerance, + conflict=conflict, + group=group, + missing_dim=missing_dim, + **kwargs, + ) @compose_docstring(doc=BaseSpool.chunk.__doc__) def chunk( @@ -694,263 +749,326 @@ def chunk( snap_coords: bool = True, tolerance: float = 1.5, conflict: Literal["drop", "raise", "keep_first"] = "raise", + group: str | Sequence[str] | None = None, + missing_dim: Literal["raise", "drop"] = "raise", **kwargs, ) -> Self: """{doc}""" - df = self._source_df.drop(columns=list(self._drop_columns), errors="ignore") - chunker = ChunkManager( + from dascore.io.index.planned import derived_catalog + from dascore.utils.chunk_plan import build_chunk_plan + + source_rows, working = self._plan_frames(next(iter(kwargs), None)) + plan = build_chunk_plan( + working, overlap=overlap, keep_partial=keep_partial, snap_coords=snap_coords, - group_columns=self._group_columns, tolerance=tolerance, conflict=conflict, + group=group, + missing_dim=missing_dim, **kwargs, ) - in_df, out_df = chunker.chunk(df) - if df.empty: - instructions = None - else: - instructions = chunker.get_instruction_df(in_df, out_df) - return self.new_from_df( - out_df, - source_df=self._source_df, - instruction_df=instructions, - merge_kwargs={"conflicts": conflict}, + merge_kwargs = { + "conflicts": conflict, + "snap_coords": snap_coords, + "tolerance": tolerance, + } + catalog = derived_catalog( + source_rows=source_rows, + plan=plan, + parent=self._catalog, + merge_kwargs=merge_kwargs, + mode="chunk", + origin_path=self.spool_path, ) + return self._new_from_catalog(catalog) - def new_from_df( - self, - df, - source_df=None, - instruction_df=None, - select_kwargs=None, - merge_kwargs=None, - ): - """Create a new instance from dataframes.""" - new = self.__class__(self) - if source_df is None or instruction_df is None: - _, source_, inst_ = self._get_dummy_dataframes(df) - source_df = source_df if source_df is not None else source_ - instruction_df = instruction_df if instruction_df is not None else inst_ - new._df = df - new._source_df = source_df - new._instruction_df = instruction_df - # Discard stale instruction indices (eg from copied caches). - new._cache.pop("_instruction_indices", None) - new._select_kwargs = dict(self._select_kwargs) - new._select_kwargs.update(select_kwargs or {}) - new._merge_kwargs = dict(self._merge_kwargs) - new._merge_kwargs.update(merge_kwargs or {}) - return new - - @compose_docstring(doc=BaseSpool.select.__doc__) - def select(self, **kwargs) -> Self: - """{doc}.""" - _, _, extra_kwargs = split_df_query(kwargs, self._df, ignore_bad_kwargs=True) - filtered_df = adjust_segments(self._df, ignore_bad_kwargs=True, **kwargs) - inst = adjust_segments( - self._instruction_df, - ignore_bad_kwargs=True, - **kwargs, - ).loc[lambda x: x["current_index"].isin(filtered_df.index)] - source = adjust_segments( - self._source_df.loc[inst.index], ignore_bad_kwargs=True, **kwargs - ) - out = self.new_from_df( - filtered_df, - # Drop rows that are no longer needed. - source_df=source, - instruction_df=inst, - select_kwargs=extra_kwargs, - ) - return out + @compose_docstring(desc=concatenate_patches.__doc__) + def concatenate(self, check_behavior: WARN_LEVELS = "warn", **kwargs) -> Self: + """{desc}""" + from dascore.io.index.planned import derived_catalog + from dascore.utils.chunk_plan import ChunkPlan - @compose_docstring(doc=BaseSpool.sort.__doc__) - def sort(self, attribute) -> Self: - """{doc}.""" - df = self._df - inst_df = self._instruction_df - - # make sure a suitable attribute is entered - attrs = set(df.columns) - if attribute not in attrs: - # make sure we can also cover coordinate names instead of the attribute - if f"{attribute}_min" in attrs: - attribute = f"{attribute}_min" - else: - msg = "Invalid attribute. Please use a valid attribute such as: 'time'" - raise IndexError(msg) - - # get a mapping from the old current index to the sorted ones - sorted_df = df.sort_values(attribute) - sorted_original_indices = sorted_df.index - sorted_df = sorted_df.reset_index(drop=True) - mapper = pd.Series(np.arange(len(sorted_df)), index=sorted_original_indices) - # swap out all the old values with new ones - new_current_index = inst_df["current_index"].map(mapper) - new_instruction_df = inst_df.assign(current_index=new_current_index) - # create new spool from new dataframes - return self.new_from_df( - df=sorted_df, - source_df=self._source_df, - instruction_df=new_instruction_df, + if len(kwargs) != 1: + msg = ( + "concatenate requires exactly one dimension keyword, " + f"got {sorted(kwargs)}" + ) + raise ParameterError(msg) + ((dim, value),) = kwargs.items() + value = None if value is Ellipsis else value + source_rows, working = self._plan_frames(dim) + # a dim absent from the metadata envelopes is legal: concatenate + # can stack patches along a brand-new dimension + has_envelope = f"{dim}_min" in working.columns + count = len(working) if value in (None,) else int(value) + count = max(count, 1) + rows = working.reset_index(drop=True) + member_frames = [] + output_rows = [] + for output_id, start in enumerate(range(0, len(rows), count)): + group_rows = rows.iloc[start : start + count] + members = pd.DataFrame( + { + "output_id": output_id, + "_patch_id": group_rows["_patch_id"].values, + "_modified": False, + } + ) + member_frames.append(members) + first = group_rows.iloc[0].to_dict() + if has_envelope: + first[f"{dim}_min"] = group_rows[f"{dim}_min"].min() + first[f"{dim}_max"] = group_rows[f"{dim}_max"].max() + first["output_id"] = output_id + first.pop("_patch_id", None) + output_rows.append(first) + outputs = pd.DataFrame(output_rows) + if member_frames: + members = pd.concat(member_frames, ignore_index=True) + else: # nothing to concatenate: an empty spool stays empty + members = pd.DataFrame( + { + "output_id": pd.Series(dtype=np.int64), + "_patch_id": pd.Series(dtype=object), + "_modified": pd.Series(dtype=bool), + } + ) + plan = ChunkPlan(outputs, members, dim, None, {}) + catalog = derived_catalog( + source_rows=source_rows, + plan=plan, + parent=self._catalog, + merge_kwargs={}, + mode="concat", + check_behavior=check_behavior, + origin_path=self.spool_path, ) + return self._new_from_catalog(catalog) - @compose_docstring(doc=BaseSpool.split.__doc__) - def split( - self, - size: int | None = None, - count: int | None = None, - ) -> Generator[Self, None, None]: - """{doc}.""" - if not ((count is not None) ^ (size is not None)): - msg = "Spool.split requires either spool_count or spool_size." - raise ParameterError(msg) - start = 0 - step = int(np.ceil(len(self) / count if count else size)) - while start < len(self): - yield self[start : start + step] - start += step + # --- construction -------------------------------------------------- - @compose_docstring(doc=BaseSpool.get_contents.__doc__) - def get_contents(self) -> pd.DataFrame: - """{doc}.""" - return self._df[filter_df(self._df, **self._select_kwargs)] + @classmethod + def from_directory(cls, path, index_path=None) -> Self: + """ + Create a spool over a directory of fiber files. - get_patch_names = get_patch_names + The directory's index (created/updated via ``update()``) backs + the catalog; ``path`` may also be an existing directory indexer. + """ + from dascore.io.index.catalog import FileResolver, PatchCatalog + from dascore.io.indexer import AbstractIndexer + + out = cls() + if isinstance(path, AbstractIndexer): + from dascore.io.index.catalog import _DIRECTORY_ORDER + + out._catalog = PatchCatalog( + backend=path._backend, + resolver=FileResolver(root=path.path), + syncer=path, + default_order=_DIRECTORY_ORDER, + ) + else: + out._catalog = PatchCatalog.from_directory(path, index_path=index_path) + return out + @classmethod + def from_file( + cls, + path, + file_format: str | None = None, + file_version: str | None = None, + ) -> Self: + """ + Create a spool over a single (multi-patch capable) fiber file. -class MemorySpool(DataFrameSpool): - """ - A Spool for storing patches in memory. + The file is scanned once; patches load lazily per row. + """ + path = path if isinstance(path, UPath) else Path(path) + if not path.exists() or path.is_dir(): + msg = f"{path} does not exist or is a directory" + raise FileNotFoundError(msg) + from dascore.io.index.catalog import PatchCatalog + + _format, _version = dc.get_format(path, file_format, file_version) + out = cls() + out._catalog = PatchCatalog.from_file( + path, file_format=_format, file_version=_version + ) + out._file_path = path + out._file_format = _format + out._file_version = _version + return out - When created from patches, the managing dataframes are built lazily - (on first access by an operation which needs them, such as chunk or - select) and simple operations (len, integer access, iteration) are - served straight from the patch tuple. This makes creating a spool - from patches nearly free, which matters when reading many files. - """ + # --- capabilities -------------------------------------------------- - def __init__(self, data: PatchType | Sequence[PatchType] | None = None): - super().__init__() - self._patches: tuple[PatchType, ...] | None = None - self._data = None - if data is not None: - if isinstance(data, dc.Patch): - self._patches = (data,) - elif isinstance(data, Sequence) and all( - isinstance(x, dc.Patch) for x in data - ): - self._patches = tuple(data) - else: # eg a spool or dataframe; needs the dataframe machinery. - self._data = data + @property + def indexer(self): + """The directory syncer, or None for non-directory spools.""" + return None if self._catalog is None else self._catalog._syncer - def _get_df(self): - """Build the managing dataframes from the input patches.""" - data = self._patches if self._patches is not None else self._data - if data is None: - return None - df, source, instruction = self._get_dummy_dataframes(patches_to_df(data)) - self._source_df = source - self._instruction_df = instruction - return df - - def _get_source_df(self): - """Build the source df (happens as part of building current df).""" - _ = self._df - return self._cache.get("_source_df") - - def _get_instruction_df(self): - """Build the instruction df (happens as part of building current df).""" - _ = self._df - return self._cache.get("_instruction_df") + @property + def spool_path(self): + """The directory or file path this spool derives from, or None.""" + indexer = self.indexer + if indexer is not None: + return indexer.path + if self._file_path is not None: + return self._file_path + return getattr(self._catalog.resolver, "origin_path", None) - def __len__(self) -> int: - if self._patches is not None: - return len(self._patches) - return super().__len__() + @property + def has_live_patches(self) -> bool: + """True when any of this spool's patches live in memory.""" + catalog = self._catalog + return catalog is not None and bool(catalog.resolver.live_entries()) - def __getitem__(self, item) -> PatchType | BaseSpool: - # Fast path: a spool created directly from patches (which never has - # select kwargs) can serve integer requests from the patch list. - patches = self._patches - if ( - patches is not None - and not self._select_kwargs - and isinstance(item, int | np.integer) - ): - try: - return patches[item] - except IndexError: - msg = f"index of [{item}] is out of bounds for spool." - raise IndexError(msg) from None - return super().__getitem__(item) + @compose_docstring(doc=BaseSpool.update.__doc__) + def update(self, progress: PROGRESS_LEVELS = "standard") -> Self: + """ + {doc} + + Update is allowed only on a root spool — one no operation has + been applied to. Directory roots re-index their directory, + single-file roots rescan the file, and purely in-memory roots + are trivially current (no-op). Any derived spool (the result of + select, slicing, sort, chunk, concatenate, or combining spools) + raises: update the root and re-apply the operations. + """ + from dascore.io.index.catalog import LiveResolver + + catalog = self._catalog + derived_msg = ( + "update() is only allowed on a root spool; this spool is the " + "result of an operation (select/slice/sort/chunk/combine). " + "Update the root spool and re-apply the operations, e.g. " + "root = root.update(); view = root.select(...)." + ) + if catalog is None or catalog.is_view: + raise InvalidSpoolError(derived_msg) + if catalog._syncer is not None: + catalog.update(progress=progress) + return self._new_from_catalog(catalog) + if self._file_path is not None: + from dascore.io.core import FiberIO + + formatter = FiberIO.manager.get_fiberio( + format=self._file_format, version=self._file_version + ) + getattr(formatter, "index", lambda _: None)(self._file_path) + return self.from_file( + self._file_path, self._file_format, self._file_version + ) + if isinstance(catalog.resolver, LiveResolver): + return self # in-memory contents are trivially current + # composite/plan roots are computed spools (unions, chunks) + raise InvalidSpoolError(derived_msg) - def __iter__(self) -> PatchType: - patches = self._patches - if patches is not None and not self._select_kwargs: - yield from patches - else: - yield from super().__iter__() + # --- equality ------------------------------------------------------ def __eq__(self, other) -> bool: """ - Equality check which ignores the state of the lazy dataframes. + Equality check which ignores the state of lazy realization. - The managed dataframes are built and compared directly so that - equality does not depend on whether they were constructed yet. + The flat relations are built and compared directly so that + equality does not depend on whether they were realized yet. """ if self is other: return True - if not isinstance(other, MemorySpool): + if not isinstance(other, Spool): return super().__eq__(other) - return deep_equality_check(self._eq_dict(), other._eq_dict()) - - def _eq_dict(self) -> dict: - """Get a dict for equality checks, normalizing lazy state.""" - out = dict(self.__dict__) - # Build (if needed) and compare the dataframes; drop the inputs - # they were built from, whose form can differ for equal contents. - out["_cache"] = { - "_df": self._df, - "_source_df": self._source_df, - "_instruction_df": self._instruction_df, - } - out.pop("_patches", None) - out.pop("_data", None) - return out + # views over the same catalog state are equal without realizing + # the relations (a 200k-row archive must not materialize for ==) + mine, theirs = self._catalog, other._catalog + if ( + mine is not None + and theirs is not None + and ( + mine is theirs + or (mine._backend is not None and mine._backend is theirs._backend) + ) + and mine._queries == theirs._queries + and mine._residuals == theirs._residuals + and mine._order == theirs._order + and mine._ids == theirs._ids + ): + return True + return deep_equality_check(self._eq_state(), other._eq_state()) + + def _eq_state(self) -> dict: + """ + The spool's semantic state, explicitly enumerated for equality. + + Equality is over *effective* rows, never backends or + representation: same length and order of patch rows, row-wise + equal semantic columns (source identity like paths and + live-vs-file backing stripped), with pending residual + selections folded into the envelopes — a trimmed view equals + its materialized twin, and spools differing only by a samples + trim differ in their adjusted envelopes. Whether rows come from + a live registry, an index file, or a plan is invisible; data + arrays are never compared (metadata-level, like everything + here). Because the state is enumerated — never ``__dict__`` — + new instance attributes cannot silently join equality. + """ + + def _strip_identity(df): + # synthetic per-catalog identities (memory:// paths, ids) and + # backend provenance (format/version) are not content; equal + # spools must compare equal without them, and column order + # (a construction artifact) must not matter. Coordinate def + # keys are representation artifacts too: a residual-trimmed + # view cannot know its trimmed fingerprint without loading, + # and data values are never compared here anyway. + drop = [ + "path", + "_patch_id", + "source_patch_id", + "file_format", + "file_version", + "_modified", + *[c for c in df.columns if str(c).endswith("_def_key")], + ] + out = df.drop(columns=drop, errors="ignore") + return out[sorted(out.columns)] + + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + catalog = self._catalog + rows = self._df + # value residuals already trim the presented envelopes (to_df); + # samples residuals fold in here. Presented-but-empty rows stay: + # a spool exposing an emptied patch is not equal to one without. + if catalog is not None and catalog._residuals: + rows = samples_adjusted_envelopes( + rows, catalog._residuals, drop_empty=False + ) + return {"rows": _strip_identity(rows)} def __rich__(self): base = super().__rich__() - df = self._df - if len(df): - t1 = df["time_min"].min() if "time_min" in df.columns else "" - t2 = df["time_min"].max() if "time_min" in df.columns else "" - tmin = get_nice_text(t1) - tmax = get_nice_text(t2) - if t1 != "" and t2 != "": - duration = get_nice_text(t2 - t1) - else: - duration = "" - base += Text(f"\n Time Span: <{duration}> {tmin} to {tmax}") + path = self.spool_path + if path is not None: + base += Text(f"\n Path: {path}") + # Only render a time span when realization is cheap: live + # contents, single files, and derived catalogs are in memory; a + # huge directory index is not realized for a repr. + cheap = self.indexer is None + if cheap: + df = self._df + if df is not None and len(df) and "time_min" in df.columns: + t1, t2 = df["time_min"].min(), df["time_min"].max() + if pd.notna(t1) and pd.notna(t2): + duration = get_nice_text(t2 - t1) + base += Text( + f"\n Time Span: <{duration}> " + f"{get_nice_text(t1)} to {get_nice_text(t2)}" + ) return base - def _load_patch(self, kwargs) -> Self: - """Load the patch into memory.""" - return kwargs["patch"] - - @compose_docstring(doc=DataFrameSpool.new_from_df.__doc__) - def new_from_df(self, *args, **kwargs): - """{doc}.""" - new = super().new_from_df(*args, **kwargs) - # The provided dataframes fully define the new spool; drop the - # construction input so derived spools don't retain their parents. - new._data = None - return new - - # Add specific implementation of concatenate patches. - concatenate = _spool_up(concatenate_patches) + get_patch_names = get_patch_names @singledispatch @@ -992,24 +1110,19 @@ def spool(obj: path_types | BaseSpool | Sequence[PatchType], **kwargs) -> BaseSp def _spool_from_str(path, **kwargs): """Get a spool from a path.""" path = coerce_to_upath(path) - # A directory was passed, create Directory Spool + # A directory was passed; index it. if path.is_dir(): requires_local_directory(path, label="Directory spool") - from dascore.clients.dirspool import DirectorySpool - - return DirectorySpool(path, **kwargs) - # A single file was passed. If the file format supports quick scanning - # Return a FileSpool (lazy file reader), else return DirectorySpool. + return Spool.from_directory(path, **kwargs) + # A single file was passed. If the file format supports quick + # scanning build a lazy file-backed spool, else read it into memory. elif path.exists(): # a single file path was passed. _format, _version = dc.get_format(path, **kwargs) formatter = dc.io.FiberIO.manager.get_fiberio(format=_format, version=_version) if formatter.implements_scan: - from dascore.clients.filespool import FileSpool - - return FileSpool(path, _format, _version) - + return Spool.from_file(path, _format, _version) else: - return MemorySpool(dc.read(path, _format, _version)) + return Spool(dc.read(path, _format, _version)) else: msg = ( f"could not get spool from argument: {path}. " @@ -1028,10 +1141,10 @@ def _spool_from_spool(spool, **kwargs): @spool.register(tuple) def _spool_from_patch_list(patch_list, **kwargs): """Return a spool from a sequence of patches.""" - return MemorySpool(patch_list) + return Spool(patch_list) @spool.register(dc.Patch) def _spool_from_patch(patch): """Get a spool from a single patch.""" - return MemorySpool([patch]) + return Spool([patch]) diff --git a/dascore/core/summary.py b/dascore/core/summary.py index 1f96952a6..a71be9540 100644 --- a/dascore/core/summary.py +++ b/dascore/core/summary.py @@ -10,6 +10,7 @@ from typing import Any import numpy as np +import pandas as pd from pydantic import ConfigDict, Field, model_validator import dascore as dc @@ -20,6 +21,27 @@ from dascore.utils.paths import coerce_to_upath, is_pathlike +def normalize_source_patch_id(value: Any) -> str: + """ + Return a source patch id as a clean string ("" when missing). + + Missing ids arrive as None, empty strings, pandas NaN/NaT, or numpy + scalars. pandas NaN is truthy, so a plain ``value or ""`` does not + normalize it — every conversion site must go through this helper to + avoid the NaN-truthiness bug the catalog resolver already had to fix. + """ + if value is None or value == "": + return "" + try: + if pd.isnull(value): + return "" + except (TypeError, ValueError): + pass # non-scalar (e.g. an array): fall through to str() + if hasattr(value, "item"): # numpy scalar -> python scalar + value = value.item() + return str(value) + + def _to_coord_summary(value: Any, dims: tuple[str, ...] = ()) -> CoordSummary: """Normalize a coordinate summary input.""" # Summary inputs can already be normalized, coord-like objects, or exact @@ -169,10 +191,8 @@ def _normalize_source_patch_id( attrs: PatchAttrs, source_patch_id: Any = "" ) -> tuple[PatchAttrs, str]: """Normalize summary and private attr source ids to one value.""" - summary_source_patch_id = ( - "" if source_patch_id in (None, "") else str(source_patch_id) - ) - attrs_source_patch_id = str(attrs.get("_source_patch_id", "") or "") + summary_source_patch_id = normalize_source_patch_id(source_patch_id) + attrs_source_patch_id = normalize_source_patch_id(attrs.get("_source_patch_id", "")) normalized = summary_source_patch_id or attrs_source_patch_id if normalized: attrs = attrs.update(_source_patch_id=normalized) diff --git a/dascore/examples.py b/dascore/examples.py index 0ca2019d8..1eb6ba01b 100644 --- a/dascore/examples.py +++ b/dascore/examples.py @@ -28,7 +28,7 @@ def _load_example_patch_from_file(path: str | Path) -> dc.Patch: - """Load the first patch from an example file without FileSpool indirection.""" + """Load the first patch from an example file without spool indirection.""" with set_config(allow_dasdae_format_unpickle=True): return dc.read(path)[0] diff --git a/dascore/exceptions.py b/dascore/exceptions.py index d687136cf..ea32e5dbf 100644 --- a/dascore/exceptions.py +++ b/dascore/exceptions.py @@ -31,6 +31,10 @@ class ParameterError(ValueError, DASCoreError): """Raised when something is wrong with an input parameter.""" +class InvalidSpoolQueryError(ParameterError): + """Raised when a spool query references unknown names or bad values.""" + + class PatchError(DASCoreError): """Parent class for more specific Patch Errors.""" @@ -105,7 +109,11 @@ class InvalidFileHandlerError(TypeError, DASCoreError): """Raised when a writable file handler is requested from a read handle.""" -class InvalidIndexVersionError(ValueError, DASCoreError): +class InvalidIndexError(ValueError, DASCoreError): + """Raised when a persisted index is invalid or incompatible.""" + + +class InvalidIndexVersionError(InvalidIndexError): """Raised when a version mismatch occurs in index.""" diff --git a/dascore/io/core.py b/dascore/io/core.py index 6b9c08785..e9d12805b 100644 --- a/dascore/io/core.py +++ b/dascore/io/core.py @@ -31,13 +31,14 @@ ) from dascore.core.attrs import PatchAttrs, str_validator from dascore.core.coordmanager import CoordManager -from dascore.core.spool import DataFrameSpool -from dascore.core.summary import PatchSummary +from dascore.core.spool import Spool +from dascore.core.summary import PatchSummary, normalize_source_patch_id from dascore.exceptions import ( DependencyError, InvalidFiberFileError, InvalidFiberIOError, MissingOptionalDependencyError, + MissingPatchError, ParameterError, PatchAttributeError, RemoteCacheError, @@ -135,9 +136,7 @@ def _make_scan_payload( "dims": tuple(dims), "shape": tuple(shape), "dtype": str(dtype), - "source_patch_id": "" - if source_patch_id in (None, "") - else str(source_patch_id), + "source_patch_id": normalize_source_patch_id(source_patch_id), } @@ -168,7 +167,10 @@ def _scan_payload_to_summary( source_path=source_path, source_format=source_format, source_version=source_version, - source_patch_id=source_patch_id or payload.get("source_patch_id") or "", + source_patch_id=( + normalize_source_patch_id(source_patch_id) + or normalize_source_patch_id(payload.get("source_patch_id")) + ), ) @@ -189,9 +191,7 @@ def _scan_result_to_summary( normalized_source_path = "" if source_path in (None, "") else source_path normalized_source_format = "" if source_format in (None, "") else source_format normalized_source_version = "" if source_version in (None, "") else source_version - summary_source_patch_id = ( - "" if source_patch_id in (None, "") else str(source_patch_id) - ) + summary_source_patch_id = normalize_source_patch_id(source_patch_id) if isinstance(patch_summary, Mapping): return _scan_payload_to_summary( patch_summary, @@ -255,25 +255,50 @@ def _patch_to_scan_payload(patch: dc.Patch) -> ScanPayload: ) -def _select_patch_from_spool(spool, source_patch_id: object = "") -> dc.Patch: - """Select one loaded patch from a spool using source identity.""" +def _resolve_read_spool(spool, source_patch_id: object = "") -> dc.Patch: + """ + Resolve one patch from a read result by source identity. - def _matches_patch_name(patch: dc.Patch, source_id: str) -> bool: - """Return True when a generated patch name matches the source id.""" - return patch.get_patch_name() == source_id + Readers that consume source_patch_id may return the single matching + patch without preserving that reload metadata on it; only trust that + when the patch doesn't claim a different identity. + """ + source_patch_id = normalize_source_patch_id(source_patch_id) + if source_patch_id and len(spool) == 1: + found = normalize_source_patch_id(spool[0].attrs.get("_source_patch_id", "")) + if found == source_patch_id or (not found and not source_patch_id.isdigit()): + return spool[0] + return _select_patch_from_spool(spool, source_patch_id=source_patch_id) + +def _select_patch_from_spool(spool, source_patch_id: object = "") -> dc.Patch: + """Select one loaded patch from a spool using source identity.""" if len(spool) == 0: - msg = "index of [0] is out of bounds for spool." - raise IndexError(msg) + # Iteration skips these with a warning, see #583. + msg = ( + "No patch remained after applying load filters; the requested " + "range may have trimmed it to nothing." + ) + raise MissingPatchError(msg) if source_patch_id not in (None, ""): source_patch_id = str(source_patch_id) + # Native source ids are preserved on patch attrs by their readers. + matches = [ + patch + for patch in spool + if normalize_source_patch_id(patch.attrs.get("_source_patch_id", "")) + == source_patch_id + ] + if len(matches) == 1: + return matches[0] + # Synthesized ids are positional within the full source read. try: index = int(source_patch_id) except (TypeError, ValueError): index = None if index is not None and 0 <= index < len(spool): return spool[index] - if len(spool) == 1 and _matches_patch_name(spool[0], source_patch_id): + if len(spool) == 1 and spool[0].get_patch_name() == source_patch_id: return spool[0] msg = "Patch could not be uniquely resolved after applying load filters." raise PatchAttributeError(msg) @@ -615,7 +640,10 @@ def _get_format( # raise, in which case the format doesn't belong. func_input = man.get_resource(required_type) format_version = func(func_input, _pre_cast=True) - except RemoteCacheError: + except RemoteCacheError: # pragma: no cover -- remote fetch only + # A remote fetch failure is a real error, not a "wrong + # format" signal, so it must propagate rather than be + # swallowed by the robustness handler below. raise # For robustness, we need to catch everything else here. except Exception: @@ -980,7 +1008,7 @@ def scan_to_df( """ if isinstance(path, pd.DataFrame): return path - if isinstance(path, DataFrameSpool): + if isinstance(path, Spool): return path.get_contents() info = scan( path=path, @@ -1297,14 +1325,48 @@ def get_format( return out +def is_directory_format(path) -> bool: + """ + Return True if a directory is itself one FiberIO scan unit. + + A directory-format source (e.g. XMLBinary) is read as a whole rather + than by traversing its members. This is the single definition of that + condition; dc.scan's traversal skips such a directory's contents and + the directory indexer treats it as one stat unit. + """ + if not Path(path).is_dir(): + return False + try: + get_format(path) + except Exception: + return False + return True + + +def _resolves_assembled_patches(spool) -> bool: + """ + Return True when the spool can produce patches that are not literal + persisted file reads (live patches or plan-assembled outputs). + + Persisted patches are always contiguous, so purely file-backed + spools skip gap inspection; plan resolvers can assemble several + sources across a real gap into a segmented coordinate. + """ + if getattr(spool, "has_live_patches", False): + return True + catalog = getattr(spool, "_catalog", None) + resolver = getattr(catalog, "resolver", None) + return bool(getattr(resolver, "plan_entries", dict)()) + + def _maybe_split_gapped_patches(spool, fiber_io, split): """Handle patches whose dimensional coords contain gaps before writing.""" from dascore.core.coords import CoordSegmented - from dascore.core.spool import MemorySpool - # Only in-memory patches are inspected; file-backed patches always have - # contiguous coordinates (gapped patches are never persisted). - if not isinstance(spool, MemorySpool): + # Gap inspection depends on what the spool resolves, not on where + # its ultimate members live: only literal file reads are always + # contiguous (gapped patches are never persisted). + if not _resolves_assembled_patches(spool): return spool def _has_gaps(patch): diff --git a/dascore/io/h5simple/utils.py b/dascore/io/h5simple/utils.py index 9e0a5a4a9..8e1308484 100644 --- a/dascore/io/h5simple/utils.py +++ b/dascore/io/h5simple/utils.py @@ -18,32 +18,6 @@ DEFAULT_ATTRS = frozenset(("CLASS", "PYTABLES_FORMAT_VERSION", "TITLE", "VERSION")) -def _get_root_attrs(h5): - """Return a mapping-like object for root attrs for either HDF5 backend.""" - if hasattr(h5, "root"): - return h5.root._v_attrs - return h5.attrs - - -def _iter_root_arrays(h5): - """Yield ``(name, node)`` pairs for array-like nodes at the HDF5 root.""" - if hasattr(h5, "list_nodes"): - for node in h5.list_nodes("/"): - if hasattr(node, "shape"): - yield node.name, node - return - for name, node in h5.items(): - if hasattr(node, "shape"): - yield name, node - - -def _get_attr_names(attrs): - """Return the set of attribute names from either backend.""" - if hasattr(attrs, "_v_attrnames"): - return set(attrs._v_attrnames) - return set(attrs) - - def _maybe_trim_data(cm, data, kwargs): """Maybe use kwargs to trim data array.""" new_cm, new_data = cm.select(array=data, **kwargs) @@ -52,12 +26,9 @@ def _maybe_trim_data(cm, data, kwargs): def _get_attrs_coords_and_data(h5, snap, fiber_io): """Return attrs, coordinate manager, and data node.""" - attrs = _get_root_attrs(h5) - attr_names = _get_attr_names(attrs) - DEFAULT_ATTRS - attr_dict = { - x: unbyte(attrs[x] if not hasattr(attrs, "_v_attrnames") else getattr(attrs, x)) - for x in attr_names - } + attrs = h5.attrs + attr_names = set(attrs) - DEFAULT_ATTRS + attr_dict = {x: unbyte(attrs[x]) for x in attr_names} attr_dict["file_version"] = fiber_io.version attr_dict["file_format"] = fiber_io.name cm, data = _get_cm_and_data(h5, snap, dims=attr_dict.get("dims")) @@ -124,7 +95,7 @@ def _get_coords_and_dims(data_node, time_node, other_nodes, snap=True, dims=None def _get_cm_and_data(h5, snap=False, dims=None): """Extract coordinate manager and data node.""" - root_nodes = dict(_iter_root_arrays(h5)) + root_nodes = {name: node for name, node in h5.items() if hasattr(node, "shape")} array_names = set(root_nodes) data_node_name = array_names & DATA_ARRAY_NAMES time_node_name = array_names & TIME_ARRAY_NAMES diff --git a/dascore/io/index/__init__.py b/dascore/io/index/__init__.py new file mode 100644 index 000000000..23f51ee03 --- /dev/null +++ b/dascore/io/index/__init__.py @@ -0,0 +1,30 @@ +""" +SQLite spool index package. + +Provides a normalized, summary-only index of patch metadata (sources, +patches, attrs, and coordinates. PatchCatalog is the spool-facing metadata +engine; the remaining exports support its internal index implementation. +""" + +from __future__ import annotations + +from dascore.io.index.backend import AbstractIndexBackend, get_backend +from dascore.io.index.catalog import ( + FileResolver, + LiveResolver, + PatchCatalog, + PatchResolver, +) +from dascore.io.index.ingest import summaries_to_records +from dascore.io.index.query import Query + +__all__ = [ + "AbstractIndexBackend", + "FileResolver", + "LiveResolver", + "PatchCatalog", + "PatchResolver", + "Query", + "get_backend", + "summaries_to_records", +] diff --git a/dascore/io/index/backend.py b/dascore/io/index/backend.py new file mode 100644 index 000000000..84b68b756 --- /dev/null +++ b/dascore/io/index/backend.py @@ -0,0 +1,1070 @@ +""" +Index backend interface and SQLite SQL implementation. + +The backend persists the seven-table schema and answers flat-relation queries. +Storage hooks remain separate from the write/query logic so the index contract +has a clear boundary. +""" + +from __future__ import annotations + +import abc +import time +import warnings +from contextlib import contextmanager, nullcontext, suppress +from pathlib import Path + +import numpy as np +import pandas as pd + +import dascore as dc +from dascore.exceptions import ( + InvalidIndexError, + InvalidIndexVersionError, + ParameterError, + UnitError, +) +from dascore.io.index.dialect import BaseDialect +from dascore.io.index.ingest import SourceRecord, attr_column_name +from dascore.io.index.query import ( + Query, + _as_query_list, + apply_residuals, + build_sql, +) +from dascore.io.index.schema import ( + COORD_DEFS, + INDEX_VERSION, + INDEXES, + KIND_STORAGE, + PATCH_COORDS, + PATCHES, + SOURCES, + TABLE_CONSTRAINTS, + TABLES, + WHAT_IS_THIS, +) +from dascore.units import convert_units +from dascore.utils.pd import resolve_selector_namespaces + +# Structural columns whose ns-integer storage maps to pandas time types. +_TIME_COLS = {"time_min": "datetime", "time_max": "datetime", "time_step": "timedelta"} + + +def _ns_to_time(series: pd.Series, flavor: str) -> pd.Series: + """ + Convert nullable ns-integer columns to datetime64/timedelta64 exactly. + + Never goes through float64: ns epochs exceed float64's 2**53 integer + range, and the resulting ~100 ns corruption breaks merge boundary + arithmetic downstream. Float input means precision was already lost + upstream (a fetch path rounding nullable integers through float64), + so it is rejected rather than silently converted. + """ + if series.dtype.kind == "f": + msg = ( + f"ns column {series.name!r} arrived as {series.dtype}; values " + "above 2**53 ns are already corrupted. Fetch nullable integer " + "columns exactly (e.g. pandas nullable Int64)." + ) + raise TypeError(msg) + mask = series.isna() + values = np.zeros(len(series), dtype="int64") + if (~mask).any(): + values[~mask.to_numpy()] = series[~mask].astype("int64").to_numpy() + dtype = "datetime64[ns]" if flavor == "datetime" else "timedelta64[ns]" + out = pd.Series(values.view(dtype), index=series.index) + if mask.any(): + out[mask] = pd.NaT + return out + + +def adapt_params(params) -> list: + """Convert numpy scalars (and NaN) to plain python for DB drivers.""" + out = [] + for p in params: + if hasattr(p, "item"): # numpy scalar + p = p.item() + if isinstance(p, float) and np.isnan(p): + p = None + out.append(p) + return out + + +class AbstractIndexBackend(abc.ABC): + """Interface every index backend must implement.""" + + @abc.abstractmethod + def write_sources(self, records: list[SourceRecord]) -> None: + """Insert or replace sources (and dependents) transactionally.""" + + @abc.abstractmethod + def delete_sources(self, source_paths: list[str], base_uri: str = "") -> None: + """Remove sources (identified by base_uri + path) and dependents.""" + + @abc.abstractmethod + def query(self, query: Query, order_by=None, patch_ids=None) -> pd.DataFrame: + """Return the flat patch-row relation matching a query.""" + + @abc.abstractmethod + def query_ids(self, query: Query, order_by=None, patch_ids=None) -> list[int]: + """Return matching patch ids in presentation order.""" + + @abc.abstractmethod + def count(self, query: Query, patch_ids=None) -> int: + """Return how many patches match a query, without projecting rows.""" + + @abc.abstractmethod + def export_records(self, patch_ids=None) -> list: + """Reconstruct source records (optionally for a subset of patches).""" + + @abc.abstractmethod + def get_sources(self) -> pd.DataFrame: + """Return the sources table.""" + + @abc.abstractmethod + def source_stats(self) -> pd.DataFrame: + """Return only (source_path, mtime_ns, size_bytes) for change checks.""" + + @abc.abstractmethod + def get_metadata(self) -> dict: + """Return index-level metadata.""" + + @abc.abstractmethod + def mark_initial_update_done(self) -> None: + """Persist that the directory index completed its first update.""" + + @abc.abstractmethod + def attr_names(self) -> set[str]: + """Return original attr names known to the index.""" + + @abc.abstractmethod + def coord_names(self) -> set[str]: + """Return coord names known to the index.""" + + @abc.abstractmethod + def close(self) -> None: + """Release resources.""" + + +class SQLIndexBackend(AbstractIndexBackend): + """Shared implementation for SQL-speaking backends.""" + + dialect: BaseDialect + + def __init__(self): + self._ensure_schema() + + # --- hooks each engine provides --------------------------------- + + @abc.abstractmethod + def _execute(self, sql: str, params=()) -> None: + """Execute one statement.""" + + @abc.abstractmethod + def _executemany(self, sql: str, seq_of_params) -> None: + """Execute one statement for many parameter tuples.""" + + @abc.abstractmethod + def _fetch_df(self, sql: str, params=()) -> pd.DataFrame: + """ + Execute a SELECT and return a dataframe. + + Contract: nullable integer columns must round-trip exactly (use a + nullable integer dtype, never float64) — ns epochs exceed + float64's 2**53 integer range. + """ + + @abc.abstractmethod + def _begin(self) -> None: + """Start a transaction.""" + + @abc.abstractmethod + def _commit(self) -> None: + """Commit the open transaction.""" + + @abc.abstractmethod + def _rollback(self) -> None: + """Roll back the open transaction.""" + + @abc.abstractmethod + def _existing_tables(self) -> set[str]: + """Return persisted user table names.""" + + @abc.abstractmethod + def _table_columns(self, table: str) -> set[str]: + """Return persisted columns for one table.""" + + @contextmanager + def _transaction(self): + """ + Run the wrapped body inside one transaction. + + Commits on normal (or early-return) exit; on any error rolls back + without letting a failed rollback mask the original exception. + The backend's statement lock (when present) is held for the whole + transaction, so readers sharing the connection can never observe + a half-applied write; the reentrant lock keeps the statement + helpers inside the body working unchanged. + """ + lock = getattr(self, "_lock", None) + with lock if lock is not None else nullcontext(): + self._begin() + try: + yield + self._commit() + except Exception: + with suppress(Exception): + self._rollback() + raise + + # --- schema ------------------------------------------------------ + + def _ensure_schema(self) -> None: + tables = self._existing_tables() + if tables: + self._validate_schema(tables) + return + with self._transaction(): + # Another connection may have initialized the file while this + # writer waited for BEGIN IMMEDIATE. Re-check under the lock. + tables = self._existing_tables() + if tables: + self._validate_schema(tables) + return + for name, columns in TABLES.items(): + self._execute( + self.dialect.create_table( + name, columns, TABLE_CONSTRAINTS.get(name, ()) + ) + ) + for index_name, table, column in INDEXES: + self._execute( + f"CREATE INDEX IF NOT EXISTS {index_name} " f"ON {table} ({column})" + ) + self._execute( + "INSERT INTO meta_data VALUES (?, ?, ?, ?)", + (WHAT_IS_THIS, INDEX_VERSION, dc.__version__, 0), + ) + + def _validate_schema(self, tables: set[str]) -> None: + """Validate an existing index before issuing any DDL or mutation.""" + required = set(TABLES) + missing = required - tables + if missing: + msg = ( + "Existing spool index is incomplete; missing tables " + f"{sorted(missing)}. Delete it and rebuild the index." + ) + raise InvalidIndexError(msg) + meta = self._fetch_df("SELECT * FROM meta_data") + if len(meta) != 1 or meta["what_is_this"].iloc[0] != WHAT_IS_THIS: + msg = "File is not a valid DASCore spool index; delete it and rebuild." + raise InvalidIndexError(msg) + version = int(meta["index_version"].iloc[0]) + if version != INDEX_VERSION: + msg = ( + f"Spool index version {version} is incompatible with supported " + f"version {INDEX_VERSION}; delete it and rebuild." + ) + raise InvalidIndexVersionError(msg) + for table, expected in TABLES.items(): + actual = self._table_columns(table) + if not set(expected) <= actual: + absent = sorted(set(expected) - actual) + msg = ( + f"Spool index table {table!r} is missing columns {absent}; " + "delete it and rebuild." + ) + raise InvalidIndexError(msg) + attr_columns = self._table_columns("attrs") + meta_columns = set(self._attr_meta().get("column_name", ())) + if not meta_columns <= attr_columns: + absent = sorted(meta_columns - attr_columns) + msg = ( + f"Spool index attrs table is missing dynamic columns {absent}; " + "delete it and rebuild." + ) + raise InvalidIndexError(msg) + + def _attr_meta(self) -> pd.DataFrame: + return self._fetch_df("SELECT * FROM attr_meta") + + def _coord_meta(self, names=None) -> pd.DataFrame: + """ + Return distinct coordinate names, kinds, and canonical units, + optionally restricted to the given coord names. + """ + sql = ( + "SELECT DISTINCT pc.coord_name, cd.value_kind, cd.units, " + "cd.is_relative FROM patch_coords pc " + "JOIN coord_defs cd ON cd.coord_def_id = pc.coord_def_id" + ) + params: list = [] + if names is not None: + params = sorted(names) + sql += f" WHERE pc.coord_name IN ({self._placeholders(len(params))})" + return self._fetch_df(sql, params) + + def _next_id(self, table: str, column: str) -> int: + df = self._fetch_df(f"SELECT max({column}) AS m FROM {table}") + value = df["m"].iloc[0] + return 1 if pd.isnull(value) else int(value) + 1 + + @staticmethod + def _units_compatible(to_units: str, from_units: str) -> bool: + """True when one unit converts to the other (same dimensionality).""" + try: + convert_units(1.0, to_units=to_units, from_units=from_units) + except UnitError: + return False + return True + + def _ensure_attr_columns( + self, records: list[SourceRecord] + ) -> tuple[dict[tuple[str, str], str], set[tuple[str, str, str]]]: + """ + Lazily add typed attr columns; return the (name, kind) -> column + map and a set of (name, kind, units) values to skip. + + attr_meta is the single source of truth for column names: distinct + attr names can sanitize to the same identifier ("Shot Number" vs + "shot_number"), so collisions get a deterministic numeric suffix. + + One attr name occasionally carries dimensionally incompatible + units across files (e.g. a "resolution" in meters here, seconds + there). A single canonical unit cannot describe both, so the + incompatible values are skipped (with a warning) rather than + failing the whole index update. + """ + meta = self._attr_meta() + mapping = { + (row.attr_name, row.value_kind): row.column_name + for row in meta.itertuples() + } + stored_units = { + (row.attr_name, row.value_kind): ( + None if pd.isnull(row.units) else row.units + ) + for row in meta.itertuples() + } + taken = set(mapping.values()) + observed: dict[tuple[str, str], set[str | None]] = {} + for record in records: + for patch in record.patches: + for name, typed in patch.attrs.items(): + key = (name, typed.kind) + observed.setdefault(key, set()).add(typed.units) + needed: dict[tuple[str, str], str | None] = {} + skip_units: set[tuple[str, str, str]] = set() + for key, units_seen in observed.items(): + canonical = stored_units.get(key) if key in mapping else None + for unit in sorted(x for x in units_seen if x is not None): + if canonical is None: + canonical = unit + elif not self._units_compatible(canonical, unit): + skip_units.add((*key, unit)) + msg = ( + f"Attr {key[0]!r} has units {unit!r} incompatible " + f"with the indexed units {canonical!r}; skipping " + "these values in the index." + ) + warnings.warn(msg, UserWarning, stacklevel=2) + if key not in mapping: + needed[key] = canonical + elif stored_units.get(key) is None and canonical is not None: + self._execute( + "UPDATE attr_meta SET units = ? " + "WHERE attr_name = ? AND value_kind = ?", + (canonical, *key), + ) + for (name, kind), units in needed.items(): + column = base = attr_column_name(name, kind) + suffix = 2 + while column in taken: + column = f"{base}_{suffix}" + suffix += 1 + taken.add(column) + self._execute(self.dialect.add_column("attrs", column, KIND_STORAGE[kind])) + self._execute( + "INSERT INTO attr_meta VALUES (?, ?, ?, ?)", + (name, kind, column, units), + ) + mapping[(name, kind)] = column + return mapping, skip_units + + def _ensure_coord_defs(self, defs_needed: dict) -> dict[str, int]: + """ + Ensure unique coord definitions exist; return def_key -> id. + + Coord summaries are deduplicated across patches: identical values + (by fingerprint, or by summary content when no fingerprint is + available) share one coord_defs row. Only fingerprint-backed rows + are exposed as exact coordinate identity to merge planning. + """ + keys = list(defs_needed) + mapping: dict[str, int] = {} + for chunk, marks in self._iter_in_batches(keys): + found = self._fetch_df( + f"SELECT def_key, coord_def_id FROM coord_defs " + f"WHERE def_key IN ({marks})", + chunk, + ) + mapping.update( + zip(found["def_key"], (int(x) for x in found["coord_def_id"])) + ) + new_keys = [k for k in keys if k not in mapping] + next_id = self._next_id("coord_defs", "coord_def_id") + def_rows = [] + for key in new_keys: + c = defs_needed[key] + def_rows.append( + ( + next_id, + key, + c.coord_hash, + c.value_kind, + c.dtype, + c.length, + c.units, + c.min_num, + c.max_num, + c.step_num, + c.min_ns, + c.max_ns, + c.step_ns, + c.min_str, + c.max_str, + c.is_monotonic, + c.is_relative, + ) + ) + mapping[key] = next_id + next_id += 1 + self._bulk_insert("coord_defs", tuple(COORD_DEFS), def_rows) + return mapping + + def _bulk_insert(self, table: str, columns: tuple, rows: list) -> None: + """Insert many rows; engines override for faster bulk paths.""" + if not rows: + return + quoted = ", ".join(self.dialect.quote(c) for c in columns) + marks = self._placeholders(len(columns)) + sql = f"INSERT INTO {self.dialect.quote(table)} ({quoted}) VALUES ({marks})" + self._executemany(sql, rows) + + # --- writes ------------------------------------------------------ + + def mark_initial_update_done(self) -> None: + """Persist successful completion of a directory index's first update.""" + with self._transaction(): + self._execute( + "UPDATE meta_data SET last_indexed_ns = ?", + (time.time_ns(),), + ) + + def write_sources(self, records: list[SourceRecord]) -> None: + """ + Insert or replace sources and all dependent rows, atomically. + + Rows are batched per table (attrs grouped by column signature) so + columnar engines aren't punished by row-at-a-time inserts. + """ + with self._transaction(): + by_base: dict[str, list[str]] = {} + for record in records: + by_base.setdefault(record.base_uri or "", []).append(record.source_path) + # Ordering contract: a replaced source keeps its ordinal + # (first-occurrence position, dict-merge semantics) while new + # sources append after every existing position. Read both + # before the delete below discards them. + kept_ordinals = self._existing_ordinals(by_base) + max_df = self._fetch_df("SELECT max(ordinal) AS m FROM sources") + max_ordinal = max_df["m"].iloc[0] + next_ordinal = 0 if pd.isnull(max_ordinal) else int(max_ordinal) + 1 + for base_uri, paths in by_base.items(): + self._delete_by_paths(paths, base_uri=base_uri) + column_map, skip_units = self._ensure_attr_columns(records) + source_id = self._next_id("sources", "source_id") + patch_id = self._next_id("patches", "patch_id") + now = time.time_ns() + source_rows, patch_rows, link_rows = [], [], [] + defs_needed: dict[str, object] = {} + attr_groups: dict[tuple[str, ...], list] = {} + for record in records: + identity = (record.base_uri or "", record.source_path) + ordinal = kept_ordinals.get(identity) + if ordinal is None: + ordinal = next_ordinal + next_ordinal += 1 + source_rows.append( + ( + source_id, + record.base_uri or "", + record.source_path, + record.source_format, + record.format_version, + record.mtime_ns, + record.size_bytes, + now, + ordinal, + ) + ) + for patch in record.patches: + patch_rows.append( + ( + patch_id, + source_id, + patch.source_patch_id, + patch.n_dims, + patch.dims, + patch.shape, + patch.sample_count_total, + patch.time_min, + patch.time_max, + patch.time_step, + patch.distance_min, + patch.distance_max, + patch.distance_step, + ) + ) + attrs = patch.attrs + if skip_units: + attrs = { + name: tv + for name, tv in attrs.items() + if (name, tv.kind, tv.units) not in skip_units + } + columns = tuple( + column_map[(name, tv.kind)] for name, tv in attrs.items() + ) + attr_groups.setdefault(columns, []).append( + [patch_id, *(tv.value for tv in attrs.values())] + ) + for c in patch.coords: + key = c.def_key + defs_needed.setdefault(key, c) + link_rows.append((patch_id, c.coord_name, c.coord_dims, key)) + patch_id += 1 + source_id += 1 + self._bulk_insert("sources", tuple(SOURCES), source_rows) + self._bulk_insert("patches", tuple(PATCHES), patch_rows) + for columns, rows in attr_groups.items(): + self._bulk_insert("attrs", ("patch_id", *columns), rows) + def_ids = self._ensure_coord_defs(defs_needed) + self._bulk_insert( + "patch_coords", + tuple(PATCH_COORDS), + [(pid, name, dims, def_ids[key]) for pid, name, dims, key in link_rows], + ) + # meta_data.last_indexed_ns is the initial-update-complete + # marker; only mark_initial_update_done (after renumbering + # succeeds) may set it, or an interruption here would defeat + # the reopen recovery path. Per-source timestamps already + # live on the sources rows. + + # Batch size for IN (...) parameter lists; SQLite caps bound + # variables (32766 by default) so large replacements must chunk. + _in_clause_batch = 5000 + + @staticmethod + def _placeholders(count: int) -> str: + """Return a comma-separated run of ``count`` ``?`` bind markers.""" + return ", ".join("?" for _ in range(count)) + + def _iter_in_batches(self, items): + """ + Yield ``(chunk, marks)`` for an ``IN (...)`` list. + + Splitting on ``_in_clause_batch`` keeps each statement under + SQLite's bound-variable cap; ``marks`` is the placeholder run for + the chunk. + """ + batch = self._in_clause_batch + for start in range(0, len(items), batch): + chunk = items[start : start + batch] + yield chunk, self._placeholders(len(chunk)) + + def _existing_ordinals(self, by_base: dict[str, list[str]]) -> dict: + """Map (base_uri, source_path) -> ordinal for already-stored sources.""" + out: dict[tuple[str, str], int] = {} + for base_uri, paths in by_base.items(): + for chunk, marks in self._iter_in_batches(paths): + df = self._fetch_df( + f"SELECT source_path, ordinal FROM sources " + f"WHERE source_path IN ({marks}) AND base_uri = ?", + [*chunk, base_uri], + ) + for row in df.itertuples(): + if not pd.isnull(row.ordinal): + out[(base_uri, row.source_path)] = int(row.ordinal) + return out + + def renumber_ordinals_by_time(self) -> None: + """ + Renumber source ordinals into time order. + + The directory syncer owns its catalog's presentation order and + calls this after each sync so file archives keep their + conventional time-ordered iteration; sources are ordered by the + earliest patch time (sources without patches last), path as the + deterministic tiebreak. + """ + with self._transaction(): + # the WHERE clause skips rows whose ordinal is already + # correct, so a one-file sync of a large archive rewrites + # one row instead of churning the whole table through WAL + self._execute( + "WITH ranked AS (" + " SELECT s2.source_id AS sid, ROW_NUMBER() OVER (" + " ORDER BY t.min_time IS NULL, t.min_time, s2.source_path" + " ) - 1 AS rn" + " FROM sources s2 LEFT JOIN (" + " SELECT source_id, MIN(time_min) AS min_time" + " FROM patches GROUP BY source_id" + " ) t ON t.source_id = s2.source_id" + ")" + "UPDATE sources SET ordinal = (" + " SELECT rn FROM ranked WHERE sid = sources.source_id" + ") WHERE ordinal IS NOT (" + " SELECT rn FROM ranked WHERE sid = sources.source_id" + ")" + ) + + def _delete_by_paths(self, source_paths: list[str], base_uri: str = "") -> None: + """ + Delete sources by (base_uri, source_path) identity. + + The schema declares sources -> patches -> attrs/patch_coords with + ON DELETE CASCADE and the connection enables foreign keys, so + deleting the sources removes every dependent row. coord_defs are + intentionally left (they may orphan; a rebuild compacts them). + """ + if not source_paths: + return + for chunk, marks in self._iter_in_batches(source_paths): + self._execute( + f"DELETE FROM sources WHERE source_path IN ({marks}) AND base_uri = ?", + [*chunk, base_uri], + ) + + def delete_sources(self, source_paths: list[str], base_uri: str = "") -> None: + """Remove sources (identified by base_uri + path) and dependents.""" + with self._transaction(): + self._delete_by_paths(source_paths, base_uri=base_uri) + + # --- queries ----------------------------------------------------- + + def _query_context(self, query, order_by=None): + """ + Normalize a query (or several) and fetch the metadata SQL needs. + + Returns ``(queries, attr_meta, coord_meta)``; coord metadata is + only consulted for coord predicates and coord ordering, so the + (whole-relation DISTINCT) scan is skipped for attr-only/empty + queries. + """ + queries = _as_query_list(query if query is not None else Query()) + attr_meta = self._attr_meta() + coord_names = {name for q in queries for name in q.coords} + if order_by is not None and order_by[0] == "coord": + coord_names.add(order_by[1]) + coord_meta = self._coord_meta(coord_names) if coord_names else pd.DataFrame() + return queries, attr_meta, coord_meta + + def query(self, query=None, order_by=None, patch_ids=None) -> pd.DataFrame: + """Return the flat patch-row relation for a query (or several).""" + queries, attr_meta, coord_meta = self._query_context(query, order_by=order_by) + sql, params, residuals = build_sql( + queries, + self.dialect, + attr_meta, + coord_meta, + order_by=order_by, + patch_ids=patch_ids, + ) + df = self._fetch_df(sql, params) + df = self._flatten(df, attr_meta) + df = self._pivot_coords(df) + if residuals: + df = apply_residuals(df, residuals) + return df.reset_index(drop=True) + + def query_ids(self, query=None, order_by=None, patch_ids=None) -> list[int]: + """Return matching patch ids in presentation order (ids only).""" + queries, attr_meta, coord_meta = self._query_context(query, order_by=order_by) + sql, params, residuals = build_sql( + queries, + self.dialect, + attr_meta, + coord_meta, + order_by=order_by, + patch_ids=patch_ids, + ids_only=True, + ) + if residuals: + # regex residuals need string values; realize the relation + df = self.query(queries, order_by=order_by, patch_ids=patch_ids) + return [int(x) for x in df["patch_id"]] + return [int(x) for x in self._fetch_df(sql, params)["patch_id"]] + + def count(self, query=None, patch_ids=None) -> int: + """Count matching patches without projecting or pivoting rows.""" + queries, attr_meta, coord_meta = self._query_context(query) + sql, params, residuals = build_sql( + queries, + self.dialect, + attr_meta, + coord_meta, + count=True, + patch_ids=patch_ids, + ) + if not residuals: + return int(self._fetch_df(sql, params)["n"].iloc[0]) + # A regex residual must inspect string values, so a database count + # cannot resolve it; the full relation already applies the residual. + return len(self.query(queries, patch_ids=patch_ids)) + + def _fetch_in(self, base_sql: str, column: str, ids: list) -> pd.DataFrame: + """Fetch ``{base_sql} WHERE {column} IN ids``, batching large sets.""" + if not ids: + return self._fetch_df(f"{base_sql} WHERE 0") + frames = [] + for chunk, marks in self._iter_in_batches(ids): + frames.append( + self._fetch_df(f"{base_sql} WHERE {column} IN ({marks})", chunk) + ) + return pd.concat(frames, ignore_index=True) + + def export_records(self, patch_ids=None) -> list: + """ + Reconstruct source records, filtering by patch id in SQL. + + With patch_ids given, only those patches (and the sources, attrs, + coordinate links, and coordinate definitions they reference) are + fetched — O(selected membership), not O(total archive). The frames + are assembled into the backend-independent transfer format. + """ + from dascore.io.index.ingest import assemble_source_records + + if patch_ids is None: + sources = self._fetch_df("SELECT * FROM sources") + patches = self._fetch_df("SELECT * FROM patches") + attrs = self._fetch_df("SELECT * FROM attrs") + links = self._fetch_df("SELECT * FROM patch_coords") + defs = self._fetch_df("SELECT * FROM coord_defs") + else: + ids = [int(x) for x in patch_ids] + patches = self._fetch_in("SELECT * FROM patches", "patch_id", ids) + if patches.empty: + return [] + source_ids = [int(x) for x in patches["source_id"].unique()] + sources = self._fetch_in("SELECT * FROM sources", "source_id", source_ids) + attrs = self._fetch_in("SELECT * FROM attrs", "patch_id", ids) + links = self._fetch_in("SELECT * FROM patch_coords", "patch_id", ids) + def_ids = ( + [int(x) for x in links["coord_def_id"].unique()] + if not links.empty + else [] + ) + defs = self._fetch_in("SELECT * FROM coord_defs", "coord_def_id", def_ids) + return assemble_source_records( + sources, patches, attrs, links, defs, self._attr_meta() + ) + + def _flatten(self, df: pd.DataFrame, attr_meta: pd.DataFrame) -> pd.DataFrame: + """Post-process raw SQL output into the flat-relation contract.""" + out = df.copy() + # structural time columns: ns ints -> numpy time types (exactly) + for col, flavor in _TIME_COLS.items(): + if col in out: + out[col] = _ns_to_time(out[col], flavor) + # numeric envelopes: engines return object columns when all-NULL; + # downstream sorting needs float64 with NaN, never object None. + for col in ("distance_min", "distance_max", "distance_step"): + if col in out: + out[col] = pd.to_numeric(out[col]) + # typed attr columns -> original names (coalesce multi-kind attrs). + # Group the metadata once and drop every typed column in a single + # pass: per-name refiltering plus one-drop-per-column was ~O(A^2) + # metadata scans and frame copies for A dynamic attrs. + # Every name that could collide with a structural or envelope + # column (RESERVED_ATTR_COLUMNS and *_min/_max/_step) is refused at + # ingest, so a dynamic attr name never shadows an existing column + # here. + cols_to_drop: list[str] = [] + new_columns: dict[str, pd.Series] = {} + for name, rows in attr_meta.groupby("attr_name", sort=False): + kinds = set(rows["value_kind"]) + multi_kind = len(rows) > 1 + series = None + for row in rows.itertuples(): + if row.column_name not in out: + continue + col = out[row.column_name] + if row.value_kind == "time": + col = _ns_to_time(col, "datetime") + elif row.value_kind == "dur": + col = _ns_to_time(col, "timedelta") + elif row.value_kind == "bool": + col = col.astype("boolean") + if multi_kind: + # multi-kind attrs coalesce in object space; typed + # extension arrays refuse cross-dtype fills + col = col.astype(object).where(col.notna(), np.nan) + series = col if series is None else series.where(series.notna(), col) + cols_to_drop.append(row.column_name) + if series is not None: + if kinds == {"str"}: + # flat-contract convention: missing strings are "" + series = series.fillna("") + new_columns[name] = series + if cols_to_drop: + out = out.drop(columns=cols_to_drop) + for name, series in new_columns.items(): + out[name] = series + # flat-contract names for source columns + renames = { + "source_path": "path", + "source_format": "file_format", + "format_version": "file_version", + } + out = out.rename(columns=renames) + if "base_uri" in out: + has_base = out["base_uri"].notna() & (out["base_uri"] != "") + out.loc[has_base, "path"] = ( + out.loc[has_base, "base_uri"].str.rstrip("/") + + "/" + + out.loc[has_base, "path"] + ) + out = out.drop(columns=["base_uri"]) + return out.drop(columns=["source_id"], errors="ignore") + + @staticmethod + def _add_envelope_objects(coords: pd.DataFrame) -> pd.DataFrame: + """ + Add per-row envelope object columns (_env_min/_env_max/_env_step) + and the merge-identity _key column to the coord-link relation. + + Conversions run on whole columns: per-row scalar pd.to_datetime + calls cost ~40us each and dominated large realizations. + """ + kind = coords["value_kind"].to_numpy() + num_mask = kind == "num" + time_mask = kind == "time" + str_mask = ~(num_mask | time_mask) + # NULL means not relative; via to_numeric so object/float/int + # columns all coerce without pandas downcasting warnings. + relative = ( + pd.to_numeric(coords["is_relative"], errors="coerce") + .to_numpy(dtype="float64", na_value=0.0) + .astype(bool) + ) + + def _time_objects(ns_series: pd.Series, flavor: str) -> np.ndarray: + """Exact int-ns -> Timestamp/Timedelta objects (None for null).""" + series = _ns_to_time(ns_series, flavor) + return series.astype(object).where(series.notna(), None).to_numpy() + + fields = ( + ("_env_min", "min_num", "min_ns", "min_str"), + ("_env_max", "max_num", "max_ns", "max_str"), + ("_env_step", "step_num", "step_ns", None), + ) + for out_col, num_col, ns_col, str_col in fields: + values = np.empty(len(coords), dtype=object) + if num_mask.any(): + values[num_mask] = coords[num_col].to_numpy(dtype=object)[num_mask] + if str_col is not None and str_mask.any(): + values[str_mask] = coords[str_col].to_numpy(dtype=object)[str_mask] + # absolute times are datetimes, relative ones timedeltas; steps + # are timedeltas either way. + time_flavors = ( + ((time_mask & ~relative), "datetime"), + ((time_mask & relative), "timedelta"), + ) + if str_col is None: + time_flavors = ((time_mask, "timedelta"),) + for mask, flavor in time_flavors: + if mask.any(): + values[mask] = _time_objects(coords[ns_col][mask], flavor) + coords[out_col] = values + # Summary-only definitions are useful for indexing/dedup but cannot + # prove coordinate value identity for merge grouping. + coords["_key"] = coords["def_key"].where(coords["fingerprint"].notna(), None) + return coords + + def _pivot_coords(self, out: pd.DataFrame) -> pd.DataFrame: + """ + Add per-coord envelope columns to the flat relation. + + Emits {name}_min/{name}_max/{name}_step for every coord in the + result beyond the time/distance envelopes already cached on + patches (memory-spool parity: chunking on any dim needs these), + plus a private _{name}_def_key column for every coord — the + globally-stable coordinate identity future chunk/merge grouping + uses (private so it does not yet participate in merge + compatibility comparisons). + """ + if out.empty or "patch_id" not in out.columns: + return out + ids = out["patch_id"].tolist() + link_sql = ( + "SELECT pc.patch_id, pc.coord_name, cd.def_key, cd.fingerprint, " + "cd.value_kind, cd.is_relative, cd.units, cd.min_num, cd.max_num, " + "cd.step_num, cd.min_ns, cd.max_ns, cd.step_ns, " + "cd.min_str, cd.max_str " + "FROM patch_coords pc " + "JOIN coord_defs cd ON cd.coord_def_id = pc.coord_def_id" + ) + n_patches = self._fetch_df("SELECT count(*) AS n FROM patches")["n"].iloc[0] + if len(ids) * 4 >= n_patches: + # Most patches selected: one scan plus a pandas filter beats + # many batched IN queries and their frame concatenation. + coords = self._fetch_df(link_sql) + coords = coords[coords["patch_id"].isin(set(ids))].reset_index(drop=True) + else: + coords = self._fetch_in(link_sql, "pc.patch_id", ids) + if coords.empty: + return out + coords = self._add_envelope_objects(coords) + for name, group in coords.groupby("coord_name"): + pids = group["patch_id"] + # last row wins for duplicate patch ids, like the mapping loop + # this replaces. + keys = dict(zip(pids, group["_key"])) + mins = dict(zip(pids, group["_env_min"])) + maxs = dict(zip(pids, group["_env_max"])) + steps = dict(zip(pids, group["_env_step"])) + units = dict(zip(pids, group["units"])) + out[f"_{name}_def_key"] = out["patch_id"].map(keys) + # canonical (base) units: numeric envelopes are stored SI, so + # this is the dimensionality marker chunk partitioning needs + out[f"_{name}_units"] = out["patch_id"].map(units) + kinds = set(group["value_kind"]) + # time/distance envelopes already live on patches... + if name in ("time", "distance"): + col = f"{name}_min" + # ...but relative-time patches leave them NULL by design; + # when the whole result is relative, serve timedelta + # envelopes so chunking on relative time works (#553). + if col in out.columns and out[col].isnull().all() and mins: + out[f"{name}_min"] = out["patch_id"].map(mins) + out[f"{name}_max"] = out["patch_id"].map(maxs) + out[f"{name}_step"] = out["patch_id"].map(steps) + continue + out[f"{name}_min"] = out["patch_id"].map(mins) + out[f"{name}_max"] = out["patch_id"].map(maxs) + out[f"{name}_step"] = out["patch_id"].map(steps) + if kinds == {"num"}: # object-None -> float NaN for sorting + for suffix in ("_min", "_max", "_step"): + out[f"{name}{suffix}"] = pd.to_numeric(out[f"{name}{suffix}"]) + return out + + # --- introspection ----------------------------------------------- + + def get_sources(self) -> pd.DataFrame: + """Return the sources table.""" + return self._fetch_df("SELECT * FROM sources") + + def source_stats(self) -> pd.DataFrame: + """Return only the columns incremental change detection needs.""" + return self._fetch_df("SELECT source_path, mtime_ns, size_bytes FROM sources") + + def get_metadata(self) -> dict: + """Return index-level metadata.""" + return self._fetch_df("SELECT * FROM meta_data").iloc[0].to_dict() + + def attr_names(self) -> set[str]: + """Return original attr names known to the index.""" + return set(self._attr_meta()["attr_name"]) + + def coord_names(self) -> set[str]: + """Return coord names known to the index.""" + df = self._fetch_df("SELECT DISTINCT coord_name FROM patch_coords") + return set(df["coord_name"]) + + def coord_dims_map(self) -> dict[str, str]: + """Return each coord name's dims string (first observed wins).""" + df = self._fetch_df("SELECT DISTINCT coord_name, coord_dims FROM patch_coords") + out: dict[str, str] = {} + for name, dims in zip(df["coord_name"], df["coord_dims"]): + out.setdefault(str(name), str(dims)) + return out + + +def resolve_query( + backend: AbstractIndexBackend, _attrs=None, _coords=None, **kwargs +) -> Query: + """ + Resolve bare kwargs into a Query: attrs first, then coords. + + Implements section 1 of the selector spec; raises on unknown names or + names supplied in more than one namespace. + """ + from dascore.io.index.query import InvalidSpoolQueryError + + def _drop_noops(mapping: dict) -> dict: + """Bare None/... selectors are no-ops, matching Patch.select.""" + return {k: v for k, v in mapping.items() if v is not None and v is not Ellipsis} + + def _shape_coord_selector(name: str, value): + """ + Normalize one coordinate selector to the spec's accepted shapes. + + Coordinates select by range — a (start, stop) tuple or slice + with None/... open ends (a 2-element list is the legacy range + form) — or by a patch-local boolean mask. Scalars and value + membership have no exact patch-level meaning and raise here, + eagerly, rather than failing when a patch is materialized. + """ + if value is None or value is Ellipsis: + return value + is_bool_array = (isinstance(value, np.ndarray) and value.dtype == np.bool_) or ( + isinstance(value, list) + and value + and all(isinstance(x, bool | np.bool_) for x in value) + ) + if is_bool_array: + # A sample mask is positional/absolute, so it is only defined + # when every patch shares the coordinate's size — a guarantee + # spools never make — and it can never reduce file reads. + msg = ( + f"Coordinate {name!r} no longer accepts boolean sample " + "masks at the spool level; apply them per patch, e.g. " + "spool.map(lambda p: p.select(...)). Boolean arrays over " + "patches (spool[mask]) still select membership." + ) + raise InvalidSpoolQueryError(msg) + if isinstance(value, tuple | list): + # range-like: a (start, stop) pair (2-element list is the + # legacy range form). Wrong arity is a malformed range. + if len(value) != 2: + msg = f"Coordinate range for {name!r} must be a length 2 sequence." + raise ParameterError(msg) + # canonicalize the open-end sentinel so equivalent selections + # (None vs ...) stay equivalent downstream (e.g. spool __eq__) + return tuple(None if v is Ellipsis else v for v in value) + msg = ( + f"Coordinate {name!r} accepts range selectors (a (start, stop) " + "tuple or slice, None/... for open ends); scalar, membership, " + f"and boolean-mask values are not supported. Got {value!r}." + ) + raise InvalidSpoolQueryError(msg) + + attrs, coords = resolve_selector_namespaces( + backend.attr_names(), + backend.coord_names(), + _attrs=_attrs, + _coords=_coords, + kwargs=kwargs, + ) + coords = {k: _shape_coord_selector(k, v) for k, v in coords.items()} + return Query(attrs=_drop_noops(attrs), coords=_drop_noops(coords)) + + +def get_backend(path: str | Path) -> AbstractIndexBackend: + """Create the SQLite spool-index backend at path.""" + from dascore.io.index.lite import SQLiteBackend + + return SQLiteBackend(path) diff --git a/dascore/io/index/catalog.py b/dascore/io/index/catalog.py new file mode 100644 index 000000000..7a4d29586 --- /dev/null +++ b/dascore/io/index/catalog.py @@ -0,0 +1,978 @@ +""" +PatchCatalog: one metadata engine for every spool type. + +The catalog owns the index tables (through a backend) and the composed +selection state; a resolver turns flat-relation rows into patches +(from files via dc.read, or from a live registry for in-memory spools); +a syncer (the directory indexer) keeps directory-backed catalogs in step +with the filesystem. See the spool index design doc and discussion #648. + +Laziness contract: creating a catalog from patches does no metadata work +until the first metadata operation (select/len/iteration), because +backend bootstrap costs ~10s of ms while holding a patch list is free. +Selection composes Query predicates without running SQL; realization +(len, to_df, iteration) runs exactly one query per view. +""" + +from __future__ import annotations + +import abc +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from pathlib import Path + +import numpy as np +import pandas as pd + +import dascore as dc +from dascore.constants import PROGRESS_LEVELS +from dascore.core.summary import normalize_source_patch_id +from dascore.exceptions import MissingPatchError +from dascore.io.index.backend import get_backend, resolve_query +from dascore.io.index.ingest import SourceRecord, patch_record +from dascore.io.index.query import ( + InvalidSpoolQueryError, + Query, +) +from dascore.io.index.schema import SPOOL_HIDDEN_COLUMNS +from dascore.utils.misc import is_range +from dascore.utils.paths import is_memory_uri +from dascore.utils.pd import adjust_segments, relative_ranges_to_absolute + +# Directory archives present in per-patch time order (source ordinals +# alone cannot interleave multi-patch files); ordinal and patch id stay +# the deterministic tiebreak inside the ORDER BY. +_DIRECTORY_ORDER = ("coord", "time", True) + + +class _CanonicalRange: + """ + A numeric coordinate range resolved to canonical SI magnitudes. + + The exact per-patch re-select defers its representation until the + target patch is known: unit-bearing coordinates get quantities + (`Patch.select` converts them to native units), unitless + coordinates get the bare magnitudes. A single eager form cannot + serve both — raw numbers trim the wrong physical interval on non-SI + patches, quantities break unitless coordinates. + + ``units`` records the query's own base unit when the original + bounds carried one, so the residual preserves the query's + dimensionality instead of adopting each patch coordinate's — a + metre query must never trim a seconds coordinate as 1-2 s. + """ + + __slots__ = ("magnitudes", "units") + + def __init__(self, magnitudes: tuple, units: str | None = None): + self.magnitudes = magnitudes + self.units = units + + def __eq__(self, other) -> bool: + """Value equality so equal selections compare equal (spool __eq__).""" + if not isinstance(other, _CanonicalRange): + return NotImplemented + return (self.magnitudes, self.units) == (other.magnitudes, other.units) + + def __hash__(self) -> int: + return hash((self.magnitudes, self.units)) + + def for_patch_coord(self, coord) -> tuple: + """Return the range in the representation this coord needs.""" + from dascore.units import get_quantity + + coord_units = getattr(coord, "units", None) + if coord_units is None: + # unitless coords: bare canonical magnitudes (documented policy) + return self.magnitudes + # a unit-bearing query keeps its own dimensionality; a bare + # numeric query means canonical SI in the coord's dimension + base = ( + get_quantity(self.units) + if self.units is not None + else get_quantity(str(coord_units)).to_base_units().units + ) + return tuple(None if mag is None else mag * base for mag in self.magnitudes) + + +def _canonical_range(value) -> _CanonicalRange | None: + """Return the canonical SI form of a numeric range, or None.""" + if not is_range(value): + return None + magnitudes = [] + units = None + for bound in value: + if bound is None or bound is Ellipsis: + magnitudes.append(None) + elif hasattr(bound, "units"): # pint quantity -> SI magnitude + base = bound.to_base_units() + magnitudes.append(float(base.magnitude)) + units = str(base.units) + elif isinstance(bound, bool | np.bool_): + return None + elif isinstance(bound, int | float | np.integer | np.floating): + magnitudes.append(float(bound)) + else: # datetimes, strings: not a numeric range + return None + if all(mag is None for mag in magnitudes): + return None + return _CanonicalRange(tuple(magnitudes), units) + + +def _envelope_range(value): + """Return a range with quantity bounds as SI magnitudes. + + Stored envelope columns are canonical SI, so the presented-envelope + adjustment needs bare magnitudes; non-numeric ranges pass through. + """ + canonical = _canonical_range(value) + return value if canonical is None else canonical.magnitudes + + +def _canonical_coord_selectors(backend, coords: dict) -> tuple[dict, dict]: + """ + Split coordinate selectors into query-side and residual-side forms. + + The query side keeps the *original* values: the SQL builder coerces + quantities itself and needs their units to constrain candidacy to + dimensionally compatible coordinate definitions (a metre query must + exclude — or raise on — a seconds coordinate, never trim it). + The residual keeps the range as a `_CanonicalRange` (canonical SI + magnitudes plus the query's base unit) so each patch decides its + own representation at load time, which keeps mixed unitful/unitless + populations correct. + + Selectors on non-numeric coordinates (time ranges, string ranges) + pass through unchanged. + """ + meta = backend._coord_meta(set(coords)) + numeric = set(meta.loc[meta["value_kind"] == "num", "coord_name"]) + query_coords, residual_coords = {}, {} + for name, value in coords.items(): + canonical = _canonical_range(value) if name in numeric else None + query_coords[name] = value + residual_coords[name] = value if canonical is None else canonical + return query_coords, residual_coords + + +def _row_source_patch_id(row: Mapping) -> str: + """Return the row's source_patch_id as a normalized string.""" + return normalize_source_patch_id(row.get("source_patch_id")) + + +def apply_exact_residuals(patch: dc.Patch, residuals) -> dc.Patch: + """ + Apply a view's exact residual selections to a loaded patch. + + Shared by catalog row resolution and plan-member loading so the + two-stage select contract has exactly one implementation. + """ + for coords, samples in residuals: + coord_map = patch.coords.coord_map + usable = { + k: ( + v.for_patch_coord(coord_map[k]) if isinstance(v, _CanonicalRange) else v + ) + for k, v in coords.items() + if k in coord_map + } + if usable: + # residual bounds are already absolute (relative queries + # resolve to absolute before the residual is recorded). + patch = patch.select(**usable, samples=samples, relative=False) + return patch + + +class PatchResolver(abc.ABC): + """Turn one flat-relation row into a Patch.""" + + @abc.abstractmethod + def resolve(self, row: Mapping, **trim) -> dc.Patch: + """ + Return the patch for a row. + + Trim kwargs are hints (a slice of the plan, not the query): + implementations may use them to read less, but exact trimming is + re-applied above, so ignoring them is slower, never wrong. + """ + + def live_entries(self) -> Mapping[str, dc.Patch]: + """Return the live patches this resolver serves (path -> patch).""" + return {} + + +class LiveResolver(PatchResolver): + """ + Serve patches from an in-memory registry. + + The registry *is* the store for live catalogs: a dict from each + patch's synthetic path (`memorypatch://`) to the patch + itself. Dict construction deduplicates identical patch instances + (set semantics by lineage), and merging catalogs unions the dicts. + """ + + def __init__(self, patches: Sequence[dc.Patch] = ()): + self._registry: dict[str, dc.Patch] = { + _patch_path(patch): patch for patch in patches + } + + def live_entries(self) -> Mapping[str, dc.Patch]: + """Return the live patch registry.""" + return self._registry + + def resolve(self, row: Mapping, **trim) -> dc.Patch: + """Look the patch up; live patches ignore trim hints.""" + path = str(row["path"]) + try: + return self._registry[path] + except KeyError: + msg = ( + f"The in-memory patch for {path} is not available in this " + "session (e.g. the row came from a reopened index). " + "In-memory patches only persist by writing them to files." + ) + raise MissingPatchError(msg) from None + + +def _patch_path(patch: dc.Patch) -> str: + """Return the synthetic source path identifying a live patch.""" + return f"memorypatch://{patch._instance_id}" + + +class FileResolver(PatchResolver): + """Load patches through dc.read; remoteness is the path layer's job.""" + + def __init__(self, root: Path | str | None = None): + self._root = Path(root) if root is not None else None + + def _read(self, path, row: Mapping, trim: dict, source_patch_id: str): + """ + Read one row's patch through dc.read. + + The recorded format/version are forwarded so dc.read skips format + probing; it reads the file exactly once (an earlier fast path that + called the reader directly re-read the file whenever the reader + returned patches lazily). + """ + id_kwargs = {"source_patch_id": source_patch_id} if source_patch_id else {} + kwargs = {"path": path} + if row.get("file_format"): + kwargs["file_format"] = row["file_format"] + if row.get("file_version"): + kwargs["file_version"] = row["file_version"] + return dc.read(**kwargs, **id_kwargs, **trim) + + def resolve(self, row: Mapping, **trim) -> dc.Patch: + """Read the patch, passing range trims down as read hints.""" + from dascore.io.core import _resolve_read_spool + + path = row["path"] + # relative paths resolve against the catalog root; URIs and + # absolute paths pass through untouched. + if self._root is not None and "://" not in str(path): + if not Path(path).is_absolute(): + path = self._root / path + source_patch_id = _row_source_patch_id(row) + if source_patch_id.isdigit(): + # Positional (synthesized) ids index the full source read; a + # trimmed read would shift or drop patches and bind the wrong + # one, so these rows read the whole source. + trim = {} + spool = self._read(path, row, trim, source_patch_id) + return _resolve_read_spool(spool, source_patch_id) + + +class CompositeResolver(PatchResolver): + """ + Route rows to a live registry, a plan, or the filesystem by scheme. + + Union catalogs mix file-backed rows (absolute paths), in-memory rows + (memory:// paths), and plan-output rows (plan://token/... paths); + this resolver dispatches accordingly. + """ + + def __init__(self): + self.live = LiveResolver() + self.file = FileResolver(root=None) + # plan:/// prefix -> the PlanResolver that owns it + self.plans: dict[str, PatchResolver] = {} + + def live_entries(self) -> Mapping[str, dc.Patch]: + """Return the merged live patch registry.""" + return self.live._registry + + def plan_entries(self) -> Mapping[str, PatchResolver]: + """Return the plan-prefix routing table.""" + return self.plans + + def absorb(self, resolver: PatchResolver, paths=None) -> None: + """ + Take over another resolver's live and plan entries. + + ``paths`` restricts absorption to the given synthetic paths + (the entries a transfer actually references); None takes all. + """ + entries = resolver.live_entries() + if paths is not None: + entries = {k: v for k, v in entries.items() if k in paths} + self.live._registry.update(entries) + plans = getattr(resolver, "plan_entries", dict)() + if paths is not None: + plans = { + prefix: plan + for prefix, plan in plans.items() + if any(str(p).startswith(prefix) for p in paths) + } + self.plans.update(plans) + + def resolve(self, row: Mapping, **trim) -> dc.Patch: + """Dispatch by scheme: live registry, plan, or file reader.""" + path = str(row.get("path", "")) + if is_memory_uri(path): + return self.live.resolve(row, **trim) + for prefix, plan in self.plans.items(): + if path.startswith(prefix): + return plan.resolve(row, **trim) + return self.file.resolve(row, **trim) + + +def _live_records(registry: Mapping[str, dc.Patch]): + """Build source records for live patches keyed by their identity.""" + records = [] + for path, patch in registry.items(): + # patch.summary is a cached_property: reuse fingerprints and + # summaries the patch already computed instead of rebuilding. + record = patch_record(patch.summary) + records.append( + SourceRecord( + source_path=path, + source_format="memory", + format_version="", + patches=(record,), + ) + ) + return records + + +def _absolutize_record(record, root): + """Return a source record whose relative path is resolved against root.""" + path = record.source_path + if "://" in path or Path(path).is_absolute(): + return record + resolved = str(Path(root) / (path if path != "." else "")) + return replace(record, source_path=resolved, base_uri=None) + + +def _membership_resolver( + resolver: PatchResolver, keep: dict, paths=() +) -> PatchResolver: + """ + Return a copy of resolver whose live registry holds only `keep`. + + ``paths`` are the synthetic paths the view's rows reference: plan + routes serving any of them must survive the restriction or mixed + planned/live views lose their plan-backed rows on serialization. + """ + if isinstance(resolver, LiveResolver): + out = LiveResolver() + out._registry = dict(keep) + return out + out = CompositeResolver() + out.live._registry = dict(keep) + plans = getattr(resolver, "plan_entries", dict)() + out.plans = { + prefix: plan + for prefix, plan in plans.items() + if any(str(p).startswith(prefix) for p in paths) + } + return out + + +def _merge_source_records(existing, new): + """ + Merge two partial records for the same source. + + Union members export only their selected patches, so two members can + hold disjoint (or overlapping) slices of one multi-patch file. The + merged record unions the patch lists by source_patch_id: a patch + keeps its first-occurrence position, a duplicate identity takes the + last occurrence's metadata (dict-merge semantics, matching the + ordering contract), and the source-level metadata (mtime, size) + comes from the last record. + """ + if existing is None: + return new + patches = {p.source_patch_id: p for p in existing.patches} + patches.update({p.source_patch_id: p for p in new.patches}) + return replace(new, patches=tuple(patches.values())) + + +@dataclass +class _CatalogRevision: + """Shared mutation revision for live catalog views.""" + + value: int = 0 + + +# sentinel: _view keeps the current order/ids spec unless told otherwise +_KEEP = object() + + +class PatchCatalog: + """ + Query-composable metadata catalog over the spool index tables. + + Instances are lightweight views: `select` returns a new catalog + sharing the backend and resolver with composed predicates. Mutation + (`add`, `update`, `remove`) is only allowed on the root view. + """ + + def __init__( + self, + *, + backend=None, + resolver: PatchResolver | None = None, + syncer=None, + queries: tuple[Query, ...] = (), + residuals: tuple[tuple[dict, bool], ...] = (), + revision: _CatalogRevision | None = None, + order: tuple | None = None, + ids: tuple | None = None, + default_order: tuple | None = None, + ): + self._backend = backend + self.resolver = resolver + self._syncer = syncer + self._queries = tuple(queries) + self._residuals = tuple(residuals) + # presentation specs (D2): an order override ("attr"|"coord", + # name, ascending) and/or an ordered patch-id membership + self._order = order + # the catalog's own presentation contract when no user order is + # set (directory archives present in per-patch time order — + # source ordinals alone cannot interleave multi-patch files). + # Not view state: a root with a default order still updates. + self._default_order = default_order + self._ids = None if ids is None else tuple(int(x) for x in ids) + self._revision = revision or _CatalogRevision() + self._df_cache: pd.DataFrame | None = None + self._df_cache_revision = -1 + self._live_cache: tuple | None = None + self._live_cache_revision = -1 + # Source records for rebuilding an in-memory backend (set by + # __getstate__ so pickled catalogs survive losing the connection). + self._rebuild_records: tuple = () + + # --- construction ------------------------------------------------- + + @classmethod + def from_patches(cls, patches: Sequence[dc.Patch] = ()) -> PatchCatalog: + """ + Catalog over live patches. No backend work happens until the + first metadata operation. + + The resolver's registry is the store; identical patch instances + collapse to a single entry (set semantics by lineage). + """ + return cls(resolver=LiveResolver(patches)) + + @classmethod + def union(cls, catalogs: Sequence[PatchCatalog]) -> PatchCatalog: + """ + Materialize several catalogs into one in-memory catalog. + + Metadata rows are merged table-to-table (coord definitions + deduplicate by def key); file-backed rows get absolute paths so + members with different roots coexist, and the same source + appearing in several members keeps a single entry (last one + wins). For catalog views, only the selected patches transfer — + note this respects row membership, not range trims; re-select + on the result for exact envelopes. + """ + resolver = CompositeResolver() + out = cls(resolver=resolver) + backend = out.backend + # Collect and merge every member's records before writing: + # write_sources replaces at (base_uri, source_path) grain, so + # partial records for the same source — two members selecting + # different patches of one multi-patch file — must merge into a + # complete record or the later write would delete the earlier + # member's patches. Dict insertion order keeps first-occurrence + # position; the merge keeps last-occurrence metadata. + merged: dict[tuple, SourceRecord] = {} + for member in catalogs: + catalog, patch_ids = member if isinstance(member, tuple) else (member, None) + if patch_ids is None and catalog.is_view: + patch_ids = catalog.to_df()["_patch_id"].tolist() + records = catalog.backend.export_records(patch_ids=patch_ids) + root = getattr(catalog.resolver, "_root", None) + if root is not None: + records = [_absolutize_record(x, root) for x in records] + for record in records: + identity = (record.base_uri or "", record.source_path) + merged[identity] = _merge_source_records(merged.get(identity), record) + # only the live entries this member actually transfers ride + # along; the rest of the registry stays with its own catalog + member_paths = {record.source_path for record in records} + resolver.absorb(catalog.resolver, paths=member_paths) + backend.write_sources(list(merged.values())) + out._invalidate() + return out + + @classmethod + def from_file( + cls, + path: str | Path, + file_format: str | None = None, + file_version: str | None = None, + ) -> PatchCatalog: + """ + Catalog over a single fiber file. + + The file is scanned eagerly (one row per contained patch) into an + in-memory backend; patches load through the file resolver on + demand. There is no syncer — a changed file needs a new catalog. + """ + from dascore.io.index.ingest import summaries_to_records + + summaries = dc.scan( + path, file_format=file_format, file_version=file_version, progress=None + ) + records = summaries_to_records(summaries) + out = cls(resolver=FileResolver()) + out.backend.write_sources(records) + out._invalidate() + return out + + @classmethod + def from_directory( + cls, + path: str | Path, + index_path: str | Path | None = None, + ) -> PatchCatalog: + """Catalog over a directory of fiber files.""" + from dascore.io.index.indexer import DBDirectoryIndexer + + syncer = DBDirectoryIndexer(path, index_path=index_path) + return cls( + backend=syncer._backend, + resolver=FileResolver(root=syncer.path), + syncer=syncer, + default_order=_DIRECTORY_ORDER, + ) + + # --- internals ------------------------------------------------------ + + @property + def backend(self): + """ + The index backend, bootstrapping lazily on first use. + + Every metadata operation funnels through here, so this is also + where a brand-new directory index gets its one automatic update. + """ + if self._backend is None: + if self._syncer is not None: + # Directory catalogs re-adopt the (unpickled) syncer's + # backend; both must keep sharing one connection. + self._backend = self._syncer._backend + else: + self._backend = get_backend(":memory:") + if self._rebuild_records: + self._backend.write_sources(list(self._rebuild_records)) + self._rebuild_records = () + elif registry := getattr(self.resolver, "_registry", None): + self._backend.write_sources(_live_records(registry)) + if self._syncer is not None and self._syncer.ensure_updated(): + self._invalidate() + return self._backend + + def __getstate__(self) -> dict: + """ + Pickle without the live DB connection. + + In-memory backends (live and union catalogs) ride along as + source records and are re-ingested on next use; the resolver + registry (the store for live patches) pickles with its patches. + Directory catalogs rebuild from their index file instead. + """ + state = dict(self.__dict__) + state["_backend"] = None + # In-memory backends are rebuilt on the other side with FRESH + # patch ids, so a stored id membership would bind to the wrong + # rows. Restrict the rebuilt content to the current membership + # instead (records/registry in presentation order, so re-ingest + # ordinals preserve it) and drop the id spec; the syncer case + # reopens the same database file, where ids stay valid. + rebuilt_membership = self._syncer is None and self._ids is not None + if rebuilt_membership: + state["_ids"] = None + # Live catalogs rebuild from their registry without touching the + # connection (which may belong to another thread during pickling); + # other in-memory catalogs (e.g. unions) capture their rows. + needs_records = ( + self._backend is not None + and self._syncer is None + and not isinstance(self.resolver, LiveResolver) + ) + if needs_records: + patch_ids = self._ids if rebuilt_membership else None + state["_rebuild_records"] = tuple( + self._backend.export_records(patch_ids=patch_ids) + ) + # A view shares the root's resolver, but must not drag the whole + # live registry across the wire: keep only the entries its rows + # reference (a one-patch view of an N-patch spool serializes one + # patch, not N — the payload Spool.map ships per task) — in + # presentation order, so a rebuilt registry keeps the view's + # ordering. + if self.is_view and self.resolver.live_entries(): + df = self.to_df() + paths = list(dict.fromkeys(df["path"].astype(str))) + entries = self.resolver.live_entries() + keep = {k: entries[k] for k in paths if k in entries} + state["resolver"] = _membership_resolver(self.resolver, keep, paths) + return state + + def _view(self, queries, residuals, order=_KEEP, ids=_KEEP) -> PatchCatalog: + out = PatchCatalog( + backend=self.backend, + resolver=self.resolver, + syncer=self._syncer, + queries=queries, + residuals=residuals, + revision=self._revision, + order=self._order if order is _KEEP else order, + ids=self._ids if ids is _KEEP else ids, + default_order=self._default_order, + ) + return out + + def order_by(self, attribute: str, ascending: bool = True) -> PatchCatalog: + """ + Return a view presenting rows ordered by an attribute or coord. + + A lazy presentation spec (D2): realization adds ORDER BY with + the ordinal contract as the deterministic tiebreak; no rows are + copied and no relation is realized here. + """ + name = str(attribute) + coords = self.backend.coord_names() + if name in coords: + spec = ("coord", name, ascending) + elif name in self.backend.attr_names(): + spec = ("attr", name, ascending) + elif name.endswith("_min") and name.removesuffix("_min") in coords: + spec = ("coord", name.removesuffix("_min"), ascending) + else: + msg = "Invalid attribute. Please use a valid attribute such as: 'time'" + raise IndexError(msg) + return self._view(self._queries, self._residuals, order=spec) + + @property + def _effective_order(self) -> tuple | None: + """The presentation order: a user order spec, else the default.""" + return self._order if self._order is not None else self._default_order + + def _ordered_ids(self) -> tuple[int, ...]: + """The view's patch ids in presentation order (ids only, cheap).""" + if self._ids is not None and self._order is None: + return self._ids + return tuple( + self.backend.query_ids( + list(self._queries) or None, + order_by=self._effective_order, + patch_ids=self._ids, + ) + ) + + def window(self, item: slice) -> PatchCatalog: + """ + Return a view restricted to a slice of the presented rows. + + Membership realizes as an ordered id list (ids only — never the + flat relation); subsequent selections compose within the window + per the D2 rules. + """ + ids = self._ordered_ids()[item] + return self._view(self._queries, self._residuals, ids=tuple(ids)) + + def restrict(self, indices) -> PatchCatalog: + """ + Return a view keeping the presented rows an array selects. + + ``indices`` is a boolean mask over rows or an array of integer + positions (order-preserving; duplicate positions collapse to + one row, matching the spool's set semantics). + """ + ids = np.asarray(self._ordered_ids()) + picked = ids[np.asarray(indices)] + deduped = tuple(dict.fromkeys(int(x) for x in picked)) + return self._view(self._queries, self._residuals, ids=deduped) + + def _invalidate(self) -> None: + self._revision.value += 1 + self._df_cache = None + self._df_cache_revision = -1 + self._live_cache = None + self._live_cache_revision = -1 + + def _cold_live_values(self) -> tuple | None: + """ + The patches, in registry (construction) order, when the registry + alone defines contents — a root live catalog whose backend was + never realized. None whenever the registry is not authoritative. + + This keeps len/iteration/indexing on freshly-built patch-list + spools allocation-free: no ingest, no SQL, no flat relation. + """ + cold = ( + self._backend is None + and self._syncer is None + and not self.is_view + and not self._rebuild_records + and isinstance(self.resolver, LiveResolver) + ) + if not cold: + return None + if ( + self._live_cache is None + or self._live_cache_revision != self._revision.value + ): + self._live_cache = tuple(self.resolver.live_entries().values()) + self._live_cache_revision = self._revision.value + return self._live_cache + + def __deepcopy__(self, memo) -> PatchCatalog: + """ + Derived spools share the catalog (live registry + connection). + + Spool copies its state on select/chunk; catalog state + is read-shared, matching the single-writer model. + """ + return self + + @property + def is_view(self) -> bool: + """True when this catalog carries selection or presentation state.""" + return bool( + self._queries + or self._residuals + or self._order is not None + or self._ids is not None + ) + + def _require_root(self, operation: str) -> None: + if self.is_view: + msg = f"{operation} is only allowed on a root catalog, not a view." + raise InvalidSpoolQueryError(msg) + + # --- selection ------------------------------------------------------ + + def select( + self, + *, + _attrs: dict | None = None, + _coords: dict | None = None, + samples: bool = False, + relative: bool = False, + **kwargs, + ) -> PatchCatalog: + """ + Compose a selection; validation is eager, execution is lazy. + + samples=True selectors are patch-local (never index predicates); + relative=True bounds resolve against the current view's global + envelope, then behave as absolute ranges. + """ + query = resolve_query(self.backend, _attrs=_attrs, _coords=_coords, **kwargs) + if samples: + if query.attrs: + msg = ( + "samples=True selections are coordinate-only; got attrs " + f"{sorted(query.attrs)}." + ) + raise InvalidSpoolQueryError(msg) + residual = (dict(query.coords), True) + return self._view(self._queries, (*self._residuals, residual)) + if relative and query.coords: + query = Query( + attrs=query.attrs, + coords=self._relative_to_absolute(query.coords), + ) + # coord range predicates are re-applied exactly at patch load; + # the residual carries canonical quantities so per-patch native + # units are respected while the query side stays SI. + residuals = self._residuals + if query.coords: + si_coords, residual_coords = _canonical_coord_selectors( + self.backend, query.coords + ) + query = Query(attrs=query.attrs, coords=si_coords) + residuals = (*residuals, (residual_coords, False)) + return self._view((*self._queries, query), residuals) + + def _relative_to_absolute(self, kwargs: dict) -> dict: + """Resolve relative bounds against the view's global envelopes.""" + return relative_ranges_to_absolute(self.to_df(), kwargs) + + # --- realization ------------------------------------------------------ + + def to_df(self) -> pd.DataFrame: + """ + The spool-facing flat patch-row relation under the selection. + + Unique-per-patch structural columns (patch_id and friends) are + hidden or renamed private so chunk merge-compatibility (which + compares all non-private columns) is not spuriously blocked. + """ + if self._df_cache is None or self._df_cache_revision != self._revision.value: + df = self.backend.query( + list(self._queries) or None, + order_by=self._effective_order, + patch_ids=self._ids, + ) + if self._ids is not None and self._order is None: + # id membership presents in its own (window/array) order + position = {pid: i for i, pid in enumerate(self._ids)} + df = df.sort_values( + "patch_id", key=lambda s: s.map(position), kind="stable" + ).reset_index(drop=True) + df = df.drop(columns=list(SPOOL_HIDDEN_COLUMNS), errors="ignore").rename( + columns={"patch_id": "_patch_id"} + ) + # SQL identifies overlapping source patches. Expose the selected + # envelopes, matching spool.get_contents() and the exact trim + # applied when each patch is materialized. Each pass copies the + # frame, so disjoint-name range sets collapse into one pass. + range_dicts = [ + ranges + for query in self._queries + if ( + ranges := { + name: _envelope_range(value) + for name, value in query.coords.items() + if is_range(value) + } + ) + ] + names = [name for ranges in range_dicts for name in ranges] + if range_dicts and len(set(names)) == len(names): + range_dicts = [{k: v for d in range_dicts for k, v in d.items()}] + for ranges in range_dicts: + df = adjust_segments(df, ignore_bad_kwargs=True, **ranges) + self._df_cache = df + self._df_cache_revision = self._revision.value + return self._df_cache + + def __len__(self) -> int: + if (live := self._cold_live_values()) is not None: + return len(live) + # Count in SQL when the relation is not already realized: coord + # range residuals only drop patches the SQL candidacy already + # excludes and samples/relative residuals never drop patches, so + # the count matches len(to_df()) without projecting or pivoting. + if ( + self._df_cache is not None + and self._df_cache_revision == self._revision.value + ): + return len(self._df_cache) + return self.backend.count(list(self._queries) or None, patch_ids=self._ids) + + def get_patch(self, index: int) -> dc.Patch: + """Materialize one patch: resolve, then exact two-stage trim.""" + if (live := self._cold_live_values()) is not None: + return live[index] + row = self.to_df().iloc[index].to_dict() + return self.resolve_row(row) + + def resolve_row(self, row: Mapping, extra_trim: Mapping | None = None) -> dc.Patch: + """ + Resolve one flat-relation row and apply exact residual selects. + + extra_trim carries caller-side read hints (e.g. chunk instruction + ranges) merged over the view's own residual ranges; like all trim + hints they only reduce reading, exactness is re-applied above. + """ + trim_hint = {} + for coords, samples in self._residuals: + if not samples: + # Canonical-SI and quantity bounds stay out of reader + # hints: readers take numbers in their native units, so + # a converted-narrower hint could drop data exactness + # cannot restore. + trim_hint.update( + { + k: v + for k, v in coords.items() + if isinstance(v, tuple) + and not any(hasattr(b, "units") for b in v) + } + ) + trim_hint.update(extra_trim or {}) + patch = self.resolver.resolve(row, **trim_hint) + return apply_exact_residuals(patch, self._residuals) + + def __iter__(self): + for index in range(len(self)): + yield self.get_patch(index) + + # --- mutation (root only) ---------------------------------------------- + + def add(self, patches: Sequence[dc.Patch] | dc.Patch) -> PatchCatalog: + """Add live patches to the catalog.""" + self._require_root("add") + if not isinstance(self.resolver, LiveResolver): + msg = "add() currently supports in-memory catalogs only." + raise NotImplementedError(msg) + patches = [patches] if isinstance(patches, dc.Patch) else list(patches) + additions = {_patch_path(x): x for x in patches} + self.resolver._registry.update(additions) + # Re-adding a patch replaces its row (same identity), so this + # stays idempotent. + self.backend.write_sources(_live_records(additions)) + self._invalidate() + return self + + def update(self, progress: PROGRESS_LEVELS = "standard") -> PatchCatalog: + """Sync a directory-backed catalog with the filesystem.""" + if self._syncer is not None: + self._syncer.update(progress=progress) + self._invalidate() + return self + + def remove(self, source_paths: Sequence[str], base_uri: str = "") -> PatchCatalog: + """Remove sources (and their patches) from the catalog.""" + self._require_root("remove") + source_paths = list(source_paths) + self.backend.delete_sources(source_paths, base_uri=base_uri) + # The live registry is the store for in-memory patches; it must + # stay in step with the backend rows (pickling rebuilds from it). + registry = self.resolver.live_entries() if self.resolver else {} + for path in source_paths: + registry.pop(path, None) + self._invalidate() + return self + + # --- introspection ------------------------------------------------------- + + def attr_names(self) -> set[str]: + """Attr names known to the index.""" + return self.backend.attr_names() + + def coord_names(self) -> set[str]: + """Coord names known to the index.""" + return self.backend.coord_names() + + def sources(self) -> pd.DataFrame: + """The sources table.""" + return self.backend.get_sources() + + def get_metadata(self) -> dict: + """Index-level metadata.""" + return self.backend.get_metadata() + + def close(self) -> None: + """Close the backend (root and all views share it).""" + if self._backend is not None: + self._backend.close() diff --git a/dascore/io/index/dialect.py b/dascore/io/index/dialect.py new file mode 100644 index 000000000..598dbd285 --- /dev/null +++ b/dascore/io/index/dialect.py @@ -0,0 +1,60 @@ +""" +SQL dialect translation for index backends. + +Everything engine-specific lives here: type names, identifier quoting, +glob matching, and table DDL generation. The rest of the package emits +logical types and parametrized SQL with `?` placeholders. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import ClassVar + + +class BaseDialect: + """Shared SQL generation for engines close to the standard.""" + + # logical type -> engine type; concrete dialects define the values. + type_map: ClassVar[Mapping[str, str]] + strict_suffix: ClassVar[str] + + def quote(self, identifier: str) -> str: + """Quote an identifier.""" + return '"' + identifier.replace('"', '""') + '"' + + def create_table( + self, name: str, columns: Mapping[str, str], constraints: tuple[str, ...] = () + ) -> str: + """Return DDL for one table from logical column types.""" + definitions = [ + f"{self.quote(col)} {self.type_map[typ]}" for col, typ in columns.items() + ] + definitions.extend(constraints) + cols = ", ".join(definitions) + quoted = self.quote(name) + return f"CREATE TABLE IF NOT EXISTS {quoted} ({cols}){self.strict_suffix}" + + def add_column(self, table: str, column: str, logical_type: str) -> str: + """Return DDL to add one nullable column.""" + return ( + f"ALTER TABLE {self.quote(table)} " + f"ADD COLUMN {self.quote(column)} {self.type_map[logical_type]}" + ) + + def glob(self, column_sql: str) -> str: + """Return a parametrized unix-glob match expression.""" + return f"{column_sql} GLOB ?" + + +class SQLiteDialect(BaseDialect): + """Dialect for SQLite; STRICT tables enforce the type contract.""" + + # SQLite STRICT tables accept INTEGER/REAL/TEXT (and INT for bool). + type_map: ClassVar[Mapping[str, str]] = { + "int64": "INTEGER", + "float64": "REAL", + "str": "TEXT", + "bool": "INTEGER", + } + strict_suffix: ClassVar[str] = " STRICT" diff --git a/dascore/io/index/indexer.py b/dascore/io/index/indexer.py new file mode 100644 index 000000000..5bdcc9ad7 --- /dev/null +++ b/dascore/io/index/indexer.py @@ -0,0 +1,328 @@ +""" +A directory indexer backed by the generic spool index. + +Drop-in alternative to `dascore.io.indexer.DirectoryIndexer`: it walks a +directory, detects new/changed/removed sources by per-source +(mtime, size) comparison, scans only what changed, and answers content +queries from the index backend. +""" + +from __future__ import annotations + +import hashlib +from contextlib import suppress +from pathlib import Path + +import pandas as pd +from typing_extensions import Self + +import dascore as dc +from dascore.compat import UPath +from dascore.config import config_attr +from dascore.constants import PROGRESS_LEVELS +from dascore.exceptions import InvalidIndexVersionError +from dascore.io.index.backend import get_backend, resolve_query +from dascore.io.index.ingest import SourceRecord, summaries_to_records +from dascore.io.index.schema import SPOOL_HIDDEN_COLUMNS +from dascore.io.indexer import ( + AbstractIndexer, + _get_index_map, + _update_index_map, +) +from dascore.utils.misc import _iter_filesystem +from dascore.utils.paths import directory_writable, requires_local_directory + + +class DBDirectoryIndexer(AbstractIndexer): + """ + Index a directory of fiber files with a database backend. + + Parameters + ---------- + path + The directory to index. + index_path + Where to keep the index; defaults to a hidden entry at the top of + the data directory. + """ + + ext: str | None = None + # user-level file tracking index locations for unwritable data dirs + index_map_path: Path = config_attr("directory_index_map_path") + + def __init__( + self, + path: str | Path, + index_path: str | Path | None = None, + ): + path = UPath(path).absolute() if isinstance(path, UPath) else Path(path) + requires_local_directory(path, label="DBDirectoryIndexer") + self.path = Path(path).absolute() + self.index_path = Path(self._find_index_path(index_path)) + try: + self._backend = get_backend(self.index_path) + except InvalidIndexVersionError: + # The index is a disposable cache and the file already + # identified itself as a dascore spool index of another + # schema version; rebuild it rather than asking the user to. + self.index_path.unlink() + self._backend = get_backend(self.index_path) + # Schema creation alone is not a successful directory scan. Read the + # transactional marker so a new process retries an interrupted first + # update instead of trusting a merely nonempty SQLite file. + metadata = self._backend.get_metadata() + self._initial_update_done = bool(metadata["last_indexed_ns"]) + + @property + def _index_name(self) -> str: + return ".dascore_index.sqlite3" + + @staticmethod + def _is_legacy_or_foreign_index(path: Path) -> bool: + """Return True if an existing file is not a SQLite database. + + Older DASCore versions recorded PyTables (.h5) index locations in + the index map; passing those to sqlite3 fails with an opaque + error instead of building the replacement index. Only the file + header decides — users may legitimately choose any suffix for a + custom index path. + """ + if not path.exists(): + return False + with suppress(OSError), open(path, "rb") as fh: + header = fh.read(16) + return len(header) >= 16 and not header.startswith(b"SQLite format 3") + return False + + def _find_index_path(self, index_path=None) -> Path: + """ + Find where the index lives (or should live). + + Mirrors the historic DirectoryIndexer behavior: in-directory by + default; when the data directory is read-only the index lives in + the dascore cache and its location is recorded in the index map. + """ + map_key = str(self.path) + if index_path: + index_path = Path(index_path).absolute() + update = {map_key: str(index_path)} + _update_index_map(update, cache_path=str(self.index_map_path)) + return index_path + expected = self.path / self._index_name + with suppress(PermissionError): + if expected.exists(): + return expected + path_map = _get_index_map(cache_path=str(self.index_map_path)) + if out := path_map.get(map_key): + mapped = Path(out) + # Index-map entries from older DASCore versions can point at + # the retired PyTables (.h5) index; those are not usable and + # a fresh SQLite index is built in their place. + if not self._is_legacy_or_foreign_index(mapped): + return mapped + if not directory_writable(self.path): + name = f"_dascore_index_{abs(hash(self.path))}.sqlite3" + index_path = self.index_map_path.parent / name + _update_index_map( + {map_key: str(index_path.absolute())}, + cache_path=str(self.index_map_path), + ) + return index_path + return expected + + def ensure_updated(self) -> bool: + """Run the initial update if the index was never populated.""" + if self._initial_update_done: + return False + self.update(progress=None) + return True + + def __str__(self) -> str: + return f"{self.__class__.__name__} managing: {self.path}" + + __repr__ = __str__ + + def __deepcopy__(self, memo) -> Self: + """ + Derived spools share the indexer (and its live DB connection). + + Spool copies its state on select/chunk; the index + connection is read-shared, matching the single-writer model. + """ + return self + + def _rel(self, path: Path) -> str: + """Relative posix path of a file under the spool root.""" + return Path(path).relative_to(self.path).as_posix() + + def _directory_format(self, path: Path) -> bool: + """Return True when a directory is itself one FiberIO scan unit.""" + from dascore.io.core import is_directory_format + + return is_directory_format(path) + + @staticmethod + def _directory_signature(path: Path) -> tuple[int, int]: + """Return a stable 128-bit manifest signature as two SQLite ints.""" + members = sorted( + ( + sub + for sub in path.rglob("*") + if sub.is_file() and not sub.name.startswith(".") + ), + key=lambda sub: sub.relative_to(path).as_posix(), + ) + digest = hashlib.sha256() + for member in members: + stat = member.stat() + relative = member.relative_to(path).as_posix().encode() + digest.update(len(relative).to_bytes(8, "big")) + digest.update(relative) + digest.update(stat.st_mtime_ns.to_bytes(8, "big", signed=True)) + digest.update(stat.st_size.to_bytes(8, "big")) + fingerprint = digest.digest() + return ( + int.from_bytes(fingerprint[:8], "big", signed=True), + int.from_bytes(fingerprint[8:16], "big", signed=True), + ) + + def _walk(self) -> dict[str, tuple[int, int, Path]]: + """ + Walk the spool directory, honoring directory-format scan units. + + Maps relative path -> (mtime_ns, size, abs path) for every scan + unit. A directory-format unit (e.g. XMLBinary) appears as one + entry keyed by the directory, with a 128-bit manifest fingerprint + split across the two integer stat fields. The fingerprint covers + every member's relative path, mtime, and size, so member changes + cannot cancel each other out. Mirrors the skip protocol dc.scan + uses so members are not offered individually. + """ + files: dict[str, tuple[int, int, Path]] = {} + gen = _iter_filesystem(self.path, ext=self.ext, include_directories=True) + signal = None + while True: + try: + # send(None) is equivalent to next() and also starts it + candidate = gen.send(signal) + except StopIteration: + break + signal = None + if candidate is None: # the reply to a "skip" send + continue + path = Path(candidate) + if path.is_dir(): + if self._directory_format(path): + signal = "skip" + signature = self._directory_signature(path) + files[self._rel(path)] = (*signature, path) + continue + stat = path.stat() + files[self._rel(path)] = (stat.st_mtime_ns, stat.st_size, path) + return files + + def update(self, paths=None, progress: PROGRESS_LEVELS = "standard") -> Self: + """ + Update the index: scan new/changed sources, drop removed ones. + + Change detection compares each source's stored (mtime_ns, + size_bytes) against the filesystem — never a global watermark — + and stale-source removal is folded in (the walk is the dominant + cost; removal afterwards is nearly free). Directory-format scan + units (e.g. XMLBinary) are rescanned whole when any member file + changes. + """ + files = self._walk() + stored = { + row.source_path: ( + None + if pd.isnull(row.mtime_ns) + else (int(row.mtime_ns), int(row.size_bytes)) + ) + for row in self._backend.source_stats().itertuples() + } + stale = [path for path in stored if path not in files] + changed = [ + rel + for rel, (mtime, size, _) in files.items() + if stored.get(rel) != (mtime, size) + ] + if paths is not None: + # restrict the rescan (not stale removal) to the given paths + keep = set() + for one in paths: + one = Path(one) + rel = ( + one.relative_to(self.path).as_posix() + if one.is_absolute() + else one.as_posix() + ) + keep.add(rel) + changed = [rel for rel in changed if rel in keep] + if stale: + self._backend.delete_sources(stale) + if changed: + # Only changed paths are rescanned, so the stat maps handed to + # summaries_to_records need only cover them — not the whole + # archive (a large mostly-unchanged directory otherwise built + # full-archive mtime/size maps for a tiny update). + changed_stats = {rel: files[rel] for rel in changed} + scan_paths = [stat[2] for stat in changed_stats.values()] + summaries = dc.scan(scan_paths, progress=progress) + # scan reports absolute source paths; stat maps use them too + records = summaries_to_records( + summaries, + relative_to=str(self.path), + mtimes_ns={str(p): m for (m, _, p) in changed_stats.values()}, + sizes_bytes={str(p): s for (_, s, p) in changed_stats.values()}, + ) + # Every visited path gets a sources row, even when scanning + # produced no patches (e.g. a non-fiber file). Otherwise such + # files look "new" on every update and force perpetual + # rescans. + recorded = {rec.source_path for rec in records} + for rel in set(changed) - recorded: + mtime, size, _ = files[rel] + records.append( + SourceRecord( + source_path=rel, + source_format="", + format_version="", + mtime_ns=mtime, + size_bytes=size, + ) + ) + if records: + self._backend.write_sources(records) + if stale or changed or not self._initial_update_done: + # Directory archives present in time order; ingest assigns + # walk-order ordinals, so each sync renumbers to keep the + # contract (iterate by ordinal) aligned with time. The + # not-yet-marked-done case covers a process killed between + # write_sources committing and this renumber: the retry sees + # no stale/changed files but must still fix walk-order + # ordinals before marking the initial update complete. + self._backend.renumber_ordinals_by_time() + if not self._initial_update_done: + self._backend.mark_initial_update_done() + self._initial_update_done = True + return self + + def get_contents(self, _attrs=None, _coords=None, **kwargs) -> pd.DataFrame: + """ + Query the index, returning the spool-facing flat relation. + + Bare kwargs resolve attrs-first then coords; `_attrs`/`_coords` + disambiguate explicitly (see the selector semantics spec). + """ + self.ensure_updated() + query = resolve_query(self._backend, _attrs=_attrs, _coords=_coords, **kwargs) + df = self._backend.query(query) + df = df.drop(columns=list(SPOOL_HIDDEN_COLUMNS), errors="ignore") + return df.rename(columns={"patch_id": "_patch_id"}) + + __call__ = get_contents + + def close(self) -> None: + """Close the backend.""" + self._backend.close() diff --git a/dascore/io/index/ingest.py b/dascore/io/index/ingest.py new file mode 100644 index 000000000..bbc245388 --- /dev/null +++ b/dascore/io/index/ingest.py @@ -0,0 +1,534 @@ +""" +Convert patch summaries into normalized index records. + +This module is backend-independent: it turns `PatchSummary` objects into +plain records (dicts/dataclasses) using only the four primitive storage +types. All unit-bearing numeric values are normalized to pint base SI +units here, so cross-patch comparisons in the index are always valid. +""" + +from __future__ import annotations + +import hashlib +import re +import warnings +from dataclasses import dataclass, field, fields, replace + +import numpy as np +import pandas as pd + +from dascore.core.summary import PatchSummary, normalize_source_patch_id +from dascore.io.index.schema import KINDS, RESERVED_ATTR_COLUMNS +from dascore.units import get_quantity +from dascore.utils.time import to_datetime64, to_int, to_timedelta64 + +_SANITIZE_RE = re.compile(r"[^a-z0-9_]+") + +# Attrs handled structurally or intentionally excluded from the index. +_SKIPPED_ATTRS = frozenset({"history", "dims", "coords"}) + +# Suffixes of the per-coordinate envelope columns the flat relation +# emits ({name}_min/{name}_max/{name}_step). An attr shaped like one of +# these is reserved catalog-wide — not just against the ingesting +# patch's own coords — so a flat-relation envelope column (e.g. +# "event_time_min") can never collide with a same-named attr contributed +# by a different patch, which would make one get_contents() column's +# meaning depend on which other patches share the catalog. (units/dtype +# do not become per-coord columns, so a real attr like "data_units" +# stays queryable, matching Patch.update_attrs.) +_ENVELOPE_SUFFIXES = ("min", "max", "step") + + +def _is_envelope_shaped(name: str) -> bool: + """True if name looks like a ``{coord}_{min,max,step}`` envelope column.""" + prefix, _, suffix = name.rpartition("_") + return bool(prefix) and suffix in _ENVELOPE_SUFFIXES + + +@dataclass(frozen=True) +class TypedValue: + """A value with its index kind and (canonical) units.""" + + kind: str + value: float | int | str | bool + units: str | None = None + + def __post_init__(self): + assert self.kind in KINDS + + +@dataclass(frozen=True) +class CoordRecord: + """One patch-coord entry (typed columns split by kind).""" + + coord_name: str + value_kind: str + dtype: str + coord_dims: str + length: int | None + units: str | None + min_num: float | None = None + max_num: float | None = None + step_num: float | None = None + min_ns: int | None = None + max_ns: int | None = None + step_ns: int | None = None + min_str: str | None = None + max_str: str | None = None + is_monotonic: bool | None = None + is_relative: bool | None = None + coord_hash: str | None = None + + @property + def def_key(self) -> str: + """ + Deduplication key for the coord definition. + + The CoordSummary fingerprint when available ("fp:" prefix; exact + value identity), otherwise a hash of the stored summary fields + ("sum:" prefix; lossless for the index but too weak for + value-identity claims). Name and dims are patch-level and + excluded. + """ + if self.coord_hash: + # truncated: 128 bits is ample and key size shows up in the + # def_key index for archives with mostly-unique time coords + return f"fp:{self.coord_hash[:32]}" + fields = ( + self.value_kind, + self.dtype, + self.length, + self.units, + self.min_num, + self.max_num, + self.step_num, + self.min_ns, + self.max_ns, + self.step_ns, + self.min_str, + self.max_str, + self.is_monotonic, + self.is_relative, + ) + digest = hashlib.sha256(repr(fields).encode()).hexdigest()[:32] + return f"sum:{digest}" + + +@dataclass(frozen=True) +class PatchRecord: + """One patch: structural fields, typed attrs, coord rows.""" + + source_patch_id: str + dims: str + shape: str + n_dims: int + sample_count_total: int | None + time_min: int | None + time_max: int | None + time_step: int | None + distance_min: float | None + distance_max: float | None + distance_step: float | None + attrs: dict[str, TypedValue] = field(default_factory=dict) + coords: tuple[CoordRecord, ...] = () + + +@dataclass(frozen=True) +class SourceRecord: + """One source (scan unit) and the patches it emitted.""" + + source_path: str + source_format: str + format_version: str + base_uri: str | None = None + mtime_ns: int | None = None + size_bytes: int | None = None + patches: tuple[PatchRecord, ...] = () + + +def sanitize_attr_name(name: str) -> str: + """Return a lowercase [a-z0-9_] identifier for an attr name.""" + out = _SANITIZE_RE.sub("_", name.lower()).strip("_") + if not out or out[0].isdigit(): + out = f"a_{out}" + return out + + +def attr_column_name(name: str, kind: str) -> str: + """Return the attrs-table column for an attr name and kind.""" + return f"{sanitize_attr_name(name)}__{kind}" + + +def _base_unit_info(value, unit_str: str | None = None) -> tuple[float, str]: + """Return a value's base-unit magnitude and canonical unit string.""" + quant = value if unit_str is None else value * get_quantity(unit_str) + quant = get_quantity(quant).to_base_units() + return float(quant.magnitude), str(quant.units) + + +def _is_missing(value) -> bool: + """Return True for values that mean 'not present'.""" + if value is None or (isinstance(value, str) and value == ""): + return True + try: + return bool(pd.isnull(value)) + except (TypeError, ValueError): + return False + + +def typed_value(value) -> TypedValue | None: + """ + Classify a python/numpy scalar into a TypedValue, or None to skip. + + Unit-bearing quantities are converted to base SI units; the canonical + unit string is recorded so queries can convert consistently. + """ + if _is_missing(value): + return None + # containers/arrays are complex attrs: never indexable scalars (and + # they must not reach the datetime fallback, which accepts arrays). + if isinstance(value, np.ndarray | list | tuple | set | frozenset | dict | bytes): + return None + # bool must precede int (bool is a subclass of int). + if isinstance(value, bool | np.bool_): + return TypedValue("bool", bool(value)) + if isinstance(value, np.datetime64): + return TypedValue("time", to_int(value)) + if isinstance(value, np.timedelta64): + return TypedValue("dur", to_int(value)) + # pint scalar quantity or unit. + if hasattr(value, "units"): + magnitude = getattr(value, "magnitude", 1) + if isinstance(magnitude, np.ndarray): + return None # array quantities are not scalar attrs + magnitude, base = _base_unit_info(value) + return TypedValue("num", magnitude, units=base) + if isinstance(value, int | np.integer | float | np.floating): + return TypedValue("num", float(value)) + if isinstance(value, str): + return TypedValue("str", value) + # datetime/timedelta and anything datetime-like numpy missed. + try: + return TypedValue("time", to_int(to_datetime64(value))) + except Exception: + pass + return None # complex attrs (sequences, dicts, ...) are skipped + + +def _extract_attrs(summary: PatchSummary) -> dict[str, TypedValue]: + """Get indexable typed attrs from a patch summary.""" + raw = summary.attrs.model_dump() + out = {} + for name, value in raw.items(): + if name in _SKIPPED_ATTRS or name.startswith("_"): + continue + # Reserve structural columns and any coordinate-envelope-shaped + # name (catalog-wide, not just this patch's own coords) so the + # meaning of a flat-relation column never depends on which other + # patches share the catalog. + if sanitize_attr_name(name) in RESERVED_ATTR_COLUMNS or ( + _is_envelope_shaped(name) + ): + msg = ( + f"Skipping reserved attr name {name!r}; it collides with a " + "structural index column. The attr stays on the patch but " + "is not queryable through the spool." + ) + warnings.warn(msg, UserWarning) + continue + typed = typed_value(value) + if typed is not None: + out[name] = typed + return out + + +def _coord_record(name: str, summary) -> CoordRecord | None: + """Convert one CoordSummary into a CoordRecord.""" + fingerprint = getattr(summary, "fingerprint", None) + if fingerprint is None and getattr(summary, "is_range_like", False): + # A range summary contains its complete representation, so recover the + # same exact identity a loaded CoordRange would have produced. + fingerprint = summary.to_coord().fingerprint() + # str() on a pint Quantity is comparatively expensive; do it once and + # reuse it below (a Quantity-keyed cache is unsafe — 1 m == 100 cm with + # equal hashes but different strings). + units_str = str(summary.units) if summary.units is not None else None + common = dict( + coord_name=name, + dtype=summary.dtype, + coord_dims=",".join(summary.dims), + length=summary.len, + units=units_str, + coord_hash=fingerprint, + ) + dtype = np.dtype(summary.dtype) if summary.dtype else None + if dtype is None: + return None # unsupported coord representation: skip, per design + if dtype.kind in "mM": # datetime64 ("M") / timedelta64 ("m") + is_datetime = dtype.kind == "M" + convert = to_datetime64 if is_datetime else to_timedelta64 + step = summary.step + return CoordRecord( + value_kind="time", + is_relative=not is_datetime, + min_ns=to_int(convert(summary.min)), + max_ns=to_int(convert(summary.max)), + step_ns=None if pd.isnull(step) else to_int(to_timedelta64(step)), + **common, + ) + if np.issubdtype(dtype, np.number): + min_num = float(summary.min) + max_num = float(summary.max) + step = summary.step + step_num = None if pd.isnull(step) else float(step) + if units_str is not None: + min_num, base = _base_unit_info(summary.min, units_str) + max_num, _ = _base_unit_info(summary.max, units_str) + if step_num is not None: + step_end, _ = _base_unit_info(summary.min + summary.step, units_str) + step_num = step_end - min_num + common["units"] = base + return CoordRecord( + value_kind="num", + min_num=min_num, + max_num=max_num, + step_num=step_num, + **common, + ) + if dtype.kind in "US" or dtype == object: + return CoordRecord( + value_kind="str", + min_str=str(summary.min), + max_str=str(summary.max), + **common, + ) + return None # unsupported coord representation: skip, per design + + +def _envelope(coords: tuple[CoordRecord, ...], name: str, kind: str): + """Pull the (min, max, step) envelope for one coord if present.""" + for rec in coords: + if rec.coord_name != name or rec.value_kind != kind: + continue + if kind == "time" and not rec.is_relative: + return rec.min_ns, rec.max_ns, rec.step_ns + if kind == "num": + return rec.min_num, rec.max_num, rec.step_num + return None, None, None + + +def patch_record(summary: PatchSummary) -> PatchRecord: + """Convert one PatchSummary into a PatchRecord.""" + coords = tuple( + rec + for name, csum in summary.coords.items() + if (rec := _coord_record(name, csum)) is not None + ) + time_min, time_max, time_step = _envelope(coords, "time", "time") + dist_min, dist_max, dist_step = _envelope(coords, "distance", "num") + shape = tuple(int(x) for x in summary.shape) + return PatchRecord( + source_patch_id=normalize_source_patch_id(summary.source_patch_id), + dims=",".join(summary.dims), + shape=",".join(str(x) for x in shape), + n_dims=len(summary.dims), + sample_count_total=int(np.prod(shape)) if shape else None, + time_min=time_min, + time_max=time_max, + time_step=time_step, + distance_min=dist_min, + distance_max=dist_max, + distance_step=dist_step, + attrs=_extract_attrs(summary), + coords=coords, + ) + + +def summaries_to_records( + summaries: list[PatchSummary], + base_uri: str | None = None, + relative_to: str | None = None, + mtimes_ns: dict[str, int] | None = None, + sizes_bytes: dict[str, int] | None = None, +) -> list[SourceRecord]: + """ + Group patch summaries by source and convert to SourceRecords. + + Parameters + ---------- + summaries + Patch summaries, e.g. from `dc.scan`. + base_uri + Optional common root persisted with each source (remote spools); + source paths are stored relative to it. + relative_to + Optional local spool root: source paths are stored relative to it + but the root itself is *not* persisted (local directory spools + resolve against their current root, per the design doc). + mtimes_ns, sizes_bytes + Optional maps of source_path -> stat values. When omitted the + caller is responsible for change detection. + """ + # Group by the original (OS-native) source path so the mtimes_ns / + # sizes_bytes maps, which the caller keys by that same path, still + # resolve. Index paths themselves are stored as POSIX so comparison + # and deletion are separator-agnostic across platforms. + by_source: dict[str, list[PatchSummary]] = {} + for summary in summaries: + by_source.setdefault(str(summary.source_path), []).append(summary) + root = base_uri or relative_to + root_posix = str(root).replace("\\", "/") if root else None + root_prefix = root_posix.rstrip("/") if root_posix else None + out = [] + for path, group in by_source.items(): + first = group[0] + patches = [] + for num, summary in enumerate(group): + record = patch_record(summary) + if record.source_patch_id == "" and len(group) > 1: + # positional identity within the source, per design doc + record = replace(record, source_patch_id=str(num)) + patches.append(record) + posix_path = path.replace("\\", "/") + store_path = posix_path + if root_prefix is not None and ( + posix_path == root_prefix or posix_path.startswith(f"{root_prefix}/") + ): + # "." (not "") marks a source that IS the root (directory units) + store_path = posix_path[len(root_prefix) :].lstrip("/") or "." + out.append( + SourceRecord( + source_path=store_path, + base_uri=base_uri, + source_format=first.source_format, + format_version=first.source_version, + mtime_ns=(mtimes_ns or {}).get(path), + size_bytes=(sizes_bytes or {}).get(path), + patches=tuple(patches), + ) + ) + return out + + +# Record fields read straight off an index row of the same name. The +# remaining fields need per-field handling: a coord's name/dims are +# patch-level (they come from the link row, not the shared definition), +# its hash is stored as "fingerprint", and a patch's id/dims/shape get +# normalized below. +_COORD_DEF_FIELDS = tuple( + f.name + for f in fields(CoordRecord) + if f.name not in ("coord_name", "coord_dims", "coord_hash") +) +_PATCH_ROW_FIELDS = tuple( + f.name + for f in fields(PatchRecord) + if f.name not in ("source_patch_id", "dims", "shape", "attrs", "coords") +) + + +def _py_scalar(value): + """Convert a fetched cell to the plain python scalar records use.""" + if value is None or pd.isnull(value): + return None + if isinstance(value, np.bool_ | bool): + return bool(value) + if isinstance(value, np.integer | int): + return int(value) + if isinstance(value, np.floating | float): + return float(value) + return value + + +def assemble_source_records( + sources: pd.DataFrame, + patches: pd.DataFrame, + attrs: pd.DataFrame, + links: pd.DataFrame, + defs: pd.DataFrame, + meta: pd.DataFrame, +) -> list[SourceRecord]: + """ + Assemble source records from already-fetched index frames. + + This is the transfer format for merging catalogs: feeding the result + to another backend's `write_sources` re-ingests the metadata with + fresh ids, coord-def deduplication (def keys are preserved), and + replace-semantics on (base_uri, source_path) identity. The caller + (an index backend's export_records) is responsible for narrowing the + frames — filtering by patch id belongs in SQL, not here. + """ + if sources.empty: + return [] + # Records transfer in catalog order: re-ingesting assigns fresh + # sequential ordinals, so record order IS the ordering contract. + if "ordinal" in sources.columns: + sources = sources.sort_values(["ordinal", "source_id"]) + if "patch_id" in patches.columns: + patches = patches.sort_values("patch_id") + col_info = { + row.column_name: (row.attr_name, row.value_kind, _py_scalar(row.units)) + for row in meta.itertuples() + } + def_map = {int(row.coord_def_id): row for row in defs.itertuples()} + attr_rows = ( + {int(k): v for k, v in attrs.set_index("patch_id").to_dict("index").items()} + if not attrs.empty + else {} + ) + link_groups = ( + {int(k): v for k, v in links.groupby("patch_id")} if not links.empty else {} + ) + patches_by_source = ( + {int(k): v for k, v in patches.groupby("source_id")} + if not patches.empty + else {} + ) + out = [] + for src in sources.itertuples(): + sub = patches_by_source.get(int(src.source_id)) + if sub is None: + continue + patch_records = [] + for patch in sub.itertuples(): + pid = int(patch.patch_id) + typed = {} + for col, value in attr_rows.get(pid, {}).items(): + if col in col_info and not pd.isnull(value): + name, kind, units = col_info[col] + typed[name] = TypedValue( + kind=kind, value=_py_scalar(value), units=units + ) + coords = [] + for link in link_groups.get(pid, pd.DataFrame()).itertuples(): + cdef = def_map[int(link.coord_def_id)] + coords.append( + CoordRecord( + coord_name=link.coord_name, + coord_dims=link.coord_dims, + coord_hash=_py_scalar(cdef.fingerprint), + **{f: _py_scalar(getattr(cdef, f)) for f in _COORD_DEF_FIELDS}, + ) + ) + patch_records.append( + PatchRecord( + source_patch_id=normalize_source_patch_id(patch.source_patch_id), + dims=_py_scalar(patch.dims) or "", + shape=_py_scalar(patch.shape) or "", + attrs=typed, + coords=tuple(coords), + **{f: _py_scalar(getattr(patch, f)) for f in _PATCH_ROW_FIELDS}, + ) + ) + out.append( + SourceRecord( + source_path=_py_scalar(src.source_path) or "", + base_uri=_py_scalar(src.base_uri) or None, + source_format=_py_scalar(src.source_format) or "", + format_version=_py_scalar(src.format_version) or "", + mtime_ns=_py_scalar(src.mtime_ns), + size_bytes=_py_scalar(src.size_bytes), + patches=tuple(patch_records), + ) + ) + return out diff --git a/dascore/io/index/lite.py b/dascore/io/index/lite.py new file mode 100644 index 000000000..f565e3ea5 --- /dev/null +++ b/dascore/io/index/lite.py @@ -0,0 +1,166 @@ +"""SQLite index backend (stdlib sqlite3, STRICT tables).""" + +from __future__ import annotations + +import sqlite3 +import threading +import weakref +from contextlib import suppress +from pathlib import Path + +import numpy as np +import pandas as pd + +from dascore.io.index.backend import SQLIndexBackend, adapt_params +from dascore.io.index.dialect import SQLiteDialect + +# A serialized SQLite build (threadsafety == 3) lets one connection be +# used and closed from any thread, so cross-thread garbage collection of +# a backend is safe. On rarer non-serialized builds the connection is +# thread-bound and must stay check_same_thread. +_SQLITE_SERIALIZED = sqlite3.threadsafety == 3 + + +def _safe_close(con: sqlite3.Connection) -> None: + """ + Close a connection, tolerating cross-thread finalization. + + On a serialized build closing works from any thread. On a + thread-bound build a finalizer firing on another thread would raise + ProgrammingError; suppress it (the underlying handle is freed at + interpreter teardown) rather than emit an unraisable exception. + """ + with suppress(sqlite3.ProgrammingError): + con.close() + + +def _adapt(params): + """Convert numpy/py types sqlite3 can't bind natively.""" + return [int(p) if isinstance(p, bool) else p for p in adapt_params(params)] + + +def _classic_dtypes(df: pd.DataFrame) -> pd.DataFrame: + """ + Convert nullable extension columns back to classic numpy dtypes. + + Fetching with dtype_backend="numpy_nullable" is what keeps nullable + INTEGER columns exact (the default assembly rounds >2**53 ns values + through float64), but downstream spool code expects classic dtypes. + Only int columns that actually hold NULLs stay nullable (Int64) — + the exactness they exist for; consumers handle them via isna(). + """ + for name in df.columns: + col = df[name] + dtype = col.dtype + if not isinstance(dtype, pd.api.extensions.ExtensionDtype): + continue + if dtype.kind == "i": + if not col.isna().any(): + df[name] = col.to_numpy(dtype="int64") + elif dtype.kind == "f": + df[name] = col.to_numpy(dtype="float64", na_value=np.nan) + else: # string/boolean/... -> classic object with None for missing + df[name] = col.to_numpy(dtype=object, na_value=None) + return df + + +class SQLiteBackend(SQLIndexBackend): + """Index backend storing tables in a single SQLite file.""" + + dialect = SQLiteDialect() + + def __init__(self, path: str | Path): + self._path = str(path) + # On a serialized build, drop the thread affinity so the shared + # backend can be used (and finalized) from worker threads, e.g. + # a thread-pool Spool.map over one catalog. + self._con = sqlite3.connect( + self._path, check_same_thread=not _SQLITE_SERIALIZED + ) + # autocommit off; we manage transactions explicitly. + self._con.isolation_level = None + self._con.execute("PRAGMA foreign_keys = ON") + self._con.execute("PRAGMA busy_timeout = 30000") + # Catalog views share this backend object, so tying connection + # cleanup to *its* collection is safe (close() stays idempotent + # for explicit use). The finalizer tolerates cross-thread firing. + self._finalizer = weakref.finalize(self, _safe_close, self._con) + # Catalog views share this backend across threads (e.g. a + # thread-pool Spool.map over split windows). SQLite serializes + # individual statements, but a pandas fetch spans many cursor + # calls; interleaving them corrupts result frames, so statement + # execution is exclusive per backend. + self._lock = threading.RLock() + try: + super().__init__() + except Exception: + self._con.close() + raise + + def __getstate__(self) -> dict: + """ + Pickle by database path; the file is the durable state. + + This makes file-backed spools usable with process pools: the + receiving process reopens its own connection. In-memory backends + have no file to reopen; their owners (catalogs) serialize their + contents separately and never pickle the backend itself. + """ + if self._path == ":memory:": + msg = ( + "In-memory index backends cannot be pickled; pickle their " + "owning catalog/spool instead." + ) + raise TypeError(msg) + return {"_path": self._path} + + def __setstate__(self, state: dict) -> None: + """Reconnect to the database file.""" + self.__init__(state["_path"]) + + def _execute(self, sql: str, params=()) -> None: + with self._lock: + self._con.execute(sql, _adapt(params)) + + def _executemany(self, sql: str, seq_of_params) -> None: + # sqlite3.executemany consumes an iterator, so adapt lazily rather + # than materializing a second copy of each already-built batch. + with self._lock: + self._con.executemany(sql, (_adapt(p) for p in seq_of_params)) + + def _fetch_df(self, sql: str, params=()) -> pd.DataFrame: + # numpy_nullable assembly keeps nullable INTEGER columns exact; + # the default path rounds them through float64, corrupting ns + # epochs (>2**53). A dtype= hint does NOT prevent that: pandas + # builds float64 first and casts after. + with self._lock: + df = pd.read_sql_query( + sql, self._con, params=_adapt(params), dtype_backend="numpy_nullable" + ) + return _classic_dtypes(df) + + def _begin(self) -> None: + self._con.execute("BEGIN IMMEDIATE") + + def _commit(self) -> None: + self._con.execute("COMMIT") + + def _rollback(self) -> None: + self._con.execute("ROLLBACK") + + def _existing_tables(self) -> set[str]: + rows = self._con.execute( + "SELECT name FROM sqlite_master " + "WHERE type = 'table' AND name NOT LIKE 'sqlite_%'" + ).fetchall() + return {row[0] for row in rows} + + def _table_columns(self, table: str) -> set[str]: + sql = f"PRAGMA table_info({self.dialect.quote(table)})" + return {row[1] for row in self._con.execute(sql).fetchall()} + + def close(self) -> None: + """Close the database connection.""" + # detach the GC finalizer; closing twice is harmless but tidy. + self._finalizer.detach() + _safe_close(self._con) diff --git a/dascore/io/index/planned.py b/dascore/io/index/planned.py new file mode 100644 index 000000000..6fb6c27b9 --- /dev/null +++ b/dascore/io/index/planned.py @@ -0,0 +1,498 @@ +""" +Derived catalogs: chunk/concat plans as first-class catalog rows. + +A restructuring operation materializes the current view's membership +into a fresh in-memory catalog whose *patch rows are the plan outputs*; +a `PlanResolver` turns an output row back into a Patch by loading the +member source patches through the parent's resolver and trimming or +merging them (the existing assembly engine). Every catalog operation — +select, order, window, union, equality, serialization — then runs the +identical code path for planned and identity spools. + +Single-writer rule: derived catalogs are always fresh in-memory +databases; the on-disk index is only ever written by the directory +syncer, and views never write. +""" + +from __future__ import annotations + +import secrets +from collections.abc import Mapping + +import numpy as np +import pandas as pd + +import dascore as dc +from dascore.io.index.backend import get_backend +from dascore.io.index.catalog import ( + CompositeResolver, + PatchCatalog, + PatchResolver, + _row_source_patch_id, + apply_exact_residuals, +) +from dascore.io.index.ingest import ( + CoordRecord, + PatchRecord, + SourceRecord, + typed_value, +) +from dascore.utils.misc import is_range +from dascore.utils.pd import adjust_segments + +PLAN_SCHEME = "plan://" +# columns that are structural/positional rather than patch attributes +_NON_ATTR = {"output_id", "dims", "coord_names", "patch"} + + +def _ns(value) -> int | None: + """Convert a datetime/timedelta-like envelope value to ns int.""" + if value is None or pd.isnull(value): + return None + if isinstance(value, pd.Timedelta | np.timedelta64): + return int(pd.Timedelta(value).value) + return int(pd.Timestamp(value).value) + + +def _num(value) -> float | None: + """Convert a numeric envelope value to float.""" + if value is None or pd.isnull(value): + return None + return float(value) + + +def _coord_record_from_row( + row: Mapping, name: str, dims: tuple[str, ...] | None = None +) -> CoordRecord | None: + """ + Build the envelope coord record for one output coordinate. + + Delegates to the ingest converter through a range CoordSummary so + virtual outputs carry the same identities real patches would: a + carried ``fp:`` def key survives for non-planned dims, and the + planned dim's range fingerprint is reconstructed exactly. ``dims`` + names the dimensions the coordinate rides (itself by default). + """ + from dascore.core.coords import CoordSummary + from dascore.io.index.ingest import _coord_record + + dims = (name,) if dims is None else dims + lo, hi = row.get(f"{name}_min"), row.get(f"{name}_max") + if lo is None or (pd.isnull(lo) and pd.isnull(hi)): + return None + step = row.get(f"{name}_step") + step = None if step is None or pd.isnull(step) else step + if isinstance(lo, str): + # string coords have no range representation; store the + # lexicographic envelope directly + key = row.get(f"_{name}_def_key") + fingerprint = ( + key[3:] if isinstance(key, str) and key.startswith("fp:") else None + ) + return CoordRecord( + coord_name=name, + value_kind="str", + dtype="str", + coord_dims=",".join(dims), + length=None, + units=None, + min_str=str(lo), + max_str=None if hi is None or pd.isnull(hi) else str(hi), + coord_hash=fingerprint, + ) + if isinstance(lo, pd.Timestamp | np.datetime64): + lo, hi = pd.Timestamp(lo).to_datetime64(), pd.Timestamp(hi).to_datetime64() + dtype = "datetime64[ns]" + elif isinstance(lo, pd.Timedelta | np.timedelta64): + lo, hi = pd.Timedelta(lo).to_timedelta64(), pd.Timedelta(hi).to_timedelta64() + dtype = "timedelta64[ns]" + else: + lo, hi = float(lo), float(hi) + step = None if step is None else abs(float(step)) + dtype = "float64" + if isinstance(step, pd.Timedelta): + step = step.to_timedelta64() + if step is not None and not step: + # a degenerate (zero) step is not a range; drop it rather than + # letting range reconstruction divide by it + step = None + # numeric envelope values are stored canonical-SI, so only the + # canonical (base) unit carried by the pivot may be attached — + # ingest's re-conversion is then the identity (time kinds attach + # without conversion) + units = row.get(f"_{name}_units") + if units == "" or (units is not None and pd.isnull(units)): + units = None + length = None + if step is not None: + length = int(round((hi - lo) / step)) + 1 + key = row.get(f"_{name}_def_key") + fingerprint = None + if isinstance(key, str) and key.startswith("fp:"): + fingerprint = key[3:] + summary = CoordSummary( + dtype=dtype, + min=lo, + max=hi, + step=step, + units=units, + dims=dims, + len=length, + fingerprint=fingerprint, + ) + return _coord_record(name, summary) + + +def _aux_coord_info( + source_rows: pd.DataFrame, + members: pd.DataFrame, + plan_dim: str, + coord_dims_map: Mapping[str, str], + trimmed_dims: frozenset[str] = frozenset(), +) -> dict[int, dict[str, dict]]: + """ + Aggregate per-output envelope info for auxiliary coordinates. + + Aggregated from the *member source rows* (authoritative, unlike the + planner's carried columns). Structural identity (def key and step, + which permit fingerprint claims) is kept only when every member + shares one def key and the values provably survive assembly: a + coordinate riding the planned dimension is trimmed/merged with it, + so only a lone unmodified member keeps identity there. Envelopes + always aggregate — the catalog contract is candidacy, with exact + values re-established at load. + """ + out: dict[int, dict[str, dict]] = {} + if not len(members) or not coord_dims_map: + return out + cols = [c for c in ("output_id", "_patch_id", "_modified") if c in members.columns] + joined = members[cols].merge(source_rows, on="_patch_id", how="left") + for name, dims_str in coord_dims_map.items(): + cmin, cmax = f"{name}_min", f"{name}_max" + if cmin not in joined.columns: + continue + dims = tuple(d for d in str(dims_str).split(",") if d) + rides = plan_dim in dims + # a residual selection trims the dims it rides at load, changing + # the values of every coordinate on those dims + trimmed = bool(set(dims) & trimmed_dims) + key_col, step_col = f"_{name}_def_key", f"{name}_step" + for output_id, sub in joined.groupby("output_id"): + lo, hi = sub[cmin].min(), sub[cmax].max() + if pd.isnull(lo) and pd.isnull(hi): + continue + keys = set(sub[key_col].dropna()) if key_col in sub.columns else set() + modified = bool(sub["_modified"].any()) if "_modified" in sub else False + keep = ( + len(keys) == 1 + and not trimmed + and (not rides or (len(sub) == 1 and not modified)) + ) + steps = ( + set(sub[step_col].dropna()) + if keep and step_col in sub.columns + else set() + ) + unit_col = f"_{name}_units" + units = set(sub[unit_col].dropna()) if unit_col in sub.columns else set() + info = { + cmin: lo, + cmax: hi, + step_col: steps.pop() if len(steps) == 1 else None, + key_col: keys.pop() if keep else None, + unit_col: units.pop() if len(units) == 1 else None, + "dims": dims, + } + out.setdefault(int(output_id), {})[name] = info + return out + + +def _output_records( + outputs: pd.DataFrame, + token: str, + aux_info: Mapping[int, Mapping[str, Mapping]] | None = None, +) -> list[SourceRecord]: + """Convert plan output rows into ingestible source records.""" + records = [] + aux_info = aux_info or {} + envelope_suffixes = ("_min", "_max", "_step", "_units") + for row in outputs.to_dict("records"): + output_id = int(row["output_id"]) + dims = str(row.get("dims") or "") + dim_names = [d for d in dims.split(",") if d] + coords = [] + for name in dim_names: + record = _coord_record_from_row(row, name) + if record is not None: + coords.append(record) + # auxiliary (non-dimension) coordinates remain on the assembled + # patches, so the catalog must keep describing them + for name, info in aux_info.get(output_id, {}).items(): + if name in dim_names: + continue + record = _coord_record_from_row(info, name, dims=info["dims"]) + if record is not None: + coords.append(record) + attrs = {} + for key, value in row.items(): + if ( + key in _NON_ATTR + or key.startswith("_") + or any(key.endswith(sfx) for sfx in envelope_suffixes) + or value is None + or (np.isscalar(value) and pd.isnull(value)) + ): + continue + typed = typed_value(value) + if typed is not None: + attrs[key] = typed + patch = PatchRecord( + source_patch_id=str(output_id), + dims=dims, + shape="", + n_dims=len(dim_names), + sample_count_total=None, + time_min=_ns(row.get("time_min")), + time_max=_ns(row.get("time_max")), + time_step=_ns(row.get("time_step")), + distance_min=_num(row.get("distance_min")), + distance_max=_num(row.get("distance_max")), + distance_step=_num(row.get("distance_step")), + attrs=attrs, + coords=tuple(coords), + ) + records.append( + SourceRecord( + source_path=f"{PLAN_SCHEME}{token}/{output_id}", + source_format="plan", + format_version="", + patches=(patch,), + ) + ) + return records + + +class PlanResolver(PatchResolver): + """ + Assemble plan-output rows from their member source patches. + + ``member_rows`` carries, per output, the full source row (path, + format, identity, attrs) with the planned dimension's envelope + replaced by the member's trim range; loading goes through ``loader`` + (live registry, files, and nested plan rows), applies the parent + view's residual selections, then trims/merges via the assembly + engine ("chunk" mode) or concatenates in order ("concat" mode). + """ + + def __init__( + self, + *, + token: str, + dim: str, + member_rows: pd.DataFrame, + loader: PatchResolver, + merge_kwargs: Mapping, + parent_residuals: tuple = (), + mode: str = "chunk", + check_behavior: str = "warn", + origin_path=None, + ): + if "output_id" not in member_rows.columns: + msg = "member_rows must carry an output_id column." + raise ValueError(msg) + # plan invariant: outputs without members must never be published + self.token = token + self.dim = dim + self.member_rows = member_rows.reset_index(drop=True) + self.loader = loader + self.merge_kwargs = dict(merge_kwargs) + self.parent_residuals = tuple(parent_residuals) + self.mode = mode + self.check_behavior = check_behavior + # informational only: the directory/file the plan derived from + self.origin_path = origin_path + + def live_entries(self) -> Mapping[str, dc.Patch]: + """Expose the loader's live registry (for absorption/transfer).""" + return self.loader.live_entries() + + def plan_entries(self) -> Mapping[str, PlanResolver]: + """Route plan:// paths with this resolver's token to it.""" + nested = dict(getattr(self.loader, "plan_entries", dict)()) + nested[f"{PLAN_SCHEME}{self.token}/"] = self + return nested + + def _assembler(self): + from dascore.utils.patch_assembly import PatchAssembler + + return PatchAssembler( + load_patch=self._load_member, + merge_kwargs=self.merge_kwargs, + ) + + def _load_member(self, kwargs: Mapping) -> dc.Patch: + """Load one member source patch, applying parent residuals.""" + trim = {} + if kwargs.get("_modified"): + trim = { + k: v + for k, v in kwargs.items() + if not str(k).startswith("_") + and k not in ("path", "file_format", "file_version", "source_patch_id") + } + patch = self.loader.resolve(kwargs, **trim) + return apply_exact_residuals(patch, self.parent_residuals) + + def resolve(self, row: Mapping, **trim) -> dc.Patch: + """Assemble the output patch a plan row describes.""" + output_id = int(_row_source_patch_id(row)) + members = self.member_rows[self.member_rows["output_id"] == output_id] + assert len(members), "no plan members found for output row" + if self.mode == "identity": + # one untouched member per output; residuals apply at load + assert len(members) == 1 + return self._load_member(members.iloc[0].to_dict()) + if self.mode == "concat": + from dascore.utils.patch import concatenate_patches + + patches = [ + self._load_member(kwargs) for kwargs in members.to_dict("records") + ] + out = concatenate_patches( + patches, check_behavior=self.check_behavior, **{self.dim: None} + ) + assert len(out) == 1 + return out[0] + joined = members.assign(current_index=output_id) + patches = self._assembler()._patch_from_instruction_df(joined) + assert len(patches) == 1 + return patches[0] + + +def _residual_ranges(residuals) -> dict: + """Envelope-applicable value ranges from a residual tuple.""" + out = {} + for coords, samples in residuals: + if samples: + continue + for name, value in coords.items(): + magnitudes = getattr(value, "magnitudes", None) + if magnitudes is not None: + out[name] = magnitudes + elif is_range(value) and not any( + hasattr(b, "units") for b in value if b is not None + ): + out[name] = value + return out + + +def derived_catalog( + *, + source_rows: pd.DataFrame, + plan, + parent: PatchCatalog | None, + merge_kwargs: Mapping, + mode: str = "chunk", + check_behavior: str = "warn", + origin_path=None, +) -> PatchCatalog: + """ + Materialize a plan into a fresh in-memory catalog. + + ``source_rows`` are the full member source rows (path/format/ + identity plus envelopes and attrs) keyed by ``_patch_id`` matching + ``plan.members``; ``parent`` supplies the resolver (live registry, + file root, nested plans) and the residual selections its view + carried, which member loading re-applies. + """ + token = secrets.token_hex(8) + name = plan.dim + trims = plan.members + trim_cols = [c for c in trims.columns if c not in ("_patch_id",)] + sources = source_rows.copy(deep=False) + if "_patch_id" not in sources.columns: + from dascore.utils.chunk_plan import _ensure_patch_id + + sources = _ensure_patch_id(sources) + member_rows = trims[["_patch_id", *[c for c in trim_cols]]].merge( + sources.drop(columns=[c for c in trim_cols if c in sources], errors="ignore"), + on="_patch_id", + how="left", + ) + # the member's trimmed range replaces the source envelope for loading + member_rows = member_rows.drop(columns=["_patch_id"]) + parent_residuals = () if parent is None else parent._residuals + # resolve stored-relative paths once; the derived catalog is + # root-independent afterwards + root = getattr(parent.resolver, "_root", None) if parent is not None else None + if root is not None and "path" in member_rows.columns: + member_rows = member_rows.assign( + path=[ + str(p) + if "://" in str(p) or str(p).startswith("/") + else str(root / str(p)) + for p in member_rows["path"] + ] + ) + loader = CompositeResolver() + if parent is not None: + member_paths = set(member_rows.get("path", pd.Series(dtype=str)).astype(str)) + loader.absorb(parent.resolver, paths=member_paths) + resolver = PlanResolver( + token=token, + dim=name, + member_rows=member_rows, + loader=loader, + merge_kwargs=merge_kwargs, + parent_residuals=parent_residuals, + mode=mode, + check_behavior=check_behavior, + origin_path=origin_path, + ) + backend = get_backend(":memory:") + coord_dims_map = {} if parent is None else parent.backend.coord_dims_map() + # residual selections trim at load; identity claims (def keys) for + # coordinates on the trimmed dims would describe the untrimmed values + residual_names = {n for coords, _ in parent_residuals for n in coords} + trimmed_dims = frozenset( + d for n in residual_names for d in str(coord_dims_map.get(n, n)).split(",") if d + ) + outputs = plan.outputs + stale_keys = [ + f"_{c}_def_key" + for c, dims_str in coord_dims_map.items() + if set(str(dims_str).split(",")) & trimmed_dims + and f"_{c}_def_key" in outputs.columns + ] + if stale_keys: + outputs = outputs.drop(columns=stale_keys) + aux_info = _aux_coord_info(sources, trims, name, coord_dims_map, trimmed_dims) + backend.write_sources(_output_records(outputs, token, aux_info=aux_info)) + return PatchCatalog(backend=backend, resolver=resolver) + + +def collapse_working_df(catalog: PatchCatalog) -> pd.DataFrame | None: + """ + Return the re-planning frame for a derived catalog, or None. + + Re-planning the *same* dimension collapses: it plans over the + current view's *members* — the trimmed source rows — restricted to + outputs the view still presents, with the view's value residuals + applied to the envelopes. (Planning a different dimension must keep + the assembled boundaries, so its caller plans over the output rows + instead and never collapses.) + """ + resolver = catalog.resolver + if not isinstance(resolver, PlanResolver): + return None + members = resolver.member_rows + if catalog.is_view: + present = { + int(_row_source_patch_id(row)) for row in catalog.to_df().to_dict("records") + } + members = members[members["output_id"].isin(present)] + ranges = _residual_ranges(catalog._residuals) + working = members.drop(columns=["output_id", "_modified"], errors="ignore") + if ranges: + working = adjust_segments(working, ignore_bad_kwargs=True, **ranges) + return working.reset_index(drop=True) diff --git a/dascore/io/index/query.py b/dascore/io/index/query.py new file mode 100644 index 000000000..a7f79d8b3 --- /dev/null +++ b/dascore/io/index/query.py @@ -0,0 +1,507 @@ +""" +Query model and SQL generation for the spool index. + +Implements the selector semantics spec (see +`.scratch/selector_semantics_spec.md`): the index only produces +candidates — predicates the summary cannot evaluate exactly are the +caller's responsibility at patch-load time. Predicates SQLite cannot evaluate +exactly are applied as pandas residual filters. +""" + +from __future__ import annotations + +import re +from collections.abc import Sequence +from dataclasses import dataclass, field + +import numpy as np +import pandas as pd + +from dascore.exceptions import InvalidSpoolQueryError, ParameterError, UnitError +from dascore.io.index.dialect import BaseDialect +from dascore.io.index.ingest import typed_value +from dascore.units import convert_units +from dascore.utils.misc import is_range + +_GLOB_CHARS = frozenset("*?[") +_UNSET = object() + + +@dataclass(frozen=True) +class Query: + """ + A resolved spool query. + + Name resolution (bare kwargs -> attrs first, then coords) happens + above this layer; a Query already knows which namespace each + predicate belongs to. + """ + + attrs: dict = field(default_factory=dict) + coords: dict = field(default_factory=dict) + + +def _is_collection(value) -> bool: + """True for non-string collections (membership predicates).""" + if isinstance(value, str | bytes): + return False + if isinstance(value, np.ndarray): + return True + return isinstance(value, list | tuple | set | frozenset) + + +# Shared join skeleton for the patch relation. The attrs join is 1:1 (one +# attrs row per patch), so it is safe for both the projection and the count. +_FROM = ( + "FROM patches p " + "JOIN sources s ON s.source_id = p.source_id " + "LEFT JOIN attrs a ON a.patch_id = p.patch_id " +) + + +def _as_query_list(query: Query | Sequence[Query]) -> list[Query]: + """Normalize a single Query or a sequence of them to a list.""" + return [query] if isinstance(query, Query) else list(query) + + +def _coerce_scalar(value, target_kinds: set[str]): + """ + Coerce a query scalar to (kind, storable value). + + Follows the coercion table in the selector spec; datetime-like + strings become time queries only when the target has a time kind. + """ + typed = typed_value(value) + if typed is None: + msg = f"Cannot use {value!r} as a spool query value." + raise InvalidSpoolQueryError(msg) + if typed.kind == "str" and "time" in target_kinds: + try: + retyped = typed_value(np.datetime64(pd.Timestamp(value), "ns")) + return retyped + except (ValueError, TypeError): + pass + return typed + + +def _normalize_unit(value) -> str | None: + """Return a nullable unit string from a dataframe value.""" + return None if value is None or pd.isnull(value) else str(value) + + +def _to_target_unit(typed, target_units: str | None, name: str): + """Validate/convert a numeric query value for one stored unit.""" + if typed.kind != "num" or typed.units is None: + return typed.value + if target_units is None: + msg = f"Cannot query unitless {name!r} with units {typed.units!r}." + raise UnitError(msg) + return convert_units(typed.value, to_units=target_units, from_units=typed.units) + + +def _range_bounds( + value, + target_kinds: set[str], + target_units: str | None | object = _UNSET, + name: str = "value", +): + """ + Return (kind, lo, hi, typed_values) from a range tuple. + + Open bounds (None/Ellipsis) are skipped; typed_values carries the + coerced usable bounds so callers don't coerce twice. + """ + lo_raw, hi_raw = value + kind = None + typed_values = [] + typed_bounds = [] + for raw, side in ((lo_raw, "lo"), (hi_raw, "hi")): + if raw is None or raw is Ellipsis: + continue + typed = _coerce_scalar(raw, target_kinds) + typed_values.append(typed) + knd = typed.kind + if kind is not None and knd != kind: + msg = f"Range bounds {value!r} have mixed kinds ({kind}, {knd})." + raise InvalidSpoolQueryError(msg) + kind = knd + typed_bounds.append((side, typed)) + if kind is None: + msg = f"Range {value!r} has no usable bounds." + raise InvalidSpoolQueryError(msg) + + # Validate all bound kinds before attempting unit conversion. An + # unsupported but internally consistent kind is a valid no-match query; + # mixed kinds remain an invalid range. + lo = hi = None + for side, typed in typed_bounds: + val = ( + typed.value + if target_units is _UNSET or kind not in target_kinds + else _to_target_unit(typed, target_units, name) + ) + if side == "lo": + lo = val + else: + hi = val + if kind in target_kinds and lo is not None and hi is not None and lo > hi: + msg = f"Range {value!r} has lo > hi after coercion." + raise InvalidSpoolQueryError(msg) + return kind, lo, hi, typed_values + + +def _compatible_coord_units( + rows: pd.DataFrame, typed_values: list, name: str +) -> set[str] | None: + """ + Return stored units compatible with quantity-valued coord selectors. + + None means the query carries no units (no unit constraint at all); a + set constrains matching to those units plus NULL-unit definitions + (which can never be proven incompatible, so they stay candidates). + Raises UnitError only when every stored definition has units and none + are compatible. + """ + query_units = { + x.units for x in typed_values if x is not None and x.units is not None + } + if not query_units: + return None + first = next(iter(query_units)) + for other in query_units - {first}: + convert_units(1.0, to_units=first, from_units=other) + stored = {_normalize_unit(x) for x in rows.get("units", ())} + compatible = set() + for unit in stored - {None}: + try: + convert_units(1.0, to_units=unit, from_units=first) + except UnitError: + continue + compatible.add(unit) + if not compatible and None not in stored: + raise UnitError(f"Coordinate {name!r} has no units compatible with {first!r}.") + return compatible + + +@dataclass +class _Where: + """Accumulates WHERE clauses and parameters.""" + + clauses: list[str] = field(default_factory=list) + params: list = field(default_factory=list) + + def add(self, clause: str, *params): + self.clauses.append(clause) + self.params.extend(params) + + @property + def sql(self) -> str: + return " AND ".join(self.clauses) if self.clauses else "TRUE" + + +def build_attr_clause( + where: _Where, + dialect: BaseDialect, + attr_meta: pd.DataFrame, + name: str, + value, +) -> re.Pattern | None: + """ + Add SQL for one attr predicate; return a residual pattern if the + predicate must be re-applied in pandas (regex). + """ + rows = attr_meta[attr_meta["attr_name"] == name] + if rows.empty: + msg = f"{name!r} is not an attribute of any patch in this spool." + raise InvalidSpoolQueryError(msg) + kinds = set(rows["value_kind"]) + columns = dict(zip(rows["value_kind"], rows["column_name"])) + units = {row.value_kind: _normalize_unit(row.units) for row in rows.itertuples()} + + def col(kind): + return f"a.{dialect.quote(columns[kind])}" + + if isinstance(value, re.Pattern): + # Regex is a residual filter; SQL only requires the attr be + # present (str kind) so candidates are a superset. + if "str" not in kinds: + where.add("FALSE") + return None + where.add(f"{col('str')} IS NOT NULL") + return value + if is_range(value): + # Attr metadata has one canonical unit per typed column. + probe = next( + ( + _coerce_scalar(x, kinds) + for x in value + if x is not None and x is not Ellipsis + ), + None, + ) + target_units = units.get(probe.kind) if probe is not None else None + kind, lo, hi, _ = _range_bounds(value, kinds, target_units, name) + if kind not in kinds: + where.add("FALSE") + return None + for bound, op in ((lo, ">="), (hi, "<=")): + if bound is not None: + where.add(f"{col(kind)} {op} ?", bound) + return None + if _is_collection(value): + coerced = [_coerce_scalar(v, kinds) for v in value] + by_kind: dict[str, list] = {} + for typed in coerced: + if typed.kind not in kinds: + continue + val = _to_target_unit(typed, units.get(typed.kind), name) + by_kind.setdefault(typed.kind, []).append(val) + subclauses = [] + params = [] + for kind, vals in by_kind.items(): + marks = ", ".join("?" for _ in vals) + subclauses.append(f"{col(kind)} IN ({marks})") + params.extend(vals) + if not subclauses: + where.add("FALSE") + else: + where.add("(" + " OR ".join(subclauses) + ")", *params) + return None + if isinstance(value, str) and _GLOB_CHARS & set(value): + if "str" not in kinds: + where.add("FALSE") + return None + where.add(dialect.glob(col("str")), value) + return None + typed = _coerce_scalar(value, kinds) + kind = typed.kind + if kind not in kinds: + where.add("FALSE") + return None + val = _to_target_unit(typed, units.get(kind), name) + where.add(f"{col(kind)} = ?", val) + return None + + +def build_coord_clause( + where: _Where, + dialect: BaseDialect, + coord_meta: pd.DataFrame, + name: str, + value, +) -> None: + """ + Add a patch_coords/coord_defs semi-join clause for one coord + predicate. + + Candidacy only: envelope overlap, never false negatives. Exact + membership/boolean masks are applied at patch load, above this layer. + """ + rows = coord_meta[coord_meta["coord_name"] == name] + if isinstance(value, tuple) and len(value) != 2: + msg = f"Coordinate range for {name!r} must be a length 2 sequence." + raise ParameterError(msg) + kinds = set(rows["value_kind"]) or {"time", "num", "str"} + typed_values = [] + if is_range(value): + kind, lo, hi, typed_values = _range_bounds(value, kinds) + else: + # Scalars, value membership, and boolean masks have no exact + # patch-level meaning spool-wide; resolve_query rejects them + # before SQL composition, so only a hand-built Query can reach + # this. + msg = f"Coordinate {name!r} accepts range selectors; got {value!r}." + raise InvalidSpoolQueryError(msg) + + compatible_units = _compatible_coord_units(rows, typed_values, name) + + min_col, max_col = { + "time": ("min_ns", "max_ns"), + "dur": ("min_ns", "max_ns"), + "num": ("min_num", "max_num"), + "str": ("min_str", "max_str"), + None: (None, None), + }[kind] + conditions = ["pc.coord_name = ?"] + params: list = [name] + if kind is not None: + if kind in ("time", "dur"): + # absolute queries match absolute coords, durations relative. + conditions.append("cd.is_relative = ?") + params.append(kind == "dur") + kind_match = "time" + else: + kind_match = kind + conditions.append("cd.value_kind = ?") + params.append(kind_match) + if compatible_units is not None: + # NULL-unit defs stay candidates: IN () never matches NULL and + # unitless values cannot be proven dimensionally incompatible. + if compatible_units: + marks = ", ".join("?" for _ in compatible_units) + conditions.append(f"(cd.units IN ({marks}) OR cd.units IS NULL)") + params.extend(sorted(compatible_units)) + else: + conditions.append("cd.units IS NULL") + # lo bounds the coord max (overlap), hi bounds the coord min. + for bound, bound_col, op in ((lo, max_col, ">="), (hi, min_col, "<=")): + if bound is None: + continue + clause = f"cd.{bound_col} {op} ?" + if compatible_units is not None: + clause = f"(cd.units IS NULL OR {clause})" + conditions.append(clause) + params.append(bound) + # A semi-join the engine can evaluate once (idx_pcoords_name) beats a + # correlated EXISTS probed per patch row (~2.5x on a 200k-source index). + where.add( + "p.patch_id IN (SELECT pc.patch_id FROM patch_coords pc " + "JOIN coord_defs cd ON cd.coord_def_id = pc.coord_def_id " + "WHERE " + " AND ".join(conditions) + ")", + *params, + ) + + +def _build_where( + queries: list[Query], + dialect: BaseDialect, + attr_meta: pd.DataFrame, + coord_meta: pd.DataFrame, +) -> tuple[_Where, list[tuple[str, re.Pattern]]]: + """Compose the shared WHERE clause and any regex residuals.""" + where = _Where() + residuals: list[tuple[str, re.Pattern]] = [] + for one in queries: + for name, value in one.attrs.items(): + residual = build_attr_clause(where, dialect, attr_meta, name, value) + if residual is not None: + residuals.append((name, residual)) + for name, value in one.coords.items(): + build_coord_clause(where, dialect, coord_meta, name, value) + return where, residuals + + +# typed coord_defs envelope-minimum column per value kind +_COORD_MIN_COLUMNS = {"num": "min_num", "time": "min_ns", "str": "min_str"} +# the two conventional dims cached as columns on the patches table +_HOT_COORDS = ("time", "distance") + + +def _order_clause( + order_by, dialect: BaseDialect, attr_meta: pd.DataFrame, coord_meta: pd.DataFrame +) -> tuple[str, list]: + """ + Resolve an order spec into an ORDER BY clause and its parameters. + + ``order_by`` is ``(kind, name, ascending)`` where kind is "attr" + (an attrs-table column ordered by its typed column) or "coord" + (ordered by the coordinate's envelope minimum — the hot patches + column when cached, otherwise the linked coord_defs typed minimum). + The ordinal contract supplies the deterministic tiebreak. + """ + kind, name, ascending = order_by + direction = "ASC" if ascending else "DESC" + params: list = [] + if kind == "coord" and name in _HOT_COORDS: + column = f"p.{dialect.quote(f'{name}_min')}" + elif kind == "coord": + rows = coord_meta[coord_meta["coord_name"] == name] + # a coord observed under several kinds orders by its first kind + value_kind = str(rows["value_kind"].iloc[0]) + min_col = _COORD_MIN_COLUMNS[value_kind] + column = ( + f"(SELECT cd.{min_col} FROM patch_coords pc " + "JOIN coord_defs cd ON cd.coord_def_id = pc.coord_def_id " + "WHERE pc.patch_id = p.patch_id AND pc.coord_name = ?)" + ) + params.append(name) + else: + rows = attr_meta[attr_meta["attr_name"] == name] + columns = [dialect.quote(c) for c in rows["column_name"]] + # an attr observed under several kinds orders by its first column + column = f"a.{columns[0]}" + # rows without a value sort last regardless of direction (matching + # the ordinal renumberer's missing-time-last rule); the null key + # repeats the column expression, so its parameters repeat too + params = [*params, *params] + sql = f"ORDER BY {column} IS NULL, {column} {direction}, s.ordinal, p.patch_id" + return sql, params + + +def build_sql( + query: Query | Sequence[Query], + dialect: BaseDialect, + attr_meta: pd.DataFrame, + coord_meta: pd.DataFrame, + count: bool = False, + order_by=None, + patch_ids=None, + ids_only: bool = False, +) -> tuple[str, list, list[tuple[str, re.Pattern]]]: + """ + Build SQL for one or more AND-composed queries. + + By default this projects the flat relation; with count=True the same + WHERE is reused for a COUNT with no projection, coordinate pivot, or + ordering; with ids_only=True only ordered patch ids are projected + (the cheap realization slices/windows use). coord_meta must cover + every coordinate the queries reference (it may be empty for + attr-only queries). ``order_by`` overrides the default ordinal + ordering (see `_order_clause`); ``patch_ids`` restricts rows to an + id membership (one JSON parameter, so the SQLite bound-variable cap + does not limit membership size). + + Returns (sql, params, residuals), where residuals pairs attr names + with regex patterns that must be re-applied to the resulting + dataframe. For a count a non-empty residual means the count is not + SQL-resolvable (regex must inspect rows) and the caller must fall + back to a projected count. + """ + import json + + queries = _as_query_list(query) + where, residuals = _build_where(queries, dialect, attr_meta, coord_meta) + if patch_ids is not None: + where.add( + "p.patch_id IN (SELECT value FROM json_each(?))", + json.dumps([int(x) for x in patch_ids]), + ) + if count: + # COUNT(p.patch_id) counts patches; a WHERE may reference a.. + sql = f"SELECT COUNT(p.patch_id) AS n {_FROM}WHERE {where.sql}" + return sql, where.params, residuals + if order_by is not None: + order, order_params = _order_clause(order_by, dialect, attr_meta, coord_meta) + else: + # the ordering contract: source ordinal, then file-internal order + order, order_params = "ORDER BY s.ordinal, p.patch_id", [] + params = [*where.params, *order_params] + if ids_only: + sql = f"SELECT p.patch_id {_FROM}WHERE {where.sql} {order}" + return sql, params, residuals + # attr columns selected explicitly: `a.*` would duplicate patch_id and + # engines disagree on how to dedupe result column names. + attr_cols = "".join( + f", a.{dialect.quote(col)}" for col in attr_meta["column_name"].unique() + ) + sql = ( + "SELECT s.source_path, s.base_uri, s.source_format, s.format_version, " + f"p.*{attr_cols} " + f"{_FROM}" + f"WHERE {where.sql} " + f"{order}" + ) + return sql, params, residuals + + +def apply_residuals( + df: pd.DataFrame, residuals: list[tuple[str, re.Pattern]] +) -> pd.DataFrame: + """Apply regex residual filters to the flat relation.""" + for name, pattern in residuals: + col = df[name] + keep = col.map( + lambda x: bool(pattern.search(x)) if isinstance(x, str) else False + ) + df = df[keep] + return df diff --git a/dascore/io/index/schema.py b/dascore/io/index/schema.py new file mode 100644 index 000000000..249108568 --- /dev/null +++ b/dascore/io/index/schema.py @@ -0,0 +1,231 @@ +""" +Logical schema for the spool index. + +The schema uses four primitive storage types (int64, float64, str, bool). +Times and durations are always epoch/plain nanoseconds stored as int64, +never engine-native timestamp types. +""" + +from __future__ import annotations + +from types import MappingProxyType + +# Version of the index schema, independent of dascore's version. +INDEX_VERSION = 3 +# Identity string so any tool can sanity-check what it opened. +WHAT_IS_THIS = "dascore_spool_index" + +# Value kinds for typed attr columns and coord rows. +KINDS = ("num", "str", "bool", "time", "dur") +# Storage type (logical) backing each kind. +KIND_STORAGE = MappingProxyType( + { + "num": "float64", + "str": "str", + "bool": "bool", + "time": "int64", # epoch ns + "dur": "int64", # ns + } +) + +META_DATA = MappingProxyType( + { + "what_is_this": "str", + "index_version": "int64", + "dascore_version": "str", + "last_indexed_ns": "int64", + } +) + +SOURCES = MappingProxyType( + { + "source_id": "int64", + "base_uri": "str", + "source_path": "str", + "source_format": "str", + "format_version": "str", + "mtime_ns": "int64", + "size_bytes": "int64", + "last_indexed_ns": "int64", + # The catalog's explicit ordering contract: patch rows present in + # (ordinal, patch_id) order. Assigned at ingest (insertion + # sequence); a replaced source keeps its position while new + # sources append, so merging catalogs concatenates and + # deduplication keeps first-occurrence position with + # last-occurrence metadata (dict-merge semantics). The directory + # syncer renumbers to time order after each sync, preserving the + # conventional time-ordered presentation of file archives. + "ordinal": "int64", + } +) + +# Frozen structural table; nothing dynamic is ever added here. The +# time/distance envelopes are cached summaries of the two conventional +# dims (hot path), not attr promotion. +PATCHES = MappingProxyType( + { + "patch_id": "int64", + "source_id": "int64", + "source_patch_id": "str", + "n_dims": "int64", + "dims": "str", + "shape": "str", + "sample_count_total": "int64", + "time_min": "int64", # epoch ns; NULL for relative-time patches + "time_max": "int64", + "time_step": "int64", + "distance_min": "float64", # canonical SI (m) + "distance_max": "float64", + "distance_step": "float64", + } +) + +# attrs table starts with only the key; typed columns (`__`) +# are added lazily at ingest. +ATTRS_BASE = MappingProxyType({"patch_id": "int64"}) + +ATTR_META = MappingProxyType( + { + "attr_name": "str", # original (unsanitized) attr name + "value_kind": "str", + "column_name": "str", # sanitized column in the attrs table + "units": "str", # canonical unit for num kinds, nullable + } +) + +# Unique coordinate summaries, deduplicated across patches. Range coordinates +# use a semantic fingerprint supplied by the scan or reconstructed exactly +# from the range summary. Non-range coordinates without a fingerprint use a +# summary hash for storage deduplication, but it is not exposed as value identity. +COORD_DEFS = MappingProxyType( + { + "coord_def_id": "int64", + "def_key": "str", + "fingerprint": "str", # nullable; semantic hash from CoordSummary + "value_kind": "str", # num | time | str + "dtype": "str", + "length": "int64", + "units": "str", # original unit string; numeric values stored SI + "min_num": "float64", + "max_num": "float64", + "step_num": "float64", + "min_ns": "int64", + "max_ns": "int64", + "step_ns": "int64", + "min_str": "str", + "max_str": "str", + "is_monotonic": "bool", + "is_relative": "bool", + } +) + +# Links a patch to its coord defs; the name and dims are patch-level +# semantics (two patches can share values under different names). +PATCH_COORDS = MappingProxyType( + { + "patch_id": "int64", + "coord_name": "str", + "coord_dims": "str", + "coord_def_id": "int64", + } +) + +TABLES = MappingProxyType( + { + "meta_data": META_DATA, + "sources": SOURCES, + "patches": PATCHES, + "attrs": ATTRS_BASE, + "attr_meta": ATTR_META, + "coord_defs": COORD_DEFS, + "patch_coords": PATCH_COORDS, + } +) + +# Keeping constraints beside the logical columns makes the stored contract +# explicit and keeps dynamic attr-column DDL separate from table identity. +TABLE_CONSTRAINTS = MappingProxyType( + { + "meta_data": ( + "PRIMARY KEY (what_is_this)", + f"CHECK (what_is_this = '{WHAT_IS_THIS}')", + ), + "sources": ( + "PRIMARY KEY (source_id)", + "UNIQUE (base_uri, source_path)", + "CHECK (base_uri IS NOT NULL)", + "CHECK (source_path IS NOT NULL)", + ), + "patches": ( + "PRIMARY KEY (patch_id)", + "UNIQUE (source_id, source_patch_id)", + "FOREIGN KEY (source_id) REFERENCES sources(source_id) ON DELETE CASCADE", + ), + "attrs": ( + "PRIMARY KEY (patch_id)", + "FOREIGN KEY (patch_id) REFERENCES patches(patch_id) ON DELETE CASCADE", + ), + "attr_meta": ( + "PRIMARY KEY (attr_name, value_kind)", + "UNIQUE (column_name)", + "CHECK (value_kind IN ('num', 'str', 'bool', 'time', 'dur'))", + ), + "coord_defs": ( + "PRIMARY KEY (coord_def_id)", + "UNIQUE (def_key)", + "CHECK (value_kind IN ('num', 'time', 'str'))", + "CHECK (is_monotonic IS NULL OR is_monotonic IN (0, 1))", + "CHECK (is_relative IS NULL OR is_relative IN (0, 1))", + ), + "patch_coords": ( + "PRIMARY KEY (patch_id, coord_name)", + "FOREIGN KEY (patch_id) REFERENCES patches(patch_id) ON DELETE CASCADE", + "FOREIGN KEY (coord_def_id) REFERENCES coord_defs(coord_def_id)", + ), + } +) + +# Attr names which would collide with structural storage or flat-relation +# columns. Attrs with these (sanitized) names stay on the patch but are +# not indexed; ingest warns about them. +RESERVED_ATTR_COLUMNS = frozenset( + { + # storage tables + "patch_id", + "source_id", + "source_patch_id", + "source_path", + "source_format", + "format_version", + "base_uri", + "mtime_ns", + "size_bytes", + "n_dims", + "dims", + "shape", + "sample_count_total", + "coord_def_id", + "def_key", + # flat-relation (spool-facing) names + "path", + "file_format", + "file_version", + # spool instruction machinery + "current_index", + "source_index", + "output_id", + "patch", + } +) + +# Structural columns the spool machinery must not see: unique-per-patch +# values block chunk merge-compatibility grouping, which compares all +# non-private columns. +SPOOL_HIDDEN_COLUMNS = ("n_dims", "sample_count_total", "shape") + +# Explicit secondary indexes. Every other access path is covered by a +# PRIMARY KEY or UNIQUE autoindex above — patch_coords(patch_id, +# coord_name), sources(base_uri, source_path), patches(source_id, +# source_patch_id), coord_defs(def_key) — and duplicating them measured +# ~25% extra file size and slower writes for no query gain. +INDEXES = (("idx_pcoords_name", "patch_coords", "coord_name"),) diff --git a/dascore/io/indexer.py b/dascore/io/indexer.py index 22cf73db3..f965025c2 100644 --- a/dascore/io/indexer.py +++ b/dascore/io/indexer.py @@ -1,40 +1,22 @@ -"""An HDF5-based indexer for local file systems.""" +""" +Base indexer interface and index-location utilities. + +The concrete directory indexer lives in `dascore.io.index.indexer` +(`DBDirectoryIndexer`); this module holds the abstract interface and the +machinery for tracking index locations when the data directory itself is +not writable (e.g. read-only archives). +""" from __future__ import annotations import abc import json -import os -import time -import warnings from contextlib import suppress from functools import cache from pathlib import Path -import pandas as pd from typing_extensions import Self -import dascore as dc -from dascore.compat import UPath -from dascore.config import config_attr, get_config -from dascore.constants import PROGRESS_LEVELS -from dascore.exceptions import InvalidIndexVersionError -from dascore.utils.hdf5 import HDFPatchIndexManager -from dascore.utils.misc import iterate -from dascore.utils.paths import requires_local_directory -from dascore.utils.pd import filter_df -from dascore.utils.time import ( - get_max_min_times, - saturate_add, - saturate_subtract, - to_timedelta64, -) - -# supported read_hdf5 kwargs -READ_HDF5_KWARGS = frozenset( - {"columns", "where", "mode", "errors", "start", "stop", "key", "chunksize"} -) - @cache def _get_index_map(cache_path) -> dict: @@ -70,20 +52,6 @@ def _update_index_map(updates, cache_path) -> dict: return data -def _directory_writable(path): - """Return True if the directory is writable else False.""" - name = "._dascore_write_test_delete_me" - path = Path(path) / name - path.parent.mkdir(exist_ok=True, parents=True) - try: - open(path, "w").close() - except (PermissionError, IsADirectoryError): - return False - else: - os.remove(path) - return True - - class AbstractIndexer: """ A base class for indexers. @@ -101,256 +69,10 @@ def update(self) -> Self: Resets any previous selection. """ - -class DirectoryIndexer(AbstractIndexer): - """ - A class for indexing a directory of dascore-readable files. - - This works by crawling the directory, getting a summary about the data it - contains, then creating a small HDF index file which can be queried later - on. - - Parameters - ---------- - path - The path to a directory containing DAS files. - index_path - The path to the index. By default, the index will be created on the - top level of the data directory. If another index is - """ - - ext = "" - _namespace = "" - _index_name = ".dascore_index.h5" # name of index file - - def __init__(self, path: str | Path, cache_size: int = 5, index_path=None): - self.max_size = cache_size - self.path = ( - UPath(path).absolute() if isinstance(path, UPath) else Path(path).absolute() - ) - requires_local_directory(self.path, label="DirectoryIndexer") - self.path = Path(self.path).absolute() - self.index_path = Path(self._find_index_file(self.path, index_path)) - self._current_index = 0 - self._index_table = HDFPatchIndexManager( - self.index_path, - self._namespace, - ) - self.cache = pd.DataFrame( - index=range(cache_size), columns="t1 t2 kwargs cindex".split() - ) - - index_map_path: Path = config_attr("directory_index_map_path") - - def _find_index_file(self, data_path, index_path=None): - """Find the path to the index file.""" - data_path = Path(data_path).absolute() - # user specified index path - if index_path: - update = {str(data_path): str(Path(index_path).absolute())} - _update_index_map(update, cache_path=str(self.index_map_path)) - return index_path - # see if expected path is in data path - expected_path = data_path / self._index_name - with suppress(PermissionError): - if expected_path.exists(): - return expected_path - # else load path map and see if it knows where the index is. - path_map = _get_index_map(cache_path=str(self.index_map_path)) - if out := path_map.get(str(data_path)): - return out - # if not, set the path to either the data path, if writable, - # else the dascore cache - if not _directory_writable(data_path): - new_path = "_dascore_index_" + str(abs(hash(data_path))) + ".h5" - index_path = self.index_map_path.parent / new_path - update = {str(data_path): str(index_path.absolute())} - _update_index_map(update, cache_path=str(self.index_map_path)) - else: - index_path = data_path / self._index_name - return index_path - - def get_contents(self, buffer=None, **kwargs) -> pd.DataFrame: - """ - Get contents of directory with specific query params. - - Parameters - ---------- - buffer - A buffer to ensure enough info is returned from hdf index. - kwargs - Used to query contents. - """ - # create index if it doesn't exist - if not self.index_path.exists(): - self.update() - # if the index still doesn't exist there are no readable files, return - # empty df. - if not self.index_path.exists(): - return pd.DataFrame(columns=self._index_table.index_columns) - time_min, time_max = get_max_min_times(kwargs.pop("time", None)) - hdf5_kwargs, kwargs = self._separate_hdf5_kwargs(kwargs) - buffer = get_config().index_query_buffer if buffer is None else buffer - buffer = to_timedelta64(buffer) - # find out if the query falls within one cached times - con1 = self.cache.t1 <= time_min - con2 = self.cache.t2 >= time_max - con3 = self.cache.kwargs == self._kwargs_to_str(kwargs) - cached_index = self.cache[con1 & con2 & con3] - if not len(cached_index): # query is not cached get it from hdf5 file - index = self._index_table.get_index( - time_min=time_min, - time_max=time_max, - **hdf5_kwargs, - ) - self._set_cache(index, time_min, time_max, hdf5_kwargs) - else: - index = cached_index.iloc[0]["cindex"] - # trim down index - con1 = index["time_min"] >= saturate_add(time_max, buffer) - con2 = index["time_max"] <= saturate_subtract(time_min, buffer) - pre_filter_df = index[~(con1 | con2)] - out = pre_filter_df[ - filter_df( - pre_filter_df, - time=(time_min, time_max), - ignore_bad_kwargs=True, - **kwargs, - ) - ] - return out - - def __str__(self): - """Rep. indexer as a string.""" - msg = f"{self.__class__.__name__} managing: {self.path}" - return msg - - __repr__ = __str__ - - __call__ = get_contents - - def _separate_hdf5_kwargs(self, kwargs): - """Ensure kwargs are supported.""" - kdf_kwargs = {i: v for i, v in kwargs.items() if i in READ_HDF5_KWARGS} - kwargs = {i: v for i, v in kwargs.items() if i not in READ_HDF5_KWARGS} - return kdf_kwargs, kwargs - - def _set_cache(self, index, starttime, endtime, kwargs): - """Cache the current index.""" - ser = pd.Series( - { - "t1": starttime, - "t2": endtime, - "cindex": index, - "kwargs": self._kwargs_to_str(kwargs), - } - ) - self.cache.loc[self._get_next_index()] = ser - - def clear_cache(self): - """Removes all cached dataframes.""" - self.cache = pd.DataFrame( - index=range(self.max_size), columns="t1 t2 kwargs cindex".split() - ) - - def _get_next_index(self): - """ - Get the next index value on cache. - Note we can't use itertools.cycle here because it cant be pickled. - """ - if self._current_index == len(self.cache.index) - 1: - self._current_index = 0 - else: - self._current_index += 1 - return self.cache.index[self._current_index] - - def _kwargs_to_str(self, kwargs): - """Convert kwargs to a string.""" - keys = sorted(list(kwargs.keys())) - out = str([(item, kwargs[item]) for item in keys]) - return out - - def _get_mtime(self, only_new=True): - """Return an iterator of potential un-indexed files.""" - # get mtime, subtract a bit to avoid odd bugs - mtime = None - # getting last updated might need the db so only call once. - last_updated = self._index_table.last_updated_timestamp if only_new else None - if last_updated is not None and only_new: - mtime = last_updated - 0.001 - # get paths to iterate - return mtime - - def _get_paths(self, paths): - path = self.path - if paths is None: - paths = path - else: - paths = [ - f"{path}/{x}" if str(path) not in str(x) else str(x) - for x in iterate(paths) - ] - return paths - - def _enforce_min_version(self): - """Ensure the minimum version is met, else delete index file.""" - try: - self._index_table.validate_version() - except InvalidIndexVersionError: - msg = ( - f"The index file at {self.path} is not compatible with this" - f" version of DASCore ({dc.__last_version__}). " - f"Recreating the index now." - ) - warnings.warn(msg, UserWarning) - os.remove(self.index_path) - self.update() - - def get_index_metadata(self): - """Return a dict of metadata about the index.""" - self.update() - up_time = dc.to_datetime64(self._index_table.last_updated_timestamp) - out = { - "index_version": self._index_table._index_version, - "last_update": up_time, - } - return out - - def update(self, paths=None, progress: PROGRESS_LEVELS = "standard") -> Self: + @abc.abstractmethod + def ensure_updated(self) -> bool: """ - Updates the contents of the Indexer. - - Also resets any previous selection. + Run the initial update if the index was never populated. - Parameters - ---------- - paths - A sequence of paths to limit the updates, if None, index all - the contents of directory. - progress - The type of progress bar to use. None disables progress bar and - "basic" is best for low latency scenarios. + Return True when an update actually ran. """ - self._enforce_min_version() # delete index if schema has changed - update_time = time.time() - timestamp = self._get_mtime(only_new=True) - paths = self._get_paths(paths) - df = dc.scan_to_df( - path=paths, - timestamp=timestamp, - progress=progress, - ext=self.ext, - ) - # Put contents found into database. - if not df.empty: - # Some users were surprised the spool wasn't sorted. We still cant - # guarantee all spools will be sorted but we can make sure most are - # by sorting the contents before dumping to index. - if "time_min" in df.columns: - df = df.sort_values("time_min").reset_index(drop=True) - # ensure the base path is not in the path column - assert "path" in set(df.columns), f"{df} has no path column" - self._index_table.write_update(df, update_time, base_path=self.path) - # clear cache out when new traces are added - self.clear_cache() - return self diff --git a/dascore/utils/chunk.py b/dascore/utils/chunk.py index dc80f0759..dae35aad6 100644 --- a/dascore/utils/chunk.py +++ b/dascore/utils/chunk.py @@ -1,28 +1,15 @@ -"""Utilities for chunking dataframes.""" +"""Utilities for chunking dataframes. -from __future__ import annotations +The interval math here is consumed by the chunk planner +(`dascore.utils.chunk_plan`), which replaced the old ChunkManager. +""" -import warnings -from collections.abc import Collection -from functools import reduce -from typing import ClassVar +from __future__ import annotations -import numpy import numpy as np import pandas as pd -from dascore.constants import attr_conflict_description, numeric_types, timeable_types -from dascore.exceptions import ChunkError, CoordMergeError, ParameterError -from dascore.utils.docs import compose_docstring -from dascore.utils.misc import get_middle_value -from dascore.utils.pd import ( - _instructions_modified, - _remove_overlaps, - get_column_names_from_dim, - get_dim_names_from_columns, - get_interval_columns, - list_ser_to_str, -) +from dascore.exceptions import ChunkError, ParameterError from dascore.utils.time import ( is_datetime64, is_timedelta64, @@ -30,8 +17,6 @@ to_timedelta64, ) -_DEFAULT_TOLERANCE = 1.5 - def get_intervals( start, @@ -65,13 +50,6 @@ def get_intervals( ------- A 2D array where first column is start and second column is end. """ - # when length is null just use entire length - if pd.isnull(length): - out = np.asarray([start, stop]) - if is_datetime64(start): - out = to_datetime64(out) - return np.atleast_2d(out) - if is_datetime64(start): # need to ensure we have numpy datetimes, not pandas start, stop = to_datetime64(start), to_datetime64(stop) @@ -85,9 +63,10 @@ def get_intervals( # get variable and perform checks overlap = length * 0 if not overlap else overlap step = length * 0 if pd.isnull(step) else step - # Check for errors - if overlap > length: - msg = "Cant chunk when overlap is greater than chunk size" + # Check for errors. Overlap equal to length would produce zero-stride + # segments, so it is also rejected. + if overlap >= length: + msg = "Cant chunk when overlap is greater than or equal to chunk size" raise ParameterError(msg) # If the step is known, we need to account for it in the total duration # See 474. @@ -116,414 +95,3 @@ def get_intervals( else: ends[bad_ends] = stop return np.stack([starts, ends]).T - - -@compose_docstring(attr_conflict=attr_conflict_description) -class ChunkManager: - """ - A class for managing the chunking of data defined in a dataframe. - - The chunk manager handles both splitting and joining of contiguous, - or near-contiguous, blocks of data. - - Parameters - ---------- - overlap - The amount of overlap between each segment, starting with the end of - first row. Negative values can be used for inducing gaps. - group_columns - A sequence of column names which should be used for sorting groups. - keep_partial - If True, keep segments which are shorter than chunk size (at end of - contiguous blocks) - tolerance - The upper limit of a gap to tolerate in terms of the sampling - along the desired dimension. E.G., the default value means entities - with gaps <= 1.5 * {name}_step will be merged. - conflict - {attr_conflict} - **kawrgs - kwargs specify the column along which to chunk. The key specifies the - column along which to chunk, typically, `time` or `distance`, and the - value specifies the chunk size. A value of None means to chunk on all - available data (e.g. merge all data). - - Notes - ----- - This class is used internally by `dc.BaseSpool.chunk`. - """ - - # Coord fingerprints are stable IDs for exact coord contents. Chunk merges - # intentionally rewrite coord extents/min/max, so inherited fingerprint - # columns are no longer expected to remain equal across merge candidates. - _merge_ignored_columns: ClassVar[set[str]] = { - "dtype", - "time_fingerprint", - "distance_fingerprint", - } - - def __init__( - self, - overlap: timeable_types | numeric_types | None = None, - group_columns: Collection[str] | None = None, - keep_partial=False, - snap_coords=True, - tolerance=_DEFAULT_TOLERANCE, - conflict="raise", - **kwargs, - ): - self._overlap = overlap - self._group_columns = group_columns - self._keep_partials = keep_partial - self._snap_coords = snap_coords - self._tolerance = tolerance - self._name, self._value = self._validate_kwargs(kwargs) - self._attr_conflict = conflict - self._validate_chunker() - - def _validate_kwargs(self, kwargs): - """Ensure kwargs is len one and has a valid.""" - if not len(kwargs) == 1: - msg = ( - f"Chunking only supported along one dimension. You passed " - f"kwargs: {kwargs}" - ) - raise ParameterError(msg) - ((key, value),) = kwargs.items() - value = None if value is ... else value - return key, value - - def _validate_chunker(self): - """Ensure selected parameters are compatible.""" - # chunker is used for merging - if pd.isnull(self._value): - if self._keep_partials or self._overlap: - msg = ( - "When chunk value is None (ie Chunker is used for merging) " - "both _keep_partials and self._overlap must not be selected." - ) - raise ParameterError(msg) - return - # ensure chunk values are greater than 0 - zero = to_timedelta64(0) if is_timedelta64(self._value) else 0 - if self._value <= zero: - msg = "Chunk value must be greater than 0." - raise ParameterError(msg) - - def _get_continuity_group_number( - self, start, stop, step, tolerance=None - ) -> pd.Series: - """Return a series of ints indicating continuity group.""" - tolerance = self._tolerance if tolerance is None else tolerance - # start by sorting according to start time - # Use positional argsort to avoid pandas label/return-type changes - args = np.argsort(start.to_numpy()) - start_sorted, stop_sorted, step_sorted = ( - start.iloc[args], - stop.iloc[args], - step.iloc[args], - ) - # next get cummax of endtimes and detect gaps - stop_cum_max = stop_sorted.cummax() - end_markers = stop_cum_max.shift() + step_sorted * tolerance - has_gap = start_sorted > end_markers - group_num = has_gap.astype(np.int64).cumsum() - return group_num[start.index] - - def _get_sampling_group_num(self, step, tolerance=0.05) -> pd.Series: - """ - Because sampling can be off a little, this adds some tolerance for - how sampling affects groups. - - Tolerance affects how close samples have to be in order to count as - the same. 5% is used here. - """ - col = step.values - sort_args = np.argsort(col) - sorted_col = col[sort_args] - roll_forward = np.roll(sorted_col, shift=1) - diff = (sorted_col - roll_forward) / sorted_col - out_of_threshold = diff > tolerance - group_number = numpy.cumsum(out_of_threshold) - # undo sorting - out = pd.Series(group_number[np.argsort(sort_args)], index=step.index) - return out - - def _get_duration_overlap(self, duration, start, step, overlap=None): - """Get duration and overlap from kwargs.""" - overlap = overlap if overlap is not None else self._overlap - # cast step/overlap to timedelta if start is datetime or timedelta; - # a span of either dtype is a duration. - if is_datetime64(start) or is_timedelta64(start): - step = to_timedelta64(step) - overlap = to_timedelta64(overlap) - if pd.isnull(overlap): - overlap = np.asarray([0], dtype=step.dtype)[0] - return duration, overlap - - def _create_df(self, df, name, start_stop, gnum): - """Reconstruct the dataframe.""" - cols = f"{name}_min", f"{name}_max" - out = pd.DataFrame(start_stop, columns=list(cols)) - out[f"{name}_step"] = get_middle_value(df[f"{name}_step"].values) - merger = df.drop(columns=out.columns) - # get dims to determine which columns are still compared. Some test - # dfs don't have dims though, so it should still work without dims col. - dims = set(df.iloc[0].get("dims", "").split(",")) - # We exclude private columns for considering if merge can happen. - for col in set(x for x in merger.columns if not x.startswith("_")): - if col in self._merge_ignored_columns: - continue - prefix = col.split("_")[0] - # If we have specified to ignore or remove conflicting attrs - # we don't need to check them here, but we do still check dims. - if self._attr_conflict != "raise" and prefix not in dims: - continue - vals = merger[col].unique() - if len(vals) > 1: - msg = ( - f"Cannot merge on dim {self._name} because all values for " - f"{col} are not equal. Consider using the `conflict` " - f"argument to loosen this restriction." - ) - raise CoordMergeError(msg) - - assert len(vals) == 1, "Haven't yet implemented non-homogenous merging" - out[col] = vals[0] - if "dims" in out.columns: - # Keep dims dtype consistent with patches_to_df/list_ser_to_str. - out["dims"] = list_ser_to_str(out["dims"]) - # add the group number for getting instruction df later - out["_group"] = gnum - return out - - def _get_chunk_overlap_inds(self, src1, src2, chu1, chu2): - """Get an index mapping from source to chunk.""" - chunk_starts = np.searchsorted(src1, chu1, side="right") - 1 - chunk_ends = np.searchsorted(src2, chu2, side="left") - # Ensure no chunks run off the end of the source. - assert np.all(chunk_ends < len(src1)), "Invalid chunk range found" - # add 1 to end so it is an exclusive end range - return np.stack([chunk_starts, chunk_ends + 1], axis=1) - - def _get_source_and_chunk_inds(self, chunk2src_inds, s_index, c_index): - """Get ndarrays of chunk index, source index.""" - # get indices for sorted arrays - source_inds_ = np.concatenate( - [np.arange(x[0], x[1], dtype=np.int64) for x in chunk2src_inds] - ) - chunk_inds_ = np.concatenate( - [ - np.ones((x[1] - x[0]), dtype=np.int64) * num - for num, x in enumerate(chunk2src_inds) - ] - ) - # use pandas index to map back to actual indices - source_inds = s_index.values[source_inds_] - chunk_inds = c_index.values[chunk_inds_] - out = { - "source_sorted": source_inds_, - "source": source_inds, - "chunk_sorted": chunk_inds_, - "chunk": chunk_inds, - } - return out - - def _get_instructions(self, sub_source, sub_chunk): - """Get source mapping to chunk.""" - min_name, max_name = f"{self._name}_min", f"{self._name}_max" - # sort inputs based on start of range, as long as we don't reset index - # we should be ok. - sub_source = sub_source.sort_values(min_name) - sub_chunk = sub_chunk.sort_values(min_name) - # need to make sure we don't have overlaps in source df. This implicitly - # handles merging. - sub_source = _remove_overlaps(sub_source, self._name) - src1, src2, _src_step = get_interval_columns( - sub_source, self._name, arrays=True - ) - chu1, chu2, _chu_step = get_interval_columns(sub_chunk, self._name, arrays=True) - dims = get_dim_names_from_columns(sub_source) - cols2keep = get_column_names_from_dim(dims) - # next get index range for which chunk times belong to. - chunk2src_inds = self._get_chunk_overlap_inds(src1, src2, chu1, chu2) - # total length of source to chunk mapping - inds = self._get_source_and_chunk_inds( - chunk2src_inds, - sub_source.index, - sub_chunk.index, - ) - source_inds, chunk_inds = inds["source_sorted"], inds["chunk_sorted"] - # get potential start/stop times. - starts = np.stack([src1[source_inds], chu1[chunk_inds]], axis=1) - ends = np.stack([src2[source_inds], chu2[chunk_inds]], axis=1) - end_values = np.min(ends, axis=1) - start_values = np.max(starts, axis=1) - data_dict = { - min_name: start_values, - max_name: end_values, - "source_index": inds["source"], - "current_index": inds["chunk"], - } - out = pd.DataFrame(data_dict) - # populate the rest of the columns needed in instruction df. - for col in cols2keep: - if col in out.columns: - continue - out[col] = sub_source[col].values[source_inds] - out = out.sort_index() - out["_modified"] = _instructions_modified(out, sub_source) - return out - - def get_instruction_df(self, source_df, chunked_df): - """ - Get a dataframe connecting the chunked dataframe to its origin. - - This is used to connect source data to desired data after chunking - operation. - - Parameters - ---------- - source_df - The dataframe before chunking - chunked_df - The chunked dataframe (output of `chunk` method) - """ - # the group column should exist and the chunked groups should be subset - # of the source groups - assert "_group" in source_df.columns and "_group" in chunked_df.columns - chunked_groups = set(chunked_df["_group"]) - if not chunked_groups: - return pd.DataFrame(columns=[*list(source_df.columns), "_modified"]) - # chunk groups should be a subset of source groups - assert chunked_groups.issubset(set(source_df["_group"])) - # iterate each group and create instruction df - out = [] - for group in chunked_groups: - sub_source = source_df[source_df["_group"] == group] - sub_chunk = chunked_df[chunked_df["_group"] == group] - out.append(self._get_instructions(sub_source, sub_chunk)) - df = pd.concat(out, axis=0).reset_index(drop=True).set_index("source_index") - return df - - def _get_col_group(self, df, cont_g): - """Get group columns based on common columns.""" - cols = list(self._group_columns or []) - columns = [x for x in cols if x in df.columns] - col_g = cont_g * 0 if not columns else df.groupby(columns).ngroup() - return col_g - - def _get_final_group(self, samp_g, col_g, cont_g): - """Combine grouping components into final group labels.""" - group_series = [x.astype(str) for x in [samp_g, col_g, cont_g]] - return reduce(lambda x, y: x + "_" + y, group_series) - - def _groups_merge_default_groups(self, group, default_cont_g) -> bool: - """Return True if any final group contains multiple default groups.""" - group_codes = pd.factorize(group, sort=False)[0] - default_codes = pd.factorize(default_cont_g, sort=False)[0] - order = np.argsort(group_codes, kind="stable") - group_sorted = group_codes[order] - default_sorted = default_codes[order] - new_group = np.r_[True, group_sorted[1:] != group_sorted[:-1]] - group_starts = np.flatnonzero(new_group) - default_min = np.minimum.reduceat(default_sorted, group_starts) - default_max = np.maximum.reduceat(default_sorted, group_starts) - return bool(np.any(default_min != default_max)) - - def _get_group(self, df, start, stop, step): - """ - Get the group designation for df. This accounts for both time intervals - being consistent and group columns matching. - """ - cont_g = self._get_continuity_group_number(start, stop, step) - samp_g = self._get_sampling_group_num(step) - col_g = self._get_col_group(df, cont_g) - group = self._get_final_group(samp_g, col_g, cont_g) - - # Check for final merges that only occur because of a non-default - # tolerance. See #662. - if self._tolerance > _DEFAULT_TOLERANCE: - default_cont_g = self._get_continuity_group_number( - start, stop, step, tolerance=_DEFAULT_TOLERANCE - ) - if self._groups_merge_default_groups(group, default_cont_g): - msg = ( - f"There is a gap in the patch along dimension {self._name} " - f"but a merge tolerance of {self._tolerance} was used to force " - "merging the patches. As a result, some patches in the chunked " - "spool may be unevenly sampled, or have their sampling rate " - "increased." - ) - warnings.warn(msg, UserWarning, stacklevel=4) - return group - - def _get_group_dfs(self, group, dur, overlap, group_mins, group_maxs, df, step): - """Get the new dataframe for a given group.""" - out = [] - for gnum in group.unique(): - g_start, g_stop = group_mins[gnum], group_maxs[gnum] - current_df = df.loc[group[group == gnum].index] - # reconstruct DF - try: - new_start_stop = get_intervals( - g_start, - g_stop, - dur, - overlap=overlap, - step=step.loc[current_df.index].iloc[0], - keep_partials=self._keep_partials, - ) - except ChunkError: # this chunk is too short, skip. - continue - # create the newly chunked dataframe - sub_new_df = self._create_df(current_df, self._name, new_start_stop, gnum) - out.append(sub_new_df) - return out - - def _filter_nan_dfs(self, df, start, stop): - """Filter NaN out of dataframe if they occur in start/stop.""" - - def chunk( - self, - df: pd.DataFrame, - ) -> tuple[pd.DataFrame, pd.DataFrame]: - """ - Chunk a dataframe into new contiguous segments. - - The dataframe must have column names {key}_max, {key}_min, and {key}_step - where {key} is the key used in the kwargs. - - Parameters - ---------- - df - Input dataframe to chunk. - - Returns - ------- - A tuple of the original dataframe with added column '_group' and an - output dataframe with column '_group'. The _group column is used - to link the two dataframes together. - """ - if df.empty: # empty df, do nothing - return df.assign(_group=None), df.assign(_group=None) - # get series of start/stop along requested dimension - start, stop, step = get_interval_columns(df, self._name) - # Filter out any NaN in start or stop. - keep = ~(pd.isnull(start) | pd.isnull(stop)) - df, start, stop, step = df[keep], start[keep], stop[keep], step[keep] - if df.empty: # Need to check again since NaN can wipe out df. - return df.assign(_group=None), df.assign(_group=None) - dur, overlap = self._get_duration_overlap(self._value, start, step) - # get group numbers - group = self._get_group(df, start, stop, step) - # get max, min for each group and expand - group_mins = start.groupby(group).min() - group_maxs = stop.groupby(group).max() - # split/group dataframe into new chunks by iterating over each group. - out = self._get_group_dfs(group, dur, overlap, group_mins, group_maxs, df, step) - if not len(out): - msg = "Could not chunk. No segments with sufficient length found." - raise ChunkError(msg) - out = pd.concat(out, axis=0).reset_index(drop=True) - return df.assign(_group=group), out diff --git a/dascore/utils/chunk_plan.py b/dascore/utils/chunk_plan.py new file mode 100644 index 000000000..57952f975 --- /dev/null +++ b/dascore/utils/chunk_plan.py @@ -0,0 +1,618 @@ +""" +Chunk planning over the flat patch relation. + +Implements the "Chunking formalities" spec: the planner consumes the +catalog's flat relation (one row per patch: `{dim}_min/max/step` envelopes, +`_{dim}_def_key` structural identity, attr columns) and produces a +[`ChunkPlan`](`dascore.utils.chunk_plan.ChunkPlan`) — an outputs table (one row +per output patch) plus a members table binding each output to trimmed +slices of source patches. No patch data is touched; assembly happens later. + +Portions of the interval/instruction math were ported from the old +`ChunkManager` (now removed) with the spec's adjudicated corrections +applied; `Spool.chunk` runs on these plans. +""" + +from __future__ import annotations + +import warnings +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import pandas as pd + +import dascore as dc +from dascore.exceptions import ( + ChunkError, + CoordMergeError, + InvalidSpoolQueryError, + ParameterError, +) +from dascore.utils.chunk import get_intervals +from dascore.utils.misc import get_middle_value, is_range +from dascore.utils.pd import _remove_overlaps, get_interval_columns +from dascore.utils.time import is_datetime64, is_timedelta64, to_float, to_timedelta64 + +# Columns which never participate in conflict policing and never carry to +# outputs: source bookkeeping (outputs are not file rows). +_SOURCE_COLUMNS = ("path", "file_format", "file_version", "source_patch_id") +# The default continuity tolerance; looser values warn when they force +# merges (#662). +_DEFAULT_TOLERANCE = 1.5 + + +@dataclass(frozen=True) +class ChunkPlan: + """ + A materialization-free description of a chunk operation. + + Attributes + ---------- + outputs + One row per output patch: `{dim}_min/max/step` for the chunked + dimension, an `output_id`, and all carried columns (group attrs, + dims, structural def keys, conflict-policed attrs). + members + Instruction rows binding outputs to sources: `output_id`, + `_patch_id`, the exact `{dim}_min/max` trim for that member, and + `_modified` (False when the member loads whole). + dim + The chunked dimension. + value + The requested chunk length (None for merge mode). + params + Resolved parameters (group attrs, tolerances, overlap, + keep_partial, conflict, snap_coords, missing_dim) — recorded, not + referencing config. + """ + + outputs: pd.DataFrame + members: pd.DataFrame + dim: str + value: Any + params: dict = field(default_factory=dict) + + @property + def merge_mode(self) -> bool: + """Return True when this plan merges (no segmenting length).""" + return self.value is None + + +def _resolve_group_attrs(group, columns) -> tuple[str, ...]: + """Resolve the group attrs: per-call > config; explicit names must exist.""" + if group is not None: + group = (group,) if isinstance(group, str) else tuple(group) + if missing := [x for x in group if x not in columns]: + msg = ( + f"group attribute(s) {missing} do not exist on any patch " + "in the spool." + ) + raise InvalidSpoolQueryError(msg) + return group + # Config (and default) names are best-effort. + return tuple(x for x in dc.get_config().groupby_attrs if x in columns) + + +def samples_adjusted_envelopes( + df: pd.DataFrame, residuals, drop_empty: bool = True +) -> pd.DataFrame: + """ + Adjust envelope columns for patch-local samples residuals. + + A ``samples=True`` index window trims each patch at load, so the + planner must consume the trimmed envelopes or it publishes outputs + that lie entirely outside the selected samples (phantom empties). + Negative indices resolve per patch against the envelope-derived + sample count (rows whose count is unknown keep their envelope as a + candidacy superset — exactness is always re-applied at load). + ``drop_empty`` removes rows whose window selects nothing (planning + truth); equality comparison keeps them, since a presented-but-empty + row is still a presented row. + """ + + def _usable_index(value) -> bool: + return value is None or isinstance(value, int | np.integer) + + df = df.copy(deep=False) + for coords, samples in residuals: + if not samples: + continue + for name, value in coords.items(): + cols = [f"{name}_min", f"{name}_max", f"{name}_step"] + if not set(cols).issubset(df.columns) or not is_range(value): + continue + lo_idx, hi_idx = value + if not (_usable_index(lo_idx) and _usable_index(hi_idx)): + continue + mins, maxs, steps = (df[c] for c in cols) + # Positions are patch-local sample indices with a stop-exclusive + # hi, so the last included position is hi - 1. Sample 0 sits at + # the envelope min for ascending coords and at the max for + # descending ones. + abs_steps = steps.abs() + descending = to_float(steps.values) < 0 + with np.errstate(invalid="ignore", divide="ignore"): + ratio = to_float((maxs - mins).values) / to_float(abs_steps.values) + counts = pd.Series(np.round(ratio) + 1, index=df.index) + + def _positions(idx, counts=counts, index=df.index): + """Per-row absolute positions (Python-slice clamping).""" + if idx is None: + return None + if idx >= 0: + return pd.Series(float(idx), index=index) + return (counts + idx).clip(lower=0) + + lo_pos, hi_pos = _positions(lo_idx), _positions(hi_idx) + unresolved = pd.Series(False, index=df.index) + for pos in (lo_pos, hi_pos): + if pos is not None: + unresolved |= pos.isna() + lo_off = None if lo_pos is None else lo_pos * abs_steps + hi_off = None if hi_pos is None else (hi_pos - 1) * abs_steps + new_min = mins if lo_off is None else mins + lo_off + new_max = maxs if hi_off is None else mins + hi_off + desc_min = maxs if hi_off is None else maxs - hi_off + desc_max = maxs if lo_off is None else maxs - lo_off + new_min = new_min.where(~descending, other=desc_min) + new_max = new_max.where(~descending, other=desc_max) + # unresolvable rows keep their envelope (candidacy superset) + new_min = new_min.mask(unresolved, mins) + new_max = new_max.mask(unresolved, maxs) + # rows whose window is empty or lies entirely outside the + # patch contribute nothing; test before clipping so such + # windows are not resurrected as one-sample envelopes + keep = (new_min <= new_max) & (new_min <= maxs) & (new_max >= mins) + keep |= unresolved + df[cols[0]] = new_min.clip(lower=mins, upper=maxs) + df[cols[1]] = new_max.clip(lower=mins, upper=maxs) + if drop_empty: + df = df[keep] + return df + + +def _ensure_patch_id(df: pd.DataFrame) -> pd.DataFrame: + """Attach the positional identity fallback for plain dataframes.""" + if "_patch_id" in df.columns: + return df + return df.assign(_patch_id=np.arange(len(df))) + + +def _dim_def_key_columns(df: pd.DataFrame, name: str) -> list[str]: + """Return def-key column names for every non-chunked dimension.""" + dim_names: set[str] = set() + if "dims" in df.columns: + for dims_str in df["dims"].dropna().unique(): + dim_names.update(str(dims_str).split(",")) + dim_names.discard(name) + return [f"_{x}_def_key" for x in sorted(dim_names)] + + +def _sampling_group(step: pd.Series, tolerance: float) -> pd.Series: + """ + Label rows whose steps are within relative tolerance (spec 2.3). + + Steps group by orientation (sign) first, then by magnitude against a + stable group anchor: a group opens at its smallest magnitude and + admits members up to ``anchor * (1 + tolerance)``, so a chain of + individually-close steps can never drift a group's endpoints past + the tolerance. Unknown (NaN) steps share one group. + """ + col = to_float(step.values) + sign = np.sign(col) + mag = np.abs(col) + # orientation-major, magnitude-minor; NaNs sort to the end of both keys + order = np.lexsort((mag, sign)) + sorted_sign, sorted_mag = sign[order], mag[order] + labels = np.zeros(len(col), dtype=np.int64) + label, i, n = 0, 0, len(col) + while i < n: + if np.isnan(sorted_mag[i]): + # NaN keys sort last, so everything from here on is unknown + j = n + else: + block_end = np.searchsorted(sorted_sign, sorted_sign[i], side="right") + bound = sorted_mag[i] * (1 + tolerance) + j = np.searchsorted(sorted_mag[:block_end], bound, side="right") + j = max(j, i + 1) + labels[order[i:j]] = label + label += 1 + i = j + return pd.Series(labels, index=step.index) + + +def _continuity_group(start, stop, step, tolerance) -> pd.Series: + """Label maximal near-contiguous runs (spec 2.4).""" + args = np.argsort(start.to_numpy()) + start_sorted = start.iloc[args] + stop_sorted = stop.iloc[args] + # envelopes are value-ordered regardless of coordinate orientation, + # so the continuity margin uses the step magnitude + step_sorted = step.iloc[args].abs() + stop_cum_max = stop_sorted.cummax() + end_markers = stop_cum_max.shift() + step_sorted * tolerance + has_gap = start_sorted > end_markers + group = has_gap.astype(np.int64).cumsum() + return group[start.index] + + +def _partition( + df, name, group_attrs, tolerance, sampling_tolerance +) -> tuple[pd.Series, bool]: + """ + Return (partition labels, forced_merge): rows sharing a label may + combine (spec 2). + + Components: group attrs, dims signature, structural def keys of + non-chunked coords, sampling group, and continuity group. Continuity + is evaluated *within* each other-component cell so unrelated patches + can never bridge a gap. `forced_merge` is True when a loosened + tolerance merged patches the default would have kept apart (#662); + the caller owns warning about it. + """ + start, stop, step = get_interval_columns(df, name) + cols = [x for x in group_attrs if x in df.columns] + if "dims" in df.columns: + cols.append("dims") + # Structural identity: def keys of non-chunked *dimensions* only + # (spec 2.2). Non-dimensional coordinate conflicts are policed at + # assembly per the `conflict` argument, never partitioned on. + cols += [x for x in _dim_def_key_columns(df, name) if x in df.columns] + # The chunked dim's canonical (base) units partition too: envelopes + # are SI magnitudes, so without this a metre patch and a second + # patch with contiguous magnitudes would plan into one unmergeable + # output. Unitless (NULL) stays its own group — assembly cannot + # merge unitless with unitful coordinates either. + if (unit_col := f"_{name}_units") in df.columns: + cols.append(unit_col) + base = ( + df.groupby(cols, dropna=False, sort=False).ngroup() + if cols + else pd.Series(0, index=df.index) + ) + samp = _sampling_group(step, sampling_tolerance) + cell = base.astype(str) + "_" + samp.astype(str) + cont = pd.Series(0, index=df.index, dtype=np.int64) + forced_merge = False + for _, index in df.groupby(cell, sort=False).groups.items(): + sub = df.loc[index] + s, e, st = get_interval_columns(sub, name) + labels = _continuity_group(s, e, st, tolerance).astype(np.int64) + cont.loc[index] = labels + if tolerance > _DEFAULT_TOLERANCE and not forced_merge: + default = _continuity_group(s, e, st, _DEFAULT_TOLERANCE) + forced_merge = default.nunique() > labels.nunique() + return cell + "_" + cont.astype(str), forced_merge + + +def _user_stacklevel() -> int: + """Return the warn stacklevel pointing at the first non-dascore frame. + + Plans are built at several call depths (spool.chunk, spool.chunk_plan, + build_chunk_plan directly), so a fixed stacklevel would blame library + frames for some entries. + """ + import inspect + + # The dascore package directory, resolved from the package itself so + # this does not depend on this module's location within it. + package_dir = str(Path(dc.__file__).resolve().parent) + # Frames after this helper's own align exactly with warn's numbering: + # level 1 is the frame calling warn. + for level, frame_info in enumerate(inspect.stack()[1:], start=1): + filename = str(Path(frame_info.filename).resolve()) + if not filename.startswith(package_dir): + return level + return 1 + + +def _coerce_length_overlap(value, overlap, start_dtype): + """Coerce the chunk length/overlap to the dimension's span dtype.""" + time_like = is_datetime64(start_dtype) or is_timedelta64(start_dtype) + if time_like: + value = to_timedelta64(value) if value is not None else None + overlap = to_timedelta64(overlap) if overlap is not None else None + return value, overlap + + +def _coord_owner(col: str, coord_names: set[str]) -> str | None: + """ + Return the coordinate owning an envelope column, if any. + + Ownership is decided by matching the full name against known + coordinates with an interval suffix; splitting on the first + underscore would mis-assign columns of dims like `event_time`. + """ + for suffix in ("_min", "_max", "_step", "_units"): + if col.endswith(suffix): + base = col[: -len(suffix)] + if base in coord_names: + return base + return None + + +def _police_columns(sub: pd.DataFrame, name, conflict) -> dict: + """ + Return the carried column values for one partition (spec 2.5/6.4). + + Group attrs, dims, and def keys are single-valued by construction. + Remaining public attrs must be single-valued, policed by `conflict`. + """ + dims = set(str(sub.iloc[0].get("dims", "")).split(",")) + coord_names = dims | {name} + carried: dict[str, Any] = {} + for col in sub.columns: + if col.startswith("_") or col in _SOURCE_COLUMNS: + continue + owner = _coord_owner(col, coord_names) + if owner == name: # chunk-dim envelope columns are rebuilt + continue + values = sub[col].unique() + single = len(values) == 1 or (len(values) and pd.isnull(values).all()) + if single: + carried[col] = values[0] + continue + # Group attrs and dims are partition keys, so they are always + # single-valued above and never reach the conflict policy here. + if owner is not None or conflict == "raise": + msg = ( + f"Cannot merge on dim {name} because all values for " + f"{col} are not equal. Consider using the `conflict` " + "argument to loosen this restriction." + ) + raise CoordMergeError(msg) + if conflict == "keep_first": + carried[col] = sub[col].iloc[0] + # conflict == "drop": omit the column entirely. + # Structural (dimension) def keys carry — single-valued by partitioning. + for col in _dim_def_key_columns(sub, name): + if col in sub.columns: + carried[col] = sub[col].iloc[0] + # Canonical units carry for every dimension, the chunked one included + # (partition-constant: units are a sampling-partition component). + for coord in coord_names: + col = f"_{coord}_units" + if col in sub.columns: + carried[col] = sub[col].iloc[0] + return carried + + +def _build_members(sub: pd.DataFrame, outputs: pd.DataFrame, name) -> pd.DataFrame: + """ + Bind one partition's outputs to trimmed source slices. + + Sources are ordered by (start, _patch_id); overlapping coverage is + deduplicated so the earlier source owns the overlap (D3: complete + overlaps keep the first member, deterministically). + """ + min_name, max_name = f"{name}_min", f"{name}_max" + step_name = f"{name}_step" + sub = sub.sort_values([min_name, "_patch_id"], kind="stable") + original = sub[[min_name, max_name]].reset_index(drop=True) + sub = _remove_overlaps(sub, name) + # Fully-covered sources become degenerate after start correction; they + # contribute nothing (deterministic keep-first dedup). + keep = sub[min_name].values <= sub[max_name].values + sub = sub[keep] + original = original[keep].reset_index(drop=True) + # sub and outputs are always non-empty here: a partition too short to + # yield an interval raises ChunkError in the caller (and the earliest + # source is never fully covered), so both keep at least one row. + steps = sub[step_name].values + src1 = sub[min_name].values + src2 = sub[max_name].values + chu1 = outputs[min_name].values + chu2 = outputs[max_name].values + out_ids = outputs["output_id"].to_numpy() + patch_ids = sub["_patch_id"].to_numpy() + orig_min = original[min_name].to_numpy() + orig_max = original[max_name].to_numpy() + modified_src = ( + sub["_modified"].to_numpy() + if "_modified" in sub + else np.zeros(len(sub), dtype=bool) + ) + # Map each output onto the source rows it draws from. + starts_ind = np.searchsorted(src1, chu1, side="right") - 1 + ends_ind = np.searchsorted(src2, chu2, side="left") + rows = [] + for out_num, (a, b) in enumerate(zip(starts_ind, ends_ind)): + a = max(int(a), 0) + for src_num in range(a, int(b) + 1): + if src_num >= len(sub): + continue + lo = max(src1[src_num], chu1[out_num]) + hi = min(src2[src_num], chu2[out_num]) + if lo > hi: # pragma: no cover -- searchsorted boundary guard + # Sources within a partition are continuous (partitioning + # splits on gaps) and start-corrected, so searchsorted does + # not offer a non-overlapping source in practice; this guards + # against a boundary off-by-one rather than a reachable state. + continue + unchanged = ( + lo == orig_min[src_num] + and hi == orig_max[src_num] + and not modified_src[src_num] + ) + rows.append( + { + "output_id": out_ids[out_num], + "_patch_id": patch_ids[src_num], + min_name: lo, + max_name: hi, + step_name: steps[src_num], + "_modified": not unchanged, + } + ) + return pd.DataFrame(rows) + + +def build_chunk_plan( + df: pd.DataFrame, + *, + overlap=None, + keep_partial: bool = False, + snap_coords: bool = True, + tolerance: float = 1.5, + conflict: Literal["drop", "raise", "keep_first"] = "raise", + group=None, + missing_dim: Literal["raise", "drop"] = "raise", + **kwargs, +) -> ChunkPlan: + """ + Build a chunk plan from a flat patch relation. + + Parameters mirror `Spool.chunk` (see the chunking formalities spec); + exactly one keyword names the dimension to chunk and its length + (`None`/`...` merges). + """ + if len(kwargs) != 1: + msg = ( + "Chunking only supported along one dimension. You passed " + f"kwargs: {kwargs}" + ) + raise ParameterError(msg) + ((name, value),) = kwargs.items() + value = None if value is Ellipsis else value + merge_mode = pd.isnull(value) + if merge_mode and (keep_partial or overlap): + msg = ( + "When chunk value is None (ie chunking is used for merging) " + "keep_partial and overlap are not supported." + ) + raise ParameterError(msg) + if not merge_mode: + zero = to_timedelta64(0) if is_timedelta64(value) else 0 + if value <= zero: + msg = "Chunk value must be greater than 0." + raise ParameterError(msg) + if missing_dim not in ("raise", "drop"): + msg = f"missing_dim must be 'raise' or 'drop', got {missing_dim!r}" + raise ParameterError(msg) + if conflict not in ("drop", "raise", "keep_first"): + msg = "conflict must be 'drop', 'raise', or 'keep_first', " f"got {conflict!r}" + raise ParameterError(msg) + + min_name, max_name = f"{name}_min", f"{name}_max" + if min_name not in df.columns and not df.empty: + msg = f"No patch in the spool has a {name!r} dimension to chunk." + raise ChunkError(msg) + empty_members = pd.DataFrame( + columns=["output_id", "_patch_id", min_name, max_name, "_modified"] + ) + params = dict( + overlap=overlap, + keep_partial=keep_partial, + snap_coords=snap_coords, + tolerance=tolerance, + conflict=conflict, + missing_dim=missing_dim, + group=_resolve_group_attrs(group, set(df.columns)), + sampling_group_tolerance=dc.get_config().sampling_group_tolerance, + ) + if df.empty: + outputs = pd.DataFrame(columns=[min_name, max_name, "output_id"]) + return ChunkPlan(outputs, empty_members, name, value, params) + df = _ensure_patch_id(df) + # Missing chunk-dim envelopes, and patches carrying the name only as + # a non-dimensional coordinate (spec 7 / D2): chunking is defined on + # dimensions, so both fall under missing_dim. Envelope presence is + # not enough — auxiliary coordinates index their envelopes too, but + # their patches cannot be trimmed or merged *along* the name. + null_rows = pd.isnull(df[min_name]) | pd.isnull(df[max_name]) + if "dims" in df.columns: + dim_lists = df["dims"].fillna("").astype(str).str.split(",") + not_a_dim = ~dim_lists.map(lambda dims: name in dims) + else: + not_a_dim = pd.Series(False, index=df.index) + unusable = null_rows | not_a_dim + if unusable.any(): + if missing_dim == "raise": + bad = df.loc[unusable, "_patch_id"].tolist() + rides = int((not_a_dim & ~null_rows).sum()) + detail = ( + f" ({rides} of them carry {name!r} only as a non-dimensional " + "coordinate; chunking is defined on dimensions)" + if rides + else "" + ) + msg = ( + f"{int(unusable.sum())} patch(es) lack the chunk dimension " + f"{name!r}{detail} (patch ids {bad[:5]}...). Pass " + "missing_dim='drop' to exclude them." + ) + raise ChunkError(msg) + df = df[~unusable] + if df.empty: + outputs = pd.DataFrame(columns=[min_name, max_name, "output_id"]) + return ChunkPlan(outputs, empty_members, name, value, params) + + labels, forced_merge = _partition( + df, name, params["group"], tolerance, params["sampling_group_tolerance"] + ) + if forced_merge: + msg = ( + f"There is a gap in the patch along dimension {name} but a " + f"merge tolerance of {tolerance} was used to force merging " + "the patches. As a result, some patches in the chunked spool " + "may be unevenly sampled, or have their sampling rate increased." + ) + warnings.warn(msg, UserWarning, stacklevel=_user_stacklevel()) + value_c, overlap_c = _coerce_length_overlap(value, overlap, df[min_name].dtype) + out_frames, member_frames = [], [] + next_id = 0 + # Deterministic partition order (spec 8): by (partition min, smallest + # member patch id) — never by anything derived from input row order. + grouped = df.groupby(labels, sort=False) + stats = grouped.agg(_min=(min_name, "min"), _pid=("_patch_id", "min")) + part_order = stats.sort_values(["_min", "_pid"], kind="stable").index + groups = grouped.groups + for label in part_order: + sub = df.loc[groups[label]] + start, stop, step = get_interval_columns(sub, name) + part_step = get_middle_value(step.values) # D7: one step everywhere + g_start, g_stop = start.min(), stop.max() + if merge_mode: + start_stop = np.atleast_2d(np.asarray([g_start, g_stop])) + else: + try: + start_stop = get_intervals( + g_start, + g_stop, + value_c, + overlap=overlap_c, + # interval arithmetic is over (direction-free) envelope + # values; a descending coordinate's negative step would + # invert the final partial interval + step=abs(part_step), + keep_partials=keep_partial, + ) + except ChunkError: # partition too short; skip (D8) + continue + sub_sorted = sub.sort_values([min_name, "_patch_id"], kind="stable") + carried = _police_columns(sub_sorted, name, conflict) + outputs = pd.DataFrame(start_stop, columns=[min_name, max_name]) + outputs[f"{name}_step"] = part_step + for col, val in carried.items(): + outputs[col] = val + outputs["output_id"] = np.arange(next_id, next_id + len(outputs)) + next_id += len(outputs) + members = _build_members(sub, outputs, name) + # Plan invariant: every published output has at least one member. + # An advertised row that cannot assemble is never surfaced as a + # runtime error; it is not surfaced at all. + fed = set(members["output_id"]) if not members.empty else set() + outputs = outputs[outputs["output_id"].isin(fed)] + out_frames.append(outputs) + member_frames.append(members) + if not out_frames or all(x.empty for x in out_frames): + msg = "Could not chunk. No segments with sufficient length found." + raise ChunkError(msg) + outputs = pd.concat(out_frames, ignore_index=True) + members = pd.concat( + [x for x in member_frames if not x.empty] or [empty_members], + ignore_index=True, + ) + return ChunkPlan(outputs, members, name, value, params) diff --git a/dascore/utils/coordmanager.py b/dascore/utils/coordmanager.py index 4e6ef04e5..9c8a60a13 100644 --- a/dascore/utils/coordmanager.py +++ b/dascore/utils/coordmanager.py @@ -19,6 +19,7 @@ def merge_coord_managers( dim: str, snap_tolerance: float | None = None, drop_conflicting: bool = False, + dim_coord=None, ) -> dc.CoordManager: """ Merge coordinate managers along a specified dimension. @@ -38,6 +39,11 @@ def merge_coord_managers( drop_conflicting If True, drop conflicting (non-dimensional) coordinates, otherwise raise an exception if they occur. + dim_coord + If provided, use this coordinate for `dim` instead of + concatenating the members' values (which materializes them); + callers which already built the merged dimension coordinate + (e.g. via `concat_coords`) pass it here to avoid that cost. """ def _get_dims(managers): @@ -112,6 +118,9 @@ def _get_merged_coords(managers, coords_to_merge): """Get the merged coordinates.""" out = {} for coord_name in coords_to_merge: + if dim_coord is not None and coord_name == dim: + out[coord_name] = (managers[0].dim_map[dim], dim_coord) + continue merge_coords = [x.coord_map[dim] for x in managers] axis = managers[0].dim_map[coord_name].index(dim) if len(units := {x.units for x in merge_coords}) != 1: @@ -126,7 +135,16 @@ def _get_merged_coords(managers, coords_to_merge): data = [x.data for x in snap_coords] dims = managers[0].dim_map[dim] new_data = np.concatenate(data, axis=axis) - out[coord_name] = (dims, new_data) + # raw value concatenation loses the coord's units; reattach + # the (verified common) units so the merge stays unit-true + common_units = next(iter(units)) + if common_units is not None: + from dascore.core.coords import get_coord + + coord = get_coord(data=new_data, units=common_units) + out[coord_name] = (dims, coord) + else: + out[coord_name] = (dims, new_data) return out def _get_new_coords(managers) -> dict[str, tuple[tuple[str, ...], ArrayLike]]: diff --git a/dascore/utils/hdf5.py b/dascore/utils/hdf5.py index 4faf7476f..cd5182669 100644 --- a/dascore/utils/hdf5.py +++ b/dascore/utils/hdf5.py @@ -1,9 +1,4 @@ -""" -Utilities for working with HDF5 files. - -Pytables should only be imported in this module in case we need to switch -out the hdf5 backend in the future. -""" +"""Utilities for working with HDF5 files (h5py-based).""" from __future__ import annotations @@ -11,42 +6,23 @@ import os import shutil import tempfile -import time -import warnings from collections.abc import Sequence -from contextlib import contextmanager, suppress +from contextlib import suppress from functools import partial from pathlib import Path -from typing import Literal import numpy as np import pandas as pd -import tables from h5py import File as H5pyFile -from packaging.version import parse as get_version -from pandas.io.common import stringify_path -from tables import ClosedNodeError -from tables import File as PyTablesFile -import dascore as dc from dascore.compat import UPath -from dascore.config import config_attr, get_config -from dascore.constants import max_lens, remote_hdf5_tuned_protocols -from dascore.exceptions import InvalidFileHandlerError, InvalidIndexVersionError -from dascore.io.core import PatchFileSummary -from dascore.utils.mapping import FrozenDict +from dascore.config import get_config +from dascore.constants import remote_hdf5_tuned_protocols from dascore.utils.misc import ( _maybe_make_parent_directory, _maybe_unpack, - cached_method, - suppress_warnings, unbyte, ) -from dascore.utils.pd import ( - _remove_base_path, - fill_defaults_from_pydantic, - list_ser_to_str, -) from dascore.utils.remote_io import ( _FallbackFileObj, _get_cached_local_file, @@ -54,15 +30,11 @@ get_local_handle, is_no_range_http_error, ) -from dascore.utils.time import get_max_min_times, to_datetime64, to_int, to_timedelta64 - -HDF5ExtError = tables.HDF5ExtError -NoSuchNodeError = tables.NoSuchNodeError -NodeError = tables.NodeError ns_to_datetime = partial(pd.to_datetime, unit="ns") ns_to_timedelta = partial(pd.to_timedelta, unit="ns") + class _ManagedH5pyFile: """ DASCore's internal h5py handle wrapper with deterministic close behavior. @@ -195,414 +167,15 @@ def open_h5_resource( handle.close() raise try: - _maybe_make_parent_directory(resource) + if mode != "r": + _maybe_make_parent_directory(resource) return _ManagedH5pyFile(constructor(resource, mode=mode)) except TypeError: msg = f"Couldn't get handle from {resource} using h5py" raise NotImplementedError(msg) -class _HDF5Store(pd.HDFStore): - """ - A work-around for pandas HDF5 store not accepting - pytables.File objects. - """ - - def __init__( # pragma: no cover - self, - path, - mode: str = "a", - complevel: int | None = None, - complib=None, - fletcher32: bool = False, - **kwargs, - ) -> None: - if isinstance(path, str | Path): - self._path = stringify_path(path) - elif isinstance(path, tables.File): - self._path = stringify_path(path.filename) - self._mode = "a" if mode is None else mode - self._handle = None - self._complevel = complevel if complevel else 0 - self._complib = complib - self._fletcher32 = fletcher32 - self._filters = None - if isinstance(path, tables.File): - self._handle = path - else: - self.open(mode) - - -@contextmanager -def open_hdf5_file( - path_or_handler: Path | str | tables.File, - mode: Literal["r", "w", "a"] = "r", -) -> tables.File: - """ - A helper function for getting a `tables.file.File` object. - - If a file reference (str or Path) is passed this context manager will - close the file when it exists. - - Parameters - ---------- - path_or_handler - The input - mode - The mode in which to open the file. - - Raises - ------ - InvalidBuffer if a writable mode is requested from a read only handler. - """ - - def _validate_mode(current_mode, desired_mode): - """Ensure modes are compatible else raise.""" - if desired_mode == "r": - return - # if a or w is desired the current mode should be w - if not current_mode == "w": - msg = ( - f"A HDF5 file handler with mode 'r' was provided but " - f"mode: {desired_mode} was requested." - ) - raise InvalidFileHandlerError(msg) - - if isinstance(path_or_handler, str | Path): - # Note: We suppress DataTypeWarnings because pytables fails to read - # 8 bit enum indicating true or false written by h5py. See: - # https://github.com/PyTables/PyTables/issues/647 - with suppress_warnings(tables.DataTypeWarning): - with tables.open_file(path_or_handler, mode) as fi: - yield fi - elif isinstance(path_or_handler, tables.File): - _validate_mode(path_or_handler.mode, mode) - yield path_or_handler - - -def _get_kernel_query(starttime: int, endtime: int, buffer: int): - """ - Create a HDF5 kernel query based on start and end times. - - This is necessary because hdf5 doesn't accept inverted conditions. - A slight buffer is applied to the ranges to make sure no edge files - are excluded. - """ - t1 = starttime - buffer - t2 = endtime + buffer - con = ( - f"(time_min>{t1:d} & time_min<{t2:d}) | " - f"((time_max>{t1:d} & time_max<{t2:d}) | " - f"(time_min<{t1:d} & time_max>{t2:d}))" - ) - return con - - -class HDFPatchIndexManager: - """ - A class for writing/querying an index table of summary patch info to hdf5. - - It creates a table of patch summary info, a table of metadata and a time - stamp of the last time it was updated. - """ - - # string column sizes in hdf5 table - _min_itemsize = max_lens - # columns which should be indexed for fast querying - _query_columns = ("time_min", "time_max") - # functions applied to encode dataframe before saving to hdf5 - _column_encoders = FrozenDict( - { - "time_min": lambda x: to_int(to_datetime64(x)), - "time_max": lambda x: to_int(to_datetime64(x)), - "time_step": lambda x: to_int(to_timedelta64(x)), - "dims": list_ser_to_str, - "path": lambda x: x.astype(str), - } - ) - # functions to apply to decode dataframe after loading from hdf file - _column_decoders = FrozenDict( - { - "time_min": ns_to_datetime, - "time_max": ns_to_datetime, - "time_step": ns_to_timedelta, - } - ) - # base model which determines fields - _base_model = PatchFileSummary - # any fields to skip - _skip_fields = () - # The minimum version of dascore required to read this index. If an older - # version is used an error will be raised. - _min_version = "0.0.13" - - def __init__(self, path, namespace=""): - super().__init__() - self.namespace = namespace - self.path = path - - @property - def index_columns(self): - """Get the columns used for indexing.""" - out = set(self._base_model.model_fields) - set(self._skip_fields) - return tuple(out) - - buffer: np.timedelta64 = config_attr("index_query_buffer") - complib: str = config_attr("hdf_index_complib") - complevel: int = config_attr("hdf_index_complevel") - max_retries: int = config_attr("hdf_index_max_retries") - - # columns which should be indexed for fast querying - @property - def _time_node(self): - """The node/table where the update time information is stored.""" - return "/".join([self.namespace, "last_updated"]) - - @property - def _index_node(self): - """Return the node/table where the index information is stored.""" - return "/".join([self.namespace, "index"]) - - @property - def _meta_node(self): - """The node/table where the update metadata is stored.""" - return "/".join([self.namespace, "metadata"]) - - def encode_table(self, df, path=None): - """Encode the table for writing to hdf5.""" - # apply column encoders, make paths relative to reference path - # and drop any non-index columns. - cols = set(df.columns) - for col, func in self._column_encoders.items(): - if col not in cols: - continue - df[col] = func(df[col]) - out = ( - df.pipe(fill_defaults_from_pydantic, self._base_model) - .loc[:, list(self.index_columns)] - .assign(path=lambda x: _remove_base_path(x["path"], path)) - ) - # there shouldn't be any null values in index now - assert not out.isnull().any().any(), "null values found in index" - return out - - def decode_table(self, df): - """Decode the table from hdf5.""" - # ensure the base path is not in the path column - for col, func in self._column_decoders.items(): - df[col] = func(df[col]) - # populate index store and update metadata - # assert not df.isnull().any().any(), "null values found in index" - return df - - def get_index(self, time_min=None, time_max=None, **kwargs): - """ - Read part of the hdf5 index from path meeting time min/max reqs. - - Parameters - ---------- - time_min - The start time of the entries to read. - time_max - The end time of the entries to read. - """ - - def _get_index(where, fail_counts=0, **kwargs): - try: - df = pd.read_hdf(self.path, self._index_node, where=where, **kwargs) - except (ClosedNodeError, Exception) as e: - # Sometimes in concurrent updates the nodes need time to open/close - # so we implement a simply "wait and retry" strategy. - # This is a bit wonky but we have found it to work well in practice. - if fail_counts >= self.max_retries: - raise e - time.sleep(0.1) - return _get_index(where, fail_counts=fail_counts + 1, **kwargs) - else: - return df - - time_min, time_max = get_max_min_times((time_min, time_max)) - where = _get_kernel_query( - time_min.view(np.int64), - time_max.view(np.int64), - self.buffer.view(np.int64), - ) - df = _get_index(where, **kwargs) - return self.decode_table(df) - - def write_update( - self, - update_df, - update_time=None, - base_path: str | Path = "", - ): - """Convert updates to dataframe, then append to index table.""" - # read in dataframe and prepare for input into hdf5 index - update_time = update_time or time.time() - df = self.encode_table(update_df.copy(), path=base_path) - with _HDF5Store(self.path) as store: - try: - nrows = store.get_storer(self._index_node).nrows - except (AttributeError, KeyError): - store.append( - self._index_node, - df, - min_itemsize=self._min_itemsize, - **self.hdf_kwargs, - ) - else: - df.index += nrows - store.append(self._index_node, df, append=True, **self.hdf_kwargs) - self._update_metadata(store, update_time) - - def _update_metadata(self, store, update_time): - # update timestamp - update_time = time.time() if update_time is None else update_time - store.put(self._time_node, pd.Series(update_time)) - # make sure meta table also exists. - # Note this is here to avoid opening the store again. - if self._meta_node not in store: - meta = self._make_meta_table() - store.put(self._meta_node, meta, format="table") - - def _read_metadata(self): - """Read the metadata table.""" - try: - with _HDF5Store(self.path, "r") as store: - out = store.get(self._meta_node) - store.close() - return out - except (FileNotFoundError, ValueError, KeyError, OSError): - with suppress(UnboundLocalError): - store.close() - self._ensure_meta_table_exists() - return pd.read_hdf(self.path, self._meta_node) - - def _ensure_meta_table_exists(self): - """If the base path exists ensure it has a meta table, if not create it.""" - if not Path(self.path).exists(): - return - with _HDF5Store(self.path) as store: - # add metadata if not in store - if self._meta_node not in store: - meta = self._make_meta_table() - store.put(self._meta_node, meta, format="table") - - def _make_meta_table(self): - """Get a dataframe of meta info.""" - meta = dict( - dascore_version=dc.__last_version__, - ) - return pd.DataFrame(meta, index=[0]) - - @property - def hdf_kwargs(self) -> dict: - """A dict of hdf_kwargs to pass to PyTables.""" - return dict( - complib=self.complib, - complevel=self.complevel, - format="table", - data_columns=list(self._query_columns), - ) - - @cached_method - def validate_version(self): - """Handles issues with version mismatches.""" - # get the version from file, if the file doesnt exist then None - version = self._version_or_none - if version is not None: - # check if index is too old to be read by this version of the parser. - # If this is the case, users of this class should handle its - # re-creation. - min_version_tuple = get_version(self._min_version) - index_version = get_version(version) - if min_version_tuple > index_version: - msg = ( - f"The indexing schema has changed since {self._min_version} " - f"and must be regenerated." - ) - raise InvalidIndexVersionError(msg) - # check if index was created with newer version of dascore - dascore_version = get_version(dc.__last_version__) - if index_version > dascore_version: - msg = ( - f"The index was created with a newer version of dascore (" - f"{version}), you are running ({dc.__last_version__}), " - f"You may encounter problems, consider updating DASCore." - ) - warnings.warn(msg) - - @property - def _index_version(self) -> str: - """Get the version of dascore used to create the index.""" - return self._read_metadata()["dascore_version"].iloc[0] - - @property - def has_index(self) -> bool: - """Return True if an index table has been written.""" - expected_node = "/".join([self.namespace, "metadata"]) - with open_hdf5_file(self.path) as h5: - try: - h5.get_node(expected_node) - except NoSuchNodeError: - return False - else: - return True - - @property - def _version_or_none(self) -> str | None: - """Return the version string or None if it doesn't yet exist.""" - try: - version = self._index_version - except FileNotFoundError: - return - return version - - @property - def last_updated_timestamp(self) -> float | None: - """Return the last modified time stored in the index, else None.""" - try: - out = pd.read_hdf(self.path, self._time_node)[0] - except (OSError, IndexError, ValueError, KeyError, AttributeError): - out = None - return out - - -class PyTablesReader(PyTablesFile): - """A thin wrapper around pytables File object for reading.""" - - mode = "r" - constructor = PyTablesFile - - @classmethod - def get_handle(cls, resource): - """Get the File object from various sources.""" - if isinstance(resource, cls | PyTablesFile): - return resource - try: - _maybe_make_parent_directory(resource) - return cls.constructor(resource, mode=cls.mode) - except TypeError: - msg = f"Couldn't get handle from {resource} using {cls}" - raise NotImplementedError(msg) - - -class LocalPyTablesReader(PyTablesReader): - """A PyTables reader which first materializes remote resources locally.""" - - @classmethod - def get_handle(cls, resource): - """Get a local-file-backed PyTables handle.""" - return get_local_handle(resource, super().get_handle) - - -class PyTablesWriter(PyTablesReader): - """A thin wrapper around pytables File object for writing.""" - - mode = "a" - - -class H5Reader(PyTablesReader): +class H5Reader: """A thin wrapper around h5py for reading files. Remote UPath resources stay remote-first and transparently retry against @@ -631,7 +204,7 @@ def get_handle(cls, resource): """ Get the HDF5 handle from local paths, remote paths, or open handles. - Unlike PyTablesReader, h5py can consume a binary file object via the + h5py can consume a binary file object via the ``fileobj`` driver, so remote UPath inputs stay streaming-based here. """ if isinstance(resource, cls | _ManagedH5pyFile): @@ -741,12 +314,6 @@ def get_handle(cls, resource): return super().get_handle(resource) -# These are left here for backward compatibility, but should not be -# used in new code. -HDF5Writer = PyTablesWriter -HDF5Reader = PyTablesReader - - def unpack_scalar_h5_dataset(dataset): """ Unpack a scalar H5Py dataset. diff --git a/dascore/utils/misc.py b/dascore/utils/misc.py index dd161159c..78f0d1444 100644 --- a/dascore/utils/misc.py +++ b/dascore/utils/misc.py @@ -432,34 +432,6 @@ def iterate(obj): return obj if isinstance(obj, Iterable) else (obj,) -class CacheDescriptor: - """A descriptor for storing infor in an instance cache (mapping).""" - - def __init__(self, cache_name, func_name, args=None, kwargs=None): - self._cache_name = cache_name - self._func_name = func_name - self._args = () if args is None else args - self._kwargs = {} if kwargs is None else kwargs - - def __set_name__(self, owner, name): - """Method to set the name of the description on the instance.""" - self._name = name - - def __get__(self, instance, owner): - """Get contents of the cache.""" - cache = getattr(instance, self._cache_name) - if self._name not in cache: - func = getattr(instance, self._func_name) - out = func(*self._args, **self._kwargs) - cache[self._name] = out - return cache[self._name] - - def __set__(self, instance, value): - """Set the cache contents.""" - cache = getattr(instance, self._cache_name) - cache[self._name] = value - - def optional_import( package_name: str, on_missing: Literal["raise", "warn", "ignore"] = "raise" ) -> ModuleType | None: @@ -766,6 +738,11 @@ def _dict_list_diffs(dict_list): return sorted(out) +def is_range(value) -> bool: + """True for a 2-tuple range (a ``(start, stop)`` selector).""" + return isinstance(value, tuple) and len(value) == 2 + + def sanitize_range_param(select) -> tuple: """Given a slice or tuple, check and return slice or tuple.""" # convert ellipses or ellipses values diff --git a/dascore/utils/patch.py b/dascore/utils/patch.py index 096a45d11..c59e7d130 100644 --- a/dascore/utils/patch.py +++ b/dascore/utils/patch.py @@ -26,12 +26,13 @@ ) from dascore.exceptions import ( CoordDataError, + CoordError, IncompatiblePatchError, ParameterError, PatchAttributeError, PatchCoordinateError, ) -from dascore.units import get_quantity, is_percent +from dascore.units import convert_units, get_quantity, is_percent from dascore.utils.attrs import combine_patch_attrs from dascore.utils.coordmanager import merge_coord_managers from dascore.utils.deprecate import deprecate @@ -40,13 +41,13 @@ from dascore.utils.misc import ( _apply_union_indexers, _merge_tuples, - all_diffs_close_enough, get_middle_value, iterate, to_object_array, warn_or_raise, yield_sub_sequences, ) +from dascore.utils.paths import is_memory_uri from dascore.utils.time import to_float attr_type = dict[str, Any] | str | Sequence[str] | None @@ -338,6 +339,10 @@ def patches_to_df( # Handle spool case if hasattr(patches, "get_contents"): df = patches.get_contents() + # get_contents() carries only metadata; embed the patches so the + # flat-dump path can serve them (the "patch" column is the point). + if "patch" not in df.columns: + df = df.assign(patch=to_object_array(list(patches))) elif isinstance(patches, pd.DataFrame): df = patches else: @@ -358,6 +363,7 @@ def patches_to_df( df["patch"] = None return df + @deprecate( info=( "merge_patches is deprecated. Use spool.chunk instead. " @@ -417,25 +423,63 @@ def _get_merge_dim(df) -> str | None: return dims_vary[dims_vary].index[0] -def _maybe_expected_step(df, dim): - """Get the expected step if all steps are close, else None.""" - col = df[f"{dim}_step"].values - if all_diffs_close_enough(col): - return get_middle_value(col) - return None +def _middle_step(coords, dim, target_units): + """Return the middle member step expressed in the merged coord's units.""" + steps = [] + for manager in coords: + coord = manager.coord_map[dim] + step = coord.step + if pd.isnull(step): + continue + if target_units is not None and coord.units is not None: + step = convert_units(step, to_units=target_units, from_units=coord.units) + steps.append(step) + if not steps: + return None + return get_middle_value(np.asarray(steps)) + + +def _split_coord_merge_kwargs(merge_kwargs) -> tuple[dict, dict]: + """Split spool merge kwargs into (attr kwargs, coord kwargs).""" + merge_kwargs = dict(merge_kwargs or {}) + coord_kwargs = { + "snap_coords": merge_kwargs.pop("snap_coords", True), + "tolerance": merge_kwargs.pop("tolerance", 1.5), + } + return merge_kwargs, coord_kwargs -def _get_merged_coord(df, merge_dim, coords, drop_conflicting=False): - """Get merged coordinates, also validate anticipated sampling.""" - new_coord = merge_coord_managers( - coords, dim=merge_dim, drop_conflicting=drop_conflicting +def _get_merged_coord( + df, merge_dim, coords, drop_conflicting=False, snap_coords=True, tolerance=1.5 +): + """ + Get merged coordinates for patches combined along merge_dim. + + The merged dimension coordinate is built by truth-preserving + concatenation of the member coords (exactly contiguous members fuse to + a plain range; recorded seams otherwise), then — when `snap_coords` — + simplified with bounded error: no value moves more than + `tolerance * step`. Merges whose gaps exceed that stay segmented + (honestly non-uniform) rather than being relabeled. + """ + from dascore.core.coords import concat_coords + + try: + merged = concat_coords(*[cm.coord_map[merge_dim] for cm in coords]) + except CoordError: + # Non-monotonic (or otherwise unsegmentable) member coordinates: + # fall back to raw value concatenation of the dim coord. + return merge_coord_managers( + coords, dim=merge_dim, drop_conflicting=drop_conflicting + ) + step = _middle_step(coords, merge_dim, merged.units) + if snap_coords and step is not None: + merged = merged.simplify(tolerance * np.abs(step)) + # Passing the pre-built dim coord avoids materializing the members' + # concatenated values only to discard them. + return merge_coord_managers( + coords, dim=merge_dim, drop_conflicting=drop_conflicting, dim_coord=merged ) - expected_step = _maybe_expected_step(df, merge_dim) - if not pd.isnull(expected_step): - new_coord = new_coord.snap(merge_dim)[0] - # TODO slightly different dt can be produced, let pass for now - # need to think more about how the merging should work. - return new_coord def _force_patch_merge(patch_dict_list, merge_kwargs, **kwargs): @@ -447,7 +491,7 @@ def _force_patch_merge(patch_dict_list, merge_kwargs, **kwargs): """ df = pd.DataFrame(patch_dict_list) merge_dim = _get_merge_dim(df) - merge_kwargs = merge_kwargs if merge_kwargs is not None else {} + attr_kwargs, coord_kwargs = _split_coord_merge_kwargs(merge_kwargs) if merge_dim is None: # nothing to merge, complete overlap return [patch_dict_list[0]] dims = df["dims"].iloc[0].split(",") @@ -461,10 +505,12 @@ def _force_patch_merge(patch_dict_list, merge_kwargs, **kwargs): attrs = [x.attrs for x in patches] new_data = np.concatenate(data, axis=axis) # Determine if conflicting non-dimensional coords should be dropped. - conf = merge_kwargs.get("conflicts", None) + conf = attr_kwargs.get("conflicts", None) drop_conf_coords = True if conf in {"drop", "keep_first"} else False - new_coord = _get_merged_coord(df, merge_dim, coords, drop_conf_coords) - new_attrs = combine_patch_attrs(attrs, **merge_kwargs) + new_coord = _get_merged_coord( + df, merge_dim, coords, drop_conf_coords, **coord_kwargs + ) + new_attrs = combine_patch_attrs(attrs, **attr_kwargs) patch = dc.Patch(data=new_data, coords=new_coord, attrs=new_attrs, dims=dims) new_dict = {"patch": patch} return [new_dict] @@ -583,15 +629,17 @@ def _get_filename(path_ser, strip_extension): # Handle special cases. if "name" in col_set: return df["name"].astype(str) - if "path" in col_set and df["path"].astype(str).str.len().gt(0).any(): - return _get_filename(df["path"], strip_extension) - # Determine the requested fields and get the ones that are there. + path_ser = df["path"].astype(str) if "path" in col_set else None + if path_ser is not None: + # synthetic in-memory identities are not real file names + usable = path_ser.str.len().gt(0) & ~path_ser.map(is_memory_uri) + if usable.all(): + return _get_filename(df["path"], strip_extension) + # Determine the requested fields; absent columns render as empty so + # names don't depend on which metadata engine produced the dataframe. coord_fields = zip([f"{x}_min" for x in coords], [f"{x}_max" for x in coords]) - requested_fields = list(attrs) + list(*coord_fields) - current = set(df.columns) - fields = [x for x in requested_fields if x in current] - # Get a sub dataframe and convert any datetime things to strings. - sub = df[fields].pipe(_format_time_columns).fillna("").astype(str) + fields = list(attrs) + [field for pair in coord_fields for field in pair] + sub = df.reindex(columns=fields).pipe(_format_time_columns).fillna("").astype(str) out = f"{prefix}_{sep}" + sub[fields[0]].str.cat(sub[fields[1:]], sep=sep) return out @@ -1256,8 +1304,9 @@ def concatenate_patches( >>> spool_concat = spool.concatenate(wave_rank=None) >>> assert "wave_rank" in spool_concat[0].dims >>> - >>> # Concatenate patches in groups of 3. - >>> big_spool = dc.spool([patch] * 12) + >>> # Concatenate patches in groups of 3. Note: spools keep one + >>> # entry per patch instance, so distinct copies are needed. + >>> big_spool = dc.spool([patch.new() for _ in range(12)]) >>> spool_concat = big_spool.concatenate(time=3) >>> assert len(spool_concat) == 4 diff --git a/dascore/utils/patch_assembly.py b/dascore/utils/patch_assembly.py new file mode 100644 index 000000000..94b8814b8 --- /dev/null +++ b/dascore/utils/patch_assembly.py @@ -0,0 +1,257 @@ +""" +Execute spool views: turn member instructions into loaded patches. + +This is the consumer of the members (instruction) table that +`dascore.utils.chunk_plan` produces and every spool view carries: it +joins member rows to their source rows, loads each source patch through +a caller-supplied loader, applies exact trims, and merges multi-member +outputs (streaming into a pre-allocated buffer when the output size is +known). The spool owns *what* rows exist; this module owns *how* a row +becomes a Patch. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass + +import numpy as np +import pandas as pd + +import dascore as dc +from dascore.exceptions import CoordMergeError +from dascore.utils.attrs import combine_patch_attrs +from dascore.utils.misc import broadcast_for_index +from dascore.utils.patch import ( + _force_patch_merge, + _get_merge_dim, + _get_merged_coord, + _split_coord_merge_kwargs, +) +from dascore.utils.pd import ( + _convert_min_max_in_kwargs, + get_dim_names_from_columns, +) + + +def _get_varying_dim(df) -> str | None: + """ + Get the single dimension whose range varies across rows of df. + + Returns None when no dimension varies, several do, or the dataframe + doesn't carry range columns for the varying dimension; those cases + need the fully materialized merge to sort out. + """ + dims = get_dim_names_from_columns(df) + varying = [] + for dim in dims: + mins, maxs = df.get(f"{dim}_min"), df.get(f"{dim}_max") + if mins.nunique(dropna=False) > 1 or maxs.nunique(dropna=False) > 1: + varying.append(dim) + return varying[0] if len(varying) == 1 else None + + +def _estimate_merge_samples(df, dim) -> int | None: + """ + Estimate the total number of samples along dim of the merged rows. + + Returns None if the estimate cannot be made (eg unknown steps), in + which case streaming the merge isn't possible. + """ + if dim is None: + return None + cols = [f"{dim}_min", f"{dim}_max", f"{dim}_step"] + if not set(cols).issubset(df.columns): + return None + mins, maxs, steps = (df[x] for x in cols) + if mins.isnull().any() or maxs.isnull().any() or steps.isnull().any(): + return None + ratios = (maxs - mins) / steps + # Degenerate steps (eg 0) make the sample counts meaningless. + if not np.isfinite(ratios.astype(np.float64)).all(): + return None + counts = np.round(ratios).astype(np.int64) + 1 + if (counts < 0).any(): + return None + return int(counts.sum()) + + +def _match_merge_units(patch, merge_dim, target_units): + """ + Convert a member's merge-dim units to the first member's. + + The planner groups by SI-canonical envelopes, so one output may mix + unit spellings of one dimensionality (metres with feet); merging + requires a single spelling, and the first member's wins. Returns + (patch, target_units); incompatible or missing units pass through + for the merge itself to police. + """ + from dascore.exceptions import UnitError + + if merge_dim is None or merge_dim not in getattr(patch.coords, "coord_map", {}): + return patch, target_units + units = patch.coords.coord_map[merge_dim].units + if target_units is None: + return patch, units + if units is None or units == target_units: + return patch, target_units + try: + patch = patch.convert_units(**{merge_dim: target_units}) + except UnitError: # incompatible dimensionality: merge will raise + return patch, target_units + return patch, target_units + + +def _coord_only_kwargs(patch, kwargs) -> dict: + """Keep only the kwargs naming a dim or coordinate of patch.""" + return { + k: v + for k, v in kwargs.items() + if k in patch.dims or k in patch.coords.coord_map + } + + +@dataclass +class PatchAssembler: + """ + Assemble output patches from joined member rows. + + ``load_patch`` resolves one member row to its source patch (residual + selections included); ``merge_kwargs`` carries the merge behavior. + The plan resolver hands this the joined member frame for one output + at a time. + """ + + load_patch: Callable[[Mapping], dc.Patch] + merge_kwargs: Mapping + + def _patch_from_instruction_df(self, joined): + """Get the patches joined columns of instruction df.""" + df_dict_list = self._df_to_dict_list(joined) + expected_len = len(joined["current_index"].unique()) + merging = len(df_dict_list) > expected_len + merge_dim = _get_varying_dim(joined) if merging else None + if merging: + # Several sources merge into one patch. When the output size can + # be determined from the instructions, stream the sources into a + # pre-allocated array so they don't all need to be in memory with + # the merged output at once. + samples = _estimate_merge_samples(joined, merge_dim) + if samples is not None: + patch = self._merge_patches_streaming( + joined, df_dict_list, merge_dim, samples + ) + return [patch] + out = [] + target_units = None + for patch_kwargs in df_dict_list: + patch = self._load_trimmed_patch(patch_kwargs, joined) + patch, target_units = _match_merge_units(patch, merge_dim, target_units) + # The index doesn't carry all the dimensional info, so get what + # merging needs from the patch coords (cheaper than attr dumps). + info = patch.coords._get_dim_summary() + info["patch"] = patch + out.append(info) + if len(out) > expected_len: + out = _force_patch_merge(out, merge_kwargs=self.merge_kwargs) + return [x["patch"] for x in out] + + def _load_trimmed_patch(self, patch_kwargs, joined) -> dc.Patch: + """Load a single patch and trim it to its instruction range.""" + # convert kwargs to format understood by parser/patch.select + kwargs = _convert_min_max_in_kwargs(patch_kwargs, joined) + patch = self.load_patch(kwargs) + # If the limits of the source patch were not modified, we can just + # skip selection. This is important for missing coordinates + # (NaN values) to not get trimmed out. + source_kwargs = kwargs if kwargs.get("_modified") else {} + # attr-style entries filter rows above; only coordinate entries + # are valid patch selections. + if select_kwargs := _coord_only_kwargs(patch, source_kwargs): + patch = patch.select(**select_kwargs) + return patch + + def _merge_patches_streaming(self, joined, df_dict_list, merge_dim, samples): + """ + Merge the patches described by the instructions along merge_dim. + + Each patch is copied into a pre-allocated output array as it is + loaded, then released; this avoids holding all source patches and + the merged output in memory at the same time, as concatenating + would. + """ + buffer, offset, axis, dims = None, 0, None, None + coords, attrs, summaries = [], [], [] + target_units = None + for patch_kwargs in df_dict_list: + patch = self._load_trimmed_patch(patch_kwargs, joined) + patch, target_units = _match_merge_units(patch, merge_dim, target_units) + if dims is None: + dims = patch.dims + axis = patch.get_axis(merge_dim) + elif patch.dims != dims: + patch = patch.transpose(*dims) + data = patch.data + if buffer is None: + shape = list(data.shape) + shape[axis] = samples + buffer = np.empty(shape, dtype=data.dtype) + # Mixed dtypes upcast, mirroring np.concatenate behavior. + dtype = np.result_type(buffer.dtype, data.dtype) + if dtype != buffer.dtype: + buffer = buffer.astype(dtype) + end = offset + data.shape[axis] + if end > buffer.shape[axis]: + # The estimate came up short (eg from slightly uneven + # sampling); grow the buffer to fit. + shape = list(buffer.shape) + shape[axis] = end + new_buffer = np.empty(shape, dtype=buffer.dtype) + head = broadcast_for_index(buffer.ndim, axis, slice(0, offset)) + new_buffer[head] = buffer[head] + buffer = new_buffer + try: + index = broadcast_for_index(buffer.ndim, axis, slice(offset, end)) + buffer[index] = data + except ValueError as e: + msg = ( + f"Cannot merge patches; their shapes are incompatible " + f"along the dimensions not being merged ({merge_dim})." + ) + raise CoordMergeError(msg) from e + offset = end + coords.append(patch.coords) + attrs.append(patch.attrs) + summaries.append(patch.coords._get_dim_summary()) + if offset != buffer.shape[axis]: # over-estimated; trim excess. + buffer = buffer[broadcast_for_index(buffer.ndim, axis, slice(0, offset))] + # Ensure the loaded patches only vary along the expected dimension, + # the same requirement _force_patch_merge enforces. + summary_df = pd.DataFrame(summaries) + found_dim = _get_merge_dim(summary_df) + if found_dim != merge_dim: + msg = ( + f"Cannot merge patches; expected them to vary along " + f"{merge_dim} but found {found_dim}." + ) + raise CoordMergeError(msg) + attr_kwargs, coord_kwargs = _split_coord_merge_kwargs(self.merge_kwargs) + conf = attr_kwargs.get("conflicts", None) + drop_conflicting = conf in {"drop", "keep_first"} + new_coord = _get_merged_coord( + summary_df, merge_dim, coords, drop_conflicting, **coord_kwargs + ) + new_attrs = combine_patch_attrs(attrs, **attr_kwargs) + return dc.Patch(data=buffer, coords=new_coord, attrs=new_attrs, dims=list(dims)) + + def _df_to_dict_list(self, df): + """ + Convert the dataframe to a list of dicts for iteration. + + This is significantly faster than iterating rows. Empty strings + (missing format fields on file rows) normalize to None; stored + relative paths pass through unchanged — the catalog's resolver + owns resolving them against the spool root. + """ + df = df.copy(deep=False).replace("", None) + return df.to_dict("records") diff --git a/dascore/utils/paths.py b/dascore/utils/paths.py index 24e76374e..4708edccf 100644 --- a/dascore/utils/paths.py +++ b/dascore/utils/paths.py @@ -2,17 +2,46 @@ from __future__ import annotations +import tempfile from pathlib import Path from dascore.compat import UPath from dascore.exceptions import InvalidSpoolError +# Synthetic URI schemes for in-memory patch identities (see +# dascore.io.index.catalog); such paths dispatch to in-memory registries +# and are never treated as file names. +_MEMORY_SCHEMES = ("memorypatch://", "memory://") + def is_pathlike(resource) -> bool: """Return True if resource is supported path-like input.""" return isinstance(resource, str | Path | UPath) +def is_memory_uri(path) -> bool: + """ + Return True if a path is a synthetic in-memory patch identity. + + Matches the exact ``memorypatch://`` / ``memory://`` schemes rather + than any string beginning with "memory", so a real file or directory + named e.g. ``memory_notes.h5`` is not misclassified. + """ + return str(path).startswith(_MEMORY_SCHEMES) + + +def directory_writable(path) -> bool: + """Return True if the directory is writable else False.""" + directory = Path(path) + try: + directory.mkdir(exist_ok=True, parents=True) + with tempfile.NamedTemporaryFile(prefix="._dascore_write_test_", dir=directory): + pass + except OSError: + return False + return True + + def coerce_to_upath(resource) -> UPath: """Return a UPath for path-like resources.""" return resource if isinstance(resource, UPath) else UPath(resource) diff --git a/dascore/utils/pd.py b/dascore/utils/pd.py index 8ce9ce9e0..bac6d5b59 100644 --- a/dascore/utils/pd.py +++ b/dascore/utils/pd.py @@ -3,7 +3,6 @@ from __future__ import annotations import fnmatch -import os from collections import defaultdict from collections.abc import Collection, Mapping, Sequence from functools import cache @@ -15,8 +14,8 @@ import dascore as dc from dascore.constants import PatchType from dascore.core.attrs import PatchAttrs -from dascore.exceptions import ParameterError -from dascore.utils.misc import order_range_tuple, sanitize_range_param +from dascore.exceptions import InvalidSpoolQueryError, ParameterError +from dascore.utils.misc import is_range, order_range_tuple, sanitize_range_param from dascore.utils.time import to_datetime64, to_timedelta64 @@ -26,18 +25,149 @@ def get_regex(seed_str): return fnmatch.translate(seed_str) # translate to re -def _remove_base_path(series: pd.Series, base="") -> pd.Series: +def relative_offset(gmin, gmax, value): """ - Ensure paths stored in column name use unix style paths and have base - path removed. + Resolve one relative bound against a global [gmin, gmax] envelope. + + Positive offsets measure from the start, negative from the end; + None/Ellipsis bounds stay open. Datetime envelopes take numeric + seconds offsets. """ - assert not series.empty, "Series must be non-empty" - unix_paths = series.str.replace(os.sep, "/") - unix_base_path = (str(base) + "/").replace(os.sep, "/") - out = unix_paths.str.replace(unix_base_path, "", regex=False) + if value is None or value is Ellipsis: + return None + if isinstance(gmin, pd.Timestamp) or isinstance(gmin, np.datetime64): + delta = to_timedelta64(abs(float(value))) + return (gmin + delta) if value >= 0 else (gmax - delta) + return (gmin + value) if value >= 0 else (gmax + value) + + +def relative_ranges_to_absolute(df, kwargs: dict) -> dict: + """ + Resolve relative (start, stop) ranges against a frame's global envelopes. + + Operates only on the dataframe's `{name}_min`/`{name}_max` envelope + columns, so both the generic dataframe select path and the catalog + share one relative-select implementation without either depending on + the index query builder. + """ + out = {} + for name, value in kwargs.items(): + lo_col, hi_col = f"{name}_min", f"{name}_max" + if lo_col not in df.columns or hi_col not in df.columns or df.empty: + msg = f"Cannot use relative select on {name!r}." + raise InvalidSpoolQueryError(msg) + if not is_range(value): + # same vocabulary as the catalog path's selector shaping + msg = ( + f"relative=True accepts range selectors only (a (start, stop) " + f"tuple or slice), got {value!r}." + ) + raise InvalidSpoolQueryError(msg) + gmin, gmax = df[lo_col].min(), df[hi_col].max() + lo, hi = value + out[name] = ( + relative_offset(gmin, gmax, lo), + relative_offset(gmin, gmax, hi), + ) return out +def normalize_range_forms(value): + """ + Normalize the patch-level slice range form to a 2-tuple. + + Only slices are converted: bare None/Ellipsis keep their own errors, + and a fully-open range is rejected downstream as having no usable + bounds (per the selector spec). + """ + if isinstance(value, slice): + return sanitize_range_param(value) + return value + + +def resolve_selector_namespaces( + known_attrs: Collection[str], + known_coords: Collection[str], + _attrs: Mapping | None = None, + _coords: Mapping | None = None, + kwargs: Mapping | None = None, +) -> tuple[dict, dict]: + """ + Split selector kwargs into (attrs, coords) per the selector spec. + + Bare kwargs resolve against attributes first, then coordinates; + `_attrs`/`_coords` name their namespace explicitly and validate + against that side only. Each accepts either a mapping of + ``name -> selector`` (the fully general form — required when a name + cannot be a Python keyword, e.g. it collides with a select parameter + or is not an identifier) or a name/collection of names tagging which + *bare kwargs* to interpret in that namespace. Unknown names, and + names supplied in more than one namespace, raise (see #435). + + Both the catalog (which pushes predicates into SQL) and the generic + dataframe select path resolve names here, so the two agree on which + names are valid, what a bare name means, and which range forms are + accepted — the paths differ only in how they *apply* a predicate. + """ + + def _tag_form(spec, kwargs, label): + """Normalize a tag-form spec (names of bare kwargs) to a dict.""" + if spec is None or isinstance(spec, Mapping): + return spec, kwargs + names = [spec] if isinstance(spec, str) else list(spec) + if not all(isinstance(n, str) for n in names): + msg = ( + f"{label} must be a mapping of name -> selector, or a " + "name/collection of names tagging bare keyword arguments." + ) + raise InvalidSpoolQueryError(msg) + kwargs = dict(kwargs or {}) + out = {} + for n in names: + if n not in kwargs: + msg = f"{label}={n!r} names no bare keyword argument." + raise InvalidSpoolQueryError(msg) + out[n] = kwargs.pop(n) + return out, kwargs + + _attrs, kwargs = _tag_form(_attrs, kwargs, "_attrs") + _coords, kwargs = _tag_form(_coords, kwargs, "_coords") + known_attrs, known_coords = set(known_attrs), set(known_coords) + # A name in both explicit namespaces is a caller error whether or not + # it is valid in either, so this precedes the membership checks. + if duplicates := set(_attrs or {}) & set(_coords or {}): + names = ", ".join(repr(x) for x in sorted(duplicates)) + raise InvalidSpoolQueryError(f"{names} given in both _attrs and _coords.") + attrs: dict = {} + coords: dict = {} + for items, allowed, out, noun in ( + (_attrs, known_attrs, attrs, "an attribute"), + (_coords, known_coords, coords, "a coordinate"), + ): + for name, value in (items or {}).items(): + if name not in allowed: + msg = f"{name!r} is not {noun} of this spool." + raise InvalidSpoolQueryError(msg) + out[name] = normalize_range_forms(value) + for name, value in (kwargs or {}).items(): + if name in attrs or name in coords: + msg = f"{name!r} given as both a bare kwarg and in _attrs/_coords." + raise InvalidSpoolQueryError(msg) + value = normalize_range_forms(value) + if name in known_attrs: + attrs[name] = value + elif name in known_coords: + coords[name] = value + else: + msg = ( + f"{name!r} is neither an attribute nor a coordinate of this " + f"spool. Attributes: {sorted(known_attrs)}; " + f"coordinates: {sorted(known_coords)}." + ) + raise InvalidSpoolQueryError(msg) + return attrs, coords + + def _get_min_max_query(kwargs, df): """ Get a dict of {column_name: Optional[min_val], Optional[max_val]}. @@ -190,7 +320,7 @@ def _convert_times(df, some_dict): return some_dict -def get_interval_columns(df, name, arrays=False): +def get_interval_columns(df, name): """ Return a series of start, stop, step for columns. @@ -200,8 +330,6 @@ def get_interval_columns(df, name, arrays=False): The input dataframe. name The name of the coordinate (eg time). - arrays - If True, return output as numpy arrays, else pandas series. """ names = f"{name}_min", f"{name}_max", f"{name}_step" missing_cols = set(names) - set(df.columns) @@ -213,10 +341,7 @@ def get_interval_columns(df, name, arrays=False): ) raise ParameterError(msg) start, stop, step = df[names[0]], df[names[1]], df[names[2]] - if not arrays: - return start, stop, step - else: - return start.values, stop.values, step.values + return start, stop, step def yield_range_tuple_from_kwargs(df, kwargs) -> tuple[str, slice]: @@ -368,16 +493,6 @@ def get_dim_names_from_columns(df: pd.DataFrame) -> list[str]: return sorted(out) -def get_column_names_from_dim(dims: Sequence[str]) -> list: - """Get column names from a sequence of dimensions.""" - out = [] - for name in dims: - out.append(f"{name}_min") - out.append(f"{name}_max") - out.append(f"{name}_step") - return out - - def fill_defaults_from_pydantic(df, base_model: type[BaseModel]): """ Fill missing columns in dataframe with defaults from base_model. @@ -486,26 +601,6 @@ def _column_or_value(df, col, value): return out -def _instructions_modified(instruct_df, sub_source): - """ - Determine if the instruction df columns are the same as the source. - - This is useful for determining which patches need select arguments. - """ - # Get the source and desired output dfs broadcast together. - names = set(sub_source.columns) & set(instruct_df.columns) - source = sub_source.loc[instruct_df["source_index"].values] - # not_modified = np.ones(len(instruct_df), dtype=bool) - not_modified = ~_column_or_value(source, "_modified", False) - for name in names: - val1, val2 = source[name].values, instruct_df[name].values - eq = val1 == val2 - null = pd.isnull(val1) & pd.isnull(val2) - not_modified &= eq | null - modified = ~not_modified - return modified - - def patch_to_dataframe(patch: PatchType) -> pd.DataFrame: """ Convert a patch to a dataframe. diff --git a/dascore/utils/time.py b/dascore/utils/time.py index 38bf6f088..9f3e7ba0d 100644 --- a/dascore/utils/time.py +++ b/dascore/utils/time.py @@ -9,10 +9,8 @@ import pandas as pd from dascore.constants import ( - LARGEDT64, NUMPY_TIME_UNIT_MAPPING, ONE_SECOND, - SMALLDT64, timeable_types, ) from dascore.exceptions import TimeError @@ -469,24 +467,3 @@ def dtype_time_like(dtype_or_array) -> bool: if is_timedelta or is_datetime: return True return False - - -def get_max_min_times(kwarg_time=None): - """ - Function to get min/max times from a tuple of possible time values. - - If None, return max/min times possible. - """ - # first unpack time from tuples - assert kwarg_time is None or len(kwarg_time) == 2 - time_min, time_max = (None, None) if kwarg_time is None else kwarg_time - # get defaults if starttime or endtime is none - time_min = None if pd.isnull(time_min) else time_min - time_max = None if pd.isnull(time_max) else time_max - time_min = to_datetime64(time_min or SMALLDT64) - time_max = to_datetime64(time_max or LARGEDT64) - if time_min is not None and time_max is not None: - if time_min > time_max: - msg = "time_min cannot be greater than time_max." - raise ValueError(msg) - return time_min, time_max diff --git a/docs/changelog.qmd b/docs/changelog.qmd index 6a9bcf6c2..16900ca05 100644 --- a/docs/changelog.qmd +++ b/docs/changelog.qmd @@ -4,6 +4,22 @@ The [releases page](https://github.com/DASDAE/dascore/releases) tracks changes f ## Unreleased API Changes +- **The spool class hierarchy is collapsed into a single concrete class.** `MemorySpool`, `DirectorySpool`, and `FileSpool` are removed (no aliases), along with the `dascore.clients` package; every spool is now `dc.Spool` under the `dc.BaseSpool` ABC. Construct via `dc.spool(...)` (or `Spool.from_directory`/`Spool.from_file`); test for in-memory content with `spool.has_live_patches` instead of `isinstance` checks. The `select_kwargs` constructor parameter for directory spools is removed — use `.select(...)` after `update()`. +- **`Spool.update()` is allowed only on a root spool.** Any derived spool — the result of `select`, slicing, `sort`, `chunk`, or `+` — raises instead of silently refreshing or widening; update the root and re-apply operations (`root = root.update(); view = root.select(...)`). +- **Spools present patches in a defined order.** Patch-list spools keep construction order on every access path; directory spools present in time order (maintained across index updates); combining spools concatenates, with duplicate patches keeping their first position. Spool equality is order-sensitive, metadata-level, and compares contents rather than backing (a live spool can equal a directory spool over identical data). +- **Coordinate boolean masks are no longer accepted by `Spool.select`** (they were only well-defined on spools whose patches share sizes, and never reduced file reads). Use `spool.map(lambda p: p.select(dim=mask))` for per-patch masking; boolean arrays over *patches* (`spool[bool_array]`) still select membership, and `samples=True` index ranges still apply per patch with Python-slice clamping. +- Memory and directory spools now share a catalog-backed metadata selection path. Attribute and coordinate candidates are pushed into SQLite lazily, while exact coordinate trimming remains a patch-load operation. +- Directory indexes now use the constrained seven-table SQLite schema in `.dascore_index.sqlite3`. Experimental DuckDB and Parquet index backends and the `engine`/`index_engine` selection parameters were removed. Prototype indexes from the earlier schema must be deleted and rebuilt. +- Combining spools (`a + b`) now preserves each operand's full current contents: pending coordinate/samples trims and sort order are baked into the union (as new patch identities backed by the same lazy loading) instead of being silently dropped. Membership-style state (attribute selections, slices, patch-id arrays) still unions by rows, so identity-based deduplication keeps working there. +- `Spool.sort(...)` accepts any coordinate of the spool (including renamed and auxiliary coordinates), not just `time`/`distance`; chunk/concatenate outputs keep describing non-dimension coordinates so they remain selectable and sortable; contiguous descending-coordinate patches now merge under `chunk(dim=None)`; and `concatenate` on an empty spool returns an empty spool. +- Chunking a restructured spool along a different dimension now plans over the spool's current patches, so `chunk(time=None).chunk(distance=...)` partitions both dimensions instead of silently undoing the first operation (re-chunking the *same* dimension still re-plans from the original members). Directory spools present patches in per-patch time order even when one multi-patch file straddles a patch of another; mixed planned/live views keep their plan-backed rows through serialization and process-backed `map`; and the segmented-coordinate write guard covers plan-assembled (file-backed) spools, not only in-memory ones. +- Spool equality compares effective contents rather than representation: pending residual selections fold into the compared envelopes (a trimmed view equals its union-materialized twin; spools differing only by a trim still differ), and internal coordinate identity keys no longer participate. Negative `samples=True` indices now resolve per patch in the planner, so chunking a tail selection reports honest envelopes. +- Directory presentation order survives combining: `+` bakes the per-patch time order into the union when record-grain transfer would present rows differently (interleaved multi-patch files); ordinary archives keep record transfer and same-source deduplication. Rows without a value for the ordering key (e.g. patches without absolute time) sort last under any spool ordering, matching the ordinal renumberer's rule. +- Chunking is defined on dimensions only: a patch carrying the chunk name solely as a non-dimensional coordinate is treated like a patch missing the dimension — `missing_dim="raise"` (the default) fails eagerly with an explanatory message and `missing_dim="drop"` excludes it. Previously such patches produced plan outputs that failed with `CoordError` only when a patch was accessed. +- Chunk partitioning is unit-aware along the chunked dimension: patches whose coordinate units have different dimensionality (or unitful vs unitless) can never plan into one output — they were previously grouped by raw SI magnitude and failed only at load. Compatible unit spellings (e.g. metres and feet) now merge correctly, converted to the first member's units. +- Spool selection now validates unknown names and quantity dimensionality, composes chained regular expressions with AND, supports explicit `_attrs` and `_coords` namespaces (as `name -> selector` mappings, or as a name/collection of names tagging bare keyword arguments), and applies relative offsets only to coordinate selectors. Values indexed without units stay candidates for quantity selectors, and an attribute observed with dimensionally incompatible units across files skips the incompatible values (with a warning) instead of failing the whole index update. +- The legacy PyTables directory index has been replaced by SQLite. SQLite supports concurrent readers and one serialized writer on local filesystems; network filesystems with unreliable locking are not supported. +- With the PyTables dependency removed, `dascore.utils.hdf5` no longer provides `PyTablesReader`, `PyTablesWriter`, or their `HDF5Reader`/`HDF5Writer` aliases. Use `H5Reader`/`H5Writer` (h5py-based) instead. - DASCore file I/O now accepts `UPath` resources across the main read, scan, spool, and write workflows. Remote file backends such as `memory://` can be used directly for supported formats, and remote directory-based formats such as `XMLBinary` now work when the backend supports listing and file reads. See the file I/O and spool tutorials for examples and current limitations. - `dc.scan(...)` now returns [`PatchSummary`](`dascore.PatchSummary`) objects rather than `PatchAttrs`. - Scan results carry metadata, coordinates, and source information without loading data. File-backed summaries contain enough source information for lazy reloads. diff --git a/docs/contributing/new_format.qmd b/docs/contributing/new_format.qmd index 7cd600c5e..f4f6dcf3a 100644 --- a/docs/contributing/new_format.qmd +++ b/docs/contributing/new_format.qmd @@ -163,7 +163,7 @@ loading the full data array. ## Support for Streams/Buffers -Rather than using paths for the IO methods as shown above, it is better practice to write a `FiberIO` which supports the [python stream interface](https://docs.python.org/3/library/io.html#io.BufferedIOBase) or an opened HDF5 file in the form of a `pytables.File` or `h5py.File` object. There are a few reasons for this: +Rather than using paths for the IO methods as shown above, it is better practice to write a `FiberIO` which supports the [python stream interface](https://docs.python.org/3/library/io.html#io.BufferedIOBase) or an opened HDF5 file in the form of an `h5py.File` object. There are a few reasons for this: * More types of inputs can be supported, including steaming file contents from the web or in-memory streams like [`BytesIO`](https://docs.python.org/3/library/io.html#io.BytesIO). * It is usually more efficient since open-file handles can be automatically reused. diff --git a/docs/index.qmd b/docs/index.qmd index ae2cb4ca3..bc6dea07c 100644 --- a/docs/index.qmd +++ b/docs/index.qmd @@ -45,14 +45,12 @@ patch = spool[0] ```{python} #| output: false import dascore as dc -# Import fetch to read DASCore example files -from dascore.utils.downloader import fetch -# Fetch a sample file path from DASCore (just to get a usable path for the rest of the cell) -directory_path = fetch('terra15_das_1_trimmed.hdf5').parent -# To read a directory of DAS data stored locally on your machine, -# simply replace the above line with: +# Write example DAS files to a local directory just to get a usable path. +# To read a directory of DAS data stored on your machine, +# simply replace the line below with: # directory_path = "/path/to/data/directory/" +directory_path = dc.examples.spool_to_directory(dc.get_example_spool()) spool = ( # Create a spool to interact with directory data @@ -60,9 +58,9 @@ spool = ( # Index the directory contents .update() # Sub-select a specific time range - .select(time_min=('2020-01-01', ...)) + .select(time=('2020-01-01', ...)) # Specify chunk of the output patches - .chunk(time=60, overlap=10) + .chunk(time=2, overlap=0.5) ) ``` diff --git a/docs/notes/notes.qmd b/docs/notes/notes.qmd index 85c73de09..82413e40f 100644 --- a/docs/notes/notes.qmd +++ b/docs/notes/notes.qmd @@ -9,3 +9,6 @@ This section of the documentation provides understanding-oriented explanation fo - [Documentation Strategy](doc_strategy.qmd) - [Fourier Transforms](dft_notes.qmd) - [Velocity to Strain Rate](velocity_to_strain_rate.qmd) +- [Spool Index](spool_index.qmd) +- [Spool Selection](spool_selection.qmd) +- [Spool Chunking](spool_chunking.qmd) diff --git a/docs/notes/spool_chunking.qmd b/docs/notes/spool_chunking.qmd new file mode 100644 index 000000000..403a3816d --- /dev/null +++ b/docs/notes/spool_chunking.qmd @@ -0,0 +1,123 @@ +--- +title: Spool Chunking +--- + +`Spool.chunk` runs in two stages: a **planner** decides everything from metadata alone, and **assembly** loads, trims, and combines patch data only when a patch is requested. The code cells below execute against the real machinery, so this note fails the doc build if it drifts from the implementation. + +## Plans + +The planner consumes the spool's flat metadata relation and produces a `ChunkPlan`: an `outputs` table (one row per patch the chunked spool will contain) and a `members` table (which slice of which source patch feeds each output). `Spool.chunk_plan` exposes the plan `chunk` would execute, with the same arguments. The chunked spool *is* a fresh in-memory catalog whose patch rows are the plan outputs; a plan resolver loads each output's members through the parent's resolver and executes the assembly engine (`dascore.utils.patch_assembly`) row by row. Chunking is one dimension per call; multi-dimensional chunking chains. Re-chunking a chunked spool along the *same* dimension re-plans from the current view's members (the collapse rule); chunking a *different* dimension plans over the spool's current output rows, preserving the boundaries the earlier operation assembled — `chunk(time=2).chunk(distance=100)` partitions both dimensions. `concatenate` is the same machinery with order-based grouping instead of continuity. + +```{python} +import dascore as dc + +spool = dc.get_example_spool("random_das") +plan = spool.chunk_plan(time=3) +assert len(plan.outputs) == len(spool.chunk(time=3)) +assert {"output_id", "_patch_id", "_modified"}.issubset(plan.members.columns) +# Resolved parameters are recorded on the plan, never left in config. +assert isinstance(plan.params["group"], tuple) +assert plan.params["sampling_group_tolerance"] == dc.get_config().sampling_group_tolerance + +# Chaining partitions both dimensions: the second chunk plans over the +# first one's outputs, not the original members. +chained = dc.get_example_spool("random_das").chunk(time=None).chunk(distance=100) +assert all(p.shape[p.get_axis("distance")] <= 100 for p in chained) +``` + +Plans are deterministic: the same spool produces the same plan regardless of metadata row order, and members with `_modified=False` load whole (no per-patch selection cost). + +## Partitioning + +Patches may only combine when they agree on all of: + +1. **Group attributes** — the config option `groupby_attrs` by default (conventional categorical identity: network, station, data type/category, tag, instrument and acquisition ids), overridden per call with `group=`. Differing group values are never an error; the patches simply land in separate outputs. Explicitly passed names must exist somewhere in the spool; config names are best-effort. +2. **Structure** — the dimensions tuple, the coordinate identity of every non-chunked dimension, and the chunked dimension's canonical units (a metre patch can never plan into one output with a seconds patch, or a unitful with a unitless one; compatible spellings such as metres and feet plan together and assembly converts them to the first member's units). +3. **Sampling** — step magnitudes within the relative tolerance `config.sampling_group_tolerance` (default 5%) of the group's smallest member, with matching orientation (ascending never merges with descending; contiguous descending patches merge with each other). +4. **Continuity** — patches within `tolerance` samples of each other, evaluated within each group so unrelated patches can never bridge a gap. + +```{python} +import numpy as np + +t0 = np.datetime64("2020-01-01", "ns") +p1 = dc.get_example_patch(time_min=t0) +time = p1.get_coord("time") +p2 = dc.get_example_patch(time_min=time.max() + time.step) +p3, p4 = p1.update_attrs(station="XX2"), p2.update_attrs(station="XX2") + +# Two stations, each with two contiguous patches: two outputs, no error. +merged = dc.spool([p1, p2, p3, p4]).chunk(time=None) +assert len(merged) == 2 +``` + +Remaining (non-group, non-dimensional) attributes must be single-valued within a partition, policed by `conflict`: `"raise"` (default), `"drop"`, or `"keep_first"`. + +Chunking is defined on dimensions: a patch that lacks the chunked dimension — including one that carries the name only as a *non-dimensional coordinate*, which cannot be trimmed or merged along it — raises by default, and `missing_dim="drop"` excludes it instead. Losing patches silently would be data loss, so it requires the explicit opt-in. + +```{python} +import pytest +from dascore.exceptions import ChunkError + +no_time = [dc.get_example_patch().mean("time") for _ in range(2)] +with pytest.raises(ChunkError, match="missing_dim"): + dc.spool(no_time).chunk(time=None) +assert len(dc.spool(no_time).chunk(time=None, missing_dim="drop")) == 0 + +# A name carried only as a non-dimensional coordinate counts as missing. +base = dc.get_example_patch() +aux = base.update_coords( + sensor=("distance", np.arange(base.shape[base.get_axis("distance")], dtype=float)) +) +with pytest.raises(ChunkError, match="non-dimensional coordinate"): + dc.spool([aux]).chunk(sensor=100) +``` + +## Merged coordinates + +Assembly builds the chunked dimension's coordinate by exact concatenation of the member coordinates: contiguous members fuse to a plain evenly sampled range, and every real seam is recorded. When `snap_coords=True` (default) the result is then simplified with **bounded error** — no coordinate value moves more than `tolerance * step`. A within-tolerance gap therefore comes back as an evenly sampled range whose worst label error is about half the gap (never more than the tolerance); this is the honest replacement for the old unconditional snap, whose error was unbounded. With `snap_coords=False`, or when accumulated gaps exceed the bound, the coordinate stays segmented — exactly non-uniform, with every gap queryable. + +```{python} +from dascore.core.coords import CoordRange, CoordSegmented + +# Contiguous members: exact fuse to a range. +patch = dc.spool([p1, p2]).chunk(time=None)[0] +assert isinstance(patch.get_coord("time"), CoordRange) + +# A gap forced together by a loose tolerance warns, and with +# snap_coords=False the output keeps the exact segmented coordinate. +gap_start = time.max() + 3 * time.step +p_gap = dc.get_example_patch(time_min=gap_start) +with pytest.warns(UserWarning, match="force merging"): + forced = dc.spool([p1, p_gap]).chunk( + time=None, tolerance=5, snap_coords=False + )[0] +coord = forced.get_coord("time") +assert isinstance(coord, CoordSegmented) +assert len(coord.get_discontinuities("gaps")) == 1 + +# The default simplifies the same merge to a range with bounded error: +# every value within tolerance * step of its exact position. +with pytest.warns(UserWarning, match="force merging"): + snapped = dc.spool([p1, p_gap]).chunk(time=None, tolerance=5)[0] +snapped_coord = snapped.get_coord("time") +assert isinstance(snapped_coord, CoordRange) +deviation = abs(snapped_coord.values - coord.values).max() +assert deviation <= 5 * time.step +``` + +Segmented coordinates are an in-memory representation only: a written patch must be contiguous, so saving raises unless `split=True` (or `patch.split_gaps()` is used first) to write each contiguous section as its own patch. + +## Union spools + +`spool + spool` produces a lazy spool over the union of both spools' metadata; chunking works across the seam, so contiguous patches from different spools (even file-backed and in-memory mixed) merge into one. + +```{python} +combined = dc.spool([p1]) + dc.spool([p2]) +assert len(combined.chunk(time=None)) == 1 + +# Spools have set semantics by patch instance: the same patch (or a copy +# of it) appears once; operations mint distinct instances. +assert len(dc.spool([p1, p1])) == 1 +assert len(dc.spool([p1, p1.new()])) == 2 +assert len(dc.spool([p1]) + dc.spool([p1])) == 1 +``` diff --git a/docs/notes/spool_index.qmd b/docs/notes/spool_index.qmd new file mode 100644 index 000000000..08a91c870 --- /dev/null +++ b/docs/notes/spool_index.qmd @@ -0,0 +1,116 @@ +--- +title: Spool Index +--- + +Directory and in-memory spools use the same metadata model. The persisted directory index is one SQLite file named `.dascore_index.sqlite3`; in-memory spools use the same schema in an in-memory SQLite database. The index stores summaries and source identities, not patch data. + +## Why seven tables? + +The schema separates records with different lifetimes and cardinalities. This avoids repeating source and coordinate metadata on every patch and gives SQLite enforceable ownership boundaries. + +| Table | One row per | Purpose | +|---|---|---| +| `meta_data` | index | Identifies the file and its schema version | +| `sources` | file or directory-format source | Tracks the source path, format, size, modification time, and presentation ordinal | +| `patches` | patch within a source | Stores patch identity and common time/distance envelopes | +| `attrs` | patch | Stores typed attribute values in dynamically added columns | +| `attr_meta` | attribute name and value kind | Maps original attribute names to typed storage columns and canonical units | +| `coord_defs` | unique coordinate value definition | Stores coordinate summaries and deduplicates identical definitions | +| `patch_coords` | coordinate attached to a patch | Links patches to coordinate definitions while retaining the coordinate name and dimensions | + +The last two tables are deliberately separate. Many patches can share a distance coordinate, so `coord_defs` stores it once and `patch_coords` supplies the many-to-many attachment. Range coordinates receive an exact semantic fingerprint, reconstructed from the coordinate values when a scan did not supply one. Non-range coordinates only receive a merge-compatibility identity when the scan provides an exact fingerprint. This distinction prevents a matching envelope from being mistaken for matching coordinate values while retaining the deduplication needed by future merge planning. + +`attrs` and `attr_meta` are also complementary. Attribute names are open-ended, so the index cannot define every typed column in advance. `attr_meta` records the stable mapping and units needed to interpret the columns added to `attrs` as data is ingested. + +## Lifecycle and validation + +The index is an incrementally updated cache. A directory update scans new or changed sources, transactionally replaces their patch rows, and removes rows for deleted sources. Foreign keys cascade source deletion through patches, attributes, and patch-coordinate links. Unreferenced coordinate definitions may remain available for reuse. + +The current schema version is validated before any mutation. An unrelated or incomplete database raises an error with instructions to delete and rebuild it; DASCore does not silently repair or migrate it. A file that identifies itself as a DASCore spool index of a different schema version is rebuilt automatically by the directory indexer — the index is a disposable cache whose truth is the files. + +## Ordering + +The base ordering contract is `(sources.ordinal, patch_id)`: ordinals are assigned at ingest, a replaced source keeps its position while new sources append, so merging catalogs concatenates and duplicate sources keep their first-occurrence position with last-occurrence metadata (dict-merge semantics). Spools built from in-memory patches therefore iterate in construction order on every path. + +Directory catalogs additionally carry a **per-patch default presentation order**: rows present by time (`time IS NULL` last, then time, with ordinal and patch id as deterministic tiebreaks), because source-grain ordinals alone cannot interleave a multi-patch file whose span straddles a patch of another file. The default order is a catalog contract, not view state — a directory root still updates — and an explicit `sort(...)` replaces it. The directory indexer still renumbers source ordinals to time order (earliest patch per source, path as tiebreak) after each sync, which keeps the ordinal grain stable for replacement and union dedup. + +SQLite permits concurrent readers and serializes writers. Initialization and updates use an immediate write transaction and a 30-second busy timeout. This relies on correct local-filesystem locking; reliable operation on network filesystems with weak locking is not promised. + +## The flat relation + +Spool-facing operations consume the tables through one flat relation: a dataframe with one row per patch carrying `{dim}_min/max/step` envelopes, private structural columns (`_patch_id`, `_{name}_def_key` coordinate identities, and `_{name}_units` canonical units — the chunk planner's unit-partition key), the `dims` signature, and one column per attribute. The chunk planner and selection both operate on this relation (see the [Spool Chunking](spool_chunking.qmd) and [Spool Selection](spool_selection.qmd) notes). The cell below runs against the real catalog so this description cannot silently drift. + +```{python} +import dascore as dc +from dascore.io.index.catalog import PatchCatalog + +catalog = PatchCatalog.from_patches(list(dc.get_example_spool("random_das"))) +df = catalog.to_df() +required = { + "_patch_id", + "_time_def_key", + "_time_units", + "dims", + "time_min", + "time_max", + "time_step", +} +assert required.issubset(df.columns) +``` + +## Patch identity and spool set semantics + +A spool of in-memory patches has **set semantics by patch identity**: constructing a spool from a sequence keeps one entry per distinct patch instance, and `spool + spool` unions membership. Each patch carries an instance identity minted eagerly at construction; because patches are immutable, copies (including deep copies and unpickled patches) share that identity, while every patch operation produces a new instance with its own. The identity is the patch's synthetic `memorypatch://` source path in the index, so identity, deduplication, and resolution all agree: + +```{python} +import copy + +patch = dc.get_example_patch() + +# a sequence with the same instance twice is one spool entry +assert len(dc.spool([patch, patch])) == 1 + +# copies share identity regardless of when they are made — identity is +# eager, so there is no access-order dependence +clone = copy.deepcopy(patch) +assert clone._instance_id == patch._instance_id +assert len(dc.spool([patch, clone])) == 1 + +# any operation mints a new, distinct identity +assert len(dc.spool([patch, patch.new()])) == 2 + +# the identity is the row's synthetic path, and it resolves back to the +# very same object +cat = PatchCatalog.from_patches([patch, patch]) # one entry, not two +row = cat.to_df().iloc[0] +assert len(cat.to_df()) == 1 +assert row["path"].startswith("memorypatch://") +assert cat.resolve_row(row.to_dict()) is patch +``` + +File-backed patches are identified by `(base_uri, source_path, source_patch_id)` instead; their rows resolve through readers rather than the registry. + +## Federation + +`spool + spool` merges catalogs table-to-table: source records are reconstructed from the member backends and re-ingested, so coordinate definitions deduplicate by definition key, the same source appearing in several members keeps a single entry, and file paths are absolutized so members with different roots coexist. A composite resolver routes in-memory rows to the shared registry, plan-output rows to their plan resolvers, and everything else through file readers. + +Row-membership state (attribute predicates, slice windows, patch-id arrays) transfers as rows, preserving the set semantics above. State that only lives Python-side — coordinate/samples residual trims, sort specs, and a directory default order that record-grain transfer would actually scramble — is first **baked into an identity-plan derived catalog** (table work only; no patch data loads), so each operand contributes exactly its current contents in its current order. Baking mints new patch identities, so a trimmed operand no longer deduplicates against its source: its contents genuinely differ. Spool equality follows the same philosophy — it compares *effective contents* (residual trims folded into the envelopes, representation artifacts like def keys and backing excluded), so a trimmed view equals its union-materialized twin. + +```{python} +other = dc.get_example_patch().new() +union = dc.spool([patch]) + dc.spool([other]) +assert len(union) == 2 +assert len(dc.spool([patch]) + dc.spool([patch])) == 1 # same patch dedups + +# A trimmed operand contributes its trimmed contents (baked, not dropped), +# and equality compares those effective contents. +t = patch.get_coord("time") +trimmed = dc.spool([patch]).select(time=(t.min() + 10 * t.step, t.min() + 20 * t.step)) +combined = trimmed + dc.spool([]) +assert combined[0].shape == trimmed[0].shape +assert combined == trimmed +``` + +## Scope + +The index answers metadata selection, identifies candidate patches, and plans chunk operations — all without touching patch data. Exact coordinate selection is applied again when a patch is loaded because summary envelopes cannot prove arbitrary-coordinate membership. Materialized derived data (persisting chunked results) remains future work. diff --git a/docs/notes/spool_selection.qmd b/docs/notes/spool_selection.qmd new file mode 100644 index 000000000..97bfafff6 --- /dev/null +++ b/docs/notes/spool_selection.qmd @@ -0,0 +1,51 @@ +--- +title: Spool Selection +--- + +`Spool.select` uses one selector model for memory and directory spools. Patch-list and directory spools compose selections in a `PatchCatalog`; ordinary metadata predicates are pushed into SQLite and remain lazy until contents, length, indexing, or iteration requires rows. + +Bare selector names resolve to attributes first and then coordinates. `_attrs` and `_coords` provide explicit namespaces when needed: either a `name -> selector` mapping (the fully general form, required when a name cannot be a Python keyword) or a name/collection of names tagging which bare keyword arguments belong to that namespace (e.g. `select(sensor=(1, 10), _coords="sensor")`). Unknown names raise immediately instead of being ignored. + +Attribute equality, membership, ranges, and glob predicates are evaluated by the index. Regular expressions use a SQL candidate predicate and an exact residual filter; chained regular expressions are combined with AND. Quantities are converted to the canonical unit recorded by the index, and dimensionally incompatible queries raise rather than silently returning incorrect matches. Values stored without units can never be proven incompatible, so they remain candidates for quantity selectors rather than being silently excluded. + +Coordinate predicates select by range — a `(start, stop)` tuple or slice, with `None`/`...` for an open end. Scalar, value-membership, and boolean-sample-mask coordinate selectors have no exact patch-level meaning spool-wide and are rejected (apply masks per patch, e.g. `spool.map(lambda p: p.select(...))`; boolean arrays over patches, `spool[mask]`, still select membership). Numeric coordinate summaries are stored in canonical SI units, so bare numeric range bounds are interpreted as canonical SI regardless of a patch's native coordinate units, and quantities convert. The exact per-patch trim defers its representation until each patch is known, so a mixed archive of unit-bearing and unitless patches is handled correctly in one selection. + +Coordinate predicates first select patches whose summary envelopes can overlap the request. The loaded patch is then selected exactly. `samples=True` is always patch-local and therefore never excludes a patch at the index stage. `relative=True` resolves coordinate ranges against the current spool view's global envelope; attribute predicates in the same call remain unchanged. + +Restructuring operations that create new patch identities (chunking, concatenation) materialize a derived in-memory catalog whose rows are the plan outputs, so selection on a chunked spool runs the identical catalog engine. Sorting, slicing, and array selection never restructure: they compose lazy order and membership specs on the current catalog. Exact selections already attached to a parent view still apply when member source patches load. + +Catalog views share their source state. Adding, removing, or rescanning sources invalidates realized metadata so existing views observe the updated catalog under their composed predicates. + +The core contract, executed here so drift fails the doc build: names resolve attributes-first then coordinates, unknown names raise, and coordinate ranges are exact on the loaded patches (candidacy by envelope at the index, exactness at load): + +```{python} +import pytest + +import dascore as dc +from dascore.exceptions import InvalidSpoolQueryError + +spool = dc.get_example_spool("diverse_das") +with pytest.raises(InvalidSpoolQueryError): + spool.select(not_a_name="anything") + +selected = spool.select(network="das2") +assert all(p.attrs.network == "das2" for p in selected) + +df = spool.get_contents() +t0 = df["time_min"].min() +window = (t0, t0 + dc.to_timedelta64(1)) +for patch in spool.select(time=window): + coord = patch.get_coord("time") + assert coord.min() >= window[0] and coord.max() <= window[1] +``` + +Bare numeric coordinate bounds mean canonical SI even on a patch whose coordinate is in another unit, and scalar/membership coordinate selectors raise: + +```{python} +ft_patch = dc.get_example_patch().convert_units(distance="ft") +coord = dc.spool([ft_patch]).select(distance=(20, 60))[0].get_coord("distance") +assert float(coord.min()) >= 65 and float(coord.max()) <= 197 # 20-60 m in ft + +with pytest.raises(InvalidSpoolQueryError): + dc.spool([ft_patch]).select(distance=100) # scalar has no range meaning +``` diff --git a/docs/tutorial/file_io.qmd b/docs/tutorial/file_io.qmd index f9975dc20..ba3739fc6 100644 --- a/docs/tutorial/file_io.qmd +++ b/docs/tutorial/file_io.qmd @@ -100,9 +100,9 @@ print(loaded_patch.data.shape) `Patch.attrs` stores non-coordinate metadata only. Coordinate summaries such as `time_min`, `time_max`, and `distance_step` are accessed through [`PatchSummary.get_coord_summary(...)`](`dascore.PatchSummary.get_coord_summary`) or via `patch.summary.get_coord_summary(...)`. ::: -## DirectorySpool +## Directory spools -The [DirectorySpool](`dascore.clients.dirspool.DirectorySpool`) is used to retrieve data from a directory of dascore-readable files. It has the same interface as other spools and is created with the [`dascore.spool`](`dascore.spool`) function. +A spool over a directory of dascore-readable files is created with the [`dascore.spool`](`dascore.spool`) function. It has the same interface as every other spool; the only difference is how it was constructed. For example: @@ -135,26 +135,23 @@ The `Patch.io` namespace also includes functionality for converting `Patch` inst ## Directory Indexer -The 'DirectoryIndexer' is used to track the contents of a directory which -contains fiber data. It creates a small, hidden HDF index file at the top -of the directory which can be efficiently queried for directory contents -(it is used internally by the `DirectorySpool`). +The `DBDirectoryIndexer` tracks the contents of a directory which contains fiber data. It creates a small, hidden SQLite index named `.dascore_index.sqlite3` at the top of the directory. Directory spools use this index internally and push metadata selections into SQLite before loading patch data. See the [spool index note](../notes/spool_index.qmd) for the schema and lifecycle. ```{python} #| output: false import dascore -from dascore.io.indexer import DirectoryIndexer +from dascore.io.index.indexer import DBDirectoryIndexer from dascore import examples as ex # Get a directory with several files diverse_spool = dascore.get_example_spool('diverse_das') path = ex.spool_to_directory(diverse_spool) -# Create an indexer and update the index. This will include any new files -# with timestamps newer than the last update, or create a new HDF index file -# if one does not yet exist. -indexer = DirectoryIndexer(path).update() +# Create an indexer and update the index. This scans new or changed files +# (detected by per-file modification time and size), removes entries of +# deleted files, and creates the index if one does not yet exist. +indexer = DBDirectoryIndexer(path).update() # get the contents of the directory's files df = indexer.get_contents() diff --git a/docs/tutorial/spool.qmd b/docs/tutorial/spool.qmd index b0f26b5cb..dd080418f 100644 --- a/docs/tutorial/spool.qmd +++ b/docs/tutorial/spool.qmd @@ -170,7 +170,7 @@ print(contents) #| echo: false from IPython.display import display -display(contents.drop(columns=['patch'])) +display(contents.drop(columns=[c for c in contents.columns if c.startswith('_')])) ``` The columns returned by `get_contents()` come from the same patch summary metadata exposed by `Patch.summary`, so fields such as `time_min`, `time_max`, and `distance_step` are available without loading the underlying patch data. Source metadata such as `path`, `file_format`, and `source_patch_id` are also available for file-backed spools. @@ -219,11 +219,7 @@ merged_spool = spool.chunk(time=None) ``` # concatenate -Similar to `chunk`, [`Spool.concatenate`](`dascore.BaseSpool.concatenate`) is used to combine patches together. However, `concatenate` doesn't account for coordinate values along the concatenation axis, and can even be used to create new patch dimensions. - -:::{.callout-warning} -However, unlike [`chunk`](`dascore.BaseSpool.chunk`), not all `Spool` types implement [`concatenate`](`dascore.BaseSpool.concatenate`). -::: +Similar to `chunk`, [`Spool.concatenate`](`dascore.BaseSpool.concatenate`) is used to combine patches together. However, `concatenate` doesn't account for coordinate values along the concatenation axis, and can even be used to create new patch dimensions. Like `chunk`, it is available on every spool and produces a lazy, plan-backed result. ```python import dascore as dc diff --git a/environment.yml b/environment.yml index 52f85af2f..6c108387c 100644 --- a/environment.yml +++ b/environment.yml @@ -11,7 +11,6 @@ dependencies: - pooch>=1.2 - xarray - pre-commit - - pytables - h5py - matplotlib>=3.5 - scipy>=1.15.0 diff --git a/pyproject.toml b/pyproject.toml index be9f14a48..968363b48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,8 +53,7 @@ dependencies = [ "pooch>=1.2", "pydantic>2.1", "rich", - "tables>=3.7", - "typing_extensions>=4.12", + "typing_extensions>=4.12", "universal-pathlib", "pint>=0.24.4", "scipy>=1.15", @@ -242,8 +241,6 @@ norecursedirs = [ "worktrees", ] filterwarnings = [ - # Ignore hdf5 warnings from pytables, See pytables #1035 - 'ignore::Warning:tables:' ] markers = [ "network: tests that require network-style filesystem access", diff --git a/scripts/_templates/_quarto.yml b/scripts/_templates/_quarto.yml index 1c46ff45e..8b2cfc950 100644 --- a/scripts/_templates/_quarto.yml +++ b/scripts/_templates/_quarto.yml @@ -220,6 +220,12 @@ website: - text: Velocity to Strain Rate href: notes/velocity_to_strain_rate.qmd + - text: Spool Index + href: notes/spool_index.qmd + + - text: Spool Selection + href: notes/spool_selection.qmd + - id: API title: "API" href: api/dascore.qmd diff --git a/tests/conftest.py b/tests/conftest.py index 66689e721..7d97dc9d4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,20 +8,19 @@ from contextlib import suppress from pathlib import Path +import h5py import matplotlib import numpy as np import pandas as pd import pytest -import tables as tb -import tables.parameters import dascore as dc import dascore.examples as ex -from dascore.clients.dirspool import DirectorySpool from dascore.compat import random_state from dascore.config import set_config from dascore.constants import SpoolType from dascore.core import Patch +from dascore.core.spool import Spool from dascore.examples import get_example_patch from dascore.io.core import read from dascore.utils.coordmanager import merge_coord_managers @@ -85,9 +84,6 @@ def pytest_sessionstart(session): if os.environ.get("CI", False): matplotlib.use("Agg") - # need to set nodes to 32 to avoid crash on p3.11. See pytables#977. - tables.parameters.NODE_CACHE_SLOTS = 32 - # Test-time debug defaults are applied by fixture to avoid state leakage. @@ -427,12 +423,15 @@ def two_patch_directory(tmp_path_factory, terra15_das_example_path, random_patch @pytest.fixture(scope="class") -def diverse_spool_directory(diverse_spool): - """Save the diverse spool contents to a directory.""" - out = ex.spool_to_directory(diverse_spool) - yield out - if out.is_dir(): - shutil.rmtree(out) +def diverse_spool_directory(diverse_spool, tmp_path_factory): + """Save the diverse spool contents to a directory. + + Pytest owns the directory's lifetime: an explicit rmtree teardown + raced lazily-finalized SQLite index connections on Windows + (WinError 32), so no teardown here. + """ + out = tmp_path_factory.mktemp("diverse_spool_dir") + return ex.spool_to_directory(diverse_spool, path=out) @pytest.fixture(scope="class") @@ -501,7 +500,7 @@ def adjacent_spool_no_overlap(random_patch) -> dc.BaseSpool: @register_func(SPOOL_FIXTURES) def one_file_directory_spool(one_file_dir): """Create a directory with a single DAS file.""" - return DirectorySpool(one_file_dir).update() + return Spool.from_directory(one_file_dir).update() @pytest.fixture(scope="class") @@ -516,15 +515,18 @@ def diverse_spool(): def diverse_directory_spool(diverse_spool_directory): """Save the diverse spool contents to a directory.""" out = dc.spool(diverse_spool_directory).update() - return out + yield out + # release the SQLite index handle so Windows can clean the temp dir + out.indexer.close() @pytest.fixture(scope="class") @register_func(SPOOL_FIXTURES) def basic_file_spool(two_patch_directory): """Return a DAS bank on basic_bank_directory.""" - out = DirectorySpool(two_patch_directory).update() - return out.update() + out = Spool.from_directory(two_patch_directory).update().update() + yield out + out.indexer.close() @pytest.fixture(scope="class") @@ -609,9 +611,9 @@ def generic_hdf5(tmp_path_factory): parent.mkdir() path = parent / "simple.hdf5" - with tb.open_file(str(path), "w") as fi: - group = fi.create_group("/", "bob") - fi.create_carray(group, "data", obj=random_state.rand(10)) + with h5py.File(str(path), "w") as fi: + group = fi.create_group("bob") + group.create_dataset("data", data=random_state.rand(10)) return path diff --git a/tests/test_core/test_coord_segmented.py b/tests/test_core/test_coord_segmented.py index 60e396fbe..12c57064d 100644 --- a/tests/test_core/test_coord_segmented.py +++ b/tests/test_core/test_coord_segmented.py @@ -1021,7 +1021,7 @@ def test_write_file_backed_spool_unaffected(self, tmp_path): """Non-memory spools skip the gap inspection (never gapped).""" path1 = dc.write(dc.get_example_patch(), tmp_path / "a.h5", "dasdae") file_spool = dc.spool(path1) - assert not isinstance(file_spool, dc.core.spool.MemorySpool) + assert not file_spool.has_live_patches path2 = dc.write(file_spool, tmp_path / "b.h5", "dasdae") assert path2.exists() @@ -1115,3 +1115,44 @@ def test_units(self): values = np.array([0.0, 1, 2, 10, 11, 12]) coord = CoordSegmented.from_array(values, units="m") assert get_quantity(coord.units) == get_quantity("m") + + +class TestPlannedSpoolWriteGuard: + """The gap write guard covers plan-assembled spools (round-4 F3).""" + + @pytest.fixture() + def gapped_planned_spool(self, tmp_path): + """A file-backed planned spool whose output spans a real gap.""" + import warnings + + src = tmp_path / "src" + src.mkdir() + p1 = dc.get_example_patch() + t = p1.get_coord("time") + p2 = p1.update_coords(time_min=t.max() + 4 * t.step) + dc.write(p1, src / "a.h5", "DASDAE") + dc.write(p2, src / "b.h5", "DASDAE") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + planned = ( + dc.spool(src) + .update(progress=None) + .chunk(time=None, tolerance=5, snap_coords=False, conflict="drop") + ) + assert isinstance(planned[0].get_coord("time"), CoordSegmented) + assert not planned.has_live_patches + return planned + + def test_write_raises_without_split(self, gapped_planned_spool, tmp_path): + """Writing a gapped planned spool raises the documented error.""" + with pytest.raises(ParameterError, match="split"): + dc.write(gapped_planned_spool, tmp_path / "out.h5", "DASDAE") + + def test_write_split_true(self, gapped_planned_spool, tmp_path): + """split=True writes each contiguous section as its own patch.""" + path = tmp_path / "out.h5" + dc.write(gapped_planned_spool, path, "DASDAE", split=True) + back = dc.spool(path) + assert len(back) == 2 + for patch in back: + assert not isinstance(patch.get_coord("time"), CoordSegmented) diff --git a/tests/test_clients/test_dirspool.py b/tests/test_core/test_directory_spool.py similarity index 73% rename from tests/test_clients/test_dirspool.py rename to tests/test_core/test_directory_spool.py index 2336acbfb..f6736cc12 100644 --- a/tests/test_clients/test_dirspool.py +++ b/tests/test_core/test_directory_spool.py @@ -1,9 +1,8 @@ -"""Tests for FileSpool.""" +"""Tests for directory-backed spools.""" from __future__ import annotations from pathlib import Path -from unittest.mock import patch as upatch import numpy as np import pandas as pd @@ -11,11 +10,9 @@ import dascore as dc import dascore.examples -from dascore.clients.dirspool import DirectorySpool from dascore.constants import ONE_SECOND +from dascore.core.spool import Spool from dascore.exceptions import MissingPatchError, ParameterError -from dascore.io.core import PatchFileSummary -from dascore.utils.hdf5 import HDFPatchIndexManager from dascore.utils.misc import register_func, suppress_warnings DIRECTORY_SPOOLS = [] @@ -43,7 +40,7 @@ def dir_spool_index_out_of_order(random_spool, tmp_path_factory): @register_func(DIRECTORY_SPOOLS) def one_directory_spool(one_file_dir): """Create a directory with a single DAS file.""" - spool = DirectorySpool(one_file_dir) + spool = Spool.from_directory(one_file_dir) return spool.update() @@ -101,7 +98,7 @@ class TestDirectorySpoolBasics: def test_isinstance(self, directory_spool): """Simply ensure expected type was returned.""" - assert isinstance(directory_spool, DirectorySpool) + assert isinstance(directory_spool, Spool) def test_selected_str(self, diverse_directory_spool): """Ensure select kwargs show up in str.""" @@ -116,7 +113,7 @@ def test_sorted_multi_patch_uses_source_patch_id(self, tmp_path): patch_2 = dc.get_example_patch() patch_1 = patch_2.update_coords(time=patch_2.coords.get_array("time") + 10) dc.write(dc.spool([patch_1, patch_2]), path / "multi_patch.h5", "dasdae") - spool = DirectorySpool(path).update().sort("time") + spool = Spool.from_directory(path).update().sort("time") patch = spool[0] assert patch.get_coord("time").min() == patch_2.get_coord("time").min() @@ -147,66 +144,131 @@ def test_merge(self, multi_patch_file_spool): class TestLoadPatchFastPath: - """Tests for the direct FiberIO read path used by _load_patch.""" + """FileResolver reads each row's patch through a single dc.read call.""" - def test_requires_concrete_format_and_version(self, one_directory_spool): - """ - Without a concrete format and version the fast path must defer to - dc.read, which detects them from the file; get_fiberio with a None - version would return the newest reader, not the file's version. - """ - spool = one_directory_spool - assert spool._read_patches({"file_format": "", "file_version": ""}) is None - assert spool._read_patches({"file_format": "DASDAE"}) is None - assert spool._read_patches({"file_version": "1"}) is None - - def test_unusual_fiberio_spool_defers_to_generic_read( + def test_forwards_recorded_format_and_version( self, one_directory_spool, monkeypatch ): - """Fast path should defer if the reader returns a non-memory spool.""" - - class _Reader: - def read(self, *args, **kwargs): - return () - - monkeypatch.setattr( - dc.io.FiberIO.manager, - "get_fiberio", - lambda format, version: _Reader(), - ) - kwargs = { - "path": one_directory_spool.get_contents()["path"].iloc[0], + """The recorded format/version are forwarded so dc.read skips probing.""" + resolver = one_directory_spool._catalog.resolver + calls = [] + + def _fake_read(**kwargs): + calls.append(kwargs) + return object() + + monkeypatch.setattr("dascore.io.index.catalog.dc.read", _fake_read) + resolver._read("path", {"file_format": "DASDAE", "file_version": "1"}, {}, "") + assert calls[-1]["file_format"] == "DASDAE" + assert calls[-1]["file_version"] == "1" + # empty format/version are simply omitted (dc.read detects them) + resolver._read("path", {"file_format": ""}, {}, "") + assert "file_format" not in calls[-1] + + def test_reads_file_once(self, one_directory_spool, monkeypatch): + """A row's patch is read exactly once regardless of the reader's return.""" + calls = [] + + def _fake_read(**kwargs): + calls.append(kwargs) + return () # an unusual (empty) reader return + + monkeypatch.setattr("dascore.io.index.catalog.dc.read", _fake_read) + row = {"file_format": "DASDAE", "file_version": "1"} + resolver = one_directory_spool._catalog.resolver + resolver._read("path", row, {}, "") + assert len(calls) == 1 + + def test_multi_patch_resolves_identity_with_single_read( + self, one_directory_spool, random_patch, monkeypatch + ): + """Multi-patch reads resolve source identity from one dc.read call.""" + patch_1 = random_patch.update_attrs(_source_patch_id="first") + patch_2 = random_patch.update_attrs(_source_patch_id="second") + reads = [] + + def _fake_read(**kwargs): + reads.append(kwargs) + return dc.spool([patch_1, patch_2]) + + monkeypatch.setattr("dascore.io.index.catalog.dc.read", _fake_read) + row = { + "path": "path", "file_format": "DASDAE", "file_version": "1", + "source_patch_id": "second", } - assert one_directory_spool._read_patches(kwargs) is None + resolver = one_directory_spool._catalog.resolver + patch = resolver.resolve(row) + assert patch.attrs["_source_patch_id"] == "second" + assert len(reads) == 1 # the file is read exactly once - def test_multi_patch_selection_defers_to_generic_read( + def test_positional_id_reads_whole_source( self, one_directory_spool, random_patch, monkeypatch ): - """Fast path should not choose the first patch from multi-patch reads.""" + """Positional ids must ignore trim hints; a trimmed read would shift them.""" + patch_2 = random_patch.update_attrs(tag="second") - class _Reader: - def read(self, *args, **kwargs): - patch_2 = random_patch.update_attrs(tag="second") - return dc.spool([random_patch, patch_2]) + def _fake_read(**kwargs): + assert "time" not in kwargs, "positional ids must read untrimmed" + return dc.spool([random_patch, patch_2]) - monkeypatch.setattr( - dc.io.FiberIO.manager, - "get_fiberio", - lambda format, version: _Reader(), - ) - path = one_directory_spool.get_contents()["path"].iloc[0] - monkeypatch.setattr(one_directory_spool, "_select_kwargs", {"tag": "second"}) - kwargs = { - "path": path, + monkeypatch.setattr("dascore.io.index.catalog.dc.read", _fake_read) + row = { + "path": "path", "file_format": "DASDAE", "file_version": "1", - "_modified": True, - "tag": "second", + "source_patch_id": "1", } + resolver = one_directory_spool._catalog.resolver + patch = resolver.resolve(row, time=(None, None)) + assert patch.attrs["tag"] == "second" + + +class TestSelectedDirectorySpools: + """Selection on directory spools (select_kwargs constructor removed).""" + + @pytest.fixture(scope="class") + def spool_dir(self, random_spool, tmp_path_factory): + """A directory holding the random spool, one file per patch.""" + path = tmp_path_factory.mktemp("select_kwargs_dir") + for num, patch in enumerate(random_spool): + patch.io.write(path / f"patch_{num}.h5", "dasdae") + return path + + @pytest.fixture(scope="class") + def first_patch_range(self, random_spool): + """The time range of the chronologically first patch.""" + patch = sorted(random_spool, key=lambda x: x.get_coord("time").min())[0] + time = patch.get_coord("time") + return (time.min(), time.max()) + + def test_contents_restricted(self, spool_dir, random_spool, first_patch_range): + """Rows outside the requested range must not appear (regression).""" + spool = Spool.from_directory(spool_dir).update().select(time=first_patch_range) + assert 1 <= len(spool) < len(random_spool) + contents = spool.get_contents() + assert (contents["time_min"] <= first_patch_range[1]).all() + assert (contents["time_max"] >= first_patch_range[0]).all() + for patch in spool: + time = patch.get_coord("time") + assert time.min() >= first_patch_range[0] + assert time.max() <= first_patch_range[1] + + def test_selected_spool_refuses_update( + self, spool_dir, random_spool, first_patch_range + ): + """D1: any operation severs update().""" + from dascore.exceptions import InvalidSpoolError + + spool = Spool.from_directory(spool_dir).update().select(time=first_patch_range) + with pytest.raises(InvalidSpoolError, match="root spool"): + spool.update() - assert one_directory_spool._read_patches(kwargs) is None + def test_select_kwargs_parameter_removed(self, spool_dir): + """The constructor no longer accepts select_kwargs.""" + with pytest.raises(TypeError, match="select_kwargs"): + Spool.from_directory(spool_dir, select_kwargs={"tag": "x"}) class TestDirectoryIndex: @@ -222,19 +284,36 @@ def test_index_exists(self, basic_file_spool): """An index should be returned.""" assert basic_file_spool.indexer.index_path.exists() - def test_index_len(self, basic_index_df, two_patch_directory): - """An index should be returned.""" - spool = dc.spool(two_patch_directory) + def test_index_len(self, random_patch, tmp_path): + """Deleting and rebuilding the index reproduces the contents.""" + # own directory so no other spool holds the index file open + dc.write(random_patch, tmp_path / "a.hdf5", "dasdae") + dc.write(random_patch.update_attrs(tag="b"), tmp_path / "b.hdf5", "dasdae") + spool = dc.spool(tmp_path) + spool.get_contents() # build the index + # close the connection so the index file can be replaced (Windows + # cannot delete a file with an open handle), then rebuild fresh. + spool.indexer.close() spool.indexer.index_path.unlink() - df = spool.update().get_contents() - bank_paths = list(Path(two_patch_directory).rglob("*hdf5")) + rebuilt = dc.spool(tmp_path).update() + df = rebuilt.get_contents() + rebuilt.indexer.close() + bank_paths = list(Path(tmp_path).rglob("*hdf5")) assert isinstance(df, pd.DataFrame) assert len(bank_paths) == len(df) def test_index_columns(self, basic_index_df): """Ensure expected columns show up in the index.""" - schema_fields = list(PatchFileSummary.model_fields) - assert set(basic_index_df).issuperset(schema_fields) + expected = { + "path", + "file_format", + "file_version", + "dims", + "time_min", + "time_max", + "time_step", + } + assert set(basic_index_df).issuperset(expected) def test_patches_extracted(self, basic_file_spool): """Ensure the patches can be extracted.""" @@ -483,10 +562,10 @@ class TestGetContents: """Tests for getting the contents of the spool.""" def test_str_columns_in_dataframe(self, diverse_directory_spool): - """Ensure all the string columns are in index.""" + """Ensure the conventional string columns are in the index.""" df = diverse_directory_spool.get_contents() - expected = HDFPatchIndexManager._min_itemsize - assert set(df.columns).issuperset(set(expected)) + expected = {"path", "file_format", "file_version", "dims", "station"} + assert set(df.columns).issuperset(expected) class TestIndexing: @@ -535,7 +614,7 @@ def test_sorted_chunked_selected_spool_can_load_patches( assert all(isinstance(patch, dc.Patch) for patch in chunked) -class TestFileSpoolIntegrations: +class TestFileBackedSpoolIntegrations: """Small integration tests for the file spool.""" @pytest.fixture(scope="class") @@ -611,13 +690,15 @@ def test_select_non_distance(self, non_distance_dir_spool): assert coord.max() <= depth_tup[1] def test_differing_distances(self, dist_differ_spool): - """Ensure iteration still works with conditions described in #583.""" + """Iteration works cleanly under the conditions described in #583. + + The generic index stores per-file distance ranges, so the select + already excluded the short-distance file; no patch needs the + historic skip-with-warning workaround. + """ assert len(dist_differ_spool) - # #583 would raise on iterating. Verify that a warning is issued when - # a patch is skipped due to coordinate mismatch. - with pytest.warns(UserWarning, match="Skipping patch at index.*#583"): - for patch in dist_differ_spool: - assert isinstance(patch, dc.Patch) + for patch in dist_differ_spool: + assert isinstance(patch, dc.Patch) def test_missing_patch_error_catchable_as_index_error(self, dist_differ_spool): """ @@ -632,26 +713,42 @@ def test_missing_patch_error_catchable_as_index_error(self, dist_differ_spool): except IndexError: pass - @pytest.mark.xfail() def test_selected_out_distance_shortens_spool(self, dist_differ_spool): - """Selecting outside of distance range should reduce spool length.""" - # Need to implement new indexing before this will pass. - with suppress_warnings(UserWarning): - assert len(dist_differ_spool) == 1 + """Selecting outside of distance range reduces spool length (#583).""" + assert len(dist_differ_spool) == 1 - def test_iteration_unexpected_index_error(self, basic_file_spool): - """ - Ensure unexpected IndexErrors (not #583) are re-raised during iteration. - """ - # TODO this can be deleted once the new indexing is implemented. - # Mock _get_patches_from_index to raise an IndexError with unexpected - # message - with upatch.object( - basic_file_spool, - "_get_patches_from_index", - side_effect=IndexError("unexpected error from pandas"), - ): - # The iteration should re-raise the unexpected IndexError - with pytest.raises(IndexError, match="unexpected error from pandas"): - for _ in basic_file_spool: - pass + +def _patch_shape(patch): + """Module-level helper (process pools need picklable functions).""" + return patch.shape + + +class TestDirectorySpoolSerialization: + """Directory spools must pickle (process-backed map depends on it).""" + + def test_pickle_round_trip(self, basic_file_spool): + """A directory spool pickles and reopens its own connection.""" + import pickle + + loaded = pickle.loads(pickle.dumps(basic_file_spool)) + assert len(loaded) == len(basic_file_spool) + assert loaded[0].shape == basic_file_spool[0].shape + + def test_pickle_selected_view(self, basic_file_spool): + """Selected views keep their selection through pickling.""" + import pickle + + df = basic_file_spool.get_contents() + sub = basic_file_spool.select(time=(df["time_min"].min(), None)) + loaded = pickle.loads(pickle.dumps(sub)) + assert len(loaded) == len(sub) + + def test_process_pool_map(self, basic_file_spool): + """Spool.map works with a process pool executor.""" + from concurrent.futures import ProcessPoolExecutor + + with ProcessPoolExecutor(max_workers=1) as client: + out = list( + basic_file_spool.map(_patch_shape, client=client, progress=False) + ) + assert len(out) == len(basic_file_spool) diff --git a/tests/test_clients/test_filespool.py b/tests/test_core/test_file_spool.py similarity index 80% rename from tests/test_clients/test_filespool.py rename to tests/test_core/test_file_spool.py index 046c4ba6e..6aa3d836c 100644 --- a/tests/test_clients/test_filespool.py +++ b/tests/test_core/test_file_spool.py @@ -6,7 +6,7 @@ from upath import UPath import dascore as dc -from dascore.clients.filespool import FileSpool +from dascore.core.spool import Spool from dascore.exceptions import PatchAttributeError @@ -15,7 +15,7 @@ class TestBasic: def test_type(self, terra15_file_spool, terra15_v5_path): """Ensure a file spool was returned.""" - assert isinstance(terra15_file_spool, FileSpool) + assert isinstance(terra15_file_spool, Spool) assert len(terra15_file_spool) == len(dc.scan_to_df(terra15_v5_path)) def test_get_patch(self, terra15_file_spool): @@ -24,14 +24,14 @@ def test_get_patch(self, terra15_file_spool): assert isinstance(patch, dc.Patch) def test_init_from_filespool(self, terra15_file_spool): - """Ensure FileSpool can init from FileSPool.""" - new = FileSpool(terra15_file_spool) - assert isinstance(new, FileSpool) + """Ensure a spool can copy-construct from another spool.""" + new = Spool(terra15_file_spool) + assert isinstance(new, Spool) def test_str(self, terra15_file_spool): """Ensure file spool works.""" out = str(terra15_file_spool) - assert "FileSpool" in out + assert "Spool" in out def test_update(self, tmp_path_factory, random_patch): """Update should preserve contents even when a format index hook is a no-op.""" @@ -46,16 +46,16 @@ def test_update(self, tmp_path_factory, random_patch): def test_raises_bad_file(self): """Simply ensures a bad file will raise.""" with pytest.raises(FileNotFoundError, match="does not exist"): - FileSpool("/not/a/directory") + Spool.from_file("/not/a/directory") def test_local_upath_file(self, terra15_v5_path): - """Ensure FileSpool accepts local UPath inputs.""" - spool = FileSpool(UPath(terra15_v5_path)) - assert isinstance(spool, FileSpool) + """Ensure from_file accepts local UPath inputs.""" + spool = Spool.from_file(UPath(terra15_v5_path)) + assert isinstance(spool, Spool) assert len(spool) def test_chunk(self, terra15_file_spool): - """Ensure chunking along time axis works with FileSpool.""" + """Ensure chunking along time axis works on a file spool.""" spool = terra15_file_spool time_coord = spool[0].get_coord("time") duration = time_coord.max() - time_coord.min() @@ -71,7 +71,7 @@ def test_sorted_multi_patch_uses_source_patch_id(self, tmp_path): patch_2 = dc.get_example_patch() patch_1 = patch_2.update_coords(time=patch_2.coords.get_array("time") + 10) dc.write(dc.spool([patch_1, patch_2]), path, "dasdae", file_version="1") - spool = FileSpool(path).sort("time") + spool = Spool.from_file(path).sort("time") loaded_patch = spool[0] assert loaded_patch.get_coord("time").min() == patch_2.get_coord("time").min() @@ -80,7 +80,7 @@ def test_multi_patch_without_source_patch_id_raises(self, tmp_path): path = tmp_path / "multi_patch.h5" spool = dc.examples.get_example_spool("random_das", length=2) dc.write(spool, path, "dasdae", file_version="1") - file_spool = FileSpool(path) + file_spool = Spool.from_file(path) kwargs = {"path": str(path), "file_format": "DASDAE", "file_version": "1"} with pytest.raises(PatchAttributeError, match="uniquely resolved"): - file_spool._load_patch(kwargs) + file_spool._catalog.resolver.resolve(kwargs) diff --git a/tests/test_core/test_patch.py b/tests/test_core/test_patch.py index 00a2ec286..4e4e5a012 100644 --- a/tests/test_core/test_patch.py +++ b/tests/test_core/test_patch.py @@ -486,6 +486,16 @@ def test_flat_dump_dim_tuple_and_exclude(self): assert "distance_min" not in out assert np.isnan(out["time_step"]) + def test_flat_dump_null_datetime_step(self, random_patch): + """A datetime coord without a clean step flat-dumps a NaT sentinel.""" + time = random_patch.coords.get_array("time").copy() + time[-1] += np.timedelta64(1, "s") # break uniformity + patch = random_patch.update_coords(time=time) + out = patch.summary.flat_dump() + step = out["time_step"] + assert isinstance(step, np.timedelta64) + assert pd.isnull(step) + def test_select_from_spool_by_integer_source_patch_id(self, random_patch): """Integer-like source ids should fall back to positional selection.""" spool = dc.spool( diff --git a/tests/test_core/test_patch_chunk.py b/tests/test_core/test_patch_chunk.py index 869f3fc45..278e40f04 100644 --- a/tests/test_core/test_patch_chunk.py +++ b/tests/test_core/test_patch_chunk.py @@ -81,8 +81,11 @@ def test_patches_match_df_contents(self, random_spool): new_content = new_spool.get_contents() # these should be (nearly) identical. common = set(chunk_df.columns) & set(new_content.columns) - # len fields may differ by ±1 between summary-based and data-based counts - skip = {"history"} | {c for c in common if c.endswith("_len")} + # len fields may differ by ±1 between summary-based and data-based + # counts; identity/provenance columns legitimately differ between + # plan rows and re-scanned live patches + skip = {"history", "path", "file_format", "file_version", "source_patch_id"} + skip |= {c for c in common if c.endswith("_len")} cols = sorted(common - skip) comp1, comp2 = chunk_df[cols], new_content[cols] equal_cols = (comp1 == comp2) | (pd.isnull(comp1) & pd.isnull(comp2)) @@ -170,10 +173,13 @@ def test_too_big_partial(self, diverse_spool): assert spool1 == spool2 def test_too_big_overlap_raises(self, diverse_spool): - """Overlap > chunk an error should raise.""" - msg = "overlap is greater than chunk size" + """Overlap >= chunk size should raise a clear error.""" + msg = "overlap is greater than or equal to chunk size" with pytest.raises(ParameterError, match=msg): diverse_spool.chunk(time=10, overlap=11) + # Equal overlap would mean zero-stride segments; also rejected. + with pytest.raises(ParameterError, match=msg): + diverse_spool.chunk(time=10, overlap=10) def test_issue_474(self, random_spool): """Ensure spools can be chunked with the duration reported by coord.""" @@ -304,9 +310,20 @@ def patches_conflicting_private_coord(self, random_patch): return p1, p2 def test_merge_unequal_other(self, distance_adjacent): - """When distance values are not equal time shouldn't be merge-able.""" - with pytest.raises(CoordMergeError): - distance_adjacent.chunk(time=...) + """Unequal distance coords partition rather than raise (0.2 change). + + Patches whose non-chunked dimension coordinates differ are never + combined; they simply land in separate output patches. + """ + out = distance_adjacent.chunk(time=...) + assert len(out) == len(distance_adjacent) + + # the differing distance envelopes are preserved, not merged/duplicated + def _distance_envelopes(spool): + df = spool.get_contents() + return sorted(zip(df["distance_min"], df["distance_max"], strict=True)) + + assert _distance_envelopes(out) == _distance_envelopes(distance_adjacent) def test_merge_adjacent(self, adjacent_spool_no_overlap): """Test simple merge of patches.""" @@ -401,6 +418,34 @@ def test_merge_distance(self, distance_adjacent): assert old_df["distance_min"].min() == new_df["distance_min"].min() assert old_df["distance_max"].max() == new_df["distance_max"].max() + def test_non_si_merge_tolerance_uses_coord_units(self, random_patch): + """Canonical index steps are not interpreted in native coord units.""" + from dascore.utils.patch import _get_merged_coord + + size = len(random_patch.get_coord("distance")) + first = dc.get_coord(data=np.arange(size, dtype=float), units="km") + second = dc.get_coord( + data=np.arange(size, dtype=float) + size + 8, + units="km", + ) + patches = [ + random_patch.update_coords(distance=first), + random_patch.update_coords(distance=second), + ] + # Index summaries store numeric dimension steps in canonical SI. + summaries = pd.DataFrame({"distance_step": [1000.0, 1000.0]}) + manager = _get_merged_coord( + summaries, + "distance", + [patch.coords for patch in patches], + tolerance=1.5, + ) + merged = manager.coord_map["distance"] + assert not merged.evenly_sampled + assert np.array_equal( + merged.values, np.concatenate([first.values, second.values]) + ) + def test_merge_distance_no_order(self, distance_adjacent_no_order): """Ensure distance can be merged with unsorted coords.""" sp = distance_adjacent_no_order.chunk(distance=...) @@ -499,8 +544,11 @@ def test_chunk_patches_with_non_coord(self, random_patch): """Tests for chunking when some patches have non coordinate dimensions.""" patches = [random_patch.mean("time") for _ in range(3)] spool = dc.spool(patches) - chunked = spool.chunk(time=None) - # Since the time dims are NaN, this can't work. + # Losing patches silently would be data loss; this raises by default + # (0.2 change) with missing_dim="drop" restoring the old behavior. + with pytest.raises(ChunkError, match="missing_dim"): + spool.chunk(time=None) + chunked = spool.chunk(time=None, missing_dim="drop") assert not len(chunked) def test_merge_with_conflicting_private_coords( @@ -642,6 +690,13 @@ def test_chunk_non_adjacent_within_tolerance_warns(self, random_patch): assert len(out) == 1 +def _bare_assembler(): + """An assembler with no frames, for direct streaming-merge tests.""" + from dascore.utils.patch_assembly import PatchAssembler + + return PatchAssembler(load_patch=None, merge_kwargs={}) + + class TestStreamingMerge: """ Tests for the streaming merge path, which copies each patch into a @@ -650,28 +705,28 @@ class TestStreamingMerge: def test_streaming_path_used(self, adjacent_spool_no_overlap, monkeypatch): """Ensure simple merges take the streaming path.""" - from dascore.core.spool import DataFrameSpool + from dascore.utils.patch_assembly import PatchAssembler called = [] - original = DataFrameSpool._merge_patches_streaming + original = PatchAssembler._merge_patches_streaming def wrapper(self, *args, **kwargs): called.append(True) return original(self, *args, **kwargs) - monkeypatch.setattr(DataFrameSpool, "_merge_patches_streaming", wrapper) + monkeypatch.setattr(PatchAssembler, "_merge_patches_streaming", wrapper) merged = adjacent_spool_no_overlap.chunk(time=None) assert isinstance(merged[0], dc.Patch) assert called def test_matches_materialized_merge(self, adjacent_spool_no_overlap, monkeypatch): """Streaming and concatenating merges must produce identical patches.""" - import dascore.core.spool as spool_module + import dascore.utils.patch_assembly as assembly_module streamed = adjacent_spool_no_overlap.chunk(time=None)[0] # Disabling the sample estimate forces the materialized path. monkeypatch.setattr( - spool_module, "_estimate_merge_samples", lambda df, dim: None + assembly_module, "_estimate_merge_samples", lambda df, dim: None ) materialized = adjacent_spool_no_overlap.chunk(time=None)[0] assert np.array_equal(streamed.data, materialized.data) @@ -692,24 +747,24 @@ def test_mixed_dtype_upcast(self, random_patch, small_first): def test_transposes_patch_to_first_patch_dims(self, random_patch, monkeypatch): """Streaming merge should tolerate patches with the same dims reordered.""" - spool = dc.spool([]) + assembler = _bare_assembler() p1 = random_patch.update_attrs(history=[]) time = p1.get_coord("time") p2 = p1.update_coords(time_min=time.max() + time.step).update_attrs(history=[]) p2 = p2.transpose(*reversed(p2.dims)) patches = iter([p1, p2]) monkeypatch.setattr( - spool, "_load_trimmed_patch", lambda patch_kwargs, joined: next(patches) + assembler, "_load_trimmed_patch", lambda patch_kwargs, joined: next(patches) ) time_axis = p1.get_axis("time") samples = p1.data.shape[time_axis] * 2 - out = spool._merge_patches_streaming(None, [{}, {}], "time", samples) + out = assembler._merge_patches_streaming(None, [{}, {}], "time", samples) assert out.dims == p1.dims assert out.data.shape[time_axis] == samples def test_incompatible_shapes_raise_merge_error(self, random_patch, monkeypatch): """Streaming merge should wrap non-merge-dimension shape mismatches.""" - spool = dc.spool([]) + assembler = _bare_assembler() p1 = random_patch.update_attrs(history=[]) time = p1.get_coord("time") p2 = p1.update_coords(time_min=time.max() + time.step).update_attrs(history=[]) @@ -717,16 +772,16 @@ def test_incompatible_shapes_raise_merge_error(self, random_patch, monkeypatch): p2 = p2.select(distance=(None, distance.max() - distance.step)) patches = iter([p1, p2]) monkeypatch.setattr( - spool, "_load_trimmed_patch", lambda patch_kwargs, joined: next(patches) + assembler, "_load_trimmed_patch", lambda patch_kwargs, joined: next(patches) ) msg = "their shapes are incompatible" with pytest.raises(CoordMergeError, match=msg): samples = p1.data.shape[p1.get_axis("time")] * 2 - spool._merge_patches_streaming(None, [{}, {}], "time", samples) + assembler._merge_patches_streaming(None, [{}, {}], "time", samples) def test_unexpected_merge_dimension_raises(self, random_patch, monkeypatch): """Streaming merge should validate the actual varying dimension.""" - spool = dc.spool([]) + assembler = _bare_assembler() p1 = random_patch.update_attrs(history=[]) dist = p1.get_coord("distance") p2 = p1.update_coords(distance_min=dist.max() + dist.step).update_attrs( @@ -734,9 +789,143 @@ def test_unexpected_merge_dimension_raises(self, random_patch, monkeypatch): ) patches = iter([p1, p2]) monkeypatch.setattr( - spool, "_load_trimmed_patch", lambda patch_kwargs, joined: next(patches) + assembler, "_load_trimmed_patch", lambda patch_kwargs, joined: next(patches) ) msg = "expected them to vary along time" with pytest.raises(CoordMergeError, match=msg): samples = p1.data.shape[p1.get_axis("time")] * 2 - spool._merge_patches_streaming(None, [{}, {}], "time", samples) + assembler._merge_patches_streaming(None, [{}, {}], "time", samples) + + +class TestDescendingChunk: + """Public chunk behavior for descending coordinates (2026-07-18 F5).""" + + def test_contiguous_descending_patches_merge(self): + """Two contiguous descending patches chunk into one patch.""" + p = dc.get_example_patch() + flipped = p.flip("time") + t = p.get_coord("time") + span = t.max() - t.min() + t.step + shifted = flipped.update_coords(time=flipped.get_coord("time").data + span) + merged = dc.spool([shifted, flipped]).chunk(time=None, conflict="drop") + assert len(merged) == 1 + patch = merged[0] + time = patch.get_coord("time") + assert time.reverse_sorted + n_time = p.shape[p.get_axis("time")] + assert patch.shape[patch.get_axis("time")] == 2 * n_time + assert time.min() == t.min() + + +class TestMixedUnitChunk: + """Chunk partitioning and merging across unit differences.""" + + @staticmethod + def _shifted(patch, units=None): + """The example patch shifted to be distance-contiguous, in units.""" + d = patch.get_coord("distance") + span = d.max() - d.min() + d.step + values = d.data + span + if units == "ft": + values = values / 0.3048 + out = patch.update_coords(distance=values) + return out.set_units(distance=units) if units else out + + def test_incompatible_dimensionality_splits(self): + """Metre and second patches with contiguous SI magnitudes stay apart.""" + p = dc.get_example_patch() + pm = p.set_units(distance="m") + ps = self._shifted(p, "s") + sp = dc.spool([pm, ps]) + plan = sp.chunk_plan(distance=None) + assert len(plan.outputs) == 2 + out = sp.chunk(distance=None, conflict="drop") + assert {str(x.get_coord("distance").units) for x in out} == {"1 m", "1 s"} + + def test_unitless_and_unitful_split(self): + """A unitless patch never merges with a unitful one.""" + p = dc.get_example_patch() + sp = dc.spool([p.set_units(distance="m"), self._shifted(p)]) + assert len(sp.chunk(distance=None, conflict="drop")) == 2 + + def test_compatible_units_convert_and_merge(self): + """Metres and feet (one dimensionality) merge, converted, unit-true.""" + p = dc.get_example_patch() + pm = p.set_units(distance="m") + pf = self._shifted(p, "ft") + out = dc.spool([pm, pf]).chunk(distance=None, conflict="drop") + assert len(out) == 1 + patch = out[0] + coord = patch.get_coord("distance") + assert str(coord.units) == "1 m" + n = p.shape[p.get_axis("distance")] + assert patch.shape[patch.get_axis("distance")] == 2 * n + assert float(coord.max()) == pytest.approx(2 * n - 1) + + def test_same_units_unchanged(self): + """The ordinary same-unit merge keeps its behavior and units.""" + p = dc.get_example_patch() + sp = dc.spool([p.set_units(distance="m"), self._shifted(p, "m")]) + out = sp.chunk(distance=None, conflict="drop") + assert len(out) == 1 + assert str(out[0].get_coord("distance").units) == "1 m" + + +class TestChainedChunk: + """Chunking a derived spool along another dimension (round-4 F1).""" + + def test_other_dim_keeps_prior_boundaries(self): + """Re-chunking distance must not undo a time concatenation.""" + p1 = dc.get_example_patch() + t = p1.get_coord("time") + p2 = p1.update_coords(time_min=t.max() + t.step) + merged = dc.spool([p1, p2]).chunk(time=None, conflict="drop") + current = merged[0] + d = current.get_coord("distance") + size = (d.max() - d.min()) / 2 + actual = merged.chunk(distance=size, keep_partial=True, conflict="drop") + expected = dc.spool([current]).chunk( + distance=size, keep_partial=True, conflict="drop" + ) + assert sorted(x.shape for x in actual) == sorted(x.shape for x in expected) + got = { + (str(x.get_coord("time").min()), str(x.get_coord("time").max())) + for x in actual + } + want = { + (str(x.get_coord("time").min()), str(x.get_coord("time").max())) + for x in expected + } + assert got == want + + def test_segment_then_segment(self): + """chunk(time=...) then chunk(distance=...) partitions both dims.""" + p = dc.get_example_patch() # (300, 2000), 8 s + out = dc.spool([p]).chunk(time=2).chunk(distance=100) + assert len(out) == 12 + assert {x.shape for x in out} == {(100, 500)} + + def test_same_dim_rechunk_still_collapses(self): + """Re-chunking the same dim re-plans from members (no nesting).""" + p1 = dc.get_example_patch() + t = p1.get_coord("time") + p2 = p1.update_coords(time_min=t.max() + t.step) + merged = dc.spool([p1, p2]).chunk(time=None, conflict="drop") + rechunk = merged.chunk(time=2) + assert len(rechunk) == 8 + assert {x.shape for x in rechunk} == {(300, 500)} + + +class TestMatchMergeUnits: + """The member unit normalizer's defensive paths.""" + + def test_incompatible_units_pass_through(self): + """Dimensionality mismatches pass through for the merge to police.""" + from dascore.units import get_quantity + from dascore.utils.patch_assembly import _match_merge_units + + patch = dc.get_example_patch().set_units(distance="m") + target = get_quantity("s").units + out, kept = _match_merge_units(patch, "distance", target) + assert out is patch # unconverted + assert kept == target diff --git a/tests/test_core/test_spool.py b/tests/test_core/test_spool.py index e75a0c8fc..4c2c4e57b 100644 --- a/tests/test_core/test_spool.py +++ b/tests/test_core/test_spool.py @@ -11,19 +11,15 @@ import pytest import dascore as dc -from dascore.clients.filespool import FileSpool -from dascore.core.spool import ( - BaseSpool, - MemorySpool, - _estimate_merge_samples, - _get_varying_dim, -) +from dascore.core.spool import BaseSpool, Spool from dascore.exceptions import ( InvalidSpoolError, MissingOptionalDependencyError, + MissingPatchError, ParameterError, ) from dascore.utils.downloader import fetch +from dascore.utils.patch_assembly import _estimate_merge_samples, _get_varying_dim from dascore.utils.time import to_datetime64, to_timedelta64 @@ -82,6 +78,47 @@ def test_base_concat_raises(self, random_spool): with pytest.raises(NotImplementedError, match=msg): BaseSpool.concatenate(random_spool, time=2) + def test_base_update_returns_self(self, random_spool): + """The BaseSpool update default (for third-party spools) no-ops.""" + assert BaseSpool.update(random_spool) is random_spool + + def test_invalid_input_raises(self): + """A non-patch, non-spool input raises a clear error.""" + with pytest.raises(InvalidSpoolError, match="accepts a Patch"): + Spool(42) + + def test_wraps_third_party_spool(self, random_spool): + """A non-Spool BaseSpool input realizes into the registry.""" + + class MiniSpool(BaseSpool): + """Minimal third-party spool over a patch list.""" + + def __init__(self, patches): + self._patches = list(patches) + + def __getitem__(self, item): + return self._patches[item] + + def __iter__(self): + return iter(self._patches) + + def __len__(self): + return len(self._patches) + + def chunk(self, **kwargs): + raise NotImplementedError + + def select(self, **kwargs): + raise NotImplementedError + + def get_contents(self): + raise NotImplementedError + + patches = list(random_spool) + wrapped = Spool(MiniSpool(patches)) + assert isinstance(wrapped, Spool) + assert list(wrapped) == patches + def test_viz_raises(self, random_spool): """Ensure Spool.viz raises AttributeError.""" msg = "Apply 'viz' on a Patch object" @@ -89,7 +126,7 @@ def test_viz_raises(self, random_spool): random_spool.viz.waterfall(random_spool) -class TestMemorySpoolLazy: +class TestLiveSpoolLazy: """ Tests for lazy behavior of in-memory spools. @@ -109,7 +146,7 @@ def test_simple_access_builds_no_dataframes(self, patch_list): assert spool[0] == patch_list[0] assert spool[-1] == patch_list[-1] assert list(spool) == patch_list - assert "_df" not in spool._cache + assert spool._catalog._backend is None def test_out_of_bounds_raises(self, patch_list): """The fast path must raise the same IndexError as the df path.""" @@ -117,14 +154,14 @@ def test_out_of_bounds_raises(self, patch_list): match = "out of bounds for spool" with pytest.raises(IndexError, match=match): _ = spool[len(patch_list)] - assert "_df" not in spool._cache + assert spool._catalog._backend is None def test_access_unchanged_after_df_built(self, patch_list): """Patch access must return the same thing before/after df built.""" spool = dc.spool(patch_list) lazy_patches = list(spool) - _ = spool.get_contents() # forces the dataframes to build - assert "_df" in spool._cache + _ = spool.get_contents() # forces the flat relation to build + assert spool._catalog._backend is not None assert list(spool) == lazy_patches assert spool[0] == lazy_patches[0] @@ -141,8 +178,7 @@ def test_input_mutation_does_not_change_spool(self, patch_list): spool = dc.spool(data) data.pop() assert len(spool) == len(patch_list) - # The snapshot itself is immutable. - assert isinstance(spool._patches, tuple) + assert list(spool) == patch_list def test_derived_spools_use_df_machinery(self, patch_list): """Chunked/selected spools must go through the instruction dfs.""" @@ -150,33 +186,38 @@ def test_derived_spools_use_df_machinery(self, patch_list): merged = spool.chunk(time=None) assert len(merged) == 1 time_coord = merged[0].get_coord("time") - expected_min = min( - x.summary.get_coord_summary("time").min for x in patch_list - ) + expected_min = min(x.summary.get_coord_summary("time").min for x in patch_list) assert time_coord.min() == expected_min - def test_derived_spool_does_not_retain_parent(self, patch_list): - """Derived spools must not hold a reference to their parent.""" + def test_derived_spool_is_own_catalog(self, patch_list): + """A chunked spool is a fresh derived catalog sharing patches.""" + from dascore.io.index.planned import PlanResolver + spool = dc.spool(patch_list) chunked = spool.chunk(time=1) - assert chunked._data is None - assert chunked._patches is None + assert chunked._catalog is not spool._catalog + assert isinstance(chunked._catalog.resolver, PlanResolver) + # member loading shares the parent's live patches, not copies + registry = chunked._catalog.resolver.live_entries() + assert {id(p) for p in registry.values()} <= {id(p) for p in patch_list} + # no other patch containers exist on the instance + assert "_patches" not in chunked.__dict__ + assert "_data" not in chunked.__dict__ def test_single_patch_input_uses_lazy_storage(self, random_patch): - """A single patch should be stored lazily just like a patch sequence.""" - spool = MemorySpool(random_patch) - assert spool._patches == (random_patch,) + """A single patch lands in the registry without realizing tables.""" + spool = Spool(random_patch) assert len(spool) == 1 + registry = spool._catalog.resolver.live_entries() + assert tuple(registry.values()) == (random_patch,) + # simple access never bootstrapped the index backend + assert spool._catalog._backend is None - def test_empty_memory_spool_has_no_dataframe(self): - """An empty MemorySpool should report no managing dataframe.""" - spool = MemorySpool() - assert spool._get_df() is None - - def test_instruction_df_builds_from_lazy_patches(self, patch_list): - """Lazy patch input should still build instruction dataframes on demand.""" - spool = dc.spool(patch_list) - assert len(spool._get_instruction_df()) == len(patch_list) + def test_empty_memory_spool(self): + """An empty Spool is a valid, iterable, zero-length spool.""" + spool = Spool() + assert len(spool) == 0 + assert list(spool) == [] class TestSpoolHelpers: @@ -226,21 +267,13 @@ def test_eq_self(self, random_spool): """A spool should always eq itself.""" assert random_spool == random_spool - def test_unequal_attr(self, random_spool): - """Simulate some attribute which isn't equal.""" + def test_foreign_attrs_do_not_join_equality(self, random_spool): + """Equality state is enumerated; stray instance attrs are ignored.""" new1 = copy.deepcopy(random_spool) new1.__dict__["bad_attr"] = 1 new2 = copy.deepcopy(random_spool) new2.__dict__["bad_attr"] = 2 - assert new1 != new2 - - def test_unequal_dicts(self, random_spool): - """Simulate some dicts which don't have the same values.""" - new1 = copy.deepcopy(random_spool) - new1.__dict__["bad_attr"] = {1: 2} - new2 = copy.deepcopy(random_spool) - new2.__dict__["bad_attr"] = {2: 3} - assert new1 != new2 + assert new1 == new2 class TestIndexing: @@ -321,8 +354,8 @@ def test_bool_some_true(self, random_spool): bool_array[1] = False out = random_spool[bool_array] assert len(out) == sum(bool_array) - df1 = out.get_contents() - df2 = random_spool.get_contents()[bool_array] + df1 = out.get_contents().reset_index(drop=True) + df2 = random_spool.get_contents()[bool_array].reset_index(drop=True) assert df1.equals(df2) @@ -674,9 +707,9 @@ def test_non_supported_type_raises(self): def test_file_spool(self, random_spool, tmp_path_factory): """ Tests for getting a file spool vs in-memory spool. Basically, - if a format supports scanning a FileSpool is returned. If it doesn't, - all the file contents have to be loaded into memory to scan so a - MemorySpool is just returned. + if a format supports scanning, a lazy file-backed spool is + returned. If it doesn't, all the file contents have to be loaded + into memory, so the spool holds live patches. """ path = tmp_path_factory.mktemp("file_spoolin") dasdae_path = path / "patch.h5" @@ -685,10 +718,10 @@ def test_file_spool(self, random_spool, tmp_path_factory): dc.write(random_spool, pickle_path, "pickle") dasdae_spool = dc.spool(dasdae_path) - assert isinstance(dasdae_spool, FileSpool) + assert not dasdae_spool.has_live_patches pickle_spool = dc.spool(pickle_path) - assert isinstance(pickle_spool, MemorySpool) + assert pickle_spool.has_live_patches class TestSpoolBehaviorOptionalImports: @@ -802,137 +835,170 @@ def test_dft_patch_access(self, random_dft_patch): assert isinstance(patch, dc.Patch) -class TestSpoolEquality: - """Tests for spool equality comparisons to ensure 100% coverage.""" - - def test_spool_equality_non_dict_comparison(self, random_spool): - """Test line 107: non-dict comparison in _vals_equal.""" - spool1 = copy.deepcopy(random_spool) - spool2 = copy.deepcopy(random_spool) - - # Add non-dict values to test the non-dict comparison path - spool1._test_string = "hello" - spool2._test_string = "hello" - - # This should be equal - assert spool1 == spool2 - - # Now make them different to test the comparison - spool2._test_string = "world" - - # This should be False - assert spool1 != spool2 +class TestDeepEqualityCheck: + """Coverage for deep_equality_check branches (formerly via spool attrs).""" - def test_spool_equality_with_objects_having_dict(self, random_spool): - """Test line 127: objects with __dict__ that are not equal.""" + def test_non_dict_comparison(self): + """Plain value comparison inside dicts.""" + from dascore.utils.misc import deep_equality_check - class TestObject: - def __init__(self, value): - self.value = value - - spool1 = copy.deepcopy(random_spool) - spool2 = copy.deepcopy(random_spool) + assert deep_equality_check({"a": "hello"}, {"a": "hello"}) + assert not deep_equality_check({"a": "hello"}, {"a": "world"}) - # Add objects with __dict__ that have different values - spool1._test_obj = TestObject(1) - spool2._test_obj = TestObject(2) # Different data - - # This should hit line 127 and return False - assert spool1 != spool2 - - def test_spool_equality_with_objects_having_dict_equal(self, random_spool): - """Test objects with __dict__ that are equal via recursive comparison.""" + def test_objects_with_dict(self): + """Objects compare via recursive __dict__ comparison.""" + from dascore.utils.misc import deep_equality_check class TestObject: def __init__(self, value): self.value = value - spool1 = random_spool - spool2 = copy.deepcopy(random_spool) - - # Add objects with __dict__ that have same internal state - spool1._test_obj = TestObject(42) - spool2._test_obj = TestObject(42) + assert deep_equality_check({"o": TestObject(42)}, {"o": TestObject(42)}) + assert not deep_equality_check({"o": TestObject(1)}, {"o": TestObject(2)}) - # This should be equal via recursive __dict__ comparison - assert spool1 == spool2 - - def test_spool_equality_mixed_types(self): - """Test equality with various mixed data types.""" - # Create simple spools to avoid cache issues - patch = dc.get_example_patch() - spool1 = dc.spool([patch]) - spool2 = dc.spool([patch]) + def test_mixed_types(self): + """Ints, lists, and numpy arrays compare by value.""" + from dascore.utils.misc import deep_equality_check - # Test with integers (non-dict) - spool1._int_val = 42 - spool2._int_val = 42 - assert spool1 == spool2 + d1 = {"i": 42, "l": [1, 2, 3], "a": np.array([1, 2, 3])} + d2 = {"i": 42, "l": [1, 2, 3], "a": np.array([1, 2, 3])} + assert deep_equality_check(d1, d2) + d2["a"] = np.array([1, 2, 4]) + assert not deep_equality_check(d1, d2) - # Test with lists (non-dict) - spool1._list_val = [1, 2, 3] - spool2._list_val = [1, 2, 3] - assert spool1 == spool2 - - # Test with numpy arrays (non-dict) - spool1._array_val = np.array([1, 2, 3]) - spool2._array_val = np.array([1, 2, 3]) - assert spool1 == spool2 + def test_dataframes(self): + """DataFrames compare via .equals.""" + from dascore.utils.misc import deep_equality_check - # Test arrays with different values - spool1._array_val = np.array([1, 2, 3]) - spool2._array_val = np.array([1, 2, 4]) - assert spool1 != spool2 - - def test_spool_equality_with_dataframes(self): - """Test equality with pandas DataFrames (has equals method).""" - # Create simple spools to avoid cache issues - patch = dc.get_example_patch() - spool1 = dc.spool([patch]) - spool2 = dc.spool([patch]) - - # Add DataFrames that should be equal df1 = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) df2 = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + assert deep_equality_check({"df": df1}, {"df": df2}) + df3 = pd.DataFrame({"a": [1, 2, 4], "b": [4, 5, 6]}) + assert not deep_equality_check({"df": df1}, {"df": df3}) - spool1._test_df = df1 - spool2._test_df = df2 + def test_unequal_sub_dicts(self): + """Nested dicts with different values are unequal.""" + from dascore.utils.misc import deep_equality_check - # Should be equal via df.equals() - assert spool1 == spool2 + assert not deep_equality_check({"d": {1: 2}}, {"d": {2: 3}}) - # Now test with different DataFrames - df3 = pd.DataFrame({"a": [1, 2, 4], "b": [4, 5, 6]}) # Different data - spool2._test_df = df3 - # Should not be equal - assert spool1 != spool2 +class TestSpoolCoverageEdges: + """Cover remaining spool-machinery branches with real operations.""" - def test_specific_coverage_lines(self): - """Test to specifically cover lines 107 and 127.""" - # Create minimal spools + @pytest.fixture(scope="class") + def many_contiguous(self): + """Twelve contiguous patches (for >10-row merge handling).""" + t0 = np.datetime64("2020-01-01", "ns") + patch = dc.get_example_patch(time_min=t0) + step = patch.get_coord("time").step + out = [patch] + for _ in range(11): + nxt = dc.get_example_patch(time_min=out[-1].get_coord("time").max() + step) + out.append(nxt) + return out + + def test_equality_and_repr(self): + """Spool equality strips synthetic identity; repr shows a time span.""" patch = dc.get_example_patch() - spool1 = dc.spool([patch]) - spool2 = dc.spool([patch]) + left, right = dc.spool([patch]), dc.spool([patch]) + left.get_contents() # realize so equality compares built frames + right.get_contents() + assert left == right + assert "Time Span" in left.__rich__().__str__() + + def test_equality_of_empty_spools(self): + """Empty spools (None frames) compare equal via the None-strip path.""" + assert Spool() == Spool() + + def test_repr_without_time_coordinate(self): + """A spool whose patches have no time coord omits the time-span line.""" + data = np.random.default_rng().random((6, 4)) + coords = {"distance": np.arange(6), "frequency": np.arange(4.0)} + patch = dc.Patch(data=data, coords=coords, dims=("distance", "frequency")) + rendered = dc.spool([patch]).__rich__().__str__() + assert "Spool" in rendered + assert "Time Span" not in rendered # no time coordinate to summarize + + def test_repr_with_time_coordinate(self): + """A normal spool renders its time span.""" + assert "Time Span" in dc.spool([dc.get_example_patch()]).__rich__().__str__() + + def test_large_merge_dedups(self, many_contiguous): + """Merging >10 sources into one patch exercises the de-dup branch.""" + merged = dc.spool(many_contiguous).chunk(time=None) + assert len(merged) == 1 + # 12 contiguous patches merge into one continuous coordinate. + assert merged[0].get_coord("time").size == sum( + p.get_coord("time").size for p in many_contiguous + ) - # Line 107: Non-dict comparison - spool1._string_test = "hello" - spool2._string_test = "hello" - assert spool1 == spool2 + def test_union_of_scanless_spool(self, tmp_path): + """A scanless (pickle) spool wraps its read patches in a live + catalog; union shares them like any in-memory member. + """ + dc.get_example_patch().io.write(tmp_path / "a.pkl", "pickle") + pickle_spool = dc.spool(tmp_path / "a.pkl") + combined = pickle_spool + dc.spool([dc.get_example_patch(tag="other")]) + assert len(combined) == 2 + + def test_union_of_chunked_spool(self, many_contiguous): + """A chunked spool is a derived catalog; unions compose it.""" + from dascore.io.index.planned import PlanResolver + + chunked = dc.spool(many_contiguous).chunk(time=None) + assert isinstance(chunked._catalog.resolver, PlanResolver) + combined = chunked + dc.spool([dc.get_example_patch(tag="other")]) + assert len(combined) == 2 + assert all(isinstance(p, dc.Patch) for p in combined) + + def test_iteration_skips_unresolvable_patch(self, monkeypatch): + """A patch that fails to resolve is skipped with a #583 warning.""" + spool = dc.spool([dc.get_example_patch()]) + + def _raise(_ind): + raise MissingPatchError("not available in this session") + + monkeypatch.setattr(spool._catalog, "get_patch", _raise) + with pytest.warns(UserWarning, match="Skipping patch"): + assert list(spool) == [] + + def test_derived_negative_and_bad_index(self): + """Derived-catalog indexing handles negatives, raises out-of-bounds.""" + patches = list(dc.get_example_spool(length=2)) + derived = dc.spool(patches).concatenate(time=1) + assert derived[-1] == derived[len(patches) - 1] + with pytest.raises(IndexError, match="out of bounds"): + _ = derived[len(patches)] + with pytest.raises(IndexError, match="out of bounds"): + _ = derived[-len(patches) - 1] + + def test_merge_buffer_grows_when_estimate_short(self, many_contiguous, monkeypatch): + """An under-estimated merge buffer is grown to fit (uneven sampling).""" + import dascore.utils.patch_assembly as assembly_mod + + # Force the pre-merge sample estimate to be too small so the + # streaming buffer must grow mid-merge. + monkeypatch.setattr(assembly_mod, "_estimate_merge_samples", lambda *a, **k: 1) + merged = dc.spool(many_contiguous).chunk(time=None) + assert merged[0].get_coord("time").size == sum( + p.get_coord("time").size for p in many_contiguous + ) - # Make them different to test line 107 return False - spool2._string_test = "world" - assert spool1 != spool2 + def test_empty_memory_spool_len_iter_repr(self): + """A bare Spool() is a valid empty spool.""" + empty = Spool() + assert len(empty) == 0 + assert list(empty) == [] + assert "Spool" in str(empty) - # Line 127: Objects with __dict__ that are different - class SimpleObj: - def __init__(self, val): - self.val = val - spool1 = dc.spool([patch]) - spool2 = dc.spool([patch]) - spool1._obj = SimpleObj(1) - spool2._obj = SimpleObj(2) +class TestEmptyConcatenate: + """Concatenating an empty spool returns an empty spool (F7).""" - # This should return False at line 127 - assert spool1 != spool2 + @pytest.mark.parametrize("kwargs", [{"time": None}, {"time": 2}, {"new_dim": None}]) + def test_empty_returns_empty(self, kwargs): + """Empty in, empty out — matching chunk's behavior.""" + out = dc.spool([]).concatenate(**kwargs) + assert len(out) == 0 + assert list(out) == [] diff --git a/tests/test_core/test_spool_contracts.py b/tests/test_core/test_spool_contracts.py new file mode 100644 index 000000000..88c407418 --- /dev/null +++ b/tests/test_core/test_spool_contracts.py @@ -0,0 +1,206 @@ +""" +Tests for the unified Spool's equality and update contracts. + +Equality is over rows, never backends: order-sensitive, metadata-level, +with pending residual selections included. update() is case-based: +directory spools sync, file spools rescan, in-memory spools are +trivially current, and combined spools with file rows raise. +""" + +from __future__ import annotations + +import pytest + +import dascore as dc +from dascore.exceptions import InvalidSpoolError + + +@pytest.fixture(scope="module") +def patches(): + """Three contiguous example patches in time order.""" + return list(dc.get_example_spool("random_das")) + + +class TestEqualityContract: + """Equality: rows not backends; order-sensitive; residuals count.""" + + def test_residual_breaks_equality(self, patches): + """A samples residual changes loaded patches, so spools differ.""" + plain = dc.spool(patches) + trimmed = dc.spool(patches).select(distance=(0, 10), samples=True) + # same visible rows (samples never excludes patches) ... + assert len(plain) == len(trimmed) + # ... but not equal spools. + assert plain != trimmed + + def test_equal_selections_compare_equal(self, patches): + """The same selection on equal spools yields equal spools.""" + sel_1 = dc.spool(patches).select(distance=(0, 10)) + sel_2 = dc.spool(patches).select(distance=(0, 10)) + assert sel_1 == sel_2 + + def test_order_matters(self, patches): + """Same patches in a different order are not equal.""" + assert dc.spool(patches) != dc.spool(list(reversed(patches))) + + def test_live_vs_directory_equal(self, patches, tmp_path): + """Identical contents compare equal across backings (rows, not + backends): a live spool equals a directory spool over the same + patches written to disk. + """ + directory = dc.examples.spool_to_directory( + dc.spool(patches), path=tmp_path / "spool_dir" + ) + dir_spool = dc.spool(directory).update(progress=None) + live_spool = dc.spool(patches) + assert live_spool == dir_spool + assert dir_spool == live_spool + + +class TestUpdateContract: + """update() means sync-with-source; the cases differ by source.""" + + def test_live_only_is_noop(self, patches): + """A purely in-memory spool is trivially current.""" + spool = dc.spool(patches) + assert spool.update() is spool + + def test_union_raises(self, patches): + """Combining spools is a computation; the result cannot update.""" + combined = dc.spool(patches[:1]) + dc.spool(patches[1:]) + with pytest.raises(InvalidSpoolError, match="root spool"): + combined.update() + + def test_union_with_file_rows_raises(self, patches, tmp_path): + """A combined spool with file rows cannot update either.""" + directory = dc.examples.spool_to_directory( + dc.spool(patches), path=tmp_path / "dir_a" + ) + combined = dc.spool(directory).update(progress=None) + dc.spool(patches[:1]) + with pytest.raises(InvalidSpoolError, match="root spool"): + combined.update() + + def test_selected_spool_raises(self, patches): + """Any operation severs update: a selected spool refuses it.""" + selected = dc.spool(patches).select(tag="random") + with pytest.raises(InvalidSpoolError, match="root spool"): + selected.update() + + def test_selected_file_spool_raises(self, patches, tmp_path): + """A selected single-file spool refuses update instead of + silently widening back to the whole file (review P1). + """ + path = tmp_path / "sel_file.h5" + dc.write(patches[0], path, "dasdae") + selected = dc.spool(path).select(distance=(0, 10), samples=True) + with pytest.raises(InvalidSpoolError, match="root spool"): + selected.update() + + def test_directory_update_picks_up_new_files(self, patches, tmp_path): + """The syncer case: new files appear after update().""" + directory = tmp_path / "dir_b" + dc.write(patches[0], directory / "a.h5", "dasdae") + spool = dc.spool(directory).update(progress=None) + assert len(spool) == 1 + dc.write(patches[1], directory / "b.h5", "dasdae") + assert len(spool.update(progress=None)) == 2 + + def test_file_spool_update_refreshes(self, patches, tmp_path): + """The single-file case: update rescans the file.""" + path = tmp_path / "single.h5" + dc.write(patches[0], path, "dasdae") + spool = dc.spool(path) + new = spool.update() + assert len(new) == len(spool) + assert new == spool + + +class TestChunkPlanContract: + """Planned views: collapse semantics and derived state.""" + + def test_replan_collapses(self, patches): + """Re-chunking a planned spool re-plans from members (no nesting).""" + spool = dc.spool(patches) + hourly = spool.chunk(time=2) + merged = hourly.chunk(time=None) + assert len(merged) == 1 + time = merged[0].get_coord("time") + mins = [p.get_coord("time").min() for p in patches] + maxs = [p.get_coord("time").max() for p in patches] + assert time.min() == min(mins) + assert time.max() == max(maxs) + + def test_planned_state_is_derived(self, patches): + """A chunked spool is a fresh derived catalog, not a mode flag.""" + from dascore.io.index.planned import PlanResolver + + spool = dc.spool(patches) + chunked = spool.chunk(time=2) + assert chunked._catalog is not spool._catalog + assert isinstance(chunked._catalog.resolver, PlanResolver) + + +class TestTypeSurface: + """The collapsed hierarchy: BaseSpool ABC with one concrete Spool.""" + + def test_every_spool_is_spool(self, patches, tmp_path): + """All construction paths yield the same concrete class.""" + live = dc.spool(patches) + directory = dc.examples.spool_to_directory( + dc.spool(patches), path=tmp_path / "types_dir" + ) + dir_spool = dc.spool(directory).update(progress=None) + file_path = tmp_path / "one.h5" + dc.write(patches[0], file_path, "dasdae") + file_spool = dc.spool(file_path) + for spool in (live, dir_spool, file_spool, live.chunk(time=2)): + assert type(spool) is dc.Spool + assert isinstance(spool, dc.BaseSpool) + + def test_removed_names_gone(self): + """The old concrete class names are deleted outright.""" + import dascore.core.spool as spool_module + + for name in ("MemorySpool", "DirectorySpool", "FileSpool"): + assert not hasattr(spool_module, name) + with pytest.raises(ImportError): + from dascore.clients.dirspool import DirectorySpool # noqa + + def test_live_patch_predicate(self, patches, tmp_path): + """has_live_patches distinguishes memory content, not class.""" + assert dc.spool(patches).has_live_patches + file_path = tmp_path / "pred.h5" + dc.write(patches[0], file_path, "dasdae") + assert not dc.spool(file_path).has_live_patches + mixed = dc.spool(file_path) + dc.spool(patches[:1]) + assert mixed.has_live_patches + + +class TestEqualityOverEffectiveRows: + """Equality compares contents, not representation (2026-07-18).""" + + @pytest.fixture() + def patch(self): + """The example patch.""" + return dc.get_example_patch() + + def test_value_trimmed_view_equals_materialized(self, patch): + """A coordinate-trimmed view equals its union-materialized twin.""" + t = patch.get_coord("time") + sel = dc.spool([patch]).select( + time=(t.min() + 10 * t.step, t.min() + 20 * t.step) + ) + assert sel == sel + dc.spool([]) + + @pytest.mark.parametrize("window", [(0, 10), (-10, None)]) + def test_samples_trimmed_view_equals_materialized(self, patch, window): + """Samples-trimmed views (negative included) equal their twins.""" + sel = dc.spool([patch]).select(time=window, samples=True) + assert sel == sel + dc.spool([]) + + def test_differing_trims_stay_unequal(self, patch): + """Different windows fold to different envelopes and stay unequal.""" + a = dc.spool([patch]).select(time=(0, 10), samples=True) + b = dc.spool([patch]).select(time=(0, 11), samples=True) + assert a != b + assert a != dc.spool([patch]) diff --git a/tests/test_core/test_spool_select_spec.py b/tests/test_core/test_spool_select_spec.py new file mode 100644 index 000000000..90f7c3ac6 --- /dev/null +++ b/tests/test_core/test_spool_select_spec.py @@ -0,0 +1,556 @@ +""" +Selector-spec behavior of Spool.select (all spool types). + +These encode the hard-break semantics adopted for 0.2: unknown names +raise (#435), samples selections are patch-local (#447), relative ranges +work at spool level (#362), and _attrs/_coords disambiguate explicitly. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +import dascore as dc +from dascore.exceptions import InvalidSpoolQueryError + + +@pytest.fixture( + scope="module", + params=("memory", "directory", "memory_derived", "directory_derived"), +) +def spool(request, tmp_path_factory): + """ + The same patches served by each spool type and catalog state. + + The ``*_derived`` params run every spec test over a derived + (plan-backed) catalog via a content-preserving concatenate — the + parity net proving one selector engine serves identity and + restructured spools alike. + """ + base = dc.get_example_spool("random_das") + if request.param.startswith("memory"): + out = dc.spool(list(base)) + else: + path = dc.examples.spool_to_directory( + base, path=tmp_path_factory.mktemp("select_spec") + ) + out = dc.spool(path).update(progress=None) + if request.param.endswith("_derived"): + from dascore.io.index.planned import PlanResolver + + out = out.concatenate(time=1) + assert isinstance(out._catalog.resolver, PlanResolver) + return out + + +class TestUnknownNames: + """Unknown selector names raise eagerly (#435).""" + + def test_unknown_kwarg_raises(self, spool): + """A name that is neither attr nor coord errors.""" + with pytest.raises(InvalidSpoolQueryError, match="neither an attribute"): + spool.select(bad_dimension=(1, 2)) + + def test_unknown_attr_namespace_raises(self, spool): + """_attrs validates against attributes only.""" + with pytest.raises(InvalidSpoolQueryError, match="not an attribute"): + spool.select(_attrs={"time": (None, None)}) + + def test_unknown_coord_namespace_raises(self, spool): + """_coords validates against coordinates only.""" + with pytest.raises(InvalidSpoolQueryError, match="not a coordinate"): + spool.select(_coords={"tag": "random"}) + + def test_double_specification_raises(self, spool): + """A name can't be bare and namespaced at once.""" + with pytest.raises(InvalidSpoolQueryError, match="both"): + spool.select(tag="random", _attrs={"tag": "random"}) + + +class TestNamespaces: + """Explicit namespaces select as their bare equivalents.""" + + def test_attr_namespace(self, spool): + """_attrs behaves like the bare attr kwarg.""" + assert len(spool.select(_attrs={"tag": "random"})) == len(spool) + + def test_coord_namespace(self, spool): + """_coords behaves like the bare coord kwarg.""" + df = spool.get_contents() + t0 = df["time_min"].min() + out = spool.select(_coords={"time": (t0, t0 + np.timedelta64(2, "s"))}) + assert len(out) == 1 + + +class TestCatalogPushdown: + """Public spool selection composes a lazy SQLite query.""" + + def test_coord_predicate_reaches_backend(self, spool, monkeypatch): + """Selection does not query all rows before applying its predicate.""" + catalog = spool._catalog + backend = catalog.backend + calls = [] + original = backend.query + + def wrapped(query=None, **kwargs): + calls.append(query) + return original(query, **kwargs) + + monkeypatch.setattr(backend, "query", wrapped) + selected = spool.select(time=("2020-01-03", "2020-01-04")) + assert calls == [] + # realizing the relation (get_contents) runs the composed query; + # len() alone counts in SQL and never fetches rows. + selected.get_contents() + queries = calls[0] + assert isinstance(queries, list) + assert queries[0].coords["time"] == ("2020-01-03", "2020-01-04") + + +class TestLazySelection: + """Selection construction never realizes the flat relation (review P1). + + Cold spools must compose the selected view without an unfiltered + backend query; realization happens on first content access. The + module-scoped ``spool`` fixture is warm by then, so these tests + build their own fresh spools. + """ + + @pytest.fixture() + def forbid_realization(self, monkeypatch): + """Return a callable that makes flat realization fail loudly.""" + from dascore.io.index.catalog import PatchCatalog + + def _boom(self): + msg = "flat relation realized during selection construction" + raise AssertionError(msg) + + def _arm(): + monkeypatch.setattr(PatchCatalog, "to_df", _boom) + + return _arm + + def test_cold_directory_select(self, tmp_path_factory, forbid_realization): + """A cold directory spool selects without touching the relation.""" + path = dc.examples.spool_to_directory( + dc.get_example_spool("random_das"), + path=tmp_path_factory.mktemp("lazy_select_dir"), + ) + dc.spool(path).update(progress=None) # build the index + fresh = dc.spool(path) + forbid_realization() + selected = fresh.select(time=("2020-01-03", "2020-01-04")) + assert selected._catalog.is_view + + def test_cold_memory_select(self, forbid_realization): + """A fresh patch-list spool selects via the catalog, lazily.""" + patches = list(dc.get_example_spool("random_das")) + forbid_realization() + fresh = dc.spool(patches) + selected = fresh.select(tag="random") + assert selected._catalog.is_view + + +class TestSamples: + """samples=True never excludes patches; trims on load (#447).""" + + def test_length_preserved(self, spool): + """The spool keeps every patch.""" + out = spool.select(distance=(0, 10), samples=True) + assert len(out) == len(spool) + + def test_patch_trimmed_on_load(self, spool): + """Loaded patches carry the sample trim.""" + out = spool.select(distance=(0, 10), samples=True) + patch = out[0] + assert len(patch.get_coord("distance")) == 10 + + def test_survives_chunk(self, spool): + """Post-selects propagate through derived spools.""" + out = spool.select(distance=(0, 10), samples=True).chunk(time=None) + patch = out[0] + assert len(patch.get_coord("distance")) == 10 + + def test_non_coord_raises(self, spool): + """Samples selections must name coordinates.""" + with pytest.raises(InvalidSpoolQueryError, match="coordinate-only"): + spool.select(tag="random", samples=True) + + def test_samples_on_materialized_spool(self, spool): + """Samples select works after chunk (the dataframe select path).""" + materialized = spool.chunk(time=None) + out = materialized.select(distance=(0, 10), samples=True) + assert len(out) == len(materialized) + assert len(out[0].get_coord("distance")) == 10 + + def test_non_coord_on_materialized_raises(self, spool): + """The coordinate-only rule also holds on the dataframe path.""" + materialized = spool.chunk(time=None) + with pytest.raises(InvalidSpoolQueryError, match="coordinate-only"): + materialized.select(tag="random", samples=True) + + +class TestRelative: + """relative=True resolves against the spool envelope (#362).""" + + def test_trims_both_ends(self, spool): + """One second off each end of the spool.""" + df = spool.get_contents() + gmin = df["time_min"].min() + gmax = df["time_max"].max() + out = spool.select(time=(1, -1), relative=True) + merged = out.chunk(time=None)[0] + time = merged.get_coord("time") + assert time.min() >= np.datetime64(gmin) + np.timedelta64(1, "s") + assert time.max() <= np.datetime64(gmax) - np.timedelta64(1, "s") + + def test_requires_range(self, spool): + """Scalars are rejected with a clear message.""" + with pytest.raises(InvalidSpoolQueryError, match="range selectors"): + spool.select(time=5, relative=True) + + def test_namespaced_coord_with_attr(self, spool): + """Only the coordinate range is converted to relative offsets.""" + out = spool.select(_coords={"time": (1, -1)}, tag="random", relative=True) + assert len(out) + assert set(out.get_contents()["tag"]) == {"random"} + + def test_relative_on_materialized_spool(self, spool): + """Relative select works after chunk (the dataframe select path).""" + materialized = spool.chunk(time=None) + gmin = materialized.get_contents()["time_min"].min() + gmax = materialized.get_contents()["time_max"].max() + out = materialized.select(time=(1, -1), relative=True) + merged = out.chunk(time=None)[0] + time = merged.get_coord("time") + assert time.min() >= np.datetime64(gmin) + np.timedelta64(1, "s") + assert time.max() <= np.datetime64(gmax) - np.timedelta64(1, "s") + + +class TestMaterializedNamespaces: + """_attrs/_coords validation on the dataframe (materialized) path.""" + + def test_namespaces_and_unknown_names(self, spool): + """Namespaced selects and unknown-name errors on a chunked spool.""" + materialized = spool.chunk(time=None) + assert len(materialized.select(_attrs={"tag": "random"})) == len(materialized) + # a valid _coords range narrows the materialized spool + df = materialized.get_contents() + t0 = df["time_min"].min() + narrowed = materialized.select( + _coords={"time": (t0, t0 + np.timedelta64(2, "s"))} + ) + assert len(narrowed) <= len(materialized) + with pytest.raises(InvalidSpoolQueryError, match="not an attribute"): + materialized.select(_attrs={"distance": (0, 10)}) + with pytest.raises(InvalidSpoolQueryError, match="not a coordinate"): + materialized.select(_coords={"tag": "random"}) + with pytest.raises(InvalidSpoolQueryError, match="both"): + materialized.select(tag="random", _attrs={"tag": "random"}) + with pytest.raises(InvalidSpoolQueryError, match="neither an attribute"): + materialized.select(nope=1) + + def test_duplicate_namespace_raises(self, spool): + """A name in both explicit namespaces raises on either path.""" + materialized = spool.chunk(time=None) + for target in (spool, materialized): + with pytest.raises(InvalidSpoolQueryError, match="both _attrs and _coords"): + target.select(_attrs={"time": (None, None)}, _coords={"time": (1, 2)}) + + def test_slice_range_form(self, spool): + """Slice selectors resolve the same on either path (#435 spec).""" + materialized = spool.chunk(time=None) + t0 = spool.get_contents()["time_min"].min() + window = slice(t0, t0 + np.timedelta64(2, "s")) + for target in (spool, materialized): + sliced = target.select(time=window) + tupled = target.select(time=(window.start, window.stop)) + assert len(sliced) == len(tupled) + assert len(sliced) >= 1 + + +class TestExistingBehaviorKept: + """The conventional selections still work.""" + + def test_attr_glob(self, spool): + """Unix-style attr matching.""" + assert len(spool.select(tag="rand*")) == len(spool) + + def test_time_range_narrows(self, spool): + """Plain time range selection.""" + df = spool.get_contents() + t0 = df["time_min"].min() + out = spool.select(time=(t0, t0 + np.timedelta64(2, "s"))) + assert len(out) == 1 + + +class TestUnitCanonicalSelection: + """Coordinate selection on non-SI patches (review P1). + + The index stores numeric coordinate summaries in canonical SI, so + range bounds are interpreted as SI end to end: bare numbers mean + canonical SI, quantities convert, and the exact per-patch residual + converts back to each patch's native units. + """ + + @pytest.fixture(scope="class") + def ft_patch(self): + """An example patch with distance in feet (0..~984 ft).""" + return dc.get_example_patch().convert_units(distance="ft") + + def test_bare_numbers_are_canonical_si(self, ft_patch): + """(20, 60) means 20-60 m even on a feet-coordinate patch.""" + coord = dc.spool([ft_patch]).select(distance=(20, 60))[0].get_coord("distance") + assert float(coord.min()) >= 65 # 20 m == 65.6 ft + assert float(coord.max()) <= 197 # 60 m == 196.9 ft + + def test_quantity_selector(self, ft_patch): + """Quantity bounds select the same physical interval.""" + from dascore.units import m + + selected = dc.spool([ft_patch]).select(distance=(20 * m, 60 * m)) + assert len(selected.get_contents()) == 1 # no DimensionalityError + coord = selected[0].get_coord("distance") + assert float(coord.min()) >= 65 + assert float(coord.max()) <= 197 + + def test_quantity_in_native_units(self, ft_patch): + """Quantities in the coordinate's own units also work.""" + from dascore.units import get_quantity + + ft = get_quantity("ft") + coord = ( + dc.spool([ft_patch]) + .select(distance=(100 * ft, 200 * ft))[0] + .get_coord("distance") + ) + assert float(coord.min()) >= 99 + assert float(coord.max()) <= 201 + + def test_unitless_coords_unchanged(self): + """Coordinates without units keep plain numeric semantics.""" + patch = dc.get_example_patch() + coord = dc.spool([patch]).select(distance=(20, 60))[0].get_coord("distance") + assert 20 <= float(coord.min()) and float(coord.max()) <= 60 + + def test_directory_spool(self, ft_patch, tmp_path): + """The same semantics hold for file-backed spools.""" + from dascore.units import m + + dc.write(ft_patch, tmp_path / "ft.h5", "dasdae") + spool = dc.spool(tmp_path).update() + coord = spool.select(distance=(20 * m, 60 * m))[0].get_coord("distance") + assert float(coord.min()) >= 65 + assert float(coord.max()) <= 197 + + def test_mixed_unitless_and_feet_bare(self, ft_patch): + """A bare range trims each patch in its own units (SI meaning).""" + plain = dc.get_example_patch() # unitless distance 0..~300 + plain = plain.update_coords( + distance=plain.get_coord("distance").set_units(None) + ) + got = dc.spool([plain, ft_patch]).select(distance=(20, 60)) + materialized = [p.get_coord("distance") for p in got] + assert len(materialized) == 2 # both overlap 20..60 m + by_units = {str(c.units): c for c in materialized} + # unitless patch: bare magnitudes applied directly + assert float(by_units["None"].min()) >= 20 + assert float(by_units["None"].max()) <= 60 + # feet patch: 20..60 m == 65.6..196.9 ft + assert float(by_units["1 ft"].min()) >= 65 + assert float(by_units["1 ft"].max()) <= 197 + + def test_mixed_unitless_and_feet_quantity(self, ft_patch): + """A metre quantity range works across a mixed population.""" + from dascore.units import m + + plain = dc.get_example_patch() + plain = plain.update_coords( + distance=plain.get_coord("distance").set_units(None) + ) + got = dc.spool([plain, ft_patch]).select(distance=(20 * m, 60 * m)) + assert len(got.get_contents()) == 2 # no UnitError on the unitless row + + def test_scalar_coord_rejected(self, ft_patch): + """A scalar coordinate selector is rejected eagerly, clearly.""" + with pytest.raises(InvalidSpoolQueryError, match="range selectors"): + dc.spool([ft_patch]).select(distance=100) + + def test_value_membership_rejected(self, ft_patch): + """A wrong-arity list is reported as a malformed range.""" + from dascore.exceptions import ParameterError + + with pytest.raises(ParameterError, match="length 2 sequence"): + dc.spool([ft_patch]).select(distance=[10, 20, 50]) + + def test_chained_views(self, ft_patch): + """Canonicalization holds across chained selections.""" + from dascore.units import m + + coord = ( + dc.spool([ft_patch]) + .select(distance=(0 * m, 90 * m)) + .select(distance=(20, 60))[0] + .get_coord("distance") + ) + assert float(coord.min()) >= 65 + assert float(coord.max()) <= 197 + + def test_boolean_mask_selectors_rejected(self, ft_patch): + """Sample masks are patch-level only; the spool points at map().""" + coord = ft_patch.get_coord("distance") + mask = np.zeros(len(coord), dtype=bool) + mask[:5] = True + with pytest.raises(InvalidSpoolQueryError, match="boolean sample"): + dc.spool([ft_patch]).select(distance=mask) + with pytest.raises(InvalidSpoolQueryError, match="boolean sample"): + dc.spool([ft_patch]).select(distance=list(mask)) + # the per-patch escape hatch still works + got = ft_patch.select(distance=mask) + assert len(got.get_coord("distance")) == 5 + + +class TestQuantityDimensionality: + """Quantity queries keep their dimensionality end to end (review P1).""" + + @pytest.fixture() + def mixed_unit_spool(self): + """Two patches whose distance coords are metres and seconds.""" + p_m = dc.get_example_patch() + p_s = p_m.update_coords( + distance=p_m.get_coord("distance").set_units("s") + ).update_attrs(history=[]) + return dc.spool([p_m, p_s]) + + def test_incompatible_coord_excluded(self, mixed_unit_spool): + """A metre query never returns (or trims) a seconds coordinate.""" + from dascore.units import get_quantity, m + + out = mixed_unit_spool.select(_coords={"distance": (1 * m, 2 * m)}) + patches = list(out) + assert len(patches) == 1 + units = get_quantity(str(patches[0].get_coord("distance").units)) + assert units.dimensionality == m.dimensionality + + def test_all_incompatible_raises(self): + """A query incompatible with every stored unit raises UnitError.""" + from dascore.exceptions import UnitError + from dascore.units import m + + p_s = dc.get_example_patch().update_coords( + distance=dc.get_example_patch().get_coord("distance").set_units("s") + ) + spool = dc.spool([p_s]) + with pytest.raises(UnitError, match="no units compatible"): + spool.select(_coords={"distance": (1 * m, 2 * m)}).get_contents() + + +class TestLazyOrderAndWindow: + """sort/slice/array selection are lazy Selection specs (D2).""" + + def test_sort_is_lazy_and_ordered(self, tmp_path_factory): + """Sorting composes a spec; realization returns ordered rows.""" + patches = list(dc.get_example_spool("random_das")) + spool = dc.spool(list(reversed(patches))) + out = spool.sort("time") + assert not out._catalog.is_view or out._catalog._order is not None + df = out.get_contents() + assert df["time_min"].is_monotonic_increasing + assert list(out) == patches + + def test_slice_is_lazy_window(self): + """Slicing keeps the catalog state and correct membership.""" + patches = list(dc.get_example_spool("random_das")) + spool = dc.spool(patches) + part = spool[1:] + assert part._catalog._ids is not None # lazy id membership + assert len(part) == len(patches) - 1 + assert list(part) == patches[1:] + + def test_select_after_slice_filters_within_window(self): + """D2 composition: predicates apply inside the window.""" + patches = list(dc.get_example_spool("random_das")) + spool = dc.spool(patches) + t0 = patches[0].get_coord("time").min() + # the first patch is outside the window, so selecting its time + # range inside the window matches nothing + windowed = spool[1:] + assert len(windowed.select(time=(None, t0 + np.timedelta64(1, "s")))) == 0 + + def test_slice_of_slice_composes(self): + """Windows compose arithmetically.""" + patches = list(dc.get_example_spool("random_das")) + part = dc.spool(patches)[1:][1:] + assert list(part) == patches[2:] + + def test_sorted_spool_slice(self): + """A slice of a sorted view respects the sort order.""" + patches = list(dc.get_example_spool("random_das")) + spool = dc.spool(list(reversed(patches))) + first = spool.sort("time")[0:1] + assert list(first) == patches[0:1] + + def test_split_parts_pickle_small(self): + """split() windows keep map() payloads at member size.""" + import pickle + + base = dc.get_example_patch() + rng = np.random.default_rng(0) + patches = [base.new(data=rng.random(base.shape)) for _ in range(5)] + spool = dc.spool(patches) + parts = list(spool.split(size=1)) + assert len(parts) == 5 + payload = len(pickle.dumps(parts[0])) + baseline = len(pickle.dumps(dc.spool([patches[0]]))) + assert payload < 2 * baseline + + +class TestNamespaceTagForm: + """_attrs/_coords accept names of bare kwargs (tag form).""" + + @pytest.fixture() + def sensor_spool(self): + """One patch with an aux coord whose name is not an attr.""" + patch = dc.get_example_patch() + n = patch.shape[patch.get_axis("distance")] + return dc.spool( + [patch.update_coords(sensor=("distance", np.arange(n, dtype=float)))] + ) + + def test_coords_tag_string(self, sensor_spool): + """A single name tags one bare kwarg as a coordinate.""" + out = sensor_spool.select(sensor=(10, 20), _coords="sensor") + coord = out[0].get_coord("sensor") + assert coord.min() == 10.0 + assert coord.max() == 20.0 + + def test_coords_tag_collection(self, sensor_spool): + """A collection of names tags several bare kwargs.""" + out = sensor_spool.select(sensor=(10, 20), _coords=["sensor"]) + assert len(out) == 1 + + def test_attrs_tag_string(self, sensor_spool): + """The attr side accepts the same tag form.""" + out = sensor_spool.select(tag="random", _attrs="tag") + assert len(out) == 1 + + def test_dict_form_unchanged(self, sensor_spool): + """The general mapping form keeps working.""" + out = sensor_spool.select(_coords={"sensor": (10, 20)}) + assert len(out) == 1 + + def test_tag_without_kwarg_raises(self, sensor_spool): + """Tagging a name with no matching bare kwarg is an error.""" + with pytest.raises(InvalidSpoolQueryError, match="names no bare keyword"): + sensor_spool.select(_coords="sensor") + + def test_non_string_tag_raises(self, sensor_spool): + """Tag collections must contain strings.""" + with pytest.raises(InvalidSpoolQueryError, match="mapping of name"): + sensor_spool.select(sensor=(1, 2), _coords=[3]) + + def test_tagged_name_validates_namespace(self, sensor_spool): + """A tagged name must belong to the claimed namespace.""" + with pytest.raises(InvalidSpoolQueryError, match="not an attribute"): + sensor_spool.select(sensor=(1, 2), _attrs="sensor") diff --git a/tests/test_examples.py b/tests/test_examples.py index 3131de126..e2e0cae52 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -42,7 +42,7 @@ def test_load_example_patch(self, name): assert isinstance(patch, dc.Patch) def test_file_backed_examples_use_direct_read(self, monkeypatch): - """File-backed examples should not depend on FileSpool patch resolution.""" + """File-backed examples should not depend on file-spool patch resolution.""" patch = dc.get_example_patch() def _fetch(_name): diff --git a/tests/test_io/test_dasdae/test_dasdae.py b/tests/test_io/test_dasdae/test_dasdae.py index 1d3494b38..847501536 100644 --- a/tests/test_io/test_dasdae/test_dasdae.py +++ b/tests/test_io/test_dasdae/test_dasdae.py @@ -11,7 +11,6 @@ import numpy as np import pandas as pd import pytest -import tables import dascore as dc from dascore.compat import random_state diff --git a/tests/test_io/test_febus/test_febusg1.py b/tests/test_io/test_febus/test_febusg1.py index 0031ecf32..410e40757 100644 --- a/tests/test_io/test_febus/test_febusg1.py +++ b/tests/test_io/test_febus/test_febusg1.py @@ -231,12 +231,19 @@ def g1_two_file_directory(self, tmp_path_factory): return out_dir def test_chunk_all_time_merges_to_single_patch(self, g1_two_file_directory): - """Ensure chunk(time=None) merges both g1 files into one patch.""" + """Ensure chunk(time=None) merges both g1 files into one patch. + + The files carry per-file attrs (temperature, freqoffset) that + differ, so conflict="keep_first" is required — matching what an + in-memory spool of the same patches requires. (The old HDF5 index + silently dropped these attrs, which masked the conflict for + directory spools.) + """ spool = dc.spool(g1_two_file_directory) # These weren't directly adjacent files so we adjust the tolerance. match = "There is a gap in the patch along dimension time" with pytest.warns(UserWarning, match=match): - merged = spool.chunk(time=None, tolerance=3) + merged = spool.chunk(time=None, tolerance=3, conflict="keep_first") assert len(merged) == 1 def test_mtx_read_raises(self, g1_mtx_buffer): diff --git a/tests/test_io/test_h5simple/test_h5simple.py b/tests/test_io/test_h5simple/test_h5simple.py index 382e1e94d..c86d372d7 100644 --- a/tests/test_io/test_h5simple/test_h5simple.py +++ b/tests/test_io/test_h5simple/test_h5simple.py @@ -5,16 +5,9 @@ import shutil import h5py -import numpy as np import pytest -import tables import dascore as dc -from dascore.io.h5simple.utils import ( - _get_attr_names, - _get_root_attrs, - _iter_root_arrays, -) from dascore.utils.downloader import fetch @@ -33,8 +26,8 @@ def h5simple_with_dim_attrs_path(self, tmp_path_factory): new_path = tmp_path_factory.mktemp("h5simple_dim_attrs") / "simple.h5" shutil.copy2(basic_path, new_path) - with tables.open_file(new_path, "a") as h5: - h5.root._v_attrs["dims"] = "distance,time" + with h5py.File(new_path, "a") as h5: + h5.attrs["dims"] = "distance,time" return new_path def test_no_snap(self, h5simple_path): @@ -46,55 +39,3 @@ def test_dims_in_attrs(self, h5simple_with_dim_attrs_path): """Ensure if 'dims' is in attrs it gets used.""" patch = dc.spool(h5simple_with_dim_attrs_path, file_format="h5simple")[0] assert isinstance(patch, dc.Patch) - - -class TestH5SimpleInternalHelpers: - """Direct tests for helper branches that still support PyTables fixtures.""" - - def test_get_root_attrs_supports_pytables(self, tmp_path): - """PyTables handles should expose root attrs through the helper.""" - path = tmp_path / "root_attrs.h5" - with tables.open_file(path, "w") as h5: - h5.root._v_attrs["dims"] = "distance,time" - attrs = _get_root_attrs(h5) - assert attrs.dims == "distance,time" - - def test_iter_root_arrays_supports_pytables(self, tmp_path): - """PyTables root arrays should still be discoverable by helper code.""" - path = tmp_path / "root_arrays.h5" - with tables.open_file(path, "w") as h5: - h5.create_array("/", "data", obj=np.arange(3)) - names = [name for name, _node in _iter_root_arrays(h5)] - assert names == ["data"] - - def test_get_attr_names_supports_pytables_attrs(self, tmp_path): - """PyTables attr containers should still expose their stored keys.""" - path = tmp_path / "attr_names.h5" - with tables.open_file(path, "w") as h5: - h5.root._v_attrs["dims"] = "distance,time" - out = _get_attr_names(h5.root._v_attrs) - assert "dims" in out - - def test_get_root_attrs_supports_h5py(self, tmp_path): - """h5py files should continue to use the attrs mapping directly.""" - path = tmp_path / "h5py_attrs.h5" - with h5py.File(path, "w") as h5: - h5.attrs["dims"] = "distance,time" - attrs = _get_root_attrs(h5) - assert attrs["dims"] == "distance,time" - - def test_iter_root_arrays_supports_h5py(self, tmp_path): - """h5py root arrays should still be discoverable by helper code.""" - path = tmp_path / "h5py_root_arrays.h5" - with h5py.File(path, "w") as h5: - h5.create_dataset("data", data=np.arange(3)) - names = [name for name, _node in _iter_root_arrays(h5)] - assert names == ["data"] - - def test_get_attr_names_supports_h5py_attrs(self, tmp_path): - """h5py attr containers should still expose their stored keys.""" - path = tmp_path / "h5py_attr_names.h5" - with h5py.File(path, "w") as h5: - h5.attrs["dims"] = "distance,time" - out = _get_attr_names(h5.attrs) - assert "dims" in out diff --git a/tests/test_io/test_index/test_catalog.py b/tests/test_io/test_index/test_catalog.py new file mode 100644 index 000000000..1ce795fec --- /dev/null +++ b/tests/test_io/test_index/test_catalog.py @@ -0,0 +1,308 @@ +"""Tests for PatchCatalog: the unified spool metadata engine.""" + +from __future__ import annotations + +import numpy as np +import pytest + +import dascore as dc +from dascore.io.index import PatchCatalog +from dascore.io.index.query import InvalidSpoolQueryError + + +@pytest.fixture(scope="class") +def patches(): + """Patches from the random example spool.""" + return tuple(dc.get_example_spool("random_das")) + + +@pytest.fixture() +def live_catalog(patches): + """A catalog over live patches.""" + return PatchCatalog.from_patches(patches) + + +class TestLaziness: + """Catalog construction does no metadata work.""" + + def test_no_backend_until_needed(self, patches): + """from_patches must not bootstrap a backend.""" + catalog = PatchCatalog.from_patches(patches) + assert catalog._backend is None + + def test_len_serves_from_registry(self, live_catalog, patches): + """Len (and patch access) never bootstrap a backend.""" + assert len(live_catalog) == len(patches) + assert live_catalog.get_patch(0) is patches[0] + assert live_catalog._backend is None + + def test_first_relation_op_bootstraps(self, live_catalog, patches): + """Realizing the flat relation creates the backend and ingests.""" + df = live_catalog.to_df() + assert len(df) == len(patches) + assert live_catalog._backend is not None + + +class TestLiveRoundtrip: + """Live patches come back identical.""" + + def test_iteration_returns_same_patches(self, live_catalog, patches): + """Iterated patches are the registered objects (construction order).""" + out = list(live_catalog) + assert len(out) == len(patches) + starts = [p.get_coord("time").min() for p in out] + assert starts == sorted(starts) + assert {id(p) for p in out} == {id(p) for p in patches} + + def test_get_patch_by_index(self, live_catalog): + """Integer access works.""" + patch = live_catalog.get_patch(0) + assert isinstance(patch, dc.Patch) + + def test_add_more_patches(self, patches): + """add() ingests additional live patches.""" + catalog = PatchCatalog.from_patches(patches[:1]) + assert len(catalog) == 1 + catalog.add(patches[1]) + assert len(catalog) == 2 + + +class TestSelectComposition: + """select composes lazily with eager validation.""" + + def test_select_narrows(self, live_catalog, patches): + """A time range select excludes non-overlapping patches.""" + t0 = patches[0].get_coord("time").min() + t1 = patches[0].get_coord("time").max() + view = live_catalog.select(time=(t0, t1)) + assert len(view) == 1 + + def test_chained_selects_and(self, live_catalog, patches): + """Chained selects AND together.""" + t0 = patches[0].get_coord("time").min() + view = live_catalog.select(time=(t0, None)).select( + time=(None, t0 + np.timedelta64(1, "s")) + ) + assert len(view) == 1 + + def test_two_stage_exact_trim(self, live_catalog, patches): + """Coord range selects trim the loaded patch exactly.""" + t0 = patches[0].get_coord("time").min() + np.timedelta64(2, "s") + view = live_catalog.select(time=(t0, None)) + patch = view.get_patch(0) + assert patch.get_coord("time").min() >= t0 + + def test_unknown_name_raises_at_select(self, live_catalog): + """Validation is eager (#435).""" + with pytest.raises(InvalidSpoolQueryError, match="neither an attribute"): + live_catalog.select(bad_dim=(1, 2)) + + def test_no_sql_at_select(self, live_catalog): + """Selection composes without realizing the dataframe.""" + view = live_catalog.select(distance=(0, 10)) + assert view._df_cache is None + + def test_views_cannot_mutate(self, live_catalog, patches): + """Mutation only on the root.""" + view = live_catalog.select(distance=(0, 10)) + with pytest.raises(InvalidSpoolQueryError, match="root catalog"): + view.add(patches[0]) + + +class TestResidualSelects: + """samples/relative are patch-local (two-stage).""" + + def test_samples_never_excludes(self, live_catalog, patches): + """samples=True keeps every patch, trims on load (#447).""" + view = live_catalog.select(distance=(0, 10), samples=True) + assert len(view) == len(patches) + patch = view.get_patch(0) + # patch-level samples semantics are authoritative (0..9) + assert len(patch.get_coord("distance")) == 10 + + def test_samples_unknown_coord_raises(self, live_catalog): + """Samples selections validate coord names.""" + with pytest.raises(InvalidSpoolQueryError, match="neither an attribute"): + live_catalog.select(wavelength=(0, 10), samples=True) + + def test_samples_attr_raises(self, live_catalog): + """Samples selections reject attribute names.""" + with pytest.raises(InvalidSpoolQueryError, match="coordinate-only"): + live_catalog.select(tag="test", samples=True) + + def test_relative_select(self, live_catalog): + """Relative bounds resolve against the global envelope (#362).""" + full = live_catalog.to_df() + span = (full["time_max"].max() - full["time_min"].min()).total_seconds() + view = live_catalog.select(time=(1, -1), relative=True) + patch = view.get_patch(0) + got_span = ( + patch.get_coord("time").max() - patch.get_coord("time").min() + ) / np.timedelta64(1, "s") + assert got_span <= span - 1 + + +class TestDirectoryCatalog: + """Directory-backed catalogs share machinery with live ones.""" + + @pytest.fixture(scope="class") + def spool_dir(self, tmp_path_factory): + """A directory of example files.""" + spool = dc.get_example_spool("random_das") + return dc.examples.spool_to_directory( + spool, path=tmp_path_factory.mktemp("catalog_dir") + ) + + def test_roundtrip(self, spool_dir): + """Directory catalog serves the same patches.""" + catalog = PatchCatalog.from_directory(spool_dir).update(progress=None) + patches = list(catalog) + assert len(patches) == 3 + assert all(isinstance(p, dc.Patch) for p in patches) + catalog.close() + + def test_select_and_trim(self, spool_dir): + """Two-stage select works through files too.""" + catalog = PatchCatalog.from_directory(spool_dir).update(progress=None) + df = catalog.to_df() + t0 = df["time_min"].min().to_datetime64() + np.timedelta64(2, "s") + view = catalog.select(time=(t0, None)) + patch = view.get_patch(0) + assert patch.get_coord("time").min() >= t0 + catalog.close() + + +class TestCatalogEdges: + """Remaining branches: errors, passthroughs, offsets.""" + + def test_relative_on_unknown_coord_raises(self, live_catalog): + """Relative select against an absent coord errors clearly.""" + with pytest.raises(InvalidSpoolQueryError, match="neither an attribute"): + live_catalog.select(wavelength=(1, -1), relative=True) + + def test_relative_requires_range(self, live_catalog): + """Relative selects take (start, stop) tuples only.""" + with pytest.raises(InvalidSpoolQueryError, match="range selectors"): + live_catalog.select(time=5, relative=True) + + def test_add_on_file_catalog_not_implemented(self, tmp_path, patches): + """add() is memory-only for now.""" + path = dc.examples.spool_to_directory( + dc.spool(list(patches)), path=tmp_path / "d" + ) + catalog = PatchCatalog.from_directory(path).update(progress=None) + with pytest.raises(NotImplementedError, match="in-memory"): + catalog.add(patches[0]) + catalog.close() + + def test_remove(self, patches): + """remove() drops sources by identity.""" + catalog = PatchCatalog.from_patches(patches) + target = catalog.sources()["source_path"].iloc[0] + catalog.remove([target]) + assert len(catalog) == len(patches) - 1 + + +class TestCount: + """len(catalog) counts in SQL and agrees with the realized relation.""" + + @pytest.fixture() + def diverse_catalog(self): + """A catalog with heterogeneous attrs and coords.""" + return PatchCatalog.from_patches(list(dc.get_example_spool("diverse_das"))) + + def _selections(self, catalog): + """Views spanning attr, coord range, regex, and chained forms.""" + import re + + df = catalog.to_df() + t0 = df["time_min"].min() + window = (t0, t0 + dc.to_timedelta64(1)) + return [ + catalog, + catalog.select(tag="random"), + catalog.select(time=window), + catalog.select(distance=(0, 50)), + catalog.select(tag=re.compile("rand.*")), # regex residual path + catalog.select(tag="random").select(time=window), + ] + + def test_count_matches_realization(self, diverse_catalog): + """Every view's len equals len(to_df()) (fresh, uncached).""" + for view in self._selections(diverse_catalog): + # a fresh view has no cached relation, so len() counts in SQL + expected = len(view.to_df()) + fresh = view._view(view._queries, view._residuals) + assert len(fresh) == expected + + def test_len_does_not_realize(self, diverse_catalog, monkeypatch): + """A cold len() must not pivot coordinates or fetch the relation.""" + catalog = diverse_catalog.select(network="das2") + fresh = catalog._view(catalog._queries, catalog._residuals) + + def _boom(self): + raise AssertionError("flat relation realized during len()") + + monkeypatch.setattr(type(fresh), "to_df", _boom) + assert isinstance(len(fresh), int) + + def test_introspection(self, live_catalog): + """Names, sources, and metadata pass through.""" + assert "time" in live_catalog.coord_names() + assert "tag" in live_catalog.attr_names() + assert len(live_catalog.sources()) == 3 + assert live_catalog.get_metadata()["what_is_this"] == "dascore_spool_index" + live_catalog.close() + + def test_open_relative_bound(self, live_catalog): + """Ellipsis/None bounds stay open through relative resolution.""" + view = live_catalog.select(time=(1, None), relative=True) + assert len(view) >= 1 + + def test_numeric_relative_offset(self, live_catalog): + """Relative selects work on numeric coords too.""" + view = live_catalog.select(distance=(5, -5), relative=True) + patch = view.get_patch(0) + assert patch.get_coord("distance").min() >= 5 + + +class TestCanonicalRange: + """Value semantics of the deferred canonical-SI range.""" + + def test_eq_and_hash(self): + """Equal magnitudes compare and hash equal; other types don't.""" + from dascore.io.index.catalog import _CanonicalRange + + r1, r2 = _CanonicalRange((1.0, 2.0)), _CanonicalRange((1.0, 2.0)) + assert r1 == r2 + assert hash(r1) == hash(r2) + assert r1 != _CanonicalRange((1.0, 3.0)) + assert r1 != (1.0, 2.0) # non-CanonicalRange comparand + + +class TestViewSerialization: + """Views serialize only the live entries their rows reference.""" + + def test_view_pickles_membership_only(self): + """A one-patch view of an N-patch live spool ships one patch.""" + import pickle + + base = dc.get_example_patch() + patches = [ + base.update_attrs(tag=str(i)).new( + data=np.random.default_rng(i).random(base.shape) + ) + for i in range(5) + ] + spool = dc.spool(patches) + view = spool.select(tag="0") + assert len(view) == 1 + payload = pickle.dumps(view) + baseline = pickle.dumps(dc.spool([patches[0]])) + assert len(payload) < 2 * len(baseline) + # the round trip serves the right patch from a one-entry registry + loaded = pickle.loads(payload) + assert len(loaded._catalog.resolver.live_entries()) == 1 + assert loaded[0].attrs["tag"] == "0" + # and the root spool's registry is untouched + assert len(spool._catalog.resolver.live_entries()) == 5 diff --git a/tests/test_io/test_index/test_db_dirspool.py b/tests/test_io/test_index/test_db_dirspool.py new file mode 100644 index 000000000..53a1fc6b8 --- /dev/null +++ b/tests/test_io/test_index/test_db_dirspool.py @@ -0,0 +1,121 @@ +""" +Integration tests: directory spools running on the database index. + +Exercises the full path — directory walk, scan, ingest, query, patch +loading, and chunk against real files. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +import dascore as dc +from dascore.core.spool import Spool +from dascore.examples import spool_to_directory + + +@pytest.fixture(scope="class") +def spool_directory(tmp_path_factory): + """A directory of DASDAE files from the random example spool.""" + spool = dc.get_example_spool("random_das") + return spool_to_directory(spool, path=tmp_path_factory.mktemp("db_spool")) + + +@pytest.fixture() +def db_spool(spool_directory): + """A directory spool using its SQLite index.""" + spool = Spool.from_directory(spool_directory) + out = spool.update(progress=None) + yield out + out.indexer.close() + + +class TestDBDirectorySpools: + """Directory spools wired to the database index.""" + + def test_length(self, db_spool): + """One entry per patch in the source spool.""" + assert len(db_spool) == 3 + + def test_contents_columns(self, db_spool): + """The contents df carries what the spool machinery needs.""" + df = db_spool.get_contents() + for col in ("path", "time_min", "time_max", "time_step", "dims"): + assert col in df.columns + + def test_load_patches(self, db_spool): + """Every indexed patch is loadable and matches the source data.""" + source = dc.get_example_spool("random_das") + source_starts = {patch.get_coord("time").min() for patch in source} + loaded_starts = set() + for patch in db_spool: + assert patch.data.size > 0 + loaded_starts.add(patch.get_coord("time").min()) + assert loaded_starts == source_starts + + def test_select_time(self, db_spool): + """Time select narrows the spool.""" + df = db_spool.get_contents() + t0 = df["time_min"].min().to_datetime64() + sub = db_spool.select(time=(t0, t0 + np.timedelta64(2, "s"))) + assert len(sub) >= 1 + patch = sub[0] + assert patch.get_coord("time").max() <= t0 + np.timedelta64(2, "s") + + def test_chunk_merge(self, db_spool): + """Contiguous patches merge with chunk(time=None).""" + merged = db_spool.chunk(time=None) + assert len(merged) == 1 + patch = merged[0] + assert patch.data.ndim == 2 + + def test_update_is_incremental(self, db_spool): + """A second update with no changes rescans nothing.""" + indexer = db_spool.indexer + before = indexer._backend.get_sources()["last_indexed_ns"].max() + db_spool.update(progress=None) + after = indexer._backend.get_sources()["last_indexed_ns"].max() + assert before == after + + +class TestUpdateLifecycle: + """New, modified, and deleted files are tracked per source.""" + + @pytest.fixture() + def fresh(self, tmp_path): + """A modifiable spool directory and database spool.""" + spool = dc.get_example_spool("random_das") + path = spool_to_directory(spool, path=tmp_path / "data") + out = Spool.from_directory(path).update(progress=None) + yield path, out + out.indexer.close() + + def test_new_file_found(self, fresh): + """A file added after indexing appears on the next update.""" + path, spool = fresh + patch = dc.get_example_patch() + patch.io.write(path / "new_file.hdf5", "dasdae") + updated = spool.update(progress=None) + assert len(updated) == 4 + + def test_deleted_file_removed(self, fresh): + """A deleted file's rows are dropped on the next update.""" + path, spool = fresh + target = next(iter(path.glob("*.hdf5"))) + target.unlink() + updated = spool.update(progress=None) + assert len(updated) == 2 + + def test_no_change_update_uses_narrow_projection(self, fresh, monkeypatch): + """A no-op update reads only the stat columns, not the wide table.""" + path, spool = fresh + backend = spool.indexer._backend + + def _boom(self): + raise AssertionError("wide get_sources() during no-change update") + + # A no-change update must not fetch the full sources table. + monkeypatch.setattr(type(backend), "get_sources", _boom) + reupdated = spool.update(progress=None) + assert len(reupdated) == 3 diff --git a/tests/test_io/test_index/test_heterogeneity_stress.py b/tests/test_io/test_index/test_heterogeneity_stress.py new file mode 100644 index 000000000..c721844dc --- /dev/null +++ b/tests/test_io/test_index/test_heterogeneity_stress.py @@ -0,0 +1,259 @@ +""" +Randomized heterogeneity stress test for the spool index. + +Generates hundreds of patch summaries with randomized dimension names, +coord dtypes/units, attr names/kinds (including hostile names that +collide after sanitization), and verifies ingest, counts, and the +no-false-negative query contract on SQLite. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from dascore.core.summary import PatchSummary +from dascore.io.index import Query, get_backend, summaries_to_records +from dascore.units import get_quantity + +N_PATCHES = 300 + +_DIM_POOL = ( + "time", + "distance", + "lag_time", + "frequency", + "velocity", + "depth", + "channel_number", + "offset", + "azimuth", + "Gauge Length (m)", # hostile coord name +) +_NUM_UNITS = (None, "m", "ft", "Hz", "1/m", "km") +_ATTR_NAMES = ( + "station", + "tag", + "shot_number", + "Shot Number", # sanitizes identically to shot_number + "GaugeLength", + "über_attr", + "9lives", + "data quality!", + "processing_level", +) + + +def _random_coord(rng, name): + """Build one random coord summary dict.""" + kind = rng.choice(["datetime", "timedelta", "float", "int", "str"]) + length = int(rng.integers(2, 5000)) + if kind == "datetime": + start = np.datetime64("2024-01-01", "ns") + np.timedelta64( + int(rng.integers(0, 10**7)), "s" + ) + span = np.timedelta64(int(rng.integers(1, 10**5)), "s") + return { + "dtype": "datetime64", + "min": start, + "max": start + span, + "step": span / max(length - 1, 1) if rng.random() > 0.3 else None, + "dims": (name,), + "len": length, + } + if kind == "timedelta": + lo = np.timedelta64(int(rng.integers(-(10**6), 0)), "ms") + hi = np.timedelta64(int(rng.integers(1, 10**6)), "ms") + return { + "dtype": "timedelta64", + "min": lo, + "max": hi, + "step": None, + "dims": (name,), + "len": length, + } + if kind == "str": + lo, hi = sorted( + [f"A{rng.integers(0, 100):03d}", f"Z{rng.integers(0, 100):03d}"] + ) + return {"dtype": "str", "min": lo, "max": hi, "dims": (name,), "len": length} + lo = float(rng.uniform(-1000, 1000)) + hi = lo + float(rng.uniform(0.001, 5000)) + units = rng.choice(_NUM_UNITS) + out = { + "dtype": "float64" if kind == "float" else "int64", + "min": lo if kind == "float" else int(lo), + "max": hi if kind == "float" else int(hi) + 1, + "step": None, + "dims": (name,), + "len": length, + } + if units is not None: + out["units"] = units + return out + + +def _random_attrs(rng) -> dict: + """Build a random attr dict with mixed kinds and hostile names.""" + out = {} + for name in _ATTR_NAMES: + roll = rng.random() + if roll < 0.35: + continue # attr missing on this patch + if name in ("station", "tag"): + # typed str fields on PatchAttrs; only extra attrs vary kind + out[name] = f"v{rng.integers(0, 50)}" + continue + kind = rng.choice(["str", "int", "float", "bool", "time", "quantity"]) + if kind == "str": + out[name] = f"v{rng.integers(0, 50)}" + elif kind == "int": + out[name] = int(rng.integers(-1000, 1000)) + elif kind == "float": + out[name] = float(rng.uniform(-1000, 1000)) + elif kind == "bool": + out[name] = bool(rng.random() > 0.5) + elif kind == "time": + out[name] = np.datetime64("2024-01-01") + np.timedelta64( + int(rng.integers(0, 10**6)), "s" + ) + else: + # "s" is dimensionally incompatible with "m"/"ft": ingest keeps + # the first-seen dimension and skips the rest with a warning. + out[name] = float(rng.uniform(0, 100)) * get_quantity( + str(rng.choice(["m", "ft", "s"])) + ) + return out + + +def make_random_summaries(n: int, seed: int = 0) -> list[PatchSummary]: + """Generate n randomized heterogeneous patch summaries.""" + rng = np.random.default_rng(seed) + out = [] + for i in range(n): + n_dims = int(rng.integers(1, 4)) + dims = tuple(rng.choice(_DIM_POOL, size=n_dims, replace=False)) + coords = {name: _random_coord(rng, name) for name in dims} + shape = tuple(int(coords[d]["len"]) for d in dims) + out.append( + PatchSummary( + attrs=_random_attrs(rng), + coords=coords, + dims=dims, + shape=shape, + dtype="float32", + source_path=f"stress/file_{i:05d}.h5", + source_format="DASDAE", + source_version="1", + ) + ) + return out + + +@pytest.fixture(scope="module") +def summaries(): + """The randomized summary population (deterministic seed).""" + return make_random_summaries(N_PATCHES, seed=42) + + +@pytest.fixture() +def backend(tmp_path_factory, summaries): + """A SQLite backend ingesting the random population.""" + path = tmp_path_factory.mktemp("stress") / "index.sqlite3" + back = get_backend(path) + back.write_sources(summaries_to_records(summaries)) + yield back + back.close() + + +class TestStressIngest: + """Every random population ingests completely.""" + + def test_all_patches_indexed(self, backend): + """One flat row per summary.""" + assert len(backend.query()) == N_PATCHES + + def test_all_coords_present(self, backend, summaries): + """Every generated coord name is known to the index.""" + expected = {name for s in summaries for name in s.coords} + assert expected <= backend.coord_names() + + def test_sanitize_collision_attrs_distinct(self, backend, summaries): + """Attrs whose names sanitize identically stay distinct.""" + assert "shot_number" in backend.attr_names() + assert "Shot Number" in backend.attr_names() + df = backend.query() + assert "shot_number" in df.columns and "Shot Number" in df.columns + # every generated attr that carried a value is indexed + expected = { + name + for s in summaries + for name, value in s.attrs.model_dump().items() + if name in _ATTR_NAMES and value not in (None, "") + } + assert expected <= backend.attr_names() + + +class TestStressNoFalseNegatives: + """Random range queries never miss a matching patch.""" + + def test_random_numeric_coord_ranges(self, backend, summaries): + """Numeric coord queries: reference computed from raw summaries.""" + rng = np.random.default_rng(1) + num_coords = ["distance", "frequency", "velocity", "depth", "offset"] + for _ in range(20): + name = str(rng.choice(num_coords)) + lo = float(rng.uniform(-500, 500)) + hi = lo + float(rng.uniform(1, 2000)) + got = set(backend.query(Query(coords={name: (lo, hi)}))["path"]) + for summary in summaries: + csum = summary.coords.get(name) + if csum is None: + continue + # numeric queries must NOT match datetime/timedelta coords + # (kind dispatch) — and numpy's type hierarchy makes + # timedelta64 a signedinteger subtype, so check dtype.kind. + if np.dtype(csum.dtype).kind not in "iuf": + continue + factor = 1.0 + if csum.units is not None: + factor = float( + get_quantity(str(csum.units)).to_base_units().magnitude + ) + if float(csum.min) * factor <= hi and float(csum.max) * factor >= lo: + assert str(summary.source_path).replace("\\", "/") in got, ( + name, + lo, + hi, + ) + + def test_random_time_ranges(self, backend, summaries): + """Absolute time queries against datetime coords.""" + rng = np.random.default_rng(2) + base = np.datetime64("2024-01-01", "ns") + for _ in range(20): + lo = base + np.timedelta64(int(rng.integers(0, 10**7)), "s") + hi = lo + np.timedelta64(int(rng.integers(1, 10**6)), "s") + got = set(backend.query(Query(coords={"time": (lo, hi)}))["path"]) + for summary in summaries: + csum = summary.coords.get("time") + if csum is None or "datetime" not in str(csum.dtype): + continue + if csum.min <= hi and csum.max >= lo: + assert str(summary.source_path).replace("\\", "/") in got + + def test_attr_equality_roundtrip(self, backend, summaries): + """Str attr equality returns every patch carrying that value.""" + rng = np.random.default_rng(3) + for _ in range(10): + summary = summaries[int(rng.integers(0, len(summaries)))] + attrs = { + k: v + for k, v in summary.attrs.model_dump().items() + if isinstance(v, str) and v and k in set(_ATTR_NAMES) + } + if not attrs: + continue + name, value = next(iter(attrs.items())) + got = set(backend.query(Query(attrs={name: value}))["path"]) + assert str(summary.source_path).replace("\\", "/") in got diff --git a/tests/test_io/test_index/test_index_contract.py b/tests/test_io/test_index/test_index_contract.py new file mode 100644 index 000000000..e5321be45 --- /dev/null +++ b/tests/test_io/test_index/test_index_contract.py @@ -0,0 +1,495 @@ +""" +Contract tests for the spool index backend. + +The SQLite backend must pass this suite; it encodes the selector +semantics spec and the summary-only/no-false-negatives contract from the +index design doc (see discussion #648). +""" + +from __future__ import annotations + +import re + +import numpy as np +import pandas as pd +import pytest + +from dascore.core.summary import PatchSummary +from dascore.exceptions import UnitError +from dascore.io.index import Query, get_backend, summaries_to_records +from dascore.io.index.backend import resolve_query +from dascore.io.index.query import InvalidSpoolQueryError +from dascore.io.index.schema import INDEX_VERSION +from dascore.units import get_quantity + + +def _time_coord(t0: str, seconds: float, step_s: float = 0.004): + """Make an absolute time coord summary dict.""" + start = np.datetime64(t0, "ns") + return { + "dtype": "datetime64", + "min": start, + "max": start + np.timedelta64(int(seconds * 1e9), "ns"), + "step": np.timedelta64(int(step_s * 1e9), "ns"), + "units": "s", + "dims": ("time",), + "len": int(seconds / step_s), + } + + +def _distance_coord(d0: float, d1: float, step: float, units="m"): + """Make a numeric distance coord summary dict.""" + return { + "dtype": "float64", + "min": d0, + "max": d1, + "step": step, + "units": units, + "dims": ("distance",), + "len": int((d1 - d0) / step) + 1, + } + + +def make_summaries() -> list[PatchSummary]: + """A deliberately heterogeneous set of patch summaries.""" + das1 = PatchSummary( + attrs={ + "station": "STA1", + "network": "NW", + "tag": "raw", + "data_type": "strain_rate", + "gauge_length": 10, + }, + coords={ + "time": _time_coord("2024-01-01T00:00:00", 60), + "distance": _distance_coord(0, 1000, 1), + }, + dims=("time", "distance"), + shape=(15000, 1001), + dtype="float32", + source_path="das/file_1.h5", + source_format="PRODML", + source_version="2.1", + ) + das2 = PatchSummary( + attrs={ + "station": "STA2", + "network": "NW", + "tag": "raw", + "data_type": "strain_rate", + "gauge_length": 10.0, + }, + coords={ + "time": _time_coord("2024-01-01T00:01:00", 60), + "distance": _distance_coord(0, 1000, 1), + }, + dims=("time", "distance"), + shape=(15000, 1001), + dtype="float32", + source_path="das/file_2.h5", + source_format="PRODML", + source_version="2.1", + ) + # correlogram: relative (timedelta) lag_time coord, shot_number attr + correlogram = PatchSummary( + attrs={"tag": "corr", "shot_number": 42, "data_type": ""}, + coords={ + "lag_time": { + "dtype": "timedelta64", + "min": np.timedelta64(-5_000_000_000, "ns"), + "max": np.timedelta64(5_000_000_000, "ns"), + "step": np.timedelta64(10_000_000, "ns"), + "dims": ("lag_time",), + "len": 1001, + }, + "distance": _distance_coord(0, 500, 5), + }, + dims=("lag_time", "distance"), + shape=(1001, 101), + dtype="float64", + source_path="products/corr_1.h5", + source_format="DASDAE", + source_version="1", + ) + # PSD-like product with distance in feet (tests SI normalization) + psd = PatchSummary( + attrs={"tag": "psd", "shot_number": "unknown"}, + coords={ + "frequency": { + "dtype": "float64", + "min": 0.0, + "max": 500.0, + "step": 0.5, + "units": "Hz", + "dims": ("frequency",), + "len": 1001, + }, + "distance": _distance_coord(0.0, 3280.0, 3.28, units="ft"), + }, + dims=("frequency", "distance"), + shape=(1001, 1001), + dtype="float64", + source_path="products/psd_1.h5", + source_format="DASDAE", + source_version="1", + ) + return [das1, das2, correlogram, psd] + + +@pytest.fixture(scope="function") +def backend(tmp_path): + """A freshly ingested SQLite index backend.""" + path = tmp_path / "index.sqlite3" + back = get_backend(path) + back._test_path = path + back.write_sources(summaries_to_records(make_summaries())) + yield back + back.close() + + +class TestFlatRelation: + """The flat patch-row relation contract.""" + + def test_row_per_patch(self, backend): + """Row per patch.""" + df = backend.query() + assert len(df) == 4 + + def test_structural_columns(self, backend): + """Structural columns.""" + df = backend.query() + for col in ("path", "file_format", "file_version", "dims", "shape"): + assert col in df.columns + assert pd.api.types.is_datetime64_dtype(df["time_min"]) + assert pd.api.types.is_timedelta64_dtype(df["time_step"]) + + def test_attr_columns_use_original_names(self, backend): + """Attr columns use original names.""" + df = backend.query() + assert "station" in df.columns + assert "gauge_length" in df.columns + assert set(df["station"].replace("", None).dropna()) == {"STA1", "STA2"} + + def test_missing_str_attrs_are_empty_string(self, backend): + """Missing str attrs are empty string.""" + df = backend.query() + corr = df[df["tag"] == "corr"] + assert (corr["station"] == "").all() + + def test_relative_time_patches_have_null_time_min(self, backend): + """Relative time patches have null time min.""" + df = backend.query() + corr = df[df["tag"] == "corr"] + assert corr["time_min"].isnull().all() + + def test_ordering_deterministic(self, backend): + """Ordering deterministic.""" + df1, df2 = backend.query(), backend.query() + pd.testing.assert_frame_equal(df1, df2) + # NULLS LAST: relative-time patches sort after absolute ones. + nulls = df1["time_min"].isnull().to_numpy() + assert not nulls[: (~nulls).sum()].any() + + +class TestAttrPredicates: + """Attr predicates are exact at the index.""" + + def test_equality(self, backend): + """Equality.""" + df = backend.query(Query(attrs={"station": "STA1"})) + assert len(df) == 1 + assert df["station"].iloc[0] == "STA1" + + def test_glob(self, backend): + """Glob.""" + df = backend.query(Query(attrs={"station": "STA*"})) + assert len(df) == 2 + + def test_regex(self, backend): + """Regex.""" + df = backend.query(Query(attrs={"station": re.compile(r"STA\d")})) + assert len(df) == 2 + + def test_membership(self, backend): + """Membership.""" + df = backend.query(Query(attrs={"station": ["STA1", "STA2", "NOPE"]})) + assert len(df) == 2 + + def test_int_matches_float_storage(self, backend): + """Int matches float storage.""" + # gauge_length stored from int 10 and float 10.0; int query hits both + df = backend.query(Query(attrs={"gauge_length": 10})) + assert len(df) == 2 + + def test_range(self, backend): + """Range.""" + df = backend.query(Query(attrs={"gauge_length": (5, 15)})) + assert len(df) == 2 + + def test_open_range(self, backend): + """Open range.""" + df = backend.query(Query(attrs={"gauge_length": (5, None)})) + assert len(df) == 2 + + def test_kind_mismatch_matches_nothing(self, backend): + """Kind mismatch matches nothing.""" + # station is a str attr; numeric query is valid but matches nothing + df = backend.query(Query(attrs={"station": 5})) + assert df.empty + + @pytest.mark.parametrize( + "value", + [ + get_quantity("1 m"), + [get_quantity("1 m"), get_quantity("2 m")], + (get_quantity("900 m"), get_quantity("1 km")), + ], + ) + def test_quantity_kind_mismatch_matches_nothing(self, backend, value): + """Quantity forms do not convert against a string-only attribute.""" + df = backend.query(Query(attrs={"station": value})) + assert df.empty + + def test_mixed_kind_attr(self, backend): + """Mixed kind attr.""" + # shot_number exists as num (42) and str ("unknown") + num = backend.query(Query(attrs={"shot_number": 42})) + assert list(num["tag"]) == ["corr"] + txt = backend.query(Query(attrs={"shot_number": "unknown"})) + assert list(txt["tag"]) == ["psd"] + + +class TestCoordPredicates: + """Coord predicates: envelope candidacy, never false negatives.""" + + def test_time_range(self, backend): + """Time range.""" + t = (np.datetime64("2024-01-01T00:00:30"), np.datetime64("2024-01-01T00:00:40")) + df = backend.query(Query(coords={"time": t})) + assert list(df["station"]) == ["STA1"] + + def test_time_range_overlap_both(self, backend): + """Time range overlap both.""" + t = (np.datetime64("2024-01-01T00:00:30"), np.datetime64("2024-01-01T00:01:30")) + df = backend.query(Query(coords={"time": t})) + assert set(df["station"]) == {"STA1", "STA2"} + + def test_absolute_time_excludes_relative(self, backend): + """Absolute time excludes relative.""" + t = (np.datetime64("1990-01-01"), np.datetime64("2100-01-01")) + df = backend.query(Query(coords={"time": t})) + assert "corr" not in set(df["tag"]) + + def test_relative_time_coord(self, backend): + """Relative time coord.""" + lag = (np.timedelta64(0, "s"), np.timedelta64(2, "s")) + df = backend.query(Query(coords={"lag_time": lag})) + assert list(df["tag"]) == ["corr"] + + def test_numeric_coord_si_normalized(self, backend): + """Numeric coord si normalized.""" + # psd distance is 0-3280 ft = 0-999.7 m; a 900-950 m query hits it + df = backend.query(Query(coords={"distance": (900, 950)})) + assert "psd" in set(df["tag"]) + + def test_quantity_coord_converts_units(self, backend): + """Quantity selectors convert to the coordinate's canonical units.""" + meter = get_quantity("m") + df = backend.query(Query(coords={"distance": (900 * meter, 950 * meter)})) + assert "psd" in set(df["tag"]) + + def test_incompatible_quantity_coord_raises(self, backend): + """A time quantity cannot query a length coordinate.""" + second = get_quantity("s") + with pytest.raises(UnitError): + backend.query(Query(coords={"distance": (1 * second, 2 * second)})) + + def test_scalar_coord_rejected(self, backend): + """Scalar coord predicates have no exact patch meaning; rejected.""" + with pytest.raises(InvalidSpoolQueryError, match="range selectors"): + backend.query(Query(coords={"frequency": 100})) + + def test_array_membership_rejected(self, backend): + """Numeric value membership on a coord is rejected, not candidacy.""" + values = np.array([10.0, 20.0, 480.0]) + with pytest.raises(InvalidSpoolQueryError, match="range selectors"): + backend.query(Query(coords={"distance": values})) + + def test_coord_missing_excludes_patch(self, backend): + """Coord missing excludes patch.""" + df = backend.query(Query(coords={"frequency": (0, 1000)})) + assert list(df["tag"]) == ["psd"] + + +class TestNameResolution: + """Bare kwargs resolve attrs first, then coords; unknown raises.""" + + def test_attr_wins(self, backend): + """Attr wins.""" + query = resolve_query(backend, station="STA1") + assert "station" in query.attrs + + def test_coord_fallback(self, backend): + """Coord fallback.""" + query = resolve_query(backend, lag_time=(0, 1)) + assert "lag_time" in query.coords + + def test_unknown_raises(self, backend): + """Unknown raises.""" + with pytest.raises(InvalidSpoolQueryError, match="neither an attribute"): + resolve_query(backend, wavelength=(1, 2)) + + def test_double_specification_raises(self, backend): + """Double specification raises.""" + with pytest.raises(InvalidSpoolQueryError, match="both"): + resolve_query(backend, station="STA1", _attrs={"station": "STA2"}) + + def test_explicit_namespaces(self, backend): + """Explicit namespaces.""" + query = resolve_query( + backend, _attrs={"tag": "raw"}, _coords={"distance": (0, 10)} + ) + assert query.attrs == {"tag": "raw"} and "distance" in query.coords + + +class TestNoFalseNegatives: + """Property: reference-matching patches always appear in results.""" + + def test_random_time_ranges(self, backend): + """Random time ranges.""" + rng = np.random.default_rng(42) + summaries = make_summaries() + base = np.datetime64("2024-01-01T00:00:00").astype("datetime64[ns]") + for _ in range(25): + lo = base + np.timedelta64(int(rng.integers(-60, 180)), "s") + hi = lo + np.timedelta64(int(rng.integers(1, 120)), "s") + result_paths = set(backend.query(Query(coords={"time": (lo, hi)}))["path"]) + for summary in summaries: + tcoord = summary.coords.get("time") + if tcoord is None or "datetime" not in str(tcoord.dtype): + continue + overlaps = tcoord.min <= hi and tcoord.max >= lo + if overlaps: + assert str(summary.source_path).replace("\\", "/") in result_paths + + def test_random_numeric_ranges(self, backend): + """Random numeric ranges.""" + rng = np.random.default_rng(7) + factor = {"m": 1.0, "ft": 0.3048} + summaries = make_summaries() + for _ in range(25): + lo = float(rng.uniform(-100, 1000)) + hi = lo + float(rng.uniform(1, 500)) + result_paths = set( + backend.query(Query(coords={"distance": (lo, hi)}))["path"] + ) + for summary in summaries: + dcoord = summary.coords.get("distance") + if dcoord is None: + continue + scale = factor[str(dcoord.units.units)] if dcoord.units else 1.0 + if dcoord.min * scale <= hi and dcoord.max * scale >= lo: + assert str(summary.source_path).replace("\\", "/") in result_paths + + +class TestSourceLifecycle: + """Source-scoped transactional replacement and deletion.""" + + def test_replace_source_drops_stale_rows(self, backend): + """Replace source drops stale rows.""" + summaries = [s for s in make_summaries() if "file_1" in str(s.source_path)] + structured = summaries[0].dump_structured() + structured["attrs"] = {"station": "NEW1"} + modified = PatchSummary(**structured) + backend.write_sources(summaries_to_records([modified])) + df = backend.query() + assert len(df) == 4 # still one row for that source + assert "STA1" not in set(df["station"]) + assert "NEW1" in set(df["station"]) + + def test_delete_cascades(self, backend): + """Delete cascades to patches, attrs, and coord links via the FK.""" + before = backend._fetch_df("SELECT patch_id FROM patches") + gone = backend._fetch_df( + "SELECT p.patch_id FROM patches p JOIN sources s " + "ON s.source_id = p.source_id WHERE s.source_path = 'das/file_1.h5'" + )["patch_id"].tolist() + assert gone # the source had patches to cascade-delete + backend.delete_sources(["das/file_1.h5"]) + df = backend.query() + assert len(df) == 3 + assert "das/file_1.h5" not in set(df["path"]) + # the deleted source's patches (and their dependents) are gone, + # not merely filtered out of the query. + remaining = set(backend._fetch_df("SELECT patch_id FROM patches")["patch_id"]) + assert remaining == set(before["patch_id"]) - set(gone) + for table in ("attrs", "patch_coords"): + ids = set(backend._fetch_df(f"SELECT patch_id FROM {table}")["patch_id"]) + assert not (ids & set(gone)) + + def test_reopen_persists(self, backend, tmp_path): + """Reopen persists.""" + path = backend._test_path + backend.close() + reopened = get_backend(path) + try: + assert len(reopened.query()) == 4 + finally: + reopened.close() + # reopen once more so fixture teardown close() has a live handle + reopened_again = get_backend(path) + backend.__dict__.update(reopened_again.__dict__) + + +class TestMetadata: + """Index metadata and introspection.""" + + def test_metadata(self, backend): + """Metadata.""" + meta = backend.get_metadata() + assert meta["what_is_this"] == "dascore_spool_index" + assert meta["index_version"] == INDEX_VERSION + + def test_names(self, backend): + """Names.""" + assert {"station", "tag", "shot_number"} <= backend.attr_names() + assert {"time", "distance", "lag_time", "frequency"} <= backend.coord_names() + + def test_sources(self, backend): + """Sources.""" + sources = backend.get_sources() + assert len(sources) == 4 + assert set(sources["source_format"]) == {"PRODML", "DASDAE"} + + +class TestCoordPivot: + """Per-coord envelope columns in the flat relation.""" + + def test_generic_coord_envelopes_present(self, backend): + """Non-conventional dims get {name}_min/max/step columns.""" + df = backend.query() + for col in ("lag_time_min", "lag_time_max", "frequency_min"): + assert col in df.columns + corr = df[df["tag"] == "corr"] + assert pd.api.types.is_timedelta64_dtype(corr["lag_time_min"].dtype) or ( + corr["lag_time_min"].map(lambda x: hasattr(x, "total_seconds")).all() + ) + + def test_time_distance_envelopes_not_duplicated(self, backend): + """patches-level envelopes are authoritative; pivot skips them.""" + df = backend.query() + assert pd.api.types.is_datetime64_dtype(df["time_min"]) + assert df.columns.tolist().count("distance_min") == 1 + + def test_def_key_columns_private_and_shared(self, backend): + """_{name}_def_key exists for every coord; shared coords share keys.""" + df = backend.query() + assert "_distance_def_key" in df.columns + das = df[df["station"].isin(["STA1", "STA2"])] + assert das["_distance_def_key"].nunique() == 1 + + def test_pivot_respects_query(self, backend): + """A filtered result only pivots the rows it contains.""" + df = backend.query(Query(attrs={"tag": "corr"})) + assert df["lag_time_min"].notna().all() + assert "frequency_min" not in df.columns or df["frequency_min"].isna().all() diff --git a/tests/test_io/test_index/test_index_edge_cases.py b/tests/test_io/test_index/test_index_edge_cases.py new file mode 100644 index 000000000..9aed8946b --- /dev/null +++ b/tests/test_io/test_index/test_index_edge_cases.py @@ -0,0 +1,1344 @@ +""" +Edge-case and error-path tests for the index package. + +Complements the contract suite: exercises failure branches, kind +mismatches, rollbacks, and the directory-format walk so the package has +full line coverage. +""" + +from __future__ import annotations + +import os +import re +import sqlite3 + +import numpy as np +import pandas as pd +import pytest +from test_index_contract import make_summaries + +import dascore as dc +from dascore.core.summary import PatchSummary +from dascore.exceptions import UnitError +from dascore.io.index import Query, get_backend, summaries_to_records +from dascore.io.index.backend import _ns_to_time, adapt_params, resolve_query +from dascore.io.index.indexer import DBDirectoryIndexer +from dascore.io.index.ingest import ( + SourceRecord, + _coord_record, + typed_value, +) +from dascore.io.index.ingest import ( + summaries_to_records as s2r, +) +from dascore.io.index.query import InvalidSpoolQueryError +from dascore.units import get_quantity, m + + +class TestIndexCoverageEdges: + """Remaining index branches covered with real backends/spools.""" + + def test_live_resolver_missing_patch(self): + """A row for a patch absent from the registry raises MissingPatchError.""" + from dascore.exceptions import MissingPatchError + from dascore.io.index.catalog import LiveResolver + + resolver = LiveResolver([dc.get_example_patch()]) + with pytest.raises(MissingPatchError, match="not available"): + resolver.resolve({"path": "memorypatch://not-a-real-id"}) + + def test_mixed_compatible_unit_range(self): + """A coord range mixing compatible units resolves.""" + from dascore.units import get_quantity + + cm = get_quantity("cm") + out = dc.spool([dc.get_example_patch()]).select(distance=(1 * m, 200 * cm)) + assert len(out.get_contents()) == 1 + + def test_backend_range_incompatible_units_raise(self, backend): + """A hand-built coord range mixing incompatible units is rejected. + + The catalog canonicalizes units before the backend, so only a + direct Query reaches the multi-unit compatibility check. + """ + from dascore.units import s + + with pytest.raises(UnitError, match="Cannot convert"): + backend.query(Query(coords={"distance": (1 * m, 2 * s)})) + + def test_export_skips_source_without_patches(self, tmp_path): + """A non-fiber file gets a sources row with no patches; export skips it.""" + dc.get_example_patch().io.write(tmp_path / "a.h5", "dasdae") + (tmp_path / "junk.txt").write_text("not a fiber file") + spool = dc.spool(tmp_path).update(progress=None) + backend = spool._catalog.backend + assert len(backend.get_sources()) == 2 # the h5 and the junk file + records = backend.export_records() + assert sum(len(r.patches) for r in records) == 1 # only the real patch + + def test_attr_meta_units_backfilled(self, tmp_path): + """An attr first seen unitless gets its unit backfilled by a later write.""" + from dascore.core.summary import PatchSummary + + def _summary(gain, path): + return PatchSummary( + attrs={"tag": "t", "gain": gain}, + coords={ + "distance": { + "dtype": "float64", + "min": 0.0, + "max": 10.0, + "step": 1.0, + "units": "m", + "dims": ("distance",), + "len": 11, + } + }, + dims=("distance",), + shape=(11,), + dtype="float64", + source_path=path, + source_format="X", + source_version="1", + ) + + back = get_backend(tmp_path / "i.sqlite3") + try: + back.write_sources(s2r([_summary(5.0, "a.h5")])) # unitless + meta = back._attr_meta() + assert list(meta.loc[meta["attr_name"] == "gain", "units"]) == [None] + back.write_sources(s2r([_summary(5.0 * m, "b.h5")])) # units -> backfill + meta = back._attr_meta() + assert list(meta.loc[meta["attr_name"] == "gain", "units"]) == ["m"] + finally: + back.close() + + def test_reopen_missing_meta_row(self, tmp_path): + """An index whose meta row was lost is rejected on reopen.""" + import sqlite3 + + from dascore.exceptions import InvalidIndexError + + path = tmp_path / "i.sqlite3" + get_backend(path).close() + con = sqlite3.connect(path) + con.execute("DELETE FROM meta_data") + con.commit() + con.close() + with pytest.raises(InvalidIndexError, match="not a valid"): + get_backend(path) + + def test_reopen_table_missing_column(self, tmp_path): + """An index whose table lost a column is rejected on reopen.""" + import sqlite3 + + from dascore.exceptions import InvalidIndexError + + path = tmp_path / "i.sqlite3" + get_backend(path).close() + con = sqlite3.connect(path) + con.execute("ALTER TABLE patches DROP COLUMN n_dims") + con.commit() + con.close() + with pytest.raises(InvalidIndexError, match="missing columns"): + get_backend(path) + + def test_legacy_index_check_unopenable_path(self, tmp_path): + """A path that exists but cannot be opened as a file is not an index.""" + idx = DBDirectoryIndexer(tmp_path) + # opening a directory raises IsADirectoryError (an OSError), which the + # header probe suppresses before concluding it is not a legacy index. + sub = tmp_path / "adir" + sub.mkdir() + assert idx._is_legacy_or_foreign_index(sub) is False + + def test_schema_creation_rolls_back_on_failure(self, tmp_path): + """A failure while creating the schema rolls back and re-raises.""" + from dascore.io.index.lite import SQLiteBackend + + class _BoomBackend(SQLiteBackend): + def _execute(self, sql, params=()): + if "INSERT INTO meta_data" in sql: + raise RuntimeError("boom during schema init") + return super()._execute(sql, params) + + with pytest.raises(RuntimeError, match="boom during schema init"): + _BoomBackend(tmp_path / "i.sqlite3") + + def test_schema_commit_failure_rolls_back(self, tmp_path): + """A schema commit failure follows the protected rollback path.""" + from dascore.io.index.lite import SQLiteBackend + + class _CommitBoomBackend(SQLiteBackend): + rolled_back = False + + def _commit(self): + raise RuntimeError("boom during schema commit") + + def _rollback(self): + type(self).rolled_back = True + return super()._rollback() + + with pytest.raises(RuntimeError, match="schema commit"): + _CommitBoomBackend(tmp_path / "commit.sqlite3") + assert _CommitBoomBackend.rolled_back + + def test_reopen_missing_dynamic_attr_column(self, tmp_path): + """attr_meta referencing an absent attrs column is rejected on reopen.""" + import sqlite3 + + from dascore.exceptions import InvalidIndexError + + path = tmp_path / "i.sqlite3" + get_backend(path).close() + con = sqlite3.connect(path) + con.execute( + "INSERT INTO attr_meta (attr_name, value_kind, column_name, units) " + "VALUES ('ghost', 'num', 'ghost__num', NULL)" + ) + con.commit() + con.close() + with pytest.raises(InvalidIndexError, match="missing dynamic columns"): + get_backend(path) + + +class TestPureHelpers: + """Small pure helpers covered directly (least-contrived form).""" + + def test_is_directory_format_on_file(self, tmp_path): + """A plain file is never a directory scan unit.""" + from dascore.io.core import is_directory_format + + f = tmp_path / "a.txt" + f.write_text("x") + assert is_directory_format(f) is False + + def test_memory_backend_refuses_pickle(self): + """An in-memory backend cannot be pickled (owners serialize rows).""" + import pickle + + from dascore.io.index.lite import SQLiteBackend + + back = SQLiteBackend(":memory:") + try: + with pytest.raises(TypeError, match="cannot be pickled"): + pickle.dumps(back) + finally: + back.close() + + def test_py_scalar_bool_and_int(self): + """_py_scalar unwraps numpy bool/int to plain python scalars.""" + from dascore.io.index.ingest import _py_scalar + + assert _py_scalar(np.bool_(True)) is True + assert _py_scalar(np.int64(5)) == 5 + assert isinstance(_py_scalar(np.int64(5)), int) + + def test_assemble_records_empty_sources(self): + """No sources yields no records.""" + from dascore.io.index.ingest import assemble_source_records + + empty = pd.DataFrame() + assert assemble_source_records(empty, empty, empty, empty, empty, empty) == [] + + def test_units_compatible(self): + """_units_compatible is True for same dimensionality, False otherwise.""" + from dascore.io.index.backend import SQLIndexBackend + + assert SQLIndexBackend._units_compatible("m", "ft") is True + assert SQLIndexBackend._units_compatible("m", "s") is False + + def test_legacy_index_check_missing_path(self, tmp_path): + """A path that does not exist is not a legacy/foreign index.""" + idx = DBDirectoryIndexer(tmp_path) + assert idx._is_legacy_or_foreign_index(tmp_path / "nope.h5") is False + + def test_wrong_arity_coord_query_raises(self, backend): + """A hand-built coord range of the wrong length is rejected.""" + from dascore.exceptions import ParameterError + + with pytest.raises(ParameterError, match="length 2 sequence"): + backend.query(Query(coords={"distance": (1, 2, 3)})) + + def test_to_target_unit_paths(self): + """Quantity on a unitless target raises; on a unit target it converts.""" + from dascore.io.index.query import _to_target_unit + from dascore.units import get_quantity as _gq + + typed = typed_value(5 * m) # a numeric TypedValue carrying units + with pytest.raises(UnitError, match="unitless"): + _to_target_unit(typed, None, "distance") + # converting to a compatible unit returns a plain magnitude + out = _to_target_unit(typed, str(_gq("m").units), "distance") + assert out == pytest.approx(5.0) + + +class TestNormalizeSourcePatchId: + """The single source-patch-id normalizer handles every missing form.""" + + def test_missing_forms_become_empty(self): + """None, empty string, and pandas NaN/NaT all normalize to ''.""" + from dascore.core.summary import normalize_source_patch_id + + assert normalize_source_patch_id(None) == "" + assert normalize_source_patch_id("") == "" + assert normalize_source_patch_id(float("nan")) == "" + assert normalize_source_patch_id(np.nan) == "" + assert normalize_source_patch_id(pd.NaT) == "" + + def test_numpy_scalar_becomes_plain_string(self): + """A numpy scalar is unwrapped before stringifying.""" + from dascore.core.summary import normalize_source_patch_id + + assert normalize_source_patch_id(np.int64(42)) == "42" + + def test_plain_values_stringify(self): + """Ordinary ids pass through as strings.""" + from dascore.core.summary import normalize_source_patch_id + + assert normalize_source_patch_id("abc") == "abc" + assert normalize_source_patch_id(7) == "7" + + def test_non_scalar_falls_through(self): + """A value pd.isnull cannot evaluate as a scalar still stringifies.""" + from dascore.core.summary import normalize_source_patch_id + + # pd.isnull on a list returns an array (truth value is ambiguous), + # so the helper must swallow that and fall through to str(). + assert normalize_source_patch_id([1, 2]) == "[1, 2]" + + +class TestCanonicalRange: + """_canonical_range recognizes only numeric ranges.""" + + def test_bare_and_quantity_bounds(self): + """Bare numbers and quantities become SI magnitudes.""" + from dascore.io.index.catalog import _canonical_range + + assert _canonical_range((20, 60)).magnitudes == (20.0, 60.0) + # 20 m .. 60 m -> SI metres + assert _canonical_range((20 * m, 60 * m)).magnitudes == (20.0, 60.0) + + def test_open_bounds_kept(self): + """A half-open numeric range keeps its open end as None.""" + from dascore.io.index.catalog import _canonical_range + + assert _canonical_range((None, 60)).magnitudes == (None, 60.0) + + @pytest.mark.parametrize( + "value", + [ + np.array([True, False]), # boolean mask, not a range + (None, None), # fully open: no numeric content + (True, False), # bool bounds are not numeric ranges + ("a", "b"), # string bounds are not numeric ranges + (1, 2, 3), # wrong arity + ], + ) + def test_non_numeric_ranges_return_none(self, value): + """Anything that is not a bounded numeric range yields None.""" + from dascore.io.index.catalog import _canonical_range + + assert _canonical_range(value) is None + + +@pytest.fixture(scope="module") +def backend(tmp_path_factory): + """One SQLite backend with the contract summaries plus extras.""" + extra = PatchSummary( + attrs={ + "tag": "extra", + "trigger_time": np.datetime64("2024-06-01T00:00:00", "ns"), + "window": np.timedelta64(10, "s"), + }, + coords={ + "time": { + "dtype": "datetime64", + "min": np.datetime64("2024-06-01T00:00:00", "ns"), + "max": np.datetime64("2024-06-01T00:01:00", "ns"), + "dims": ("time",), + "len": 100, + }, + }, + dims=("time",), + shape=(100,), + dtype="float32", + source_path="extras/trigger.h5", + source_format="DASDAE", + source_version="1", + ) + path = tmp_path_factory.mktemp("edge") / "index.sqlite3" + back = get_backend(path) + back.write_sources(summaries_to_records([*make_summaries(), extra])) + yield back + back.close() + + +class TestAdaptAndBackendBasics: + """Small helpers and backend plumbing.""" + + def test_adapt_params_nan_becomes_none(self): + """NaN floats bind as NULL.""" + assert adapt_params([float("nan"), 1])[0] is None + + def test_bulk_insert_empty_rows_noop(self, tmp_path): + """Empty bulk inserts are no-ops.""" + back = get_backend(tmp_path / "insert.sqlite3") + back._bulk_insert("attr_meta", ("attr_name",), []) + back._executemany( + "INSERT INTO attr_meta VALUES (?, ?, ?, ?)", + [("a", "num", "a__num", None)], + ) + assert len(back._attr_meta()) == 1 + back.close() + + def test_write_failure_rolls_back(self, tmp_path): + """A failing write leaves the index unchanged.""" + back = get_backend(tmp_path / "rollback.sqlite3") + records = summaries_to_records(make_summaries()) + back.write_sources(records[:1]) + before = len(back.query()) + + def boom(*args, **kwargs): + raise RuntimeError("simulated failure") + + back._bulk_insert = boom + with pytest.raises(RuntimeError, match="simulated"): + back.write_sources(records[1:]) + del back.__dict__["_bulk_insert"] + assert len(back.query()) == before + back.close() + + @pytest.mark.parametrize("operation", ["write", "delete"]) + def test_commit_failure_rolls_back(self, tmp_path, monkeypatch, operation): + """Write and delete commit failures both release their transaction.""" + back = get_backend(tmp_path / f"{operation}.sqlite3") + records = summaries_to_records(make_summaries()) + initial = records[:1] if operation == "write" else records + back.write_sources(initial) + before = len(back.query()) + original_rollback = back._rollback + rolled_back = [] + + def rollback(): + rolled_back.append(True) + original_rollback() + + def commit_failure(): + raise RuntimeError("simulated commit failure") + + monkeypatch.setattr(back, "_rollback", rollback) + monkeypatch.setattr(back, "_commit", commit_failure) + with pytest.raises(RuntimeError, match="commit failure"): + if operation == "write": + back.write_sources(records[1:2]) + else: + back.delete_sources([records[0].source_path]) + assert rolled_back + assert len(back.query()) == before + back.close() + + def test_marker_commit_failure_rolls_back(self, tmp_path, monkeypatch): + """A marker commit failure restores the initial metadata value.""" + back = get_backend(tmp_path / "marker.sqlite3") + before = back.get_metadata()["last_indexed_ns"] + original_rollback = back._rollback + rolled_back = [] + + def rollback(): + rolled_back.append(True) + original_rollback() + + def commit_failure(): + raise RuntimeError("simulated commit failure") + + monkeypatch.setattr(back, "_rollback", rollback) + monkeypatch.setattr(back, "_commit", commit_failure) + with pytest.raises(RuntimeError, match="commit failure"): + back.mark_initial_update_done() + assert rolled_back + assert back.get_metadata()["last_indexed_ns"] == before + back.close() + + def test_delete_failure_rolls_back(self, tmp_path): + """A failing delete leaves the index unchanged.""" + back = get_backend(tmp_path / "delete.sqlite3") + back.write_sources(summaries_to_records(make_summaries())) + before = len(back.query()) + + def boom(paths, base_uri=""): + raise RuntimeError("simulated failure") + + back._delete_by_paths = boom + with pytest.raises(RuntimeError, match="simulated"): + back.delete_sources(["das/file_1.h5"]) + del back.__dict__["_delete_by_paths"] + assert len(back.query()) == before + back.close() + + def test_delete_no_paths_noop(self, backend): + """Deleting an empty path list does nothing.""" + before = len(backend.query()) + backend.delete_sources([]) + assert len(backend.query()) == before + + def test_flatten_skips_absent_columns(self, backend): + """attr_meta rows without a matching result column are skipped.""" + df = backend._fetch_df("SELECT patch_id FROM patches LIMIT 2") + out = backend._flatten(df, backend._attr_meta()) + assert len(out) == 2 + + def test_duration_attr_roundtrip(self, backend): + """dur-kind attrs come back as timedeltas.""" + df = backend.query(Query(attrs={"window": np.timedelta64(10, "s")})) + assert len(df) == 1 + assert pd.api.types.is_timedelta64_dtype(df["window"]) + + +class TestResolveQueryErrors: + """Explicit-namespace validation.""" + + def test_duplicate_explicit_namespace_raises(self, backend): + """A name cannot be supplied in both explicit namespaces.""" + with pytest.raises(InvalidSpoolQueryError, match="both _attrs and _coords"): + resolve_query( + backend, + _attrs={"distance": (0, 1)}, + _coords={"distance": (0, 1)}, + ) + + def test_unknown_attr_in_explicit_namespace(self, backend): + """Unknown key in _attrs raises.""" + with pytest.raises(InvalidSpoolQueryError, match="not an attribute"): + resolve_query(backend, _attrs={"nope": 1}) + + def test_unknown_coord_in_explicit_namespace(self, backend): + """Unknown key in _coords raises.""" + with pytest.raises(InvalidSpoolQueryError, match="not a coordinate"): + resolve_query(backend, _coords={"nope": (1, 2)}) + + +class TestQueryValueEdges: + """Coercion, kind-mismatch, and malformed-value behavior.""" + + def test_none_value_raises(self, backend): + """None is not a valid predicate value.""" + with pytest.raises(InvalidSpoolQueryError, match="Cannot use"): + backend.query(Query(attrs={"station": None})) + + def test_datetime_string_matches_time_attr(self, backend): + """A datetime-like string queries a time-kind attr.""" + df = backend.query(Query(attrs={"trigger_time": "2024-06-01T00:00:00"})) + assert list(df["tag"]) == ["extra"] + + def test_mixed_kind_range_raises(self, backend): + """Range bounds of different kinds raise.""" + with pytest.raises(InvalidSpoolQueryError, match="mixed kinds"): + backend.query(Query(attrs={"gauge_length": ("a", 5)})) + + def test_fully_open_range_raises(self, backend): + """A range with no usable bounds raises.""" + with pytest.raises(InvalidSpoolQueryError, match="no usable bounds"): + backend.query(Query(attrs={"gauge_length": (None, ...)})) + + def test_inverted_range_raises(self, backend): + """Lo > hi raises.""" + with pytest.raises(InvalidSpoolQueryError, match="lo > hi"): + backend.query(Query(attrs={"gauge_length": (5, 1)})) + + def test_unknown_attr_in_query_raises(self, backend): + """A Query naming an unknown attr raises at SQL build.""" + with pytest.raises(InvalidSpoolQueryError, match="not an attribute"): + backend.query(Query(attrs={"nope": 1})) + + def test_regex_on_non_str_attr_empty(self, backend): + """Regex against a numeric-only attr matches nothing.""" + df = backend.query(Query(attrs={"gauge_length": re.compile("x")})) + assert df.empty + + def test_range_kind_mismatch_empty(self, backend): + """A numeric range on a str-only attr matches nothing.""" + df = backend.query(Query(attrs={"station": (1, 2)})) + assert df.empty + + def test_membership_mixed_kinds(self, backend): + """Wrong-kind members are ignored; right-kind ones match.""" + df = backend.query(Query(attrs={"station": ["STA1", 5]})) + assert list(df["station"]) == ["STA1"] + + def test_membership_all_wrong_kind_empty(self, backend): + """All-wrong-kind membership matches nothing.""" + df = backend.query(Query(attrs={"station": [1, 2]})) + assert df.empty + + def test_glob_on_non_str_attr_empty(self, backend): + """Glob against a numeric-only attr matches nothing.""" + df = backend.query(Query(attrs={"gauge_length": "1*"})) + assert df.empty + + def test_boolean_array_coord_rejected(self, backend): + """Boolean sample masks are no longer index predicates.""" + mask = np.array([True, False, True]) + with pytest.raises(InvalidSpoolQueryError, match="range selectors"): + backend.query(Query(coords={"distance": mask})) + + def test_slice_range_form(self, backend): + """Slices resolve to the same range tuples patch selects accept.""" + lo = np.datetime64("2024-06-01T00:00:00", "ns") + query = resolve_query(backend, time=slice(lo, None)) + assert query.coords["time"] == (lo, None) + df = backend.query(query) + assert list(df["tag"]) == ["extra"] + + +class TestUnitHeterogeneity: + """Mixed unit populations across sources.""" + + @staticmethod + def _summary(path, attrs=None, coord_units=None): + """One summary with a 100-200 distance coord.""" + coord = { + "dtype": "float64", + "min": 100.0, + "max": 200.0, + "dims": ("distance",), + "len": 10, + } + if coord_units is not None: + coord["units"] = coord_units + return PatchSummary( + attrs=attrs or {"tag": "units"}, + coords={"distance": coord}, + dims=("distance",), + shape=(10,), + dtype="float32", + source_path=path, + source_format="DASDAE", + source_version="1", + ) + + def test_null_unit_coord_defs_stay_candidates(self, tmp_path): + """Quantity queries must not drop unitless coord defs (candidacy).""" + back = get_backend(tmp_path / "units.sqlite3") + back.write_sources( + summaries_to_records( + [ + self._summary("with_units.h5", coord_units="m"), + self._summary("no_units.h5"), + ] + ) + ) + meters = get_quantity("m") + df = back.query(Query(coords={"distance": (150 * meters, 300 * meters)})) + assert set(df["path"]) == {"with_units.h5", "no_units.h5"} + back.close() + + def test_nonoverlapping_unitless_coord_stays_candidate(self, tmp_path): + """Numeric envelopes cannot exclude a unitless quantity candidate.""" + back = get_backend(tmp_path / "nonoverlap-unitless.sqlite3") + back.write_sources( + summaries_to_records( + [ + self._summary("with_units.h5", coord_units="m"), + self._summary("no_units.h5"), + ] + ) + ) + meters = get_quantity("m") + df = back.query(Query(coords={"distance": (250 * meters, 300 * meters)})) + assert list(df["path"]) == ["no_units.h5"] + back.close() + + def test_all_unitless_quantity_query_keeps_candidates(self, tmp_path): + """Unitless defs cannot be proven incompatible; they stay candidates.""" + back = get_backend(tmp_path / "unitless.sqlite3") + back.write_sources(summaries_to_records([self._summary("no_units.h5")])) + meters = get_quantity("m") + df = back.query(Query(coords={"distance": (150 * meters, 300 * meters)})) + assert list(df["path"]) == ["no_units.h5"] + back.close() + + def test_incompatible_units_only_raises(self, tmp_path): + """When every def carries units and none are compatible, raise.""" + back = get_backend(tmp_path / "incompat.sqlite3") + back.write_sources( + summaries_to_records([self._summary("s.h5", coord_units="s")]) + ) + meters = get_quantity("m") + with pytest.raises(UnitError, match="no units compatible"): + back.query(Query(coords={"distance": (150 * meters, 300 * meters)})) + back.close() + + def test_incompatible_attr_units_warn_not_fail(self, tmp_path): + """One rogue attr dimension must not abort the whole index update.""" + summaries = [ + self._summary("a.h5", attrs={"resolution": 1.0 * get_quantity("m")}), + self._summary("b.h5", attrs={"resolution": 1.0 * get_quantity("s")}), + ] + back = get_backend(tmp_path / "attr_units.sqlite3") + with pytest.warns(UserWarning, match="incompatible"): + back.write_sources(summaries_to_records(summaries)) + # both patches indexed; only the incompatible value is skipped + df = back.query() + assert set(df["path"]) == {"a.h5", "b.h5"} + got = back.query(Query(attrs={"resolution": (0.5, 2.0)})) + assert list(got["path"]) == ["a.h5"] + back.close() + + +class TestExactNsFetch: + """Nullable ns-integer columns must never round through float64.""" + + # an epoch-ns value float64 rounds to ...768: exactness is observable + NS = 1_752_244_251_123_456_789 + + def _summary(self, path, coords, dims): + return PatchSummary( + attrs={"tag": "ns"}, + coords=coords, + dims=dims, + shape=tuple(10 for _ in dims), + dtype="float32", + source_path=path, + source_format="DASDAE", + source_version="1", + ) + + def test_time_envelopes_exact_when_column_nullable(self, tmp_path): + """NULLs from other kinds/patches must not degrade ns columns.""" + t0 = np.datetime64(self.NS, "ns") + with_time = self._summary( + "abs.h5", + { + "event_time": { + "dtype": "datetime64", + "min": t0, + "max": t0 + np.timedelta64(60, "s"), + "dims": ("event_time",), + "len": 10, + } + }, + ("event_time",), + ) + # a numeric coord puts NULL min_ns rows in the same link fetch, + # and no time coord leaves patches.time_min NULL for this patch. + numeric_only = self._summary( + "num.h5", + { + "distance": { + "dtype": "float64", + "min": 0.0, + "max": 10.0, + "dims": ("distance",), + "len": 10, + } + }, + ("distance",), + ) + back = get_backend(tmp_path / "exact.sqlite3") + back.write_sources(summaries_to_records([with_time, numeric_only])) + df = back.query().set_index("path") + got = df.loc["abs.h5", "event_time_min"] + assert pd.Timestamp(got).value == self.NS + back.close() + + def test_mtime_ns_exact_with_null_row(self, tmp_path): + """A single NULL mtime row must not corrupt the others (rescans).""" + back = get_backend(tmp_path / "mtime.sqlite3") + records = [ + SourceRecord( + source_path="a.h5", + source_format="", + format_version="", + mtime_ns=self.NS, + size_bytes=1, + ), + SourceRecord(source_path="b.h5", source_format="", format_version=""), + ] + back.write_sources(records) + sources = back.get_sources().set_index("source_path") + assert int(sources.loc["a.h5", "mtime_ns"]) == self.NS + back.close() + + def test_float_ns_column_rejected(self): + """The conversion helper refuses already-corrupted float input.""" + series = pd.Series([1.5e18, np.nan], name="min_ns") + with pytest.raises(TypeError, match="already corrupted"): + _ns_to_time(series, "datetime") + + +class TestIngestEdges: + """typed_value and record-building edge cases.""" + + def test_offset_quantity_attr_uses_full_conversion(self): + """Affine quantity attrs are converted with their offset.""" + out = typed_value(get_quantity("0 degC")) + assert out is not None + assert out.value == pytest.approx(273.15) + assert out.units == "K" + + def test_offset_coord_uses_full_conversion(self): + """Affine coord bounds use offsets while steps remain deltas.""" + summary = PatchSummary( + attrs={"tag": "temperature"}, + coords={ + "temperature": { + "dtype": "float64", + "min": 0.0, + "max": 100.0, + "step": 1.0, + "units": "degC", + "dims": ("temperature",), + "len": 101, + } + }, + dims=("temperature",), + shape=(101,), + dtype="float32", + source_path="temperature.h5", + source_format="DASDAE", + source_version="1", + ) + out = _coord_record("temperature", summary.coords["temperature"]) + assert out is not None + assert out.min_num == pytest.approx(273.15) + assert out.max_num == pytest.approx(373.15) + assert out.step_num == pytest.approx(1.0) + assert out.units == "K" + + def test_relative_root_requires_path_boundary(self): + """A similarly prefixed path is not made relative to the root.""" + data = make_summaries()[0].dump_structured() + data["source_path"] = "/data/foobar/file.h5" + record = s2r([PatchSummary(**data)], relative_to="/data/foo")[0] + assert record.source_path == "/data/foobar/file.h5" + + def test_plain_array_skipped(self): + """Arrays are complex attrs; skipped.""" + assert typed_value(np.array([1, 2])) is None + + def test_array_quantity_skipped(self): + """Array-valued quantities are skipped.""" + assert typed_value(np.array([1.0, 2.0]) * get_quantity("m")) is None + + def test_reserved_attr_name_warns(self): + """An attr named patch_id is skipped with a warning.""" + summary = PatchSummary( + attrs={"patch_id": 5, "tag": "x"}, + coords={ + "distance": { + "dtype": "float64", + "min": 0.0, + "max": 1.0, + "dims": ("distance",), + "len": 2, + } + }, + dims=("distance",), + shape=(2,), + dtype="float32", + source_path="a.h5", + source_format="DASDAE", + source_version="1", + ) + with pytest.warns(UserWarning, match="reserved attr name"): + records = s2r([summary]) + assert "patch_id" not in records[0].patches[0].attrs + + @pytest.mark.parametrize("dtype", ["", np.dtype(bool)]) + def test_unsupported_coord_dtype_skipped(self, dtype): + """A coord with a missing or unsupported dtype produces no record.""" + + class _Stub: + dims = ("x",) + len = 2 + units = None + fingerprint = None + min = 0 + max = 1 + step = None + + # "" exercises the missing-dtype guard; a bool dtype is a real + # dtype that none of the value-kind branches handle. + stub = _Stub() + stub.dtype = dtype + assert _coord_record("x", stub) is None + + def test_multipatch_source_gets_positional_ids(self): + """Multi-patch sources get positional source_patch_ids.""" + base = make_summaries()[0].dump_structured() + one = PatchSummary(**base) + two = PatchSummary(**{**base, "attrs": {"station": "STA9"}}) + records = s2r([one, two]) + assert len(records) == 1 + ids = [p.source_patch_id for p in records[0].patches] + assert ids == ["0", "1"] + + +class TestIndexerEdges: + """DBDirectoryIndexer edge behavior.""" + + def test_failed_initial_update_can_retry(self, tmp_path, monkeypatch): + """A new process retries when the first automatic update failed.""" + indexer = DBDirectoryIndexer(tmp_path) + index_path = indexer.index_path + + def fail_walk(): + raise OSError("simulated walk failure") + + monkeypatch.setattr(indexer, "_walk", fail_walk) + with pytest.raises(OSError, match="walk failure"): + indexer.ensure_updated() + assert not indexer._initial_update_done + indexer.close() + + retry = DBDirectoryIndexer(tmp_path, index_path=index_path) + assert not retry._initial_update_done + assert retry.ensure_updated() + assert retry._initial_update_done + retry.close() + + reopened = DBDirectoryIndexer(tmp_path, index_path=index_path) + assert reopened._initial_update_done + reopened.close() + + def test_directory_manifest_detects_equal_stat_name_swap( + self, tmp_path, monkeypatch + ): + """Equal-size, equal-mtime member replacement changes the signature.""" + unit = tmp_path / "unit" + unit.mkdir() + old = unit / "old.raw" + old.write_bytes(b"same-size") + stamp = 1_700_000_000_000_000_000 + os.utime(old, ns=(stamp, stamp)) + indexer = DBDirectoryIndexer(tmp_path) + monkeypatch.setattr(indexer, "_directory_format", lambda path: path == unit) + monkeypatch.setattr(dc, "scan", lambda *_args, **_kwargs: []) + indexer.update(progress=None) + before = tuple( + indexer._backend.source_stats().loc[0, ["mtime_ns", "size_bytes"]] + ) + + old.unlink() + new = unit / "new.raw" + new.write_bytes(b"same-size") + os.utime(new, ns=(stamp, stamp)) + indexer.update(progress=None) + after = tuple( + indexer._backend.source_stats().loc[0, ["mtime_ns", "size_bytes"]] + ) + assert before != after + indexer.close() + + def test_auto_update_on_first_query(self, tmp_path, random_patch): + """A brand-new index triggers one update on first query.""" + random_patch.io.write(tmp_path / "one.hdf5", "dasdae") + indexer = DBDirectoryIndexer(tmp_path) + assert len(indexer()) == 1 # no explicit update() call + + def test_empty_index_file_updates_on_first_query(self, tmp_path, random_patch): + """A pre-created empty SQLite path is still a new index.""" + random_patch.io.write(tmp_path / "one.hdf5", "dasdae") + index_path = tmp_path / "empty.sqlite3" + index_path.touch() + indexer = DBDirectoryIndexer(tmp_path, index_path=index_path) + assert len(indexer()) == 1 + + def test_directory_format_unit(self, tmp_path): + """Directory-format sources (xml binary) group as one scan unit.""" + import sys + + sys.path.insert(0, "tests/test_io/test_xml_binary") + from test_xml_binary import metadata + + sub = tmp_path / "unit" + sub.mkdir() + (sub / "metadata.xml").write_text(metadata) + # hidden files and subdirectories inside a unit are ignored + (sub / ".hidden_state").write_text("x") + (sub / "logs").mkdir() + rand = np.random.default_rng(0).random((5000, 10)).astype("float32") + for name in ( + "DAS_20240530T011500_000000Z.raw", + "DAS_20240530T011530_000000Z.raw", + ): + with (sub / name).open("wb") as fi: + rand.tofile(fi) + indexer = DBDirectoryIndexer(tmp_path).update(progress=None) + df = indexer() + assert len(df) == 2 + # unchanged: second update rescans nothing + before = indexer._backend.get_sources()["last_indexed_ns"].max() + indexer.update(progress=None) + after = indexer._backend.get_sources()["last_indexed_ns"].max() + assert before == after + indexer.close() + + +class TestDirSpoolPassthrough: + """Directory spools accept a prebuilt indexer.""" + + def test_spool_from_indexer(self, tmp_path, random_patch): + """Passing an indexer instance to from_directory works.""" + from dascore.core.spool import Spool + + random_patch.io.write(tmp_path / "one.hdf5", "dasdae") + indexer = DBDirectoryIndexer(tmp_path) + spool = Spool.from_directory(indexer).update(progress=None) + assert len(spool) == 1 + + +class TestFinalCoverage: + """Remaining edge branches.""" + + def test_datetime_object_becomes_time(self): + """A python datetime routes through the datetime fallback.""" + import datetime + + out = typed_value(datetime.datetime(2024, 1, 1)) + assert out is not None and out.kind == "time" + + def test_arbitrary_object_skipped(self): + """Unclassifiable objects are skipped.""" + + class _Odd: + """Not datetime-convertible, not a scalar.""" + + assert typed_value(_Odd()) is None + + +class TestCoordDeduplication: + """Coord summaries are stored once per unique definition.""" + + def test_shared_coord_stored_once(self, tmp_path): + """Identical distance coords across patches share one def row.""" + back = get_backend(tmp_path / "dedup.sqlite3") + back.write_sources(summaries_to_records(make_summaries())) + links = back._fetch_df("SELECT * FROM patch_coords") + defs = back._fetch_df("SELECT * FROM coord_defs") + assert len(defs) < len(links) + # das1 and das2 share an identical distance coord: one def, two links + dist_links = links[links["coord_name"] == "distance"] + das_defs = dist_links["coord_def_id"].value_counts() + assert (das_defs >= 2).any() + back.close() + + def test_defs_reused_across_writes(self, tmp_path): + """A second write with known coords creates no new defs.""" + back = get_backend(tmp_path / "reuse.sqlite3") + summaries = make_summaries() + back.write_sources(summaries_to_records(summaries[:1])) + n_defs = len(back._fetch_df("SELECT * FROM coord_defs")) + # das2 shares the distance def with das1; only time is new + back.write_sources(summaries_to_records(summaries[1:2])) + n_defs_after = len(back._fetch_df("SELECT * FROM coord_defs")) + assert n_defs_after == n_defs + 1 + back.close() + + def test_fingerprint_backed_defs(self, tmp_path): + """Summaries from real patches carry fingerprints into defs.""" + summary = PatchSummary.from_patch(dc.get_example_patch()) + structured = summary.dump_structured() + structured.update( + { + "source_path": "fp/one.h5", + "source_format": "DASDAE", + "source_version": "1", + } + ) + back = get_backend(tmp_path / "fp.sqlite3") + back.write_sources(summaries_to_records([PatchSummary(**structured)])) + defs = back._fetch_df("SELECT def_key, fingerprint FROM coord_defs") + assert defs["fingerprint"].notna().all() + assert defs["def_key"].str.startswith("fp:").all() + back.close() + + def test_irregular_coord_hashes_values(self): + """A non-range coordinate carries the hash of its complete array.""" + patch = dc.get_example_patch() + old = patch.get_coord("distance") + values = np.arange(len(old), dtype=float) + values[2:] += 0.5 + patch = patch.update_coords(distance=values) + summary = PatchSummary.from_patch(patch) + record = _coord_record("distance", summary.coords["distance"]) + assert record.coord_hash == patch.get_coord("distance").fingerprint() + assert record.def_key.startswith("fp:") + + def test_orphan_defs_tolerated(self, tmp_path): + """Deleting sources leaves defs behind without breaking queries.""" + back = get_backend(tmp_path / "orphan.sqlite3") + back.write_sources(summaries_to_records(make_summaries())) + n_defs = len(back._fetch_df("SELECT * FROM coord_defs")) + back.delete_sources(["das/file_1.h5", "das/file_2.h5"]) + assert len(back._fetch_df("SELECT * FROM coord_defs")) == n_defs + assert len(back.query()) == 2 + back.close() + + +class TestPivotEdge: + """Pivot with coord-less patches.""" + + def test_no_coords_patch(self, tmp_path): + """A patch with no coords pivots to nothing, without error.""" + summary = PatchSummary( + attrs={"tag": "bare"}, + coords={}, + dims=(), + shape=(), + dtype="float32", + source_path="bare.h5", + source_format="DASDAE", + source_version="1", + ) + back = get_backend(tmp_path / "bare.sqlite3") + back.write_sources(summaries_to_records([summary])) + df = back.query() + assert len(df) == 1 + assert not [c for c in df.columns if c.endswith("_def_key")] + back.close() + + +class TestCompositeSourceIdentity: + """Sources are identified by (base_uri, source_path).""" + + def test_same_path_different_base_coexist(self, tmp_path): + """Identical relative paths under different bases don't collide.""" + base = make_summaries()[0].dump_structured() + one = PatchSummary(**base) + records_a = summaries_to_records([one], base_uri="s3://bucket-a") + # base_uri strip only applies when paths share the base; set directly + records_a = [ + type(r)(**{**r.__dict__, "base_uri": "s3://bucket-a"}) for r in records_a + ] + records_b = [ + type(r)(**{**r.__dict__, "base_uri": "s3://bucket-b"}) for r in records_a + ] + back = get_backend(tmp_path / "multi.sqlite3") + back.write_sources(records_a) + back.write_sources(records_b) + df = back.query() + assert len(df) == 2 + prefixes = {p.split("/das/")[0] for p in df["path"]} + assert prefixes == {"s3://bucket-a", "s3://bucket-b"} + # deletion is base-scoped + back.delete_sources([records_a[0].source_path], base_uri="s3://bucket-a") + df = back.query() + assert len(df) == 1 + assert df["path"].iloc[0].startswith("s3://bucket-b") + back.close() + + def test_replacement_is_base_scoped(self, tmp_path): + """Rewriting a source under one base leaves the other base alone.""" + base = make_summaries()[0].dump_structured() + one = PatchSummary(**base) + rec = summaries_to_records([one])[0] + rec_a = type(rec)(**{**rec.__dict__, "base_uri": "s3://a"}) + rec_b = type(rec)(**{**rec.__dict__, "base_uri": "s3://b"}) + back = get_backend(tmp_path / "scoped.sqlite3") + back.write_sources([rec_a, rec_b]) + assert len(back.query()) == 2 + back.write_sources([rec_a]) # replace only the s3://a copy + assert len(back.query()) == 2 + back.close() + + +class TestResourceCleanup: + """Backends must not leak SQLite connections (review P2).""" + + def test_gc_closes_connection(self): + """Garbage collection closes the backend connection silently.""" + import gc + import warnings as warnings_mod + + import dascore as dc + + spool = dc.spool([dc.get_example_patch()]) + spool.get_contents() # realize the backend + with warnings_mod.catch_warnings(record=True) as caught: + warnings_mod.simplefilter("always", ResourceWarning) + del spool + gc.collect() + resource = [w for w in caught if issubclass(w.category, ResourceWarning)] + assert not resource + + def test_explicit_close_idempotent(self): + """Explicit close works and GC afterwards stays quiet.""" + import dascore as dc + from dascore.io.index.catalog import PatchCatalog + + catalog = PatchCatalog.from_patches([dc.get_example_patch()]) + catalog.to_df() + catalog.close() + + def test_finalization_from_worker_thread(self): + """Dropping the last backend reference off-thread does not raise.""" + import gc + import sys + import threading + + from dascore.io.index.lite import SQLiteBackend + + holder = [SQLiteBackend(":memory:")] + unraisable = [] + original = sys.unraisablehook + + def _drop(): + holder.clear() + gc.collect() + + sys.unraisablehook = unraisable.append + try: + worker = threading.Thread(target=_drop) + worker.start() + worker.join() + gc.collect() + finally: + sys.unraisablehook = original + errors = [u for u in unraisable if u.exc_type is sqlite3.ProgrammingError] + assert not errors + + +class TestLegacyIndexMap: + """Index-map entries pointing at retired .h5 indexes are bypassed.""" + + def test_legacy_entry_ignored(self, tmp_path): + """A mapped legacy HDF5 index does not break index creation.""" + import h5py + + import dascore as dc + from dascore.io.index.indexer import ( + _update_index_map, + ) + + # data directory with one file, plus a fake legacy index mapping. + data_dir = tmp_path / "data" + data_dir.mkdir() + dc.write(dc.get_example_patch(), data_dir / "a.h5", "dasdae") + legacy = tmp_path / "legacy_index.h5" + with h5py.File(legacy, "w") as fh: + fh.create_dataset("x", data=[1, 2, 3]) + map_path = tmp_path / "cache_paths.json" + with dc.set_config(directory_index_map_path=map_path): + _update_index_map({str(data_dir): str(legacy)}, cache_path=str(map_path)) + spool = dc.spool(data_dir).update() + assert len(spool) == 1 + + +class TestReservedAttrNames: + """Attrs colliding with structural columns are skipped with a warning.""" + + @pytest.mark.parametrize( + "name,value", + [("path", "user-path"), ("file_format", "attr-format"), ("source_id", 42)], + ) + def test_reserved_attr_warns_and_skips(self, name, value): + """Reserved names warn at ingest and never corrupt the relation.""" + import dascore as dc + + patch = dc.get_example_patch().update_attrs(**{name: value}) + with pytest.warns(UserWarning, match="reserved attr"): + df = dc.spool([patch]).get_contents() + assert not df.columns.duplicated().any() + # structural values win; the attr stays on the patch itself. + assert patch.attrs[name] == value + + def test_non_reserved_attr_round_trips(self): + """Ordinary arbitrary attrs still index and select normally.""" + import dascore as dc + + patch = dc.get_example_patch().update_attrs(experiment="exp42") + spool = dc.spool([patch]) + assert spool.get_contents()["experiment"].iloc[0] == "exp42" + assert len(spool.select(experiment="exp42")) == 1 + + def test_cross_patch_envelope_attr_reserved(self): + """An envelope-shaped attr is reserved even against another patch's coord.""" + import dascore as dc + + base = dc.get_example_patch() + # construction-time attrs bypass the update_attrs coord-name guard. + p1 = dc.Patch( + data=base.data, + coords=base.coords, + dims=base.dims, + attrs=dict(base.attrs) | {"event_time_min": 123.0}, + ) + p2 = dc.get_example_patch().rename_coords(time="event_time") + with pytest.warns(UserWarning, match="reserved attr name 'event_time_min'"): + spool = dc.spool([p1, p2]) + df = spool.get_contents() + backend = spool._catalog.backend + assert "event_time_min" not in backend.attr_names() + # the column is the coordinate envelope, not the stray attr value. + assert "event_time_min" in df.columns + assert 123.0 not in set(df["event_time_min"].dropna()) + + def test_data_units_attr_still_indexed(self): + """A real ``*_units`` attr is not an envelope column; stays queryable.""" + import dascore as dc + + patch = dc.get_example_patch().update_attrs(data_units="strain") + spool = dc.spool([patch]) + assert "data_units" in spool._catalog.backend.attr_names() + + +class TestTransactionIsolation: + """The statement lock covers whole transactions (round-4 F5).""" + + def test_reader_never_sees_half_written_replacement(self, tmp_path): + """A concurrent reader blocks during a source replacement.""" + import threading + + from dascore.io.index.backend import get_backend + from dascore.io.index.ingest import SourceRecord, patch_record + + patch = dc.get_example_patch() + record = SourceRecord( + source_path="mem://one", + source_format="mem", + format_version="", + patches=(patch_record(patch.summary),), + ) + backend = get_backend(str(tmp_path / "idx.sqlite3")) + backend.write_sources([record]) + assert len(backend.query()) == 1 + + in_delete = threading.Event() + release = threading.Event() + original = type(backend)._delete_by_paths + + def paused_delete(self, *args, **kwargs): + out = original(self, *args, **kwargs) + in_delete.set() + release.wait(timeout=10) + return out + + counts = [] + + def read(): + counts.append(len(backend.query())) + + writer = threading.Thread( + target=lambda: backend.write_sources([record]), daemon=True + ) + type(backend)._delete_by_paths = paused_delete + try: + writer.start() + assert in_delete.wait(timeout=10) + # the writer sits mid-transaction with the row deleted; a + # reader must block on the transaction lock, not observe it + reader = threading.Thread(target=read, daemon=True) + reader.start() + reader.join(timeout=0.3) + assert reader.is_alive(), "reader observed a half-written state" + assert counts == [] + release.set() + writer.join(timeout=10) + reader.join(timeout=10) + finally: + type(backend)._delete_by_paths = original + release.set() + assert counts == [1] + backend.close() diff --git a/tests/test_io/test_index/test_ordering.py b/tests/test_io/test_index/test_ordering.py new file mode 100644 index 000000000..cd14e417d --- /dev/null +++ b/tests/test_io/test_index/test_ordering.py @@ -0,0 +1,334 @@ +""" +Tests for the catalog ordering contract (source ordinals). + +Patch rows present in (ordinal, patch_id) order: live spools keep +construction order, unions concatenate (dedup keeps first-occurrence +position), and directory archives present in time order (the syncer +renumbers after each sync). +""" + +from __future__ import annotations + +import numpy as np +import pytest + +import dascore as dc + + +@pytest.fixture(scope="module") +def three_patches(): + """Three time-contiguous example patches.""" + t0 = np.datetime64("2020-01-01", "ns") + p1 = dc.get_example_patch(time_min=t0) + time = p1.get_coord("time") + p2 = dc.get_example_patch(time_min=time.max() + time.step) + p3 = dc.get_example_patch(time_min=p2.get_coord("time").max() + time.step) + return p1, p2, p3 + + +class TestLiveSpoolOrder: + """Patch-list spools keep construction order on every path.""" + + def test_out_of_time_order_kept(self, three_patches): + """Construction order wins even when it disagrees with time order.""" + p1, p2, p3 = three_patches + spool = dc.spool([p3, p1, p2]) + # the tuple fast path + loaded = list(spool) + assert loaded[0] is p3 and loaded[1] is p1 and loaded[2] is p2 + # the catalog relation presents the same order + df = spool.get_contents() + expected = [p.get_coord("time").min() for p in (p3, p1, p2)] + assert list(df["time_min"]) == expected + # and indexing after realization still agrees + assert spool[0] is p3 + + def test_selection_preserves_relative_order(self, three_patches): + """A narrowed view keeps the surviving rows in spool order.""" + p1, p2, p3 = three_patches + spool = dc.spool([p3, p1, p2]) + t3 = p3.get_coord("time").min() + t1 = p1.get_coord("time").min() + selected = spool.select(time=(min(t1, t3), None)) + df = selected.get_contents() + assert list(df["time_min"])[:2] == [t3, t1] + + +class TestUnionOrder: + """Combined spools are list concatenation, deduped dict-merge style.""" + + def test_concatenation_order(self, three_patches): + """(a + b) presents a's rows then b's.""" + p1, p2, p3 = three_patches + combined = dc.spool([p3]) + dc.spool([p1, p2]) + loaded = list(combined) + assert loaded[0] is p3 and loaded[1] is p1 and loaded[2] is p2 + + def test_dedup_keeps_first_position(self, three_patches): + """A patch in both members keeps its first position, appears once.""" + p1, p2, p3 = three_patches + combined = dc.spool([p1, p2]) + dc.spool([p2, p3]) + assert len(combined) == 3 + loaded = list(combined) + assert loaded[0] is p1 and loaded[1] is p2 and loaded[2] is p3 + + def test_union_of_union_order(self, three_patches): + """Order survives a second union (export/re-ingest round trip).""" + p1, p2, p3 = three_patches + combined = (dc.spool([p3]) + dc.spool([p2])) + dc.spool([p1]) + loaded = list(combined) + assert loaded[0] is p3 and loaded[1] is p2 and loaded[2] is p1 + + +class TestDirectoryOrder: + """File archives present in time order, maintained across syncs.""" + + def test_time_order_disagrees_with_name_order(self, tmp_path): + """Presentation follows patch time, not file names or walk order.""" + t0 = np.datetime64("2020-01-01", "ns") + early = dc.get_example_patch(time_min=t0) + late = dc.get_example_patch(time_min=t0 + np.timedelta64(3600, "s")) + dc.write(late, tmp_path / "a_late.h5", "dasdae") + dc.write(early, tmp_path / "z_early.h5", "dasdae") + spool = dc.spool(tmp_path).update(progress=None) + df = spool.get_contents() + assert df["time_min"].is_monotonic_increasing + + def test_update_interleaves_new_files_by_time(self, tmp_path): + """A later-indexed but earlier-in-time file sorts into place.""" + t0 = np.datetime64("2020-01-01", "ns") + mid = dc.get_example_patch(time_min=t0 + np.timedelta64(1800, "s")) + late = dc.get_example_patch(time_min=t0 + np.timedelta64(3600, "s")) + dc.write(mid, tmp_path / "mid.h5", "dasdae") + dc.write(late, tmp_path / "late.h5", "dasdae") + dc.spool(tmp_path).update(progress=None) # build the index + early = dc.get_example_patch(time_min=t0) + dc.write(early, tmp_path / "early.h5", "dasdae") + updated = dc.spool(tmp_path).update(progress=None) + df = updated.get_contents() + assert len(df) == 3 + assert df["time_min"].is_monotonic_increasing + assert df["time_min"].iloc[0] == early.get_coord("time").min() + + +class TestIndexVersionRebuild: + """Old-version index files rebuild automatically (disposable cache).""" + + def test_version_mismatch_rebuilds(self, tmp_path): + """An index of another schema version is replaced, not fatal.""" + import sqlite3 + + patch = dc.get_example_patch() + dc.write(patch, tmp_path / "a.h5", "dasdae") + spool = dc.spool(tmp_path).update(progress=None) + index_path = spool.indexer.index_path + spool.indexer.close() + # simulate an index written by another (older/newer) schema version; + # close the connection explicitly (the sqlite3 context manager only + # manages transactions) or Windows cannot unlink the file below. + con = sqlite3.connect(index_path) + try: + con.execute("UPDATE meta_data SET index_version = 1") + con.commit() + finally: + con.close() + reopened = dc.spool(tmp_path).update(progress=None) + assert len(reopened) == 1 + reopened.indexer.close() + + def test_indexer_deepcopy_shares_instance(self, tmp_path): + """Derived spools share the indexer (and its live DB connection).""" + import copy + + dc.write(dc.get_example_patch(), tmp_path / "a.h5", "dasdae") + spool = dc.spool(tmp_path).update(progress=None) + assert copy.deepcopy(spool.indexer) is spool.indexer + spool.indexer.close() + + +class TestSortNonHotCoords: + """Sorting by coords without patches-table columns (2026-07-18 F3).""" + + @pytest.fixture() + def renamed_spool(self): + """Two patches whose time coord is renamed (not a hot column).""" + p = dc.get_example_patch().rename_coords(time="event_time") + t = p.get_coord("event_time") + span = t.max() - t.min() + t.step + p2 = p.update_coords(event_time=t.data + span) + return dc.spool([p2, p]) # deliberately out of order + + @pytest.mark.parametrize("key", ["event_time", "event_time_min"]) + def test_sort_renamed_datetime_coord(self, renamed_spool, key): + """A renamed datetime coord sorts through coord_defs.""" + srt = renamed_spool.sort(key) + mins = [x.get_coord("event_time").min() for x in srt] + assert mins == sorted(mins) + # the realized relation agrees + contents = srt.get_contents() + assert contents["event_time_min"].is_monotonic_increasing + + def test_sort_non_hot_numeric_coord(self): + """A numeric aux coord sorts through coord_defs.""" + p = dc.get_example_patch() + n = p.shape[p.get_axis("distance")] + lo = p.update_coords(sensor=("distance", np.arange(n, dtype=float))) + hi = p.update_coords(sensor=("distance", np.arange(n, dtype=float) + 1000)) + srt = dc.spool([hi, lo]).sort("sensor") + mins = [x.get_coord("sensor").min() for x in srt] + assert mins == sorted(mins) + + def test_sort_string_coord(self): + """A string coord sorts lexicographically through coord_defs.""" + p = dc.get_example_patch() + n = p.shape[p.get_axis("distance")] + pa = p.update_coords(station=("distance", np.array(["a"] * n))) + pb = p.update_coords(station=("distance", np.array(["b"] * n))) + srt = dc.spool([pb, pa]).sort("station") + firsts = [x.get_coord("station").values[0] for x in srt] + assert firsts == ["a", "b"] + + def test_hot_coords_still_sort(self): + """time/distance keep the cached patches-column path.""" + p = dc.get_example_patch() + t = p.get_coord("time") + p2 = p.update_coords(time_min=t.max() + t.step) + srt = dc.spool([p2, p]).sort("time") + mins = [x.get_coord("time").min() for x in srt] + assert mins == sorted(mins) + + +class TestInterleavedSourceOrder: + """Directory time order across interleaved multi-patch files (F4).""" + + def test_multi_patch_file_straddles_another(self, tmp_path): + """A patch between two patches of another file presents in order.""" + p0 = dc.get_example_patch() + t = p0.get_coord("time") + span = t.max() - t.min() + t.step + p1 = p0.update_coords(time_min=t.min() + span) + p2 = p0.update_coords(time_min=t.min() + 2 * span) + dc.write(dc.spool([p0, p2]), tmp_path / "a.h5", "DASDAE") + dc.write(p1, tmp_path / "b.h5", "DASDAE") + spool = dc.spool(tmp_path).update(progress=None) + contents = spool.get_contents() + assert contents["time_min"].is_monotonic_increasing + mins = [x.get_coord("time").min() for x in spool] + assert mins == sorted(mins) + # windows and sorting stay consistent with the presentation + assert spool[1:2][0].get_coord("time").min() == mins[1] + assert spool.sort("distance").get_contents().shape[0] == 3 + + def test_default_order_is_not_view_state(self, tmp_path): + """The presentation contract does not make a root a view.""" + dc.write(dc.get_example_patch(), tmp_path / "a.h5", "DASDAE") + spool = dc.spool(tmp_path).update(progress=None) + assert spool.update() is not None # root update allowed + + +class TestInterruptedInitialUpdate: + """The initial-update marker only sets after renumbering (F6).""" + + def test_interruption_before_renumber_recovers(self, tmp_path): + """A crash after write_sources still renumbers on the retry.""" + from dascore.io.index.indexer import DBDirectoryIndexer + + p0 = dc.get_example_patch() + t = p0.get_coord("time") + span = t.max() - t.min() + t.step + late = p0.update_coords(time_min=t.min() + span) + # walk order (file names) disagrees with time order + dc.write(late, tmp_path / "a_late.h5", "DASDAE") + dc.write(p0, tmp_path / "b_early.h5", "DASDAE") + + class _InterruptedError(RuntimeError): + pass + + indexer = DBDirectoryIndexer(tmp_path) + original = type(indexer._backend).renumber_ordinals_by_time + + def _boom(self): + raise _InterruptedError + + type(indexer._backend).renumber_ordinals_by_time = _boom + try: + with pytest.raises(_InterruptedError): + indexer.ensure_updated() + finally: + type(indexer._backend).renumber_ordinals_by_time = original + del indexer # simulate the process dying after write_sources + + # a fresh open must not treat the interrupted update as done + spool = dc.spool(tmp_path).update(progress=None) + mins = list(spool.get_contents()["time_min"]) + assert mins == sorted(mins) + sources = spool._catalog.backend.get_sources() + by_ordinal = sources.sort_values("ordinal")["source_path"].tolist() + assert [p.split("/")[-1] for p in by_ordinal] == [ + "b_early.h5", + "a_late.h5", + ] + + +class TestDefaultOrderThroughUnion: + """Directory presentation order survives combining (round-5).""" + + @pytest.fixture() + def interleaved_dir_spool(self, tmp_path): + """A directory whose multi-patch file straddles another file.""" + p0 = dc.get_example_patch() + t = p0.get_coord("time") + span = t.max() - t.min() + t.step + p1 = p0.update_coords(time_min=t.min() + span) + p2 = p0.update_coords(time_min=t.min() + 2 * span) + dc.write(dc.spool([p0, p2]), tmp_path / "a.h5", "DASDAE") + dc.write(p1, tmp_path / "b.h5", "DASDAE") + return dc.spool(tmp_path).update(progress=None) + + def test_empty_union_keeps_order_and_equality(self, interleaved_dir_spool): + """Adding an empty spool preserves contents, order, and equality.""" + source = interleaved_dir_spool + combined = source + dc.spool([]) + want = [x.get_coord("time").min() for x in source] + got = [x.get_coord("time").min() for x in combined] + assert got == want + assert combined == source + + def test_live_append_keeps_directory_prefix(self, interleaved_dir_spool): + """A live operand appends after the directory's presented rows.""" + source = interleaved_dir_spool + later = dc.get_example_patch(time_min="2030-01-01") + combined = source + dc.spool([later]) + mins = [x.get_coord("time").min() for x in combined] + assert mins[:3] == [x.get_coord("time").min() for x in source] + assert len(mins) == 4 + + def test_non_interleaved_union_still_dedups(self, tmp_path): + """Ordinary archives keep record-grain transfer and dedup.""" + p0 = dc.get_example_patch() + t = p0.get_coord("time") + span = t.max() - t.min() + t.step + dc.write(p0, tmp_path / "a.h5", "DASDAE") + dc.write(p0.update_coords(time_min=t.min() + span), tmp_path / "b.h5", "DASDAE") + spool = dc.spool(tmp_path).update(progress=None) + assert len(spool + spool) == len(spool) + + +class TestMissingTimeSortsLast: + """Rows without a value sort last under any order (round-5).""" + + def test_directory_no_time_patch_presents_last(self, tmp_path): + """A distance-only patch follows every time-bearing patch.""" + timed = dc.get_example_patch().update_attrs(tag="time") + no_time = timed.mean("time").squeeze().update_attrs(tag="no_time") + dc.write(timed, tmp_path / "a_time.h5", "DASDAE") + dc.write(no_time, tmp_path / "b_no_time.h5", "DASDAE") + spool = dc.spool(tmp_path).update(progress=None) + assert list(spool.get_contents()["tag"]) == ["time", "no_time"] + + def test_sort_puts_missing_values_last(self): + """Explicit sort also presents value-less rows last.""" + timed = dc.get_example_patch().update_attrs(tag="a_time") + no_time = timed.mean("time").squeeze().update_attrs(tag="b_no_time") + spool = dc.spool([no_time, timed]).sort("time") + assert [x.attrs.tag for x in spool] == ["a_time", "b_no_time"] diff --git a/tests/test_io/test_index/test_plan.py b/tests/test_io/test_index/test_plan.py new file mode 100644 index 000000000..a8db642a3 --- /dev/null +++ b/tests/test_io/test_index/test_plan.py @@ -0,0 +1,664 @@ +"""Tests for the chunk planner (chunking formalities spec).""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +import dascore as dc +from dascore.exceptions import ( + ChunkError, + CoordMergeError, + InvalidSpoolQueryError, + ParameterError, +) +from dascore.io.index.catalog import PatchCatalog +from dascore.utils.chunk_plan import ChunkPlan, build_chunk_plan +from dascore.utils.time import to_timedelta64 + +ONE_S = np.timedelta64(1, "s") + + +def _flat(patches) -> pd.DataFrame: + """Get the flat relation for a list of patches.""" + return PatchCatalog.from_patches(list(patches)).to_df() + + +@pytest.fixture(scope="module") +def random_flat() -> pd.DataFrame: + """Flat relation of the contiguous random_das example spool.""" + return _flat(dc.get_example_spool("random_das")) + + +@pytest.fixture(scope="module") +def diverse_flat() -> pd.DataFrame: + """Flat relation of the diverse example spool.""" + return _flat(dc.get_example_spool("diverse_das")) + + +class TestValidation: + """Parameter validation per the spec errors table.""" + + def test_no_kwargs_raises(self, random_flat): + """Exactly one dimension kwarg is required.""" + with pytest.raises(ParameterError, match="one dimension"): + build_chunk_plan(random_flat) + + def test_two_kwargs_raise(self, random_flat): + """Two chunk kwargs raise.""" + with pytest.raises(ParameterError, match="one dimension"): + build_chunk_plan(random_flat, time=10, distance=10) + + def test_non_positive_value_raises(self, random_flat): + """Chunk lengths must be positive.""" + with pytest.raises(ParameterError, match="greater than 0"): + build_chunk_plan(random_flat, time=0) + + def test_merge_mode_forbids_overlap(self, random_flat): + """Merge mode does not accept overlap/keep_partial.""" + with pytest.raises(ParameterError, match="merging"): + build_chunk_plan(random_flat, time=None, overlap=1) + with pytest.raises(ParameterError, match="merging"): + build_chunk_plan(random_flat, time=..., keep_partial=True) + + def test_overlap_ge_length_raises(self, random_flat): + """D6: overlap >= length raises cleanly.""" + with pytest.raises(ParameterError, match="overlap"): + build_chunk_plan(random_flat, time=2, overlap=2) + + def test_unknown_group_raises(self, random_flat): + """Explicit group names must exist somewhere in the spool.""" + with pytest.raises(InvalidSpoolQueryError, match="bob"): + build_chunk_plan(random_flat, time=None, group=("bob",)) + + def test_unknown_dim_raises(self, random_flat): + """Chunking a dimension no patch has raises.""" + with pytest.raises(ChunkError, match="quelle"): + build_chunk_plan(random_flat, quelle=10) + + def test_bad_missing_dim_raises(self, random_flat): + """missing_dim accepts only raise/drop.""" + with pytest.raises(ParameterError, match="missing_dim"): + build_chunk_plan(random_flat, time=None, missing_dim="bob") + + +class TestMergePlan: + """Merge-mode planning on contiguous data.""" + + def test_contiguous_spool_single_output(self, random_flat): + """A contiguous spool merges to one output.""" + plan = build_chunk_plan(random_flat, time=None) + assert isinstance(plan, ChunkPlan) + assert plan.merge_mode + assert len(plan.outputs) == 1 + out = plan.outputs.iloc[0] + assert out["time_min"] == random_flat["time_min"].min() + assert out["time_max"] == random_flat["time_max"].max() + # every source patch appears exactly once as a member + assert len(plan.members) == len(random_flat) + assert set(plan.members["_patch_id"]) == set(random_flat["_patch_id"]) + + def test_members_unmodified_when_contiguous(self, random_flat): + """Contiguous members load whole (no trims).""" + plan = build_chunk_plan(random_flat, time=None) + assert not plan.members["_modified"].any() + + def test_diverse_partitions(self, diverse_flat): + """The diverse spool partitions by identity attrs, never raising.""" + plan = build_chunk_plan(diverse_flat, time=None) + assert len(plan.outputs) > 1 + # every output's members share that output's group attr values + merged = plan.members.merge( + diverse_flat[["_patch_id", "network", "station", "tag"]], + on="_patch_id", + ).merge( + plan.outputs[["output_id", "network", "station", "tag"]], + on="output_id", + suffixes=("_src", "_out"), + ) + for col in ("network", "station", "tag"): + src, out = merged[f"{col}_src"], merged[f"{col}_out"] + equal = (src == out) | (src.isnull() & out.isnull()) + assert equal.all() + + def test_gap_splits_partition(self): + """A gap larger than tolerance yields separate outputs.""" + t0 = np.datetime64("2020-01-01", "ns") + p1 = dc.get_example_patch(time_min=t0) + time = p1.get_coord("time") + dt = time.step + span = time.max() - time.min() + dt + p2 = dc.get_example_patch(time_min=time.max() + dt) # contiguous + p3 = dc.get_example_patch(time_min=time.max() + span + 10 * dt) + plan = build_chunk_plan(_flat([p1, p2, p3]), time=None) + assert len(plan.outputs) == 2 + + def test_plan_records_params(self, random_flat): + """Plans record resolved parameters, not config references.""" + with dc.set_config(sampling_group_tolerance=0.02): + plan = build_chunk_plan(random_flat, time=None) + assert plan.params["sampling_group_tolerance"] == 0.02 + assert isinstance(plan.params["group"], tuple) + + +class TestSegmentPlan: + """Segment-mode planning.""" + + def test_intervals_cover_envelope(self, random_flat): + """Chunk segments tile the envelope with the requested length.""" + plan = build_chunk_plan(random_flat, time=3) + out = plan.outputs + lengths = (out["time_max"] - out["time_min"]) + out["time_step"] + expected = to_timedelta64(3) + assert (abs(lengths - expected) <= out["time_step"]).all() + + def test_members_reference_real_patches(self, random_flat): + """All members point at rows of the input relation.""" + plan = build_chunk_plan(random_flat, time=3) + assert set(plan.members["_patch_id"]) <= set(random_flat["_patch_id"]) + # member trims stay within their output's envelope + joined = plan.members.merge( + plan.outputs[["output_id", "time_min", "time_max"]], + on="output_id", + suffixes=("", "_out"), + ) + assert (joined["time_min"] >= joined["time_min_out"]).all() + assert (joined["time_max"] <= joined["time_max_out"]).all() + + def test_too_short_partition_skipped(self): + """D8: partitions shorter than the length are skipped silently.""" + t0 = np.datetime64("2020-01-01", "ns") + p1 = dc.get_example_patch(time_min=t0) # ~8s long + time = p1.get_coord("time") + gap = time.max() - time.min() + 100 * ONE_S + p2 = dc.get_example_patch(time_min=time.min() + gap) + df = _flat([p1, p2]) + plan = build_chunk_plan(df, time=5) + assert len(plan.outputs) == 2 # one 5s chunk per 8s partition + with pytest.raises(ChunkError, match="sufficient length"): + build_chunk_plan(df, time=100) + + def test_overlap(self, random_flat): + """Overlapping chunks step by length minus overlap.""" + plan = build_chunk_plan(random_flat, time=4, overlap=2) + starts = plan.outputs["time_min"].sort_values().values + strides = np.diff(starts) + assert (abs(strides - to_timedelta64(2)) <= to_timedelta64(0.01)).all() + + def test_middle_value_step(self): + """D7: the partition step is the middle value of member steps.""" + t0 = np.datetime64("2020-01-01", "ns") + p1 = dc.get_example_patch(time_min=t0) + time = p1.get_coord("time") + p2 = dc.get_example_patch(time_min=time.max() + time.step) + plan = build_chunk_plan(_flat([p1, p2]), time=None) + assert plan.outputs["time_step"].iloc[0] == time.step + + +class TestMissingDim: + """Spec section 7 (D2): patches lacking the chunk dim.""" + + @pytest.fixture() + def flat_with_null(self, random_flat): + """A flat relation with one null time envelope.""" + df = random_flat.copy() + df.loc[df.index[0], ["time_min", "time_max"]] = (pd.NaT, pd.NaT) + return df + + def test_raise_by_default(self, flat_with_null): + """Null chunk-dim envelopes raise by default.""" + with pytest.raises(ChunkError, match="missing_dim"): + build_chunk_plan(flat_with_null, time=None) + + def test_drop_opt_in(self, flat_with_null): + """missing_dim='drop' excludes the offending rows.""" + plan = build_chunk_plan(flat_with_null, time=None, missing_dim="drop") + dropped = flat_with_null["_patch_id"].iloc[0] + assert dropped not in set(plan.members["_patch_id"]) + + +class TestConflict: + """Spec 2.5: attr policing within a partition.""" + + @pytest.fixture(scope="class") + def conflicted_patches(self): + """Two contiguous patches with a differing non-group attr.""" + t0 = np.datetime64("2020-01-01", "ns") + p1 = dc.get_example_patch(time_min=t0) + time = p1.get_coord("time") + p2 = dc.get_example_patch(time_min=time.max() + time.step) + p2 = p2.update_attrs(data_units="m/s") + return [p1, p2] + + def test_raise(self, conflicted_patches): + """Differing non-group attrs raise by default.""" + with pytest.raises(CoordMergeError, match="data_units"): + build_chunk_plan(_flat(conflicted_patches), time=None) + + def test_keep_first(self, conflicted_patches): + """keep_first carries the first member's value.""" + df = _flat(conflicted_patches) + plan = build_chunk_plan(df, time=None, conflict="keep_first") + first_id = df.sort_values("time_min")["_patch_id"].iloc[0] + expected = df.loc[df["_patch_id"] == first_id, "data_units"].iloc[0] + assert plan.outputs["data_units"].iloc[0] == expected + + def test_drop(self, conflicted_patches): + """Drop omits the conflicting attr from outputs.""" + plan = build_chunk_plan(_flat(conflicted_patches), time=None, conflict="drop") + assert "data_units" not in plan.outputs.columns + + def test_unknown_policy_raises(self, conflicted_patches): + """A misspelled conflict policy cannot silently behave like drop.""" + with pytest.raises(ParameterError, match="conflict must be"): + build_chunk_plan(_flat(conflicted_patches), time=None, conflict="keep_fist") + + +class TestGroupParameter: + """Group attrs partition instead of raising.""" + + @pytest.fixture(scope="class") + def two_station_flat(self): + """Contiguous patches from two stations.""" + t0 = np.datetime64("2020-01-01", "ns") + p1 = dc.get_example_patch(time_min=t0) + time = p1.get_coord("time") + p2 = dc.get_example_patch(time_min=time.max() + time.step) + p3 = p1.update_attrs(station="XX2") + p4 = p2.update_attrs(station="XX2") + return _flat([p1, p2, p3, p4]) + + def test_station_partitions(self, two_station_flat): + """Different stations produce separate outputs, no error.""" + plan = build_chunk_plan(two_station_flat, time=None) + assert len(plan.outputs) == 2 + assert set(plan.outputs["station"]) == set(two_station_flat["station"]) + + def test_group_override(self, two_station_flat): + """An explicit empty group means station conflicts raise.""" + with pytest.raises(CoordMergeError, match="station"): + build_chunk_plan(two_station_flat, time=None, group=()) + + def test_config_group(self, two_station_flat): + """Config groupby_attrs drives the default partitioning.""" + with dc.set_config(groupby_attrs=("network",)): + with pytest.raises(CoordMergeError, match="station"): + build_chunk_plan(two_station_flat, time=None) + + +class TestDeterminism: + """Spec section 8.""" + + def test_repeat_identical(self, diverse_flat): + """Identical inputs give identical plans.""" + p1 = build_chunk_plan(diverse_flat, time=None) + p2 = build_chunk_plan(diverse_flat, time=None) + pd.testing.assert_frame_equal(p1.outputs, p2.outputs) + pd.testing.assert_frame_equal(p1.members, p2.members) + + def test_input_order_invariant(self, diverse_flat): + """Row order of the flat relation does not change the plan.""" + shuffled = diverse_flat.sample(frac=1, random_state=0) + p1 = build_chunk_plan(diverse_flat, time=None) + p2 = build_chunk_plan(shuffled, time=None) + cols = ["time_min", "time_max"] + pd.testing.assert_frame_equal( + p1.outputs[cols].reset_index(drop=True), + p2.outputs[cols].reset_index(drop=True), + ) + + +class TestOracleParity: + """Sanity: plans describe exactly what spool.chunk produces.""" + + def test_merge_envelopes_match(self, random_flat): + """Merge-mode output envelopes match spool.chunk(time=None).""" + spool = dc.get_example_spool("random_das") + merged = spool.chunk(time=None) + contents = merged.get_contents() + plan = build_chunk_plan(random_flat, time=None) + assert len(plan.outputs) == len(contents) + assert plan.outputs["time_min"].iloc[0] == contents["time_min"].iloc[0] + assert plan.outputs["time_max"].iloc[0] == contents["time_max"].iloc[0] + + def test_segment_envelopes_match(self, random_flat): + """Segment-mode envelopes match spool.chunk(time=3).""" + spool = dc.get_example_spool("random_das") + chunked = spool.chunk(time=3) + contents = chunked.get_contents().sort_values("time_min") + plan = build_chunk_plan(random_flat, time=3) + outs = plan.outputs.sort_values("time_min") + assert len(outs) == len(contents) + assert np.array_equal(outs["time_min"].values, contents["time_min"].values) + assert np.array_equal(outs["time_max"].values, contents["time_max"].values) + + +class TestChunkPlanAccessor: + """Tests for the public spool.chunk_plan diagnostic.""" + + def test_matches_chunk(self): + """The plan describes exactly what chunk produces.""" + spool = dc.get_example_spool("random_das") + plan = spool.chunk_plan(time=3) + chunked = spool.chunk(time=3) + assert len(plan.outputs) == len(chunked) + contents = chunked.get_contents().sort_values("time_min") + outs = plan.outputs.sort_values("time_min") + assert np.array_equal(outs["time_min"].values, contents["time_min"].values) + + def test_records_params(self): + """Plans record the resolved parameters.""" + spool = dc.get_example_spool("random_das") + plan = spool.chunk_plan(time=None, tolerance=2.0) + assert plan.params["tolerance"] == 2.0 + assert isinstance(plan.params["group"], tuple) + assert plan.merge_mode + + def test_members_reference_sources(self): + """Members bind outputs to source patches without loading data.""" + spool = dc.get_example_spool("diverse_das") + plan = spool.chunk_plan(time=None) + assert set(plan.members["output_id"]) == set(plan.outputs["output_id"]) + assert len(plan.members) == len(spool) + + def test_directory_spool(self, tmp_path): + """chunk_plan works on file-backed spools.""" + patch = dc.get_example_patch() + dc.write(patch, tmp_path / "a.h5", "dasdae") + spool = dc.spool(tmp_path).update() + plan = spool.chunk_plan(time=None) + assert len(plan.outputs) == 1 + + +class TestUnderscoreDimNames: + """Dims with underscores must not confuse column ownership.""" + + def test_chunk_event_time(self): + """Adjacent patches along an underscore dim merge cleanly.""" + patch = dc.get_example_patch().rename_coords(time="event_time") + coord = patch.get_coord("event_time") + middle = coord.values[len(coord) // 2] + p1 = patch.select(event_time=(None, middle)) + p2 = patch.select(event_time=(middle + coord.step, None)) + merged = dc.spool([p1, p2]).chunk(event_time=None) + assert len(merged) == 1 + out = merged[0].get_coord("event_time") + assert out.min() == coord.min() + assert out.max() == coord.max() + + def test_segment_event_time(self): + """Segmenting along an underscore dim works too.""" + patch = dc.get_example_patch().rename_coords(time="event_time") + spool = dc.spool([patch]) + chunked = spool.chunk(event_time=2) + assert len(chunked) > 1 + for sub in chunked: + assert "event_time" in sub.dims + + +class TestChunkPlanCoverageEdges: + """Remaining chunk-planner branches.""" + + def test_partial_overlap_members(self): + """Overlapping sources drop the covered span (non-overlap skip).""" + import numpy as np + + t0 = np.datetime64("2020-01-01", "ns") + p1 = dc.get_example_patch(time_min=t0) + step = p1.get_coord("time").step + # p2 overlaps p1's tail by 100 samples. + p2 = dc.get_example_patch(time_min=p1.get_coord("time").max() - 100 * step) + merged = dc.spool([p1, p2]).chunk(time=None) + # the overlap is removed, so the merge is shorter than the naive sum. + naive = p1.get_coord("time").size + p2.get_coord("time").size + assert merged[0].get_coord("time").size < naive + + def test_user_stacklevel_fallback(self, monkeypatch): + """With no user frame in the stack, the stacklevel falls back to 1.""" + import inspect as _inspect + + import dascore.utils.chunk_plan as cp + + # Every frame reports a dascore path, so no "user" frame is found. + class _Frame: + filename = cp.__file__ + + monkeypatch.setattr(_inspect, "stack", lambda: [_Frame()] * 3) + assert cp._user_stacklevel() == 1 + + +class TestSamplingGroups: + """Sampling-group partitioning invariants (2026-07-18 F5).""" + + def test_close_steps_group(self): + """Steps within tolerance of the anchor share a group.""" + from dascore.utils.chunk_plan import _sampling_group + + labels = _sampling_group(pd.Series([1.0, 1.04]), 0.05) + assert labels.nunique() == 1 + + def test_chain_does_not_drift_past_tolerance(self): + """Adjacent-close steps cannot chain past the group anchor.""" + from dascore.utils.chunk_plan import _sampling_group + + labels = _sampling_group(pd.Series([1.00, 1.04, 1.08]), 0.05) + assert list(labels) == [0, 0, 1] + + def test_negative_steps_group_by_magnitude(self): + """Widely different negative steps never share a group.""" + from dascore.utils.chunk_plan import _sampling_group + + labels = _sampling_group(pd.Series([-1.0, -5.0]), 0.05) + assert labels.nunique() == 2 + + def test_mixed_orientation_never_groups(self): + """Equal magnitudes with opposite signs stay separate.""" + from dascore.utils.chunk_plan import _sampling_group + + labels = _sampling_group(pd.Series([1.0, -1.0]), 0.05) + assert labels.nunique() == 2 + + def test_unknown_steps_share_one_group(self): + """NaN steps keep their historical single-group behavior.""" + from dascore.utils.chunk_plan import _sampling_group + + labels = _sampling_group(pd.Series([1.0, np.nan, np.nan]), 0.05) + assert labels.nunique() == 2 + assert labels.iloc[1] == labels.iloc[2] + + def test_descending_contiguous_merges(self): + """Contiguous descending patches produce a single merge output.""" + p = dc.get_example_patch() + flipped = p.flip("time") + t = p.get_coord("time") + span = t.max() - t.min() + t.step + shifted = flipped.update_coords(time=flipped.get_coord("time").data + span) + plan = dc.spool([shifted, flipped]).chunk_plan(time=None) + assert len(plan.outputs) == 1 + assert len(plan.members) == 2 + + +class TestSamplesAdjustedEnvelopes: + """Samples residual envelope adjustment (2026-07-18 F4).""" + + @staticmethod + def _frame(): + """One ascending row: 10 samples at step 1 spanning [0, 9].""" + return pd.DataFrame({"time_min": [0.0], "time_max": [9.0], "time_step": [1.0]}) + + def test_exclusive_stop(self): + """The stop index is exclusive: (0, 5) ends at sample 4.""" + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + out = samples_adjusted_envelopes(self._frame(), (({"time": (0, 5)}, True),)) + assert out["time_max"].iloc[0] == 4.0 + + def test_empty_window_drops_row(self): + """A zero-length window contributes nothing.""" + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + out = samples_adjusted_envelopes(self._frame(), (({"time": (3, 3)}, True),)) + assert len(out) == 0 + + def test_start_past_end_drops_row(self): + """A window beyond the patch contributes nothing.""" + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + out = samples_adjusted_envelopes(self._frame(), (({"time": (50, 60)}, True),)) + assert len(out) == 0 + + def test_stop_clamps(self): + """A stop past the end clamps to the envelope max.""" + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + out = samples_adjusted_envelopes(self._frame(), (({"time": (2, 99)}, True),)) + assert out["time_min"].iloc[0] == 2.0 + assert out["time_max"].iloc[0] == 9.0 + + def test_descending_orientation(self): + """On a descending coord, sample 0 sits at the envelope max.""" + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + df = pd.DataFrame({"time_min": [0.0], "time_max": [9.0], "time_step": [-1.0]}) + out = samples_adjusted_envelopes(df, (({"time": (0, 5)}, True),)) + assert out["time_max"].iloc[0] == 9.0 + assert out["time_min"].iloc[0] == 5.0 + + def test_public_same_dim_envelope_matches_patch(self): + """The chunked catalog's envelope matches the loaded patch (F4).""" + p = dc.get_example_patch() + out = dc.spool([p]).select(time=(0, 10), samples=True).chunk(time=None) + patch = out[0] + want = pd.Timestamp(patch.get_coord("time").max()) + got = out.get_contents()["time_max"].iloc[0] + assert got == want + assert patch.shape[patch.get_axis("time")] == 10 + + +class TestChunkOnlyOnDims: + """Chunking is defined on dimensions; non-dim coords are 'missing'.""" + + @pytest.fixture() + def aux_time_patch(self): + """A patch carrying time only as a coord riding distance.""" + p = dc.get_example_patch() + t = p.get_coord("time") + base = p.mean("time").squeeze() + n = base.shape[base.get_axis("distance")] + return base.update_coords(time=("distance", t.data[:n])) + + def test_aux_only_raises_with_detail(self, aux_time_patch): + """The default error explains the name rides as a coordinate.""" + from dascore.exceptions import ChunkError + + with pytest.raises(ChunkError, match="non-dimensional coordinate"): + dc.spool([aux_time_patch]).chunk(time=None) + + def test_mixed_population_raises_by_default(self, aux_time_patch): + """Mixed dim/aux populations fail eagerly, not at patch access.""" + from dascore.exceptions import ChunkError + + sp = dc.spool([dc.get_example_patch(), aux_time_patch]) + with pytest.raises(ChunkError, match="lack the chunk dimension"): + sp.chunk(time=None) + + def test_drop_excludes_aux_patches(self, aux_time_patch): + """missing_dim='drop' keeps only patches with the real dimension.""" + sp = dc.spool([dc.get_example_patch(), aux_time_patch]) + out = sp.chunk(time=None, missing_dim="drop", conflict="drop") + assert len(out) == 1 + assert out[0].dims == ("distance", "time") + + def test_segmenting_aux_coord_raises(self): + """chunk(=value) is rejected, not accidentally served.""" + import numpy as np + + from dascore.exceptions import ChunkError + + p = dc.get_example_patch() + q = p.update_coords(sensor=("distance", np.arange(p.shape[0], dtype=float))) + with pytest.raises(ChunkError, match="non-dimensional coordinate"): + dc.spool([q]).chunk(sensor=100) + + +class TestDimlessFrames: + """Plain planner frames without a dims column still plan.""" + + def test_frame_without_dims_column(self): + """Dimension membership checks are skipped when dims is absent.""" + df = pd.DataFrame( + { + "time_min": [0.0, 10.0], + "time_max": [10.0, 20.0], + "time_step": [1.0, 1.0], + } + ) + plan = build_chunk_plan(df, time=None) + assert len(plan.outputs) == 1 + assert len(plan.members) == 2 + + +class TestNegativeSamplesEnvelopes: + """Negative samples indices resolve per patch (2026-07-18).""" + + @staticmethod + def _frame(): + """One ascending row: 10 samples at step 1 spanning [0, 9].""" + return pd.DataFrame({"time_min": [0.0], "time_max": [9.0], "time_step": [1.0]}) + + def test_negative_start_resolves(self): + """(-3, None) selects the last three samples.""" + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + out = samples_adjusted_envelopes(self._frame(), (({"time": (-3, None)}, True),)) + assert out["time_min"].iloc[0] == 7.0 + assert out["time_max"].iloc[0] == 9.0 + + def test_negative_stop_resolves(self): + """(None, -2) drops the last two samples.""" + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + out = samples_adjusted_envelopes(self._frame(), (({"time": (None, -2)}, True),)) + assert out["time_max"].iloc[0] == 7.0 + + def test_unknown_step_keeps_envelope(self): + """Rows whose count is unknown keep their candidacy envelope.""" + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + df = pd.DataFrame({"time_min": [0.0], "time_max": [9.0], "time_step": [np.nan]}) + out = samples_adjusted_envelopes(df, (({"time": (-3, None)}, True),)) + assert len(out) == 1 + assert out["time_min"].iloc[0] == 0.0 + assert out["time_max"].iloc[0] == 9.0 + + def test_drop_empty_false_keeps_rows(self): + """Equality's variant keeps presented-but-empty rows.""" + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + out = samples_adjusted_envelopes( + self._frame(), (({"time": (3, 3)}, True),), drop_empty=False + ) + assert len(out) == 1 + + def test_public_negative_window_contents_honest(self): + """Derived contents match the loaded patch for negative windows.""" + p = dc.get_example_patch() + out = dc.spool([p]).select(time=(-10, None), samples=True).chunk(time=None) + got = out.get_contents()["time_min"].iloc[0] + want = pd.Timestamp(out[0].get_coord("time").min()) + assert got == want + assert out[0].shape[out[0].get_axis("time")] == 10 + + +class TestNonIntSamplesIndices: + """Non-integer samples indices leave envelopes untouched.""" + + def test_float_index_skipped(self): + """A float index cannot adjust; the envelope stays candidacy.""" + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + df = pd.DataFrame({"time_min": [0.0], "time_max": [9.0], "time_step": [1.0]}) + out = samples_adjusted_envelopes(df, (({"time": (0.5, None)}, True),)) + assert out["time_min"].iloc[0] == 0.0 + assert out["time_max"].iloc[0] == 9.0 diff --git a/tests/test_io/test_index/test_planned.py b/tests/test_io/test_index/test_planned.py new file mode 100644 index 000000000..163017bdc --- /dev/null +++ b/tests/test_io/test_index/test_planned.py @@ -0,0 +1,328 @@ +""" +Tests for derived catalogs (plan-as-catalog) and coverage of their edges. +""" + +from __future__ import annotations + +import re + +import numpy as np +import pandas as pd +import pytest + +import dascore as dc +from dascore.exceptions import ParameterError +from dascore.io.index.planned import ( + PlanResolver, + _coord_record_from_row, + _ns, + collapse_working_df, + derived_catalog, +) + + +@pytest.fixture(scope="module") +def patches(): + """Three contiguous example patches.""" + return list(dc.get_example_spool("random_das")) + + +class TestHelpers: + """Unit coverage for the conversion helpers.""" + + def test_ns_forms(self): + """All datetime/timedelta forms convert to the same ns.""" + ts = pd.Timestamp("2020-01-01") + assert _ns(ts) == _ns(ts.to_datetime64()) == ts.value + td = pd.Timedelta(seconds=1) + assert _ns(td) == _ns(td.to_timedelta64()) == td.value + assert _ns(None) is None + + def test_coord_record_numpy_datetimes(self): + """np.datetime64 envelope values build the same record.""" + lo = np.datetime64("2020-01-01", "ns") + hi = np.datetime64("2020-01-02", "ns") + row = {"time_min": lo, "time_max": hi, "time_step": np.timedelta64(1, "s")} + record = _coord_record_from_row(row, "time") + assert record.value_kind == "time" + assert record.min_ns == _ns(lo) + + def test_coord_record_zero_step_length(self): + """A degenerate step leaves length unknown instead of raising.""" + row = {"time_min": 0.0, "time_max": 1.0, "time_step": 0.0} + record = _coord_record_from_row(row, "time") + assert record.length is None + + def test_plan_resolver_requires_output_id(self): + """member_rows without output_id is a construction error.""" + with pytest.raises(ValueError, match="output_id"): + PlanResolver( + token="x", + dim="time", + member_rows=pd.DataFrame({"path": []}), + loader=None, + merge_kwargs={}, + ) + + def test_derived_catalog_adds_patch_ids(self, patches): + """source_rows without _patch_id get positional ids.""" + from dascore.utils.chunk_plan import ChunkPlan + + spool = dc.spool(patches) + rows = spool.get_contents().drop(columns=["_patch_id"], errors="ignore") + rows = rows.reset_index(drop=True) + members = pd.DataFrame( + {"output_id": [0], "_patch_id": [0], "_modified": [False]} + ) + outputs = rows.iloc[:1].assign(output_id=0) + plan = ChunkPlan(outputs, members, "time", None, {}) + catalog = derived_catalog( + source_rows=rows, + plan=plan, + parent=spool._catalog, + merge_kwargs={}, + mode="concat", + ) + assert len(catalog) == 1 + + +class TestDerivedComposition: + """Operation-order coverage over derived catalogs.""" + + def test_collapse_with_value_residual(self, patches): + """Chunk of a selected chunked spool re-plans from trimmed members.""" + t0 = patches[0].get_coord("time").min() + t1 = patches[1].get_coord("time").max() + chunked = dc.spool(patches).chunk(time=2) + selected = chunked.select(time=(t0, t1)) + merged = selected.chunk(time=None) + assert len(merged) >= 1 + out = merged[0] + assert out.get_coord("time").min() >= t0 + + def test_sort_by_attr_on_windowed_regex_view(self, patches): + """Regex selection + window + attr sort compose through SQL.""" + tagged = [ + p.update_attrs(tag=f"t{num}", history=[]) for num, p in enumerate(patches) + ] + spool = dc.spool(tagged) + view = spool.select(tag=re.compile("t[0-9]"))[1:] + out = view.sort("tag") + tags = [p.attrs["tag"] for p in out] + assert tags == sorted(tags) + + def test_attr_membership_array(self, patches): + """Attr membership with a numpy array of values selects rows.""" + tagged = [ + p.update_attrs(tag=f"t{num}", history=[]) for num, p in enumerate(patches) + ] + spool = dc.spool(tagged) + out = spool.select(tag=np.array(["t0", "t2"])) + assert len(out) == 2 + + def test_sort_by_envelope_column_name(self, patches): + """Sort accepts the explicit `{dim}_min` column form.""" + spool = dc.spool(list(reversed(patches))) + out = spool.sort("time_min") + assert out.get_contents()["time_min"].is_monotonic_increasing + + def test_concatenate_requires_one_kwarg(self, patches): + """Concatenate validates its dimension keyword.""" + with pytest.raises(ParameterError, match="exactly one dimension"): + dc.spool(patches).concatenate(time=None, distance=None) + + def test_union_view_of_live_spools_pickles_composite(self, patches): + """A selected union pickles a membership-restricted composite.""" + import pickle + + t0 = patches[0].get_coord("time") + combined = dc.spool(patches[:2]) + dc.spool(patches[2:]) + view = combined.select(time=(None, t0.max())) + assert len(view) == 1 + loaded = pickle.loads(pickle.dumps(view)) + assert len(loaded) == 1 + assert isinstance(loaded[0], dc.Patch) + + def test_missing_live_patch_getitem(self, patches): + """A missing registry entry surfaces as MissingPatchError, not + out-of-bounds. + """ + from dascore.exceptions import MissingPatchError + + spool = dc.spool(patches[:1]) + _ = spool.get_contents() # realize rows + spool._catalog.resolver._registry.clear() + with pytest.raises(MissingPatchError, match="not available"): + spool[0] + + def test_union_with_third_party_spool(self, patches): + """The BaseSpool fallback materializes third-party members.""" + from dascore.core.spool import BaseSpool + + class MiniSpool(BaseSpool): + def __init__(self, inner): + self._inner = list(inner) + + def __getitem__(self, item): + return self._inner[item] + + def __iter__(self): + return iter(self._inner) + + def __len__(self): + return len(self._inner) + + def chunk(self, **kwargs): + raise NotImplementedError + + def select(self, **kwargs): + raise NotImplementedError + + def get_contents(self): + raise NotImplementedError + + combined = dc.spool(patches[:1]) + MiniSpool(patches[1:]) + assert len(combined) == len(patches) + + def test_samples_negative_index_skips_envelope_adjust(self, patches): + """Negative samples windows stay candidacy supersets (no crash).""" + spool = dc.spool(patches[:1]).select(time=(0, -10), samples=True) + merged = spool.chunk(time=None) + assert len(merged) == 1 + + def test_complete_overlap_merge(self): + """Two identical-envelope patches merge by keeping the first.""" + patch = dc.get_example_patch() + twin = patch.new() + merged = dc.spool([patch, twin]).chunk(time=None) + assert len(merged) == 1 + assert isinstance(merged[0], dc.Patch) + + +class TestRemainingEdges: + """Direct coverage of defensive/rare branches.""" + + def test_collapse_with_quantity_residual(self, patches): + """A quantity-selected chunked view re-chunks without applying + unit-bearing bounds to envelopes (they stay load residuals). + """ + from dascore.units import m + + chunked = dc.spool(patches).chunk(time=2) + selected = chunked.select(_coords={"distance": (0 * m, 10 * m)}) + merged = selected.chunk(time=None) + assert len(merged) == 1 + coord = merged[0].get_coord("distance") + assert float(coord.max()) <= 10 + + def test_samples_adjust_skips_missing_columns(self): + """Residuals naming absent envelope columns pass through.""" + from dascore.utils.chunk_plan import samples_adjusted_envelopes + + df = pd.DataFrame({"time_min": [0.0], "time_max": [1.0]}) + residuals = (({"depth": (0, 5)}, True),) + out = samples_adjusted_envelopes(df, residuals) + assert out.equals(df) + + +class TestAuxiliaryCoords: + """Derived catalogs keep non-dimension coords (2026-07-18 F2).""" + + @pytest.fixture() + def sensor_spool(self): + """Two contiguous patches carrying an aux coord on distance.""" + p = dc.get_example_patch() + sensor = np.arange(p.shape[p.get_axis("distance")], dtype=float) + p = p.update_coords(sensor=("distance", sensor)) + t = p.get_coord("time") + p2 = p.update_coords(time_min=t.max() + t.step) + return dc.spool([p, p2]) + + @pytest.mark.parametrize("op", ["chunk", "concatenate"]) + def test_aux_coord_survives(self, sensor_spool, op): + """Chunk and concat outputs keep describing the aux coord.""" + if op == "chunk": + derived = sensor_spool.chunk(time=None, conflict="drop") + else: + derived = sensor_spool.concatenate(time=None) + contents = derived.get_contents() + assert "sensor_min" in contents.columns + assert "sensor_max" in contents.columns + # and it stays selectable + out = derived.select(sensor=(10, 20)) + assert len(out) == 1 + coord = out[0].get_coord("sensor") + assert coord.min() == 10.0 + assert coord.max() == 20.0 + + def test_aux_identity_preserved_when_unchanged(self, sensor_spool): + """Members sharing one def key off the planned dim keep identity.""" + derived = sensor_spool.chunk(time=None, conflict="drop") + source_key = sensor_spool._catalog.to_df()["_sensor_def_key"].iloc[0] + derived_key = derived._catalog.to_df()["_sensor_def_key"].iloc[0] + assert derived_key == source_key + assert str(derived_key).startswith("fp:") + + def test_aux_identity_dropped_when_riding_trimmed_dim(self, sensor_spool): + """A residual trim on distance voids sensor's identity claim.""" + d = sensor_spool[0].get_coord("distance") + lo, hi = d.min() + 5 * d.step, d.min() + 50 * d.step + selected = sensor_spool.select(distance=(lo, hi)) + derived = selected.chunk(time=None, conflict="drop") + key = derived._catalog.to_df()["_sensor_def_key"].iloc[0] + assert not str(key).startswith("fp:") + # loading still yields the trimmed coord + assert derived[0].get_coord("sensor").min() == 5.0 + + def test_string_aux_coord(self): + """String-valued aux coords survive with a lexicographic envelope.""" + p = dc.get_example_patch() + n = p.shape[p.get_axis("distance")] + labels = np.array([f"s{i:03d}" for i in range(n)]) + p = p.update_coords(station=("distance", labels)) + derived = dc.spool([p]).chunk(time=None) + contents = derived.get_contents() + assert contents["station_min"].iloc[0] == "s000" + assert contents["station_max"].iloc[0] == f"s{n - 1:03d}" + assert "station" in derived[0].coords.coord_map + + +class TestAuxInfoEdges: + """Edge branches of the aux-coord aggregation helpers.""" + + def test_coord_record_missing_envelope_returns_none(self): + """A row without envelope values yields no coord record.""" + from dascore.io.index.planned import _coord_record_from_row + + assert _coord_record_from_row({}, "time") is None + + def test_absent_envelope_columns_skipped(self): + """A mapped coord with no envelope columns contributes nothing.""" + from dascore.io.index.planned import _aux_coord_info + + members = pd.DataFrame( + {"output_id": [0], "_patch_id": [1], "_modified": [False]} + ) + sources = pd.DataFrame({"_patch_id": [1]}) + assert _aux_coord_info(sources, members, "time", {"ghost": "distance"}) == {} + + def test_all_null_group_skipped(self): + """An output whose members carry no values for a coord is skipped.""" + from dascore.io.index.planned import _aux_coord_info + + members = pd.DataFrame( + {"output_id": [0], "_patch_id": [1], "_modified": [False]} + ) + sources = pd.DataFrame( + {"_patch_id": [1], "sensor_min": [np.nan], "sensor_max": [np.nan]} + ) + assert _aux_coord_info(sources, members, "time", {"sensor": "distance"}) == {} + + +class TestCollapseGuard: + """collapse_working_df only applies to plan-backed catalogs.""" + + def test_non_plan_catalog_returns_none(self): + """A live catalog has no plan to collapse.""" + catalog = dc.spool([dc.get_example_patch()])._catalog + assert collapse_working_df(catalog) is None diff --git a/tests/test_io/test_index/test_schema.py b/tests/test_io/test_index/test_schema.py new file mode 100644 index 000000000..d442665e7 --- /dev/null +++ b/tests/test_io/test_index/test_schema.py @@ -0,0 +1,72 @@ +"""Schema and initialization tests for the SQLite spool index.""" + +from __future__ import annotations + +import sqlite3 +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier + +import pytest + +from dascore.exceptions import InvalidIndexError, InvalidIndexVersionError +from dascore.io.index import get_backend +from dascore.io.index.schema import INDEX_VERSION, TABLES + + +class TestSchemaValidation: + """Existing files are validated without repair or implicit migration.""" + + def test_unrelated_database_rejected(self, tmp_path): + """A SQLite database belonging to another application is rejected.""" + path = tmp_path / "other.sqlite3" + con = sqlite3.connect(path) + con.execute("CREATE TABLE other_app (value TEXT)") + con.close() + with pytest.raises(InvalidIndexError, match="missing tables"): + get_backend(path) + + def test_old_version_rejected(self, tmp_path): + """Prototype schemas require an explicit delete and rebuild.""" + path = tmp_path / "old.sqlite3" + backend = get_backend(path) + backend._execute("UPDATE meta_data SET index_version = ?", (INDEX_VERSION - 1,)) + backend.close() + with pytest.raises(InvalidIndexVersionError, match="delete it and rebuild"): + get_backend(path) + + def test_schema_has_foreign_keys_and_constraints(self, tmp_path): + """SQLite enforces source/patch ownership and cascades.""" + backend = get_backend(tmp_path / "index.sqlite3") + assert backend._con.execute("PRAGMA foreign_keys").fetchone()[0] == 1 + assert backend._con.execute("PRAGMA busy_timeout").fetchone()[0] == 30_000 + tables = backend._existing_tables() + assert set(TABLES) <= tables + with pytest.raises(sqlite3.IntegrityError): + backend._execute( + "INSERT INTO patches (patch_id, source_id, source_patch_id) " + "VALUES (1, 999, '0')" + ) + backend.close() + + +class TestConcurrentInitialization: + """Only one writer initializes a new index file.""" + + def test_concurrent_open(self, tmp_path): + """Connections racing to create one index all open successfully.""" + path = tmp_path / "shared.sqlite3" + barrier = Barrier(4) + + def open_index(_): + barrier.wait() + backend = get_backend(path) + metadata = backend.get_metadata() + backend.close() + return metadata["index_version"] + + with ThreadPoolExecutor(max_workers=4) as pool: + versions = list(pool.map(open_index, range(4))) + assert versions == [INDEX_VERSION] * 4 + backend = get_backend(path) + assert len(backend._fetch_df("SELECT * FROM meta_data")) == 1 + backend.close() diff --git a/tests/test_io/test_index/test_union.py b/tests/test_io/test_index/test_union.py new file mode 100644 index 000000000..a2d0fdf5f --- /dev/null +++ b/tests/test_io/test_index/test_union.py @@ -0,0 +1,493 @@ +"""Tests for spool union (spool + spool) and catalog merging.""" + +from __future__ import annotations + +import numpy as np +import pytest + +import dascore as dc +from dascore.core.coords import CoordRange +from dascore.io.index.catalog import CompositeResolver, PatchCatalog + + +@pytest.fixture(scope="module") +def contiguous_patches(): + """Two contiguous example patches.""" + t0 = np.datetime64("2020-01-01", "ns") + p1 = dc.get_example_patch(time_min=t0) + time = p1.get_coord("time") + p2 = dc.get_example_patch(time_min=time.max() + time.step) + return p1, p2 + + +@pytest.fixture() +def dir_spool(contiguous_patches, tmp_path): + """A directory spool holding the first patch.""" + p1, _ = contiguous_patches + dc.write(p1, tmp_path / "a.h5", "dasdae") + return dc.spool(tmp_path).update() + + +class TestMemoryUnion: + """Unions of in-memory spools.""" + + def test_lengths_add(self): + """The union contains the patches of both.""" + sp1 = dc.get_example_spool("random_das") + sp2 = dc.get_example_spool("diverse_das") + combined = sp1 + sp2 + assert len(combined) == len(sp1) + len(sp2) + + def test_patches_shared_not_copied(self, contiguous_patches): + """In-memory patches resolve to the same objects.""" + p1, p2 = contiguous_patches + combined = dc.spool([p1]) + dc.spool([p2]) + loaded = list(combined) + assert any(x is p1 for x in loaded) + assert any(x is p2 for x in loaded) + + def test_union_of_materialized_member(self): + """A planned (materialized but catalog-backed) member unions by ids.""" + sp = dc.get_example_spool("random_das") + other = dc.get_example_spool("diverse_das") + # a content-preserving derived catalog (concat groups of one) + materialized = sp.concatenate(time=1) + combined = materialized + other + assert len(combined) == len(sp) + len(other) + + def test_union_of_sorted_member(self): + """A lazily sorted member unions by its ordered membership.""" + sp = dc.get_example_spool("random_das") + other = dc.get_example_spool("diverse_das") + combined = sp.sort("time") + other + assert len(combined) == len(sp) + len(other) + + def test_select_on_union(self): + """Selection works over the merged metadata.""" + sp1 = dc.get_example_spool("random_das") + sp2 = dc.get_example_spool("diverse_das") + combined = sp1 + sp2 + selected = combined.select(network="das2") + assert len(selected) + for patch in selected: + assert patch.attrs.network == "das2" + + def test_non_spool_add(self): + """Adding a non-spool returns NotImplemented semantics.""" + sp1 = dc.get_example_spool("random_das") + with pytest.raises(TypeError): + _ = sp1 + 42 + + def test_chunk_across_members(self, contiguous_patches): + """Contiguous patches from different spools merge into one.""" + p1, p2 = contiguous_patches + combined = dc.spool([p1]) + dc.spool([p2]) + merged = combined.chunk(time=None) + assert len(merged) == 1 + patch = merged[0] + time = patch.get_coord("time") + assert isinstance(time, CoordRange) + assert time.min() == p1.get_coord("time").min() + assert time.max() == p2.get_coord("time").max() + + def test_selection_carries_by_membership(self): + """A selected input contributes only its selected rows.""" + sp2 = dc.get_example_spool("diverse_das") + sub = sp2.select(network="das2") + combined = dc.get_example_spool("random_das") + sub + assert len(combined) == len(dc.get_example_spool("random_das")) + len(sub) + + +class TestFileUnion: + """Unions involving file-backed spools.""" + + def test_dir_plus_memory(self, dir_spool, contiguous_patches): + """A directory spool and memory spool combine lazily.""" + _, p2 = contiguous_patches + combined = dir_spool + dc.spool([p2]) + assert len(combined) == 2 + loaded = list(combined) + assert all(x.shape for x in loaded) + + def test_chunk_across_file_and_memory(self, dir_spool, contiguous_patches): + """The union seam merges: file patch + memory patch -> one patch.""" + p1, p2 = contiguous_patches + combined = dir_spool + dc.spool([p2]) + merged = combined.chunk(time=None) + assert len(merged) == 1 + patch = merged[0] + assert patch.shape[patch.get_axis("time")] == ( + p1.shape[p1.get_axis("time")] + p2.shape[p2.get_axis("time")] + ) + + def test_same_source_dedups(self, dir_spool): + """The same source in both members keeps a single entry.""" + combined = dir_spool + dir_spool + assert len(combined) == len(dir_spool) + + def test_constructor_select_kwargs_restrict_union(self, tmp_path): + """A selection-restricted directory spool unions only its rows.""" + base = dc.get_example_spool("random_das") + dc.examples.spool_to_directory(base, path=tmp_path) + full = dc.spool(tmp_path).update() + df = full.get_contents().sort_values("time_min") + window = (df["time_min"].iloc[0], df["time_max"].iloc[0]) # first patch + restricted = full.select(time=window) + assert 0 < len(restricted) < len(full) + combined = restricted + dc.spool([dc.get_example_patch(tag="mem")]) + # the union must not reintroduce the rows the constructor excluded + assert len(combined) == len(restricted) + 1 + + def test_union_preserves_def_keys(self, dir_spool, contiguous_patches): + """Coord definitions deduplicate by def key across members.""" + _, p2 = contiguous_patches + combined = dir_spool + dc.spool([p2]) + df = combined._catalog.to_df() + # both patches share the same distance coord identity + assert df["_distance_def_key"].nunique() == 1 + + +class TestCompositeResolver: + """Resolver dispatch for union catalogs.""" + + def test_routes_memory_rows(self, contiguous_patches): + """memory:// rows go to the live registry.""" + p1, _ = contiguous_patches + cat = PatchCatalog.union([PatchCatalog.from_patches([p1])]) + assert isinstance(cat.resolver, CompositeResolver) + row = cat.to_df().iloc[0].to_dict() + assert cat.resolve_row(row) is p1 + + def test_union_of_union(self, contiguous_patches): + """Unions compose (a union catalog can be a member).""" + p1, p2 = contiguous_patches + first = PatchCatalog.union([PatchCatalog.from_patches([p1])]) + second = PatchCatalog.union([first, PatchCatalog.from_patches([p2])]) + assert len(second.to_df()) == 2 + patches = [second.resolve_row(x) for _, x in second.to_df().iterrows()] + assert {id(x) for x in patches} == {id(p1), id(p2)} + + +class TestPatchIdentity: + """Set semantics by patch instance identity (lineage).""" + + def test_duplicate_instances_collapse(self): + """The same patch instance twice is one spool entry.""" + patch = dc.get_example_patch() + assert len(dc.spool([patch, patch])) == 1 + + def test_deepcopy_shares_identity(self): + """Copies of an immutable patch share its identity. + + Identity is minted eagerly at construction, so copies share it + regardless of when they are made (no access-order dependence). + """ + import copy + + patch = dc.get_example_patch() + clone = copy.deepcopy(patch) + assert clone._instance_id == patch._instance_id + assert len(dc.spool([patch, clone])) == 1 + + def test_new_instance_distinct(self): + """patch.new() (and any patch op) mints a distinct identity.""" + patch = dc.get_example_patch() + assert len(dc.spool([patch, patch.new()])) == 2 + + def test_ops_mint_new_identity(self): + """Operations produce instances with their own identity.""" + patch = dc.get_example_patch() + other = patch.update_attrs(tag="x") + assert patch._instance_id != other._instance_id + + def test_pickle_round_trip_preserves_content(self): + """Spools of live patches pickle and rebuild their backend.""" + import pickle + + patch = dc.get_example_patch() + spool = dc.spool([patch]) + _ = len(spool) # realize the catalog + loaded = pickle.loads(pickle.dumps(spool)) + assert len(loaded) == 1 + assert loaded[0] == patch + + def test_union_pickles(self): + """Union spools survive pickling (rows ride along as records).""" + import pickle + + p1 = dc.get_example_patch() + p2 = p1.new() + combined = dc.spool([p1]) + dc.spool([p2]) + loaded = pickle.loads(pickle.dumps(combined)) + assert len(loaded) == 2 + + def test_remove_updates_live_registry(self): + """Removing a live source removes it from the store as well.""" + import pickle + + patch = dc.get_example_patch() + catalog = PatchCatalog.from_patches([patch]) + path = catalog.to_df().iloc[0]["path"] + catalog.remove([path]) + assert len(catalog.to_df()) == 0 + # A pickled catalog rebuilds from the registry; the removed patch + # must not resurrect. + loaded = pickle.loads(pickle.dumps(catalog)) + loaded.attr_names() # bootstrap the backend + loaded._invalidate() + assert len(loaded.to_df()) == 0 + + +class TestExportPushdown: + """Selected-membership export must not scan the whole archive.""" + + def test_export_one_of_many_is_narrow(self): + """Exporting one patch fetches only its own rows, not all sources.""" + patches = [dc.get_example_patch().update_attrs(tag=f"t{i}") for i in range(40)] + catalog = PatchCatalog.from_patches(patches) + catalog.to_df() # bootstrap the backend + backend = catalog.backend + con = backend._con + + fetched_patches = [] + + def _trace(sql): + # Count how many patch rows any SELECT against patches pulls. + if "from patches" in sql.lower() and sql.lower().lstrip().startswith( + "select" + ): + fetched_patches.append(sql) + + target = int(catalog.to_df()["_patch_id"].iloc[0]) + con.set_trace_callback(_trace) + try: + records = backend.export_records(patch_ids=[target]) + finally: + con.set_trace_callback(None) + + # exactly one source/patch comes back... + assert sum(len(r.patches) for r in records) == 1 + # ...and every patches query was id-filtered (no full-table scan). + assert fetched_patches + assert all("patch_id in" in sql.lower() for sql in fetched_patches) + + def test_export_all_matches_full(self): + """export_records() with no ids returns every source, unchanged.""" + patches = [dc.get_example_patch().update_attrs(tag=f"t{i}") for i in range(5)] + catalog = PatchCatalog.from_patches(patches) + catalog.to_df() + records = catalog.backend.export_records() + assert sum(len(r.patches) for r in records) == 5 + + def test_export_empty_patch_ids(self): + """Exporting an empty id set returns no records without querying.""" + catalog = PatchCatalog.from_patches([dc.get_example_patch()]) + catalog.to_df() + assert catalog.backend.export_records(patch_ids=[]) == [] + + def test_absolutize_record_passthrough(self, tmp_path): + """A record already carrying an absolute/URI path is returned as-is.""" + from pathlib import Path + + from dascore.io.index.catalog import _absolutize_record + from dascore.io.index.ingest import SourceRecord + + # an OS-native absolute path (drive-qualified on Windows) + abs_path = str((tmp_path / "a.h5").resolve()) + assert Path(abs_path).is_absolute() + rec = SourceRecord(source_path=abs_path, source_format="X", format_version="1") + assert _absolutize_record(rec, str(tmp_path)) is rec + uri = SourceRecord( + source_path="s3://bucket/a.h5", source_format="X", format_version="1" + ) + assert _absolutize_record(uri, str(tmp_path)) is uri + + def test_dir_union_absolutizes_relative_paths(self, tmp_path): + """A directory member's relative source path is absolutized on union.""" + dc.get_example_patch().io.write(tmp_path / "a.h5", "dasdae") + dir_spool = dc.spool(tmp_path).update(progress=None) + combined = dir_spool + dc.spool([dc.get_example_patch(tag="mem")]) + assert len(combined) == 2 + # the file-backed member still loads (its path was made absolute) + contents = combined.get_contents() + file_row = contents[contents["path"].str.endswith("a.h5")] + assert len(file_row) == 1 + loaded = [p for p in combined] + assert len(loaded) == 2 + + +class TestSameFileUnion: + """Unions of members selecting patches from the same multi-patch file.""" + + @pytest.fixture() + def two_patch_file_spool(self, contiguous_patches, tmp_path): + """A file spool over one file holding two patches.""" + p1, p2 = (x.update_attrs(history=[]) for x in contiguous_patches) + path = tmp_path / "two_patch.h5" + dc.write(dc.spool([p1, p2]), path, "dasdae") + return dc.spool(path) + + def test_disjoint_selections_union(self, two_patch_file_spool): + """Two members holding different patches of one file both survive.""" + sp = two_patch_file_spool + combined = sp[:1] + sp[1:] + assert len(combined) == 2 + for patch in combined: + assert isinstance(patch, dc.Patch) + + def test_overlapping_selections_dedup(self, two_patch_file_spool): + """A patch present in both members appears once (dict-merge).""" + sp = two_patch_file_spool + combined = sp[:2] + sp[1:] + assert len(combined) == 2 + + def test_union_absorbs_only_member_registry_entries(self, contiguous_patches): + """Live entries outside a member's rows don't ride into the union.""" + p1, p2 = contiguous_patches + t1 = p1.get_coord("time") + narrowed = dc.spool([p1, p2]).select(time=(None, t1.max())) + assert len(narrowed) == 1 + other = dc.get_example_patch(time_min="2030-01-01") + combined = narrowed + dc.spool([other]) + # the trimmed operand materializes into a plan, so its live + # member rides inside the plan's loader, not the top registry + registry = combined._catalog.resolver.live_entries() + assert len(registry) == 1 # other + plans = combined._catalog.resolver.plan_entries() + nested_live = {k for p in plans.values() for k in p.live_entries()} + assert len(nested_live) == 1 # p1; p2 stayed home + assert len(combined) == 2 + + +class TestLossyStateUnion: + """Residual trims and order specs survive combining (2026-07-18 F1).""" + + def test_value_trim_survives_union(self): + """A coordinate-range selection's trim is baked in, not dropped.""" + p = dc.get_example_patch() + t = p.get_coord("time") + lo, hi = t.min() + 10 * t.step, t.min() + 20 * t.step + selected = dc.spool([p]).select(time=(lo, hi)) + combined = selected + dc.spool([]) + assert len(combined) == 1 + got, want = combined[0], selected[0] + assert got.shape == want.shape + assert got.get_coord("time").min() == want.get_coord("time").min() + assert got.get_coord("time").max() == want.get_coord("time").max() + + def test_samples_trim_survives_union(self): + """A samples window's trim is baked in, not dropped.""" + p = dc.get_example_patch() + selected = dc.spool([p]).select(time=(0, 10), samples=True) + combined = selected + dc.spool([]) + assert combined[0].shape == selected[0].shape + + def test_file_backed_value_trim_survives_union(self, dir_spool): + """The same guarantee holds for file-backed catalogs.""" + patch = dir_spool[0] + t = patch.get_coord("time") + lo, hi = t.min() + 10 * t.step, t.min() + 20 * t.step + selected = dir_spool.select(time=(lo, hi)) + combined = selected + dc.spool([]) + assert combined[0].shape == selected[0].shape + + def test_sort_order_survives_union(self): + """A sort spec bakes into ordinals instead of silently reverting.""" + p = dc.get_example_patch() + t = p.get_coord("time") + early = p.update_attrs(tag="early") + late = p.update_coords(time_min=t.max() + t.step).update_attrs(tag="late") + srt = dc.spool([late, early]).sort("time") + combined = srt + dc.spool([]) + assert [x.attrs.tag for x in combined] == [x.attrs.tag for x in srt] + + def test_membership_selections_still_dedup(self): + """Row-membership state unions as rows: identity dedup preserved.""" + p = dc.get_example_patch() + spool = dc.spool([p.update_attrs(tag=f"t{i}") for i in range(4)]) + combined = spool.select(tag="t2") + spool + assert len(combined) == 4 + + def test_windows_still_union_by_rows(self): + """Slice windows survive as membership without materializing.""" + p = dc.get_example_patch() + spool = dc.spool([p.update_attrs(tag=f"t{i}") for i in range(4)]) + combined = spool[1:3] + dc.spool([]) + assert [x.attrs.tag for x in combined] == ["t1", "t2"] + + def test_two_selected_operands(self): + """Different trims on each operand both survive as new contents.""" + p = dc.get_example_patch() + t = p.get_coord("time") + lo = t.min() + 10 * t.step + a = dc.spool([p]).select(time=(lo, t.min() + 20 * t.step)) + b = dc.spool([p]).select(time=(t.min(), lo)) + combined = a + b + assert len(combined) == 2 + assert {x.shape for x in combined} == {a[0].shape, b[0].shape} + + def test_sorted_derived_union(self): + """A derived (chunked) catalog with an order spec also survives.""" + p = dc.get_example_patch() + t = p.get_coord("time") + early = p.update_attrs(tag="z_early") + # the gap keeps two outputs; distinct tags survive unmerged + late = p.update_coords(time_min=t.max() + 10 * t.step).update_attrs( + tag="a_late" + ) + chunked = dc.spool([early, late]).chunk(time=None) + srt = chunked.sort("tag") + want = [x.attrs.tag for x in srt] + assert want == ["a_late", "z_early"] # tag order != time order + combined = srt + dc.spool([]) + assert [x.attrs.tag for x in combined] == want + + def test_combined_pickles(self): + """A union holding a materialized operand round-trips pickling.""" + import pickle + + p = dc.get_example_patch() + t = p.get_coord("time") + lo, hi = t.min() + 10 * t.step, t.min() + 20 * t.step + combined = dc.spool([p]).select(time=(lo, hi)) + dc.spool([]) + loaded = pickle.loads(pickle.dumps(combined)) + assert len(loaded) == 1 + assert loaded[0].shape == combined[0].shape + + +def _patch_shape(patch): + """Module-level shape getter (process pools need a picklable callable).""" + return patch.shape + + +class TestMixedViewPickle: + """Serialization keeps plan routes in mixed views (round-4 F2).""" + + def test_sliced_mixed_union_pickles(self): + """A sliced union of planned and live rows loads all rows back.""" + import pickle + + p = dc.get_example_patch() + t = p.get_coord("time") + trimmed = dc.spool([p]).select( + time=(t.min() + 10 * t.step, t.min() + 20 * t.step) + ) + other = p.new().update_attrs(tag="other") + view = (trimmed + dc.spool([other]))[:] + loaded = pickle.loads(pickle.dumps(view)) + shapes = {loaded[i].shape for i in range(len(loaded))} + assert shapes == {(300, 11), (300, 2000)} + + def test_mixed_union_map_processes(self): + """Process-backed map ships plan routes with each task.""" + from concurrent.futures import ProcessPoolExecutor + + p = dc.get_example_patch() + t = p.get_coord("time") + trimmed = dc.spool([p]).select( + time=(t.min() + 10 * t.step, t.min() + 20 * t.step) + ) + combined = trimmed + dc.spool([p.new().update_attrs(tag="other")]) + with ProcessPoolExecutor(2) as executor: + shapes = set(combined.map(_patch_shape, client=executor)) + assert shapes == {(300, 11), (300, 2000)} diff --git a/tests/test_io/test_indexer.py b/tests/test_io/test_indexer.py index d0b89b03f..b7451aa40 100644 --- a/tests/test_io/test_indexer.py +++ b/tests/test_io/test_indexer.py @@ -7,39 +7,31 @@ import shutil from contextlib import suppress from pathlib import Path -from unittest.mock import patch -import numpy as np import pandas as pd import pytest -from packaging.version import parse as get_version from upath import UPath -import dascore as dc from dascore.config import set_config -from dascore.examples import spool_to_directory from dascore.exceptions import InvalidSpoolError -from dascore.io.indexer import DirectoryIndexer -from dascore.utils.hdf5 import HDFPatchIndexManager +from dascore.io.index.indexer import DBDirectoryIndexer from dascore.utils.patch import get_patch_names @pytest.fixture(scope="class") def basic_indexer(two_patch_directory): - """Return and indexer on the basic spool directory.""" - return DirectoryIndexer(two_patch_directory) - - -@pytest.fixture(scope="class") -def adjacent_indexer(adjacent_spool_directory): - """Return and indexer on the basic spool directory.""" - return DirectoryIndexer(adjacent_spool_directory).update() + """Return an indexer on the basic spool directory.""" + indexer = DBDirectoryIndexer(two_patch_directory).update(progress=None) + yield indexer + indexer.close() @pytest.fixture(scope="class") def diverse_indexer(diverse_spool_directory): - """Return and indexer on the basic spool directory.""" - return DirectoryIndexer(diverse_spool_directory).update() + """Return an indexer on the diverse spool directory.""" + indexer = DBDirectoryIndexer(diverse_spool_directory).update(progress=None) + yield indexer + indexer.close() @pytest.fixture(scope="class") @@ -48,23 +40,13 @@ def diverse_df(diverse_indexer): return diverse_indexer() -@pytest.fixture() -def diverse_df_reset_cache(diverse_indexer): - """Return the indexer with a reset cache.""" - return DirectoryIndexer(diverse_indexer.path) - - -@pytest.fixture(params=[diverse_indexer, diverse_df_reset_cache]) -def diverse_ind(request): - """Aggregate the diverse indexers.""" - return request.getfixturevalue(request.param.__name__) - - @pytest.fixture() def empty_index(tmp_path_factory): """Create an index around an empty directory.""" path = tmp_path_factory.mktemp("index_created_test") - return DirectoryIndexer(path).update() + indexer = DBDirectoryIndexer(path).update(progress=None) + yield indexer + indexer.close() class TestFindIndex: @@ -73,66 +55,59 @@ class TestFindIndex: @pytest.fixture() def unwritable_directory(self, tmp_path_factory): """Return an un-writable directory.""" - # currently this doesn't work on windows so we need to skip any test - # that depend on this fixture if running on windows if "windows" in platform.system().lower(): pytest.skip("Cant run this test on windows") path = tmp_path_factory.mktemp("read_only_data_file") os.chmod(path, 0o444) - return path + yield path + os.chmod(path, 0o755) @pytest.fixture() def directory_indexer_bad_cache(self, tmp_path_factory): - """Create a subclass of indexer which has a bd index_map file.""" + """Create a bad index_map file.""" path = tmp_path_factory.mktemp("corrupt_cache_test") cache_path = path / "corrupt_cache.json" - with cache_path.open("wt") as fi: fi.write("{'bad': 'json'") return cache_path def test_directory_cant_write(self, unwritable_directory): """Ensure correct path is found when a read-only directory is used.""" - dir_index = DirectoryIndexer(unwritable_directory) + dir_index = DBDirectoryIndexer(unwritable_directory) index_path = dir_index.index_path index_map_path = dir_index.index_map_path assert index_map_path.parent == index_path.parent def test_specify_index_path(self, tmp_path_factory): - """Ensure specifying a Path works.""" + """Ensure specifying a Path works and is remembered.""" data_path = tmp_path_factory.mktemp("data_dir") - index_path = tmp_path_factory.mktemp("index_dir") / "index.h5" - dir_index = DirectoryIndexer(data_path, index_path=index_path) + index_path = tmp_path_factory.mktemp("index_dir") / "index.sqlite" + dir_index = DBDirectoryIndexer(data_path, index_path=index_path) assert dir_index.index_path == index_path - # loading a new data dir should now remember where this is. - dir_index2 = DirectoryIndexer(data_path) + # loading the same data dir should now remember where this is. + dir_index2 = DBDirectoryIndexer(data_path) assert dir_index2.index_path == index_path def test_writeable_dir_index_not_there(self, tmp_path_factory): - """Tests for when there is writeable directory.""" + """Tests for when there is a writeable directory.""" path = tmp_path_factory.mktemp("normal_indexer_test") - dir_indexer = DirectoryIndexer(path) + dir_indexer = DBDirectoryIndexer(path) assert dir_indexer.index_path.parent == path def test_writable_dir_index_exists(self, tmp_path_factory): """A test case where the index does exist.""" path = tmp_path_factory.mktemp("normal_indexer_test") - index_path = path / DirectoryIndexer._index_name - index_path.open("w").close() - dir_indexer = DirectoryIndexer(path) - assert dir_indexer.index_path == index_path - - def test_corrupt_cache( - self, - directory_indexer_bad_cache, - tmp_path_factory, - ): - """Ensure a corrupted cache doesnt crash indexing. See #508.""" + first = DBDirectoryIndexer(path) + second = DBDirectoryIndexer(path) + assert first.index_path == second.index_path + assert first.index_path.exists() + + def test_corrupt_cache(self, directory_indexer_bad_cache, tmp_path_factory): + """Ensure a corrupted cache doesn't crash indexing. See #508.""" path = tmp_path_factory.mktemp("corrupt_cache_test") - # Test passes if this doesn't raise should not raise. assert directory_indexer_bad_cache.exists() with set_config(directory_index_map_path=directory_indexer_bad_cache): - DirectoryIndexer(path) + DBDirectoryIndexer(path) assert not directory_indexer_bad_cache.exists() def test_remote_directory_not_supported(self): @@ -140,11 +115,11 @@ def test_remote_directory_not_supported(self): path = UPath("memory://dascore/indexer") (path / "file.txt").write_text("x") with pytest.raises(InvalidSpoolError, match="local filesystem"): - DirectoryIndexer(path) + DBDirectoryIndexer(path) def test_local_upath_normalized_to_path(self, tmp_path): """Local UPath inputs should normalize to pathlib.Path internally.""" - out = DirectoryIndexer(UPath(tmp_path)) + out = DBDirectoryIndexer(UPath(tmp_path)) assert isinstance(out.path, Path) assert out.path == Path(tmp_path).absolute() @@ -152,7 +127,7 @@ def test_index_map_path_comes_from_config(self, tmp_path): """Index map paths should be sourced from runtime configuration.""" index_map_path = tmp_path / "cache_paths.json" with set_config(directory_index_map_path=index_map_path): - out = DirectoryIndexer(tmp_path) + out = DBDirectoryIndexer(tmp_path) assert out.index_map_path == index_map_path @@ -164,24 +139,11 @@ def test_str_repr(self, basic_indexer): out = str(basic_indexer) assert "object at" not in out - def test_version(self, basic_indexer): - """Ensure the version written to file is correct.""" - updated = basic_indexer.update() - index_version = updated._index_table._index_version - assert index_version == dc.__last_version__ - assert get_version(index_version) > get_version("0.0.1") - - def test_update_does_not_reconstruct_patch_summary_from_flat_dicts( - self, two_patch_directory - ): - """Indexer update should not reconstruct PatchSummary from flat rows.""" - indexer = DirectoryIndexer(two_patch_directory) - with patch( - "dascore.core.summary.PatchSummary.model_validate", - side_effect=AssertionError("unexpected PatchSummary.model_validate call"), - ): - updated = indexer.update() - assert updated.index_path.exists() + def test_metadata(self, basic_indexer): + """The index records its schema version and identity.""" + meta = basic_indexer._backend.get_metadata() + assert meta["what_is_this"] == "dascore_spool_index" + assert meta["index_version"] >= 1 class TestGetContents: @@ -197,175 +159,120 @@ def test_get_contents(self, basic_indexer, two_patch_directory): names_files = {x.name for x in files} assert names_df == names_files - def test_filter_large_starttime(self, diverse_df, diverse_ind): - """Ensure the index can be filtered by end time.""" + def test_filter_time_after(self, diverse_df, diverse_indexer): + """Half-open time range keeps every file overlapping it.""" max_starttime = diverse_df["time_min"].max() - filtered = diverse_df[diverse_df["time_min"] >= max_starttime] - out = diverse_ind(time_min=max_starttime) - assert len(out) == len(filtered) + expected = diverse_df[diverse_df["time_max"] >= max_starttime] + out = diverse_indexer(time=(max_starttime, None)) + assert len(out) == len(expected) - def test_filter_small_starttime(self, diverse_df, diverse_ind): - """Ensure the index can be filtered by start time.""" + def test_filter_time_before(self, diverse_df, diverse_indexer): + """Half-open time range keeps every file overlapping it.""" min_endtime = diverse_df["time_max"].min() - filtered = diverse_df[diverse_df["time_max"] <= min_endtime] - out = diverse_ind(time_max=min_endtime) - assert len(out) == len(filtered) + expected = diverse_df[diverse_df["time_min"] <= min_endtime] + out = diverse_indexer(time=(None, min_endtime)) + assert len(out) == len(expected) - def test_filter_station_exact(self, diverse_df, diverse_ind): - """Ensure contents can be filtered on time.""" - # tests for filtering with exact station name + def test_filter_station_exact(self, diverse_df, diverse_indexer): + """Ensure contents can be filtered on an attr.""" exact_name = diverse_df["station"].unique()[0] - new_df = diverse_ind(station=exact_name) + new_df = diverse_indexer(station=exact_name) assert (new_df["station"] == exact_name).all() - def test_filter_isin(self, diverse_df, diverse_ind): - """Ensure contents can be filtered on time.""" - # tests for filtering with exact station name - exact_name = diverse_df["station"].unique()[0] - new_df = diverse_ind(station=exact_name) - assert (new_df["station"] == exact_name).all() + def test_filter_isin(self, diverse_df, diverse_indexer): + """Ensure contents can be filtered with a collection.""" + # empty strings mean "attr missing" and are not queryable (spec). + stations = [x for x in diverse_df["station"].unique() if x] + new_df = diverse_indexer(station=stations[:2]) + assert set(new_df["station"]) <= set(stations[:2]) + assert len(new_df) def test_empty_index(self, empty_index): """An empty index should return an empty dataframe.""" df = empty_index() assert df.empty - def test_default_buffer_comes_from_config(self, basic_indexer, monkeypatch): - """Configured index buffer should be used when no explicit buffer is passed.""" - seen = {} - - def _fake_to_timedelta64(value): - seen["buffer"] = value - return value - - monkeypatch.setattr("dascore.io.indexer.to_timedelta64", _fake_to_timedelta64) - with set_config(index_query_buffer=np.timedelta64(5, "s")): - basic_indexer.get_contents() - assert seen["buffer"] == np.timedelta64(5, "s") - - def test_explicit_buffer_overrides_config(self, basic_indexer): - """Explicit get_contents buffer should override config defaults.""" - seen = {} - - def _fake_to_timedelta64(value): - seen["buffer"] = value - return value - - with set_config(index_query_buffer=np.timedelta64(10, "s")): - with pytest.MonkeyPatch.context() as monkeypatch: - monkeypatch.setattr( - "dascore.io.indexer.to_timedelta64", _fake_to_timedelta64 - ) - basic_indexer.get_contents(buffer=np.timedelta64(0, "s")) - assert seen["buffer"] == np.timedelta64(0, "s") - - -class TestHDFPatchIndexManager: - """Tests for config-backed HDF index defaults.""" - - def test_buffer_comes_from_config(self, tmp_path): - """Index manager buffer property should reflect runtime config.""" - manager = HDFPatchIndexManager(tmp_path / "index.h5") - with set_config(index_query_buffer=np.timedelta64(7, "s")): - assert manager.buffer == np.timedelta64(7, "s") - class TestUpdate: - """Tests for updating index.""" - - def make_simple_index_with_version(self, path, version): - """Helper function to make a simple index with desired version.""" - patch = dc.get_example_patch() - spool_to_directory([patch], path) - # this ensure the version is set to fake version - old_version = dc.__last_version__ - # for some reason monkeypatch fixture wasnt setting version back - # so I had to manually set and revert dascore version. - setattr(dc, "__last_version__", version) - spool = dc.spool(path).update() - setattr(dc, "__last_version__", old_version) - # ensure version monkey patch worked. - meta = spool.indexer.get_index_metadata() - assert meta["index_version"] == version - return path + """Tests for updating the index.""" @pytest.fixture(scope="class") def spool_directory_with_non_das_file(self, two_patch_directory, tmp_path_factory): """Create a directory with some das files and some non-das files.""" new = tmp_path_factory.mktemp("unreadable_test") / "sub" shutil.copytree(two_patch_directory, new) - indexer = DirectoryIndexer(new) - # remove index if it exists with suppress(FileNotFoundError): - indexer.index_path.unlink() - # add a non das file + for index in Path(new).glob(".dascore_index*"): + index.unlink() with open(new / "not_das.open", "w") as fi: fi.write("cant be das, can it?") return new - @pytest.fixture() - def index_old_version(self, monkeypatch, tmp_path_factory): - """Create an index which has an old, incompatible version.""" - # cant use random_patch fixture due to scope-mismatch w/ monkeypatch. - path = tmp_path_factory.mktemp("index_old_version ") - self.make_simple_index_with_version(path, "0.0.1") - return path - - @pytest.fixture() - def index_new_version(self, monkeypatch, tmp_path_factory): - """Create an index which has an old, incompatible version.""" - # cant use random_patch fixture due to scope-mismatch w/ monkeypatch. - path = tmp_path_factory.mktemp("index_new_version ") - # a ridiculously high version - fake_version = "1000.0.1" - assert get_version(fake_version) > get_version(dc.__last_version__) - self.make_simple_index_with_version(path, fake_version) - return path - def test_add_one_patch(self, empty_index, random_patch): """Ensure a new patch added to the directory shows up.""" path = empty_index.path / get_patch_names(random_patch).iloc[0] random_patch.io.write(path, file_format="dasdae") - new_index = empty_index.update() + new_index = empty_index.update(progress=None) contents = new_index() assert len(contents) == 1 def test_index_with_bad_file(self, spool_directory_with_non_das_file): """Ensure if one file is not readable index continues.""" - indexer = DirectoryIndexer(spool_directory_with_non_das_file) - # if this doesn't fail the test passes - updated = indexer.update() - assert isinstance(updated, DirectoryIndexer) - - def test_old_index_recreated(self, index_old_version): - """Ensure the old index is recreated when update is called.""" - msg = "Recreating the index now." - with pytest.warns(UserWarning, match=msg): - dc.spool(index_old_version).update() - - def test_new_version_warnings(self, index_new_version): - """Ensure an index file with a newer version of dascore issues a warning.""" - msg = "The index was created with a newer version of dascore" - dc.spool(index_new_version) - with pytest.warns(UserWarning, match=msg): - dc.spool(index_new_version).update() + indexer = DBDirectoryIndexer(spool_directory_with_non_das_file) + updated = indexer.update(progress=None) + assert isinstance(updated, DBDirectoryIndexer) + assert len(updated()) == 2 + + def test_removed_file_dropped(self, two_patch_directory, tmp_path_factory): + """A deleted file's rows disappear on the next update.""" + new = tmp_path_factory.mktemp("removed_file_test") / "sub" + shutil.copytree(two_patch_directory, new) + for index in Path(new).glob(".dascore_index*"): + index.unlink() + indexer = DBDirectoryIndexer(new).update(progress=None) + assert len(indexer()) == 2 + next(iter(Path(new).glob("*.hdf5"))).unlink() + assert len(indexer.update(progress=None)()) == 1 + + def test_noop_update_rescans_nothing(self, basic_indexer): + """Unchanged sources are not rescanned.""" + before = basic_indexer._backend.get_sources()["last_indexed_ns"].max() + basic_indexer.update(progress=None) + after = basic_indexer._backend.get_sources()["last_indexed_ns"].max() + assert before == after def test_update_with_specific_paths(self, basic_indexer): - """Test updating with specific file paths to cover _get_paths method.""" - # Get files in the directory - files = list(basic_indexer.path.rglob("*.hdf5")) - assert len(files) > 0, "Need at least one file for testing" - - # Test update with specific paths (relative paths) - relative_paths = [f.name for f in files[:1]] # Use just first file - updated = basic_indexer.update(paths=relative_paths) - contents = updated() - - # Should have at least the file we specified - assert len(contents) >= 1 - - # Test update with absolute paths - absolute_paths = [str(f) for f in files[:1]] - updated2 = basic_indexer.update(paths=absolute_paths) - contents2 = updated2() - assert len(contents2) >= 1 + """Updating with specific paths restricts the rescan.""" + files = sorted(basic_indexer.path.rglob("*.hdf5")) + assert len(files) >= 2 + + def _indexed_times(): + sources = basic_indexer._backend.get_sources().set_index("source_path") + return sources["last_indexed_ns"].to_dict() + + before = _indexed_times() + for path in files[:2]: + stat = path.stat() + os.utime(path, ns=(stat.st_atime_ns, stat.st_mtime_ns + 1_000_000_000)) + + first, second = (basic_indexer._rel(path) for path in files[:2]) + basic_indexer.update(paths=[files[0].name], progress=None) + after_relative = _indexed_times() + assert after_relative[first] > before[first] + assert after_relative[second] == before[second] + + basic_indexer.update(paths=[str(files[1])], progress=None) + after_absolute = _indexed_times() + assert after_absolute[first] == after_relative[first] + assert after_absolute[second] > after_relative[second] + + +class TestNameResolution: + """Unknown names raise per the selector spec.""" + + def test_unknown_name_raises(self, basic_indexer): + """Names in neither namespace error clearly (#435).""" + from dascore.io.index.query import InvalidSpoolQueryError + + with pytest.raises(InvalidSpoolQueryError, match="neither an attribute"): + basic_indexer(bad_dimension=(1, 2)) diff --git a/tests/test_io/test_io_core.py b/tests/test_io/test_io_core.py index b6e47b392..222f63e17 100644 --- a/tests/test_io/test_io_core.py +++ b/tests/test_io/test_io_core.py @@ -928,3 +928,47 @@ def test_get_supported_io_table(self): # assert that the length of the DataFrame is not 0 assert len(result_df) > 0 + + +class TestIOCoreCoverageEdges: + """Remaining io.core resolution/robustness branches.""" + + def test_numeric_singleton_without_identity_not_trusted(self): + """A positional ID cannot resolve an anonymous trimmed singleton.""" + from dascore.exceptions import PatchAttributeError + from dascore.io.core import _resolve_read_spool + + spool = dc.spool([dc.get_example_patch()]) + with pytest.raises(PatchAttributeError, match="uniquely resolved"): + _resolve_read_spool(spool, source_patch_id="1") + + def test_non_unique_patch_resolution_raises(self): + """An unresolvable source id in a multi-patch read raises clearly.""" + from dascore.exceptions import PatchAttributeError + from dascore.io.core import _select_patch_from_spool + + spool = dc.spool([dc.get_example_patch(tag="a"), dc.get_example_patch(tag="b")]) + with pytest.raises(PatchAttributeError, match="uniquely resolved"): + _select_patch_from_spool(spool, source_patch_id="neither-id-nor-index") + + def test_single_patch_resolved_by_name(self): + """A one-patch read resolves when the id matches the patch name.""" + from dascore.io.core import _select_patch_from_spool + + patch = dc.get_example_patch() + spool = dc.spool([patch]) + resolved = _select_patch_from_spool( + spool, source_patch_id=str(patch.get_patch_name()) + ) + assert resolved == patch + + def test_corrupt_file_format_detection_is_robust(self, tmp_path): + """A reader raising during format detection is caught, not propagated.""" + from dascore.exceptions import UnknownFiberFormatError + + # valid HDF5 magic followed by garbage: an HDF5 reader raises while + # probing, which format detection must swallow before giving up. + bad = tmp_path / "bad.h5" + bad.write_bytes(b"\x89HDF\r\n\x1a\n" + b"\x00" * 256) + with pytest.raises(UnknownFiberFormatError): + dc.get_format(bad) diff --git a/tests/test_io/test_pickle/test_pickle.py b/tests/test_io/test_pickle/test_pickle.py index fe81fc808..835d21ba7 100644 --- a/tests/test_io/test_pickle/test_pickle.py +++ b/tests/test_io/test_pickle/test_pickle.py @@ -42,6 +42,18 @@ def test_read_pickle(self, pickle_patch_path, random_patch): assert isinstance(out[0], dc.Patch) assert random_patch == out[0] + def test_spool_from_pickle(self, pickle_patch_path, random_patch): + """dc.spool on a scanless format wraps the read spool and serves it. + + PICKLE implements read but not scan, so dc.spool routes through + Spool(dc.read(...)); the wrapped patches must load back. + """ + spool = dc.spool(pickle_patch_path) + assert len(spool) == 1 + assert len(spool.get_contents()) == 1 + assert spool[0] == random_patch + assert next(iter(spool)) == random_patch + def test_file_not_there(self): """Get format should return false if the file doesn't exist.""" parser = PickleIO() diff --git a/tests/test_io/test_prodml/test_prod_ml.py b/tests/test_io/test_prodml/test_prod_ml.py index 6f3fd519d..5e8ad374d 100644 --- a/tests/test_io/test_prodml/test_prod_ml.py +++ b/tests/test_io/test_prodml/test_prod_ml.py @@ -7,7 +7,6 @@ import h5py import pandas as pd import pytest -import tables import dascore as dc from dascore.core.coords import get_coord @@ -47,12 +46,9 @@ def issue_221_patch_path(self, tmp_path_factory): tmp_path = tmp_path_factory.mktemp("issue_221") path = dc.utils.downloader.fetch("prodml_2.0.h5") new_path = shutil.copy2(path, tmp_path / "prod_2_monkey_patched.h5") - with tables.open_file(new_path, "a") as fi: + with h5py.File(new_path, "a") as fi: # monkey patch dimensions to simulate issue. - new_dims = "time, locus" - parent_node = fi.root.Acquisition["Raw[0]"] - node = parent_node["RawData"] - node._v_attrs.Dimensions = new_dims + fi["Acquisition/Raw[0]/RawData"].attrs["Dimensions"] = "time, locus" return new_path @pytest.fixture(scope="class") diff --git a/tests/test_io/test_terra15/test_terra15.py b/tests/test_io/test_terra15/test_terra15.py index 903a05738..ab04b522b 100644 --- a/tests/test_io/test_terra15/test_terra15.py +++ b/tests/test_io/test_terra15/test_terra15.py @@ -5,10 +5,10 @@ import shutil from typing import ClassVar +import h5py import numpy as np import pandas as pd import pytest -import tables import dascore as dc from dascore.io.terra15 import Terra15FormatterV4 @@ -22,8 +22,8 @@ def missing_gps_terra15_hdf5(self, terra15_v5_path, tmp_path_factory): """Creates a terra15 file with missing GPS Time.""" new = tmp_path_factory.mktemp("missing_gps") / "missing.hdf5" shutil.copy(terra15_v5_path, new) - with tables.open_file(new, "a") as fi: - fi.root.data_product.gps_time._f_remove() + with h5py.File(new, "a") as fi: + del fi["data_product/gps_time"] return new def test_missing_gps_time(self, missing_gps_terra15_hdf5): diff --git a/tests/test_io/test_xml_binary/test_xml_binary.py b/tests/test_io/test_xml_binary/test_xml_binary.py index a3726eb52..dbbed1347 100644 --- a/tests/test_io/test_xml_binary/test_xml_binary.py +++ b/tests/test_io/test_xml_binary/test_xml_binary.py @@ -191,6 +191,14 @@ def test_mtime(self, binary_xml_directory): scan3 = dc.scan(binary_xml_directory, timestamp=mtime - 50) assert len(scan3) == 2 + def test_direct_scan_filters_all_by_mtime(self, binary_xml_directory): + """The FiberIO scan contract returns empty after filtering every file.""" + fiber = XMLBinaryV1() + newest = max( + path.stat().st_mtime for path in binary_xml_directory.glob("*.raw") + ) + assert fiber.scan(binary_xml_directory, timestamp=newest + 1) == [] + def test_remote_directory(self, remote_binary_xml_directory): """Remote XMLBinary directories should be scannable.""" fiber = XMLBinaryV1() diff --git a/tests/test_utils/test_chunk.py b/tests/test_utils/test_chunk.py index acf9b689c..3f12ada00 100644 --- a/tests/test_utils/test_chunk.py +++ b/tests/test_utils/test_chunk.py @@ -9,8 +9,9 @@ import pytest import dascore as dc -from dascore.exceptions import ParameterError -from dascore.utils.chunk import ChunkManager, get_intervals +from dascore.exceptions import ChunkError +from dascore.utils.chunk import get_intervals +from dascore.utils.chunk_plan import build_chunk_plan from dascore.utils.time import to_timedelta64 STARTTIME = np.datetime64("2020-01-03") @@ -105,27 +106,25 @@ def test_timedelta_start_numeric_length(self): assert out[-1, 1] == stop -class TestBasicChunkDF: - """Test basic DF chunking.""" +class TestChunkPlanDF: + """Dataframe-level chunk planning (ported from the old ChunkManager tests).""" @pytest.fixture() def df_different_sample_rates(self, contiguous_df): - """Tests for a df which does have overlaps but different sampling rates.""" + """Adjacent blocks with different sampling rates.""" df1 = contiguous_df.copy() df2 = contiguous_df.copy() time_span = df1["time_max"].max() - df1["time_min"].min() df2["time_min"] += time_span df2["time_max"] += time_span df2["time_step"] = df1["time_step"] * 2 - out = pd.concat([df1, df2], axis=0).reset_index(drop=True) - return out + return pd.concat([df1, df2], axis=0).reset_index(drop=True) def test_rechunk_contiguous(self, contiguous_df): """Test rechunking with no gaps.""" time_interval = (contiguous_df["time_max"] - contiguous_df["time_min"]).max() new_time_interval = time_interval / 2 - chunker = ChunkManager(time=new_time_interval) - _, out = chunker.chunk(contiguous_df) + out = build_chunk_plan(contiguous_df, time=new_time_interval).outputs assert len(out) == 2 * len(contiguous_df) time_step = out["time_step"].iloc[0] new_interval = (out["time_max"] - out["time_min"] + time_step).max() @@ -137,17 +136,15 @@ def test_rechunk_contiguous_with_sr_separation(self, contiguous_sr_spaced_df): sr = df["time_step"] time_interval = (sr + df["time_max"] - df["time_min"]).max() new_time_interval = time_interval / 2 - chunker = ChunkManager(time=new_time_interval) - _, out = chunker.chunk(df) + out = build_chunk_plan(df, time=new_time_interval).outputs assert len(out) == 2 * len(df) new_interval = (out["time_max"] - out["time_min"]).max() assert new_interval == (new_time_interval - sr.iloc[0]) def test_rechunk_different_sr(self, df_different_sample_rates): - """Ensure segments with different sample rates don't get combined.""" + """Segments with different sample rates don't get combined.""" df = df_different_sample_rates - chunker = ChunkManager(overlap=None, time=23) - _, out = chunker.chunk(df) + out = build_chunk_plan(df, time=23).outputs dt = np.sort(np.unique(out["time_step"])) assert len(dt) == 2, "both dt should remain" # the second part of the df should start at the one minute mark @@ -170,105 +167,73 @@ def test_chunk_uses_step_from_each_group(self): "time_step": [np.timedelta64(1, "s"), np.timedelta64(10, "s")], } ) - - _, chunked = ChunkManager(time=50).chunk(df) - + chunked = build_chunk_plan(df, time=50).outputs ten_second_group = chunked[chunked["time_step"] == np.timedelta64(10, "s")] first = ten_second_group.iloc[0] assert first["time_max"] - first["time_min"] == np.timedelta64(40, "s") def test_keep_leftovers(self, contiguous_df): - """Ensure leftovers show up in df.""" - chunker = ChunkManager(overlap=None, keep_partial=True, time=28) - _, out = chunker.chunk(contiguous_df) + """Ensure leftovers show up in outputs.""" + out = build_chunk_plan(contiguous_df, keep_partial=True, time=28).outputs assert len(out) == 3 assert out["time_max"].max() == contiguous_df["time_max"].max() def test_overlap(self, contiguous_df): - """Ensure overlapping segments work.""" + """Ensure overlapping segments work, with timedelta or float overlap.""" over = to_timedelta64(10) - chunker1 = ChunkManager(overlap=over, time=20) - _, out = chunker1.chunk(contiguous_df) + out = build_chunk_plan(contiguous_df, overlap=over, time=20).outputs expected = over - contiguous_df["time_step"].iloc[0] olap = out.shift()["time_max"] - out["time_min"] assert np.all(pd.isnull(olap) | (olap == expected)) - # now ensure floats work for overlap param - chunker2 = ChunkManager(overlap=10, time=20) - _, out2 = chunker2.chunk(contiguous_df) + out2 = build_chunk_plan(contiguous_df, overlap=10, time=20).outputs assert out.equals(out2) def test_chunk_on_split(self, terra15_file_spool): """Ensure chunking which creates a slice at the end time works.""" - # this spool was selected because I first observed the issue in it. df = terra15_file_spool.get_contents() dur = (df["time_max"] - df["time_min"]).iloc[0] seg_len = dur / 3 dt = df["time_step"].iloc[0] - chunker = ChunkManager(keep_partial=True, time=seg_len) - _, chunk_df = chunker.chunk(df) + chunk_df = build_chunk_plan(df, keep_partial=True, time=seg_len).outputs duration = chunk_df["time_max"] - chunk_df["time_min"] assert duration.sum() == ((seg_len - dt) * 3) assert len(duration) == 3 assert (duration > np.timedelta64(0, "s")).all() def test_nan_in_df(self, contiguous_df): - """Ensure contiguous df with nan inside still works.""" + """A null envelope row breaks continuity when dropped.""" df = contiguous_df.copy() - # Adding null values on row 3 df.loc[3, "time_min"] = dc.to_datetime64("NaT") - # Which means new time should start in row 4 because of the gap. expected_start = df.loc[4, "time_min"] - chunker = ChunkManager(keep_partial=True, time=dc.to_timedelta64(15)) - _, chunk_df = chunker.chunk(df) - assert expected_start in set(chunk_df["time_min"]) + plan = build_chunk_plan( + df, keep_partial=True, missing_dim="drop", time=dc.to_timedelta64(15) + ) + assert expected_start in set(plan.outputs["time_min"]) def test_all_nan(self, contiguous_df): - """Ensure when all NaNs are encountered the chunked df is empty.""" + """When all rows lack the dim (and are dropped) the plan is empty.""" nat = dc.to_datetime64("NaT") df = contiguous_df.assign(time_min=nat, time_max=nat) - chunker = ChunkManager(time=dc.to_timedelta64(1.2)) - _, chunk_df = chunker.chunk(df) - assert chunk_df.empty + plan = build_chunk_plan(df, missing_dim="drop", time=dc.to_timedelta64(1.2)) + assert plan.outputs.empty def test_nan_in_sample_ok(self, contiguous_df): """Ensure a NaN in the sampling rate is ok.""" df = contiguous_df.assign(time_step=dc.to_timedelta64("NaT")) dur = (df["time_max"] - df["time_min"]).iloc[0] - chunker = ChunkManager(time=dc.to_timedelta64(dur / 2)) - _, chunk_df = chunker.chunk(df) + chunk_df = build_chunk_plan(df, time=dc.to_timedelta64(dur / 2)).outputs assert isinstance(chunk_df, pd.DataFrame) assert len(chunk_df) == 2 * len(contiguous_df) assert np.all(pd.isnull(chunk_df["time_step"])) + def test_unknown_dim_raises(self, contiguous_df): + """An unknown chunk dimension raises a clear error.""" + with pytest.raises(ChunkError, match="Time"): + build_chunk_plan(contiguous_df, Time=10) -class TestChunkExceptions: - """Tests for various exceptions from the chunk manager.""" - - def test_raises_overlap_no_chunksize(self): - """Specifying an overlap and no chunk size should raise.""" - with pytest.raises(ParameterError, match="used for merging"): - ChunkManager(time=None, overlap=10) - - def test_raises_zero_length_multiple_kwargs(self): - """Ensure multiple kwargs raises nice error.""" - with pytest.raises(ParameterError, match="along one dimension"): - ChunkManager(time=10, distance=1) - def test_raises_zero_length_chunk(self): - """Ensure zero length chunk raises.""" - with pytest.raises(ParameterError, match="must be greater than 0"): - ChunkManager(time=0) - - def test_raises_invalid_key_in_kwargs(self, contiguous_df): - """Ensure an invalid key in kwargs raises an error.""" - chunk_manager = ChunkManager(Time=10) - chunk_manager.patch = type("Patch", (object,), {"dims": ["time", "distance"]})() - with pytest.raises(ParameterError, match="Cannot chunk spool or"): - chunk_manager.chunk(contiguous_df) - - -class TestChunkToMerge: - """Tests for using chunking to merge contiguous, or overlapping, data.""" +class TestChunkPlanToMerge: + """Merge-mode planning on raw dataframes.""" @pytest.fixture() def gapy_df(self, contiguous_df): @@ -285,29 +250,26 @@ def gapy_df_unordered(self, gapy_df): def test_chunk_can_merge(self, contiguous_df): """Ensure chunk can be used to merge unspecified segment lengths.""" - cm = ChunkManager(time=None) - _, out = cm.chunk(contiguous_df) + out = build_chunk_plan(contiguous_df, time=None).outputs assert len(out) == 1 assert out["time_min"].min() == contiguous_df["time_min"].min() def test_doesnt_merge_gappy_df(self, gapy_df): """Ensure the gappy dataframe doesn't get merged.""" - cm = ChunkManager(time=None) - _, out = cm.chunk(gapy_df) + out = build_chunk_plan(gapy_df, time=None).outputs assert len(gapy_df) == len(out) - expected_durations = gapy_df["time_max"] - gapy_df["time_min"] - durations = out["time_max"] - out["time_min"] - assert expected_durations.equals(durations) + expected = (gapy_df["time_max"] - gapy_df["time_min"]).sort_values() + durations = (out["time_max"] - out["time_min"]).sort_values() + assert np.array_equal(expected.values, durations.values) def test_doesnt_merge_unordered_gappy_df(self, gapy_df_unordered): - """Ensure the gappy dataframe doesn't get merged.""" + """Row order must not affect merge results.""" df = gapy_df_unordered - cm = ChunkManager(time=None) - _, out = cm.chunk(df) + out = build_chunk_plan(df, time=None).outputs assert len(df) == len(out) - expected_durations = df["time_max"] - df["time_min"] - durations = out["time_max"] - out["time_min"] - assert expected_durations.equals(durations) + expected = (df["time_max"] - df["time_min"]).sort_values() + durations = (out["time_max"] - out["time_min"]).sort_values() + assert np.array_equal(expected.values, durations.values) def test_no_warning_when_final_groups_stay_separate(self, contiguous_df): """No warning if other group components prevent final forced merge.""" @@ -317,89 +279,69 @@ def test_no_warning_when_final_groups_stay_separate(self, contiguous_df): df.loc[1, "time_min"] = df.loc[0, "time_max"] + 5 * step.iloc[0] df.loc[1, "time_max"] = df.loc[1, "time_min"] + 10 * step.iloc[1] df["station"] = ["sta1", "sta2"] - cm = ChunkManager(time=None, tolerance=10, group_columns=("station",)) - with warnings.catch_warnings(): warnings.filterwarnings("error") - _, out = cm.chunk(df) + plan = build_chunk_plan(df, time=None, tolerance=10, group=("station",)) + assert len(plan.outputs) == 2 - assert len(out) == 2 + def test_forced_merge_warns(self, contiguous_df): + """A tolerance forcing a merge across a real gap warns (#662).""" + df = contiguous_df.iloc[:2].copy() + step = df["time_step"].iloc[0] + df.loc[0, "time_max"] = df.loc[0, "time_min"] + 10 * step + df.loc[1, "time_min"] = df.loc[0, "time_max"] + 5 * step + df.loc[1, "time_max"] = df.loc[1, "time_min"] + 10 * step + with pytest.warns(UserWarning, match="force merging"): + plan = build_chunk_plan(df, time=None, tolerance=10) + assert len(plan.outputs) == 1 def test_modified_flag_after_merge(self, contiguous_df): - """Test that the modified flag shows False for simple merge.""" - cm = ChunkManager(time=None) - # Need to remove overlapping sample so these really are contiguous - # with no overlaps. - contiguous_df = contiguous_df.assign( - time_max=lambda x: x["time_max"] - x["time_step"] - ) - source, current = cm.chunk(contiguous_df) - inst_df = cm.get_instruction_df(source, current) - assert len(current) == 1 - assert current["time_min"].min() == contiguous_df["time_min"].min() - assert not inst_df["_modified"].any() - - -class TestInstructionDF: - """Sanity checks on intermediary df.""" - - def test_indices(self, contiguous_df): - """Ensure the input/output index belong to input/output df.""" - chunker = ChunkManager(overlap=0, time=10) - in_df, out_df = chunker.chunk(contiguous_df) - instruction = chunker.get_instruction_df(in_df, out_df) - # ensure the source index is set as the index of the instruction_df - assert instruction.index.name == "source_index" - assert set(instruction.index).issubset(set(contiguous_df.index)) - assert set(instruction["current_index"]).issubset(set(out_df.index)) + """The modified flag shows False for a simple contiguous merge.""" + df = contiguous_df.assign(time_max=lambda x: x["time_max"] - x["time_step"]) + plan = build_chunk_plan(df, time=None) + assert len(plan.outputs) == 1 + assert plan.outputs["time_min"].min() == df["time_min"].min() + assert not plan.members["_modified"].any() + + +class TestPlanMembers: + """Sanity checks on the members (instruction) table.""" + + def test_ids(self, contiguous_df): + """Members reference real sources and outputs.""" + plan = build_chunk_plan(contiguous_df, overlap=0, time=10) + members = plan.members + assert set(members["_patch_id"]).issubset(set(range(len(contiguous_df)))) + assert set(members["output_id"]).issubset(set(plan.outputs["output_id"])) def test_different_group_columns(self, contiguous_df_two_stations): - """Ensure instruction df honors differences in group columns.""" + """Ensure members honor differences in group columns.""" df = contiguous_df_two_stations - chunker = ChunkManager( - overlap=0, - time=10, - group_columns=("station",), - keep_partial=True, + plan = build_chunk_plan( + df, overlap=0, time=10, group=("station",), keep_partial=True + ) + joined = plan.members.merge( + df.assign(_patch_id=np.arange(len(df)))[["_patch_id", "station"]], + on="_patch_id", + ).merge( + plan.outputs[["output_id", "station"]], + on="output_id", + suffixes=("_src", "_out"), ) - in_df, out_df = chunker.chunk(df) - instruction = chunker.get_instruction_df(in_df, out_df) - # ensure each output has exactly one station. - for _current_index, sub in instruction.groupby("current_index"): - source = df.loc[sub.index] - # there should only be on station in the source for this group - unique_stations = source["station"].unique() - assert len(unique_stations) == 1 - # ensure all stations are present. - used = in_df.loc[instruction.index] - assert set(used["station"]) == set(in_df["station"]) - assert set(used["_group"]) == set(in_df["_group"]) + assert (joined["station_src"] == joined["station_out"]).all() + assert set(plan.outputs["station"]) == set(df["station"]) def test_modified_flag_if_chunked(self, contiguous_df): """Ensure the modified flag shows up for modified rows.""" - df = contiguous_df - chunker = ChunkManager( - overlap=0, - time=5, - group_columns=("station",), - keep_partial=True, - ) - in_df, out_df = chunker.chunk(df) - instruction = chunker.get_instruction_df(in_df, out_df) - assert instruction["_modified"].all() + plan = build_chunk_plan(contiguous_df, overlap=0, time=5, keep_partial=True) + assert plan.members["_modified"].all() def test_modified_flag_no_chunk(self, contiguous_df): - """Ensure the rows that don't change limits aren't modified.""" + """Rows whose limits don't change aren't modified.""" time_diff = contiguous_df["time_max"] - contiguous_df["time_min"] df = contiguous_df.assign(time_max=lambda x: (x["time_max"] - x["time_step"])) - chunker = ChunkManager( - overlap=0, - time=time_diff.iloc[0], - group_columns=("station",), - keep_partial=True, + plan = build_chunk_plan( + df, overlap=0, time=time_diff.iloc[0], keep_partial=True ) - in_df, out_df = chunker.chunk(df) - - assert (out_df[sorted(out_df.columns)]).equals(in_df[sorted(in_df.columns)]) - instruction = chunker.get_instruction_df(in_df, out_df) - assert not instruction["_modified"].any() + assert len(plan.outputs) == len(df) + assert not plan.members["_modified"].any() diff --git a/tests/test_utils/test_config.py b/tests/test_utils/test_config.py index 162f6e10b..ccf48e725 100644 --- a/tests/test_utils/test_config.py +++ b/tests/test_utils/test_config.py @@ -75,3 +75,44 @@ class _UsesConfig: assert _UsesConfig().value == get_config().display_float_precision with set_config(display_float_precision=7): assert _UsesConfig().value == 7 + + def test_groupby_attrs_default(self): + """The default group attrs are the conventional identity set.""" + expected = ( + "network", + "station", + "data_type", + "data_category", + "tag", + "instrument_id", + "acquisition_id", + ) + assert get_config().groupby_attrs == expected + + def test_groupby_attrs_override(self): + """groupby_attrs round-trips through scoped set_config.""" + previous = get_config() + with set_config(groupby_attrs=("network", "station")): + assert get_config().groupby_attrs == ("network", "station") + assert get_config() == previous + + def test_groupby_attrs_coerced_to_tuple(self): + """List inputs coerce to the immutable tuple form.""" + with set_config(groupby_attrs=["tag"]): + assert get_config().groupby_attrs == ("tag",) + + def test_sampling_group_tolerance_default(self): + """The default sampling group tolerance is 5%.""" + assert get_config().sampling_group_tolerance == 0.05 + + def test_sampling_group_tolerance_override(self): + """sampling_group_tolerance round-trips through scoped set_config.""" + previous = get_config() + with set_config(sampling_group_tolerance=0.01): + assert get_config().sampling_group_tolerance == 0.01 + assert get_config() == previous + + def test_sampling_group_tolerance_must_be_positive(self): + """Non-positive tolerances are rejected.""" + with pytest.raises(ValueError, match="sampling_group_tolerance"): + set_config(sampling_group_tolerance=0) diff --git a/tests/test_utils/test_coordmanager_utils.py b/tests/test_utils/test_coordmanager_utils.py index ced4baa4c..d00cbf66f 100644 --- a/tests/test_utils/test_coordmanager_utils.py +++ b/tests/test_utils/test_coordmanager_utils.py @@ -165,3 +165,22 @@ def test_conflicting_non_dimensional_coords(self, conflicting_non_dim_coords): with pytest.raises(CoordMergeError, match="cannot be merged"): merge_coord_managers([c1, c2], dim="time", drop_conflicting=False) + + +class TestRawMergeKeepsUnits: + """The raw-concatenation merge fallback keeps common units.""" + + def test_units_survive_value_merge(self): + """Merging value-backed coords with one common unit keeps it.""" + import numpy as np + + import dascore as dc + from dascore.utils.coordmanager import merge_coord_managers + + p1 = dc.get_example_patch().set_units(distance="m") + d = p1.get_coord("distance") + # non-uniform values force the raw concatenation path + values = np.sort(np.random.default_rng(0).uniform(400, 500, len(d.data))) + p2 = p1.update_coords(distance=values).set_units(distance="m") + merged = merge_coord_managers([p1.coords, p2.coords], dim="distance") + assert str(merged.coord_map["distance"].units) == "1 m" diff --git a/tests/test_utils/test_hdf_utils.py b/tests/test_utils/test_hdf_utils.py index ee87017d1..0187e8d67 100644 --- a/tests/test_utils/test_hdf_utils.py +++ b/tests/test_utils/test_hdf_utils.py @@ -3,27 +3,17 @@ from __future__ import annotations from contextlib import closing -from pathlib import Path import h5py -import pandas as pd import pytest -import tables -from tables.exceptions import ClosedNodeError -import dascore as dc -from dascore.config import set_config -from dascore.exceptions import InvalidFileHandlerError from dascore.utils.downloader import fetch from dascore.utils.hdf5 import ( H5Reader, - HDFPatchIndexManager, - LocalPyTablesReader, - PyTablesWriter, + H5Writer, extract_h5_attrs, get_h5py_file, h5_matches_structure, - open_hdf5_file, ) @@ -36,172 +26,31 @@ def h5_example_file(): fi.close() -class TestGetHDF5Handlder: - """Tests for opening an HDF5 file from various inputs.""" +class TestH5Readers: + """Tests for the h5py-based reader/writer handles.""" - @pytest.fixture() - def simple_hdf_path(self, tmp_path_factory): - """Create a hdf5 file in a temporary directory.""" - new = tmp_path_factory.mktemp("dummy_hdf5") / "test.h5" - with tables.open_file(str(new), mode="w") as fi: - bob = fi.create_group(fi.root, name="bob") - bob._v_attrs["lightening"] = 1 - return Path(new) - - @pytest.fixture() - def simple_hdf_file_handler_read(self, simple_hdf_path): - """Return a tables file handler in read mode.""" - with tables.open_file(simple_hdf_path, mode="r") as fi: - yield fi - - @pytest.fixture() - def simple_hdf_file_handler_append(self, simple_hdf_path): - """Return a tables file handler in append mode.""" - with tables.open_file(simple_hdf_path, mode="a") as fi: - yield fi - - def test_path_read(self, simple_hdf_path): - """Ensure passing a path works.""" - with open_hdf5_file(simple_hdf_path) as fi: - assert isinstance(fi, tables.File) - - def test_table_file_read(self, simple_hdf_file_handler_read): - """Ensure a tables file also works.""" - with open_hdf5_file(simple_hdf_file_handler_read) as fi: - assert isinstance(fi, tables.File) - - def test_read_only_filehandle_raises(self, simple_hdf_file_handler_read): - """If write is requested but read handler is provided an error should raise.""" - with pytest.raises(InvalidFileHandlerError, match="but mode"): - with open_hdf5_file(simple_hdf_file_handler_read, mode="w"): - pass - - def test_read_with_write_filehandler(self, simple_hdf_file_handler_append): - """ - Ensure a file handler is returned if read mode is requested but write - mode is provided. This works because write is a superset of read - functionality. - """ - with open_hdf5_file(simple_hdf_file_handler_append, mode="r") as fi: - assert isinstance(fi, tables.File) - - -class TestHDFPatchIndexManager: - """Tests for the HDF5 index manager.""" - - @pytest.fixture - def index_manager(self, tmp_path_factory): - """Create a new index.""" - path = Path(tmp_path_factory.mktemp("example")) / ".index" - return HDFPatchIndexManager(path) - - @pytest.fixture - def index_manager_with_content(self, index_manager, random_spool): - """Add content to the index manager.""" - spool_df = dc.scan_to_df(random_spool) - index_manager.write_update(spool_df) - return index_manager - - def test_extra_columns(self, index_manager, random_spool): - """ - Only the columns used for indexing should be kept, extras discarded. - - Here we include a column with types that can't be serialized. If the - write_update works the test passes. - """ - df = dc.scan_to_df(random_spool).assign( - bad_cols=[[] for _ in range(len(random_spool))] - ) - index_manager.write_update(df) - - def test_empty_tuple(self, index_manager, random_spool): - """Empty dims should convert to empty string.""" - df = dc.scan_to_df(random_spool).assign( - dims=[() for _ in range(len(random_spool))], - ) - index_manager.write_update(df) - - def test_has_content(self, index_manager_with_content, tmp_path): - """`has_index` should return True if data have been written else False.""" - assert index_manager_with_content.has_index - # create hdf5 file with no index - path = tmp_path / "empty.h5" - df = pd.DataFrame([1, 2, 3], columns=["first"]) - df.to_hdf(str(path), key="df") - # assert it doesn't have an index - assert not HDFPatchIndexManager(path).has_index - - def test_closed_node_error(self, index_manager_with_content, monkeypatch): - """ - Test for when the file fails to open. This is a bit contrived but the - closed node issues does happen sometimes in multiple thread environments. - """ - failed_count = 0 - old_func = pd.read_hdf - - def _new_read(*args, **kwargs): - nonlocal failed_count - if failed_count < 1: - failed_count += 1 - raise ClosedNodeError("Simulated failed node opening") - else: - return old_func(*args, **kwargs) - - monkeypatch.setattr(pd, "read_hdf", _new_read) - - df = index_manager_with_content.get_index() - assert len(df) - - # now insure the exception propagates - failed_count = 0 - with set_config(hdf_index_max_retries=0): - with pytest.raises(ClosedNodeError): - index_manager_with_content.get_index() - - def test_metadata_created(self, tmp_path_factory): - """Tests for getting info from a index that doesnt yet exist.""" - path = tmp_path_factory.mktemp("non_existent_index") / "index.hdf5" - with tables.open_file(path, "w"): - pass - index = HDFPatchIndexManager(path) - meta = index._read_metadata() - assert meta is not None - - def test_encode_table_skips_missing_encoded_columns(self, index_manager): - """Missing encoder columns should be ignored safely.""" - df = dc.scan_to_df(dc.get_example_spool()) - index_manager._column_encoders = dict(index_manager._column_encoders) | { - "missing_column": lambda x: x - } - out = index_manager.encode_table(df.copy(), path=None) - assert "path" in out.columns - - def test_hdf_kwargs_come_from_config(self, index_manager): - """Compression defaults should come from runtime configuration.""" - with set_config(hdf_index_complib="zlib", hdf_index_complevel=1): - out = index_manager.hdf_kwargs - assert out["complib"] == "zlib" - assert out["complevel"] == 1 - - -class TestHDFReaders: - """Tests for HDF5 readers.""" - - def test_get_handle(self, tmp_path_factory): - """Ensure we can get a handle with the class.""" + def test_writer_get_handle(self, tmp_path_factory): + """Ensure we can get a writable handle from a path or handle.""" path = tmp_path_factory.mktemp("hdf_handle_test") / "test_file.h5" - with closing(PyTablesWriter.get_handle(path)) as handle: - assert isinstance(handle, tables.File) - handle_2 = PyTablesWriter.get_handle(handle) - assert isinstance(handle_2, tables.File) + with closing(H5Writer.get_handle(path)) as handle: + handle.create_group("waveforms") + handle_2 = H5Writer.get_handle(handle) + assert handle_2 is handle - def test_local_pytables_reader_get_handle(self, tmp_path): - """The local-materializing reader should open local files directly.""" + def test_reader_get_handle(self, tmp_path): + """Ensure the reader opens existing files.""" path = tmp_path / "local_reader.h5" - with tables.open_file(path, mode="w") as h5: - h5.create_group("/", "waveforms") - with closing(LocalPyTablesReader.get_handle(path)) as handle: - assert isinstance(handle, tables.File) + with h5py.File(path, "w") as h5: + h5.create_group("waveforms") + with closing(H5Reader.get_handle(path)) as handle: + assert "waveforms" in handle + + def test_missing_reader_does_not_create_parent(self, tmp_path): + """Opening a missing file for reading does not mutate the filesystem.""" + path = tmp_path / "missing_parent" / "missing.h5" + with pytest.raises(OSError): + H5Reader.get_handle(path) + assert not path.parent.exists() class TestGetH5pyFile: diff --git a/tests/test_utils/test_io_utils.py b/tests/test_utils/test_io_utils.py index e89bfd6cb..36dcd687b 100644 --- a/tests/test_utils/test_io_utils.py +++ b/tests/test_utils/test_io_utils.py @@ -8,7 +8,6 @@ import h5py import pytest -from tables import File from upath import UPath import dascore as dc @@ -18,8 +17,6 @@ from dascore.utils.hdf5 import ( H5Reader, H5Writer, - HDF5Reader, - HDF5Writer, LocalH5Reader, open_h5_resource, ) @@ -149,14 +146,15 @@ def test_binary_stream_not_text_reader(self): def test_path_to_hdf5_reader(self, generic_hdf5): """Ensure we get a reader from tmp path reader.""" - with closing(get_handle_from_resource(generic_hdf5, HDF5Reader)) as handle: - assert isinstance(handle, File) + with closing(get_handle_from_resource(generic_hdf5, H5Reader)) as handle: + assert "bob" in handle # h5py-file-like def test_path_to_hdf5_writer(self, tmp_path): - """Ensure we get a reader from tmp path reader.""" + """Ensure we get a writer from tmp path.""" path = tmp_path / "test_hdf_writer.h5" - with closing(get_handle_from_resource(path, HDF5Writer)) as handle: - assert isinstance(handle, File) + with closing(get_handle_from_resource(path, H5Writer)) as handle: + handle.create_group("waveforms") + assert "waveforms" in handle def test_get_path(self, tmp_path): """Ensure we can get a path.""" @@ -446,9 +444,9 @@ def test_not_implemented(self): with pytest.raises(NotImplementedError): get_handle_from_resource(bad_instance, BinaryWriter) with pytest.raises(NotImplementedError): - get_handle_from_resource(bad_instance, HDF5Writer) + get_handle_from_resource(bad_instance, H5Writer) with pytest.raises(NotImplementedError): - get_handle_from_resource(bad_instance, HDF5Reader) + get_handle_from_resource(bad_instance, H5Reader) class TestIOResourceManager: @@ -471,13 +469,12 @@ def test_basic_context_manager(self, tmp_path): assert isinstance(path_from_hint, Path) path = man.get_resource(Path) assert isinstance(path, Path) - hf = man.get_resource(HDF5Writer) + hf = man.get_resource(H5Writer) fi = man.get_resource(BinaryWriter) - # Why didn't pytables implement the stream like pythons? - assert hf.isopen + assert not hf.closed assert not fi.closed - # after the context manager exists everything should be closed. - assert not hf.isopen + # after the context manager exits everything should be closed. + assert hf.closed assert fi.closed def test_get_none_resource_returns_source(self): diff --git a/tests/test_utils/test_misc.py b/tests/test_utils/test_misc.py index 378db4b2c..7d87e8764 100644 --- a/tests/test_utils/test_misc.py +++ b/tests/test_utils/test_misc.py @@ -384,6 +384,17 @@ def test_empty_sequence_false(self): """An empty set of diffs should not be considered close enough.""" assert not all_diffs_close_enough([]) + @pytest.mark.parametrize( + "diffs", + [ + np.array([np.nan, np.nan]), + np.array(["NaT", "NaT"], dtype="timedelta64[ns]"), + ], + ) + def test_all_null_false(self, diffs): + """Diffs containing only null values are not close enough.""" + assert not all_diffs_close_enough(diffs) + class TestOptionalImport: """Ensure the optional import works.""" diff --git a/tests/test_utils/test_patch_utils.py b/tests/test_utils/test_patch_utils.py index b3f09eb54..b19f0ece5 100644 --- a/tests/test_utils/test_patch_utils.py +++ b/tests/test_utils/test_patch_utils.py @@ -30,8 +30,8 @@ get_patch_names, get_patch_window_size, get_window_axis_step, - merge_patches, merge_compatible_coords_attrs, + merge_patches, patches_to_df, stack_patches, swap_kwargs_dim_to_axis, @@ -396,10 +396,28 @@ class TestPatchesToDF: """Test for getting metadata from patch into a dataframe.""" def test_spool_input(self, random_spool): - """A spool should return its contents.""" + """A spool should return its contents with its patches embedded.""" df = patches_to_df(random_spool) assert isinstance(df, pd.DataFrame) assert len(df) == len(random_spool) + # the "patch" column carries the actual patches, not None + assert "patch" in df.columns + assert all(isinstance(x, dc.Patch) for x in df["patch"]) + + def test_list_of_patches_input(self, random_spool): + """A plain sequence of patches is scanned and the patches embedded.""" + patches = list(random_spool) + df = patches_to_df(patches) + assert isinstance(df, pd.DataFrame) + assert len(df) == len(patches) + assert list(df["patch"]) == patches + + def test_empty_list_input(self): + """An empty sequence produces an empty frame with the right columns.""" + df = patches_to_df([]) + assert isinstance(df, pd.DataFrame) + assert len(df) == 0 + assert "patch" in df.columns def test_dataframe_input(self, random_spool): """The function should be idempotent.""" @@ -568,8 +586,13 @@ def test_different_dims_raises(self, random_patch): concatenate_patches([p1, p2], time=None) def test_duplicate_patches_existing_dim(self, random_patch): - """Ensure duplicate patches are concatenated together.""" - spool = dc.spool([random_patch, random_patch]) + """Ensure equal (but distinct) patches are concatenated together. + + Note: the same patch *instance* twice would collapse to one entry + (spools have set semantics by patch identity); patch.new() mints + a distinct instance with equal contents. + """ + spool = dc.spool([random_patch, random_patch.new()]) out = concatenate_patches(spool, time=None) assert len(out) == 1 patch = out[0] @@ -614,7 +637,8 @@ def test_concat_chunk_to_new_dimension(self, random_patch): """Ensure the new dimension can be chunked by an int value.""" # When new_dim = 1 it should only add a new dimension to each patch # and not change the original shape. - spool = dc.spool([random_patch] * 6) + # Distinct instances: identical instances would dedup to one. + spool = dc.spool([random_patch] + [random_patch.new() for _ in range(5)]) # Test for single values along new dimension new = spool.concatenate(new_dim=1) assert len(new) == len(spool) @@ -833,6 +857,22 @@ def test_path_column_leave_extension(self, random_directory_spool): names = get_patch_names(random_directory_spool, strip_extension=False) assert "." in names.iloc[0] + def test_mixed_path_sources_use_metadata(self, random_spool): + """Mixed real and memory paths consistently use metadata names.""" + df = random_spool.get_contents().iloc[:2].copy() + df["path"] = ["/tmp/real_file.h5", "memory://registry/patch"] + names = get_patch_names(df) + assert names.iloc[0] != "real_file" + assert names.iloc[1] != "patch" + + def test_multiple_coord_fields(self, random_spool): + """Naming on more than one coordinate flattens the min/max fields.""" + # drop path so the coordinate-based naming branch is exercised + df = random_spool.get_contents().drop(columns=["path"], errors="ignore") + names = get_patch_names(df, coords=("time", "distance")) + assert len(names) == len(df) + assert names.str.len().gt(0).all() + class TestSwapKwargsDimToAxis: """Tests for swap_kwargs_dim_to_axis function.""" @@ -1093,3 +1133,22 @@ def test_none_overlap_matches_default(self, random_patch): random_patch, distance=self.window * step, overlap=None ) assert out == (self.window, random_patch.get_axis("distance"), None) + + +class TestForcePatchMergeOverlap: + """_force_patch_merge tolerates complete-envelope overlap (keep first).""" + + def test_complete_overlap_keeps_first(self, random_patch): + """Identical envelopes merge to the first patch.""" + from dascore.utils.patch import _force_patch_merge + + twin = random_patch.new() + infos = [] + for patch in (random_patch, twin): + info = patch.coords._get_dim_summary() + info["patch"] = patch + info["dims"] = ",".join(patch.dims) + infos.append(info) + out = _force_patch_merge(infos, merge_kwargs={}) + assert len(out) == 1 + assert out[0]["patch"] is random_patch diff --git a/tests/test_utils/test_paths.py b/tests/test_utils/test_paths.py index 9c52086bd..e19b1931e 100644 --- a/tests/test_utils/test_paths.py +++ b/tests/test_utils/test_paths.py @@ -11,6 +11,7 @@ from dascore.utils.paths import ( coerce_to_local_path, coerce_to_upath, + directory_writable, get_path_protocol, is_local_path, is_pathlike, @@ -18,6 +19,28 @@ ) +class TestDirectoryWritable: + """directory_writable probes without leaking exceptions.""" + + def test_writable_directory(self, tmp_path): + """A normal writable directory returns True.""" + assert directory_writable(tmp_path) is True + + def test_unwritable_returns_false(self, tmp_path): + """A probe that can't create its parent returns False, not OSError.""" + # a path *under a file* makes mkdir raise NotADirectoryError (OSError) + a_file = tmp_path / "a_file" + a_file.write_text("x") + assert directory_writable(a_file / "sub") is False + + def test_existing_legacy_probe_is_preserved(self, tmp_path): + """The writability probe never truncates a predictable old sentinel.""" + sentinel = tmp_path / "._dascore_write_test_delete_me" + sentinel.write_text("keep me") + assert directory_writable(tmp_path) is True + assert sentinel.read_text() == "keep me" + + class TestIsPathlike: """Tests for ``is_pathlike``.""" diff --git a/tests/test_utils/test_pd.py b/tests/test_utils/test_pd.py index 4ceffd08c..894719e32 100644 --- a/tests/test_utils/test_pd.py +++ b/tests/test_utils/test_pd.py @@ -464,3 +464,25 @@ def test_raises(self, example_df_2): msg = "Cannot chunk spool or dataframe" with pytest.raises(ParameterError, match=msg): get_interval_columns(example_df_2, "money") + + +class TestRelativeRangesToAbsolute: + """Relative range resolution validates its envelope columns.""" + + def test_missing_max_column_raises_spool_error(self): + """A frame with the min but not the max column raises the doc'd error.""" + from dascore.exceptions import InvalidSpoolQueryError + from dascore.utils.pd import relative_ranges_to_absolute + + df = pd.DataFrame({"time_min": [0.0]}) # no time_max + with pytest.raises(InvalidSpoolQueryError, match="relative select"): + relative_ranges_to_absolute(df, {"time": (1, -1)}) + + def test_non_tuple_value_raises(self): + """A non-(start, stop) relative value raises rather than mis-resolving.""" + from dascore.exceptions import InvalidSpoolQueryError + from dascore.utils.pd import relative_ranges_to_absolute + + df = pd.DataFrame({"time_min": [0.0], "time_max": [1.0]}) + with pytest.raises(InvalidSpoolQueryError, match="range selectors"): + relative_ranges_to_absolute(df, {"time": 5}) diff --git a/tests/test_utils/test_time.py b/tests/test_utils/test_time.py index e98ce0e6f..5bc147da0 100644 --- a/tests/test_utils/test_time.py +++ b/tests/test_utils/test_time.py @@ -13,7 +13,6 @@ from dascore.compat import random_state from dascore.exceptions import TimeError from dascore.utils.time import ( - get_max_min_times, is_datetime64, is_timedelta64, saturate_add, @@ -321,6 +320,13 @@ def test_timedelta(self): assert isinstance(out, np.timedelta64) assert out == to_timedelta64(3600) + def test_series(self): + """A Series converts to timedeltas without losing its index.""" + ser = pd.Series([1.0, 2.0], index=["first", "second"]) + out = to_timedelta64(ser) + expected = pd.Series(to_timedelta64(ser.values), index=ser.index) + pd.testing.assert_series_equal(out, expected) + def test_pandas_string_array(self): """Ensure pandas StringArray converts to timedelta64[ns].""" arr = pd.array(["1s", "2s", None], dtype="string") @@ -424,6 +430,13 @@ def test_datetime64(self): out = to_int(to_datetime64("1970-01-01") + np.timedelta64(1, "ns")) assert out == 1 + def test_series(self): + """A datetime Series converts to integer ns and preserves its index.""" + ser = pd.Series(to_datetime64(["1970-01-01", "2000-01-01"])) + ser.index = ["first", "second"] + out = to_int(ser) + pd.testing.assert_series_equal(out, ser.astype(np.int64)) + def test_timedelta64_array(self): """Ensure int ns is returned for datetime64.""" array = to_datetime64(["2017-01-01", "1970-01-01", "1999-01-01"]) @@ -662,14 +675,3 @@ def test_dtype(self): d2 = np.array([1, 2]).astype("timedelta64[ms]").dtype assert not is_timedelta64(d1) assert is_timedelta64(d2) - - -class TestGetmaxMinTimes: - """Tests for max_min fetching.""" - - def test_raises_bad_value(self): - """Simple test to make sure error is raised if unordered tuple.""" - t1 = to_datetime64("2020-01-01") - t2 = to_datetime64("1994-01-01") - with pytest.raises(ValueError): - get_max_min_times((t1, t2))