diff --git a/dascore/core/annotation_loader.py b/dascore/core/annotation_loader.py index d3cb3f428..3b4321799 100644 --- a/dascore/core/annotation_loader.py +++ b/dascore/core/annotation_loader.py @@ -48,9 +48,10 @@ _VERTEX_COLUMNS, ANNOTATION_STEM, ATTRS_STEM, + DIMS_KEY, OBJECT_SUFFIXES, RESERVED_COLUMNS, - TABLE_SUFFIX, + TABLE_SUFFIXES, VERTEX_STEM, AnnotationSet, _text, @@ -59,7 +60,12 @@ from dascore.models.registry import TAG_FIELD from dascore.utils.misc import iterate, optional_import from dascore.utils.paths import quote_path -from dascore.utils.tables import parse_cell, read_table +from dascore.utils.tables import ( + parse_cell, + read_parquet, + read_parquet_metadata, + read_table, +) from dascore.utils.time import to_datetime64 # What an attrs file declares itself to be; the model writes its own tag. @@ -80,6 +86,10 @@ _COMMENT = "#" _DIMS_PRAGMA = "dims" +# The suffix which names the parquet encoding, for the branches which read +# a table rather than merely find one. +PARQUET_SUFFIX = TABLE_SUFFIXES[1] + # The name a directory of data carries its annotations under, in either # form: the directory `.annotations/`, holding a set or a directory of sets, # or the bare table `.annotations.csv`. Hidden, like the `.inventory` which @@ -231,8 +241,17 @@ def _dimension_spellings(dims: Sequence[str]) -> frozenset[str]: return frozenset(x for dim in dims for x in (dim, f"{dim}{_START}", f"{dim}{_END}")) +def _is_text(series: pd.Series) -> bool: + """Whether a column holds text, whichever way pandas is spelling it.""" + return getattr(series.dtype, "kind", "") in "OTU" + + def _read_cells( - frame: pd.DataFrame, dims: Sequence[str], path: Path, ordered: bool = False + frame: pd.DataFrame, + dims: Sequence[str], + path: Path, + ordered: bool = False, + typed: bool = False, ) -> pd.DataFrame: """ Read a table's text cells as the values each column holds. @@ -240,17 +259,37 @@ def _read_cells( Only the vertices are ordered, so only they read ``seq`` as the number it is: the annotations table does not reserve that name, and an annotation carrying its own ``seq`` means whatever it says. + + A typed table -- parquet -- states what each column holds, so only its + text columns are read further, and only as far as a dimension or a + basis: text elsewhere is text, since a format with a boolean of its own + would have used one. Stating a type is not the same as stating a usable + one, so a dimension or a vertex order which arrives as something a + coordinate cannot be is refused rather than trusted. """ spellings = _dimension_spellings(dims) out = {} for name in frame.columns: series = frame[name] - if str(name) in spellings: + if typed and not _is_text(series): + # The file stated what this column holds, so nothing here has + # to work it out from a spelling -- only check that what it + # states is a thing the column is allowed to hold. + if str(name) in spellings: + _check_kind(series, name, path, "iufMm", "numbers or times") + elif ordered and str(name) == _ORDINAL: + _check_kind(series, name, path, "iuf", "a number") + out[name] = series + elif str(name) in spellings: out[name] = _read_dimension(series, path) elif ordered and str(name) == _ORDINAL: out[name] = _read_ordinal(series, path) elif str(name) == "basis": out[name] = _read_basis(series, path) + elif typed: + # Text in a typed table is text: a cell reading 'true' in a + # format which has a boolean is the word, not the boolean. + out[name] = series elif str(name) in _TEXT_COLUMNS or series.isna().all(): # A column no row states says nothing about what it holds, and # reading its emptiness as a type would invent one. @@ -260,6 +299,17 @@ def _read_cells( return pd.DataFrame(out) +def _check_kind(series: pd.Series, name, path: Path, kinds: str, what: str) -> None: + """Refuse a typed column whose type the field it names cannot be.""" + if series.dtype.kind in kinds: + return + msg = ( + f"The column {str(name)!r} of {quote_path(path)} holds " + f"{series.dtype}, where it states {what}." + ) + raise ParameterError(msg) + + def _read_extra(cell): """ Read one cell of a column the set does not model. @@ -286,15 +336,62 @@ def _read_set_table( """ Read one of a set's tables, with its cells typed. - A table stating nothing at all is what a set of no annotations writes, - so it reads back as none rather than as a table with no columns. + A table stating nothing at all is what a set of no annotations writes + -- a blank CSV, or a parquet file with no columns -- so it reads back + as none rather than as a table which states nothing. """ + if _is_parquet(path): + frame, _ = read_parquet(path, what=what, empty=True) + if not len(frame.columns): + return None + return _read_cells(frame, dims, path, ordered=ordered, typed=True) if _is_blank(path): return None frame = read_table(path, what=what, skip=skip) return _read_cells(frame, dims, path, ordered=ordered) +def _is_parquet(path: Path) -> bool: + """Whether a table's name says it is parquet rather than CSV.""" + return path.suffix.casefold() == PARQUET_SUFFIX + + +def _read_table_dims(path: Path) -> tuple[tuple[str, ...] | None, int]: + """ + Return the dimensions a table declares for itself, and the CSV lines + above its header. + + Each encoding declares them where it can: a CSV in a comment above the + header, a parquet file in the metadata its footer holds. + """ + if not _is_parquet(path): + return _read_pragma(path) + stated = read_parquet_metadata(path) + return _stated_dims(stated.get(DIMS_KEY), path), 0 + + +def _stated_dims(document: str | None, path: Path) -> tuple[str, ...] | None: + """Read the dimensions a parquet file states in its metadata.""" + if document is None: + return None + try: + stated = json.loads(document) + except ValueError as error: + msg = ( + f"{quote_path(path)} states {DIMS_KEY} as {document!r}, which is not " + f"a JSON document: {error}." + ) + raise ParameterError(msg) from error + names = tuple(str(x) for x in iterate(stated)) + if not names: + msg = ( + f"{quote_path(path)} states {DIMS_KEY} but names none; a table which " + "declares its dimensions names them." + ) + raise ParameterError(msg) + return names + + def _read_pragma(path: Path) -> tuple[tuple[str, ...] | None, int]: """ Return the dimensions a table declares above its header, and the lines @@ -372,24 +469,31 @@ def _is_blank(path: Path) -> bool: return not stream.read().strip() +def _tables(directory: Path) -> list[Path]: + """ + Every file in a directory whose name says it is a table. + + Hidden names are not tables here, as they are not sets: an editor's + ``.annotations.csv.swp`` or a half-copied ``.annotations.csv`` is a + companion the directory keeps, and taking a directory down for one would + make a stray sync file fatal. + """ + return [ + x + for x in _entries(directory) + if not x.name.startswith(".") and x.suffix.casefold() in TABLE_SUFFIXES + ] + + def _refuse_stray_tables(directory: Path, known: Collection[str], what: str) -> None: """ Refuse a table whose name names no part of what a directory holds. A ``vertexes.csv`` beside an ``annotations.csv`` claims to participate in this convention and gets it wrong, which is worth more than being - quietly skipped. Hidden names are not tables here, as they are not sets: - an editor's ``.annotations.csv.swp`` or a half-copied ``.annotations.csv`` - is a companion the directory keeps, and taking a directory down for one - would make a stray sync file fatal. + quietly skipped. """ - stray = sorted( - x.name - for x in _entries(directory) - if not x.name.startswith(".") - and x.suffix.casefold() == TABLE_SUFFIX - and x.stem not in known - ) + stray = sorted(x.name for x in _tables(directory) if x.stem not in known) if stray: msg = f"{quote_path(directory)} holds the table(s) {', '.join(stray)}, {what}" raise ParameterError(msg) @@ -470,10 +574,15 @@ def _load_path(path: Path, dims, **kwargs) -> AnnotationSet: def _states_annotations(path: Path) -> bool: - """Whether a directory states annotations of its own; a file states none.""" + """ + Whether a directory states annotations of its own, in either encoding. + + A file states none: the scans below walk what a directory holds, and a + plain file among them is not a set which spelled itself oddly. + """ if not path.is_dir(): return False - return _one_spelling(path, ANNOTATION_STEM, (TABLE_SUFFIX,)) is not None + return _one_spelling(path, ANNOTATION_STEM, TABLE_SUFFIXES) is not None def find_annotations(directory: str | os.PathLike) -> Path | None: @@ -502,29 +611,37 @@ def find_annotations(directory: str | os.PathLike) -> Path | None: """ root = Path(directory) tree = root / BLESSED_NAME - table = tree.with_suffix(TABLE_SUFFIX) - found = [x for x in (tree, table) if x.exists()] + named = ", ".join(tree.with_suffix(x).name for x in TABLE_SUFFIXES) + # The suffix is matched without regard to case, as every other table + # this format finds is: a `.annotations.PARQUET` is the file the lower + # case name would be, and reading one and not the other would make the + # carried name mean less than the visible one. + tables = [ + x + for x in _entries(root) + if x.stem == BLESSED_NAME and x.suffix.casefold() in TABLE_SUFFIXES + ] + found = [x for x in (tree, *tables) if x.exists()] if not found: return None if len(found) > 1: + listed = ", ".join(x.name for x in found) msg = ( - f"{quote_path(root)} carries annotations twice over: " - f"{tree.name} and {table.name}. A directory states what it carries " - "once; keep the one it means." + f"{quote_path(root)} carries annotations more than once: {listed}. " + "A directory states what it carries once; keep the one it means." ) raise InvalidAnnotationError(msg) (only,) = found if only == tree and not only.is_dir(): msg = ( f"{quote_path(only)} is a file. The annotations a directory carries " - f"are the set directory {BLESSED_NAME}/, or the bare table " - f"{table.name}." + f"are the set directory {BLESSED_NAME}/, or a bare table: {named}." ) raise InvalidAnnotationError(msg) - if only == table and only.is_dir(): + if only != tree and only.is_dir(): msg = ( f"{quote_path(only)} is a directory. A set held as a directory is " - f"named {BLESSED_NAME}/; the {TABLE_SUFFIX} name spells a bare table." + f"named {BLESSED_NAME}/; a name with a suffix spells a bare table." ) raise InvalidAnnotationError(msg) return only @@ -566,7 +683,7 @@ def _child_sets(directory: Path) -> list[Path]: if _one_spelling(child, ATTRS_STEM, OBJECT_SUFFIXES) is not None: msg = ( f"{quote_path(child)} states the attributes of a set but no " - f"{ANNOTATION_STEM}{TABLE_SUFFIX}, so it states no annotations." + f"{ANNOTATION_STEM} table, so it states no annotations." ) raise ParameterError(msg) _refuse_colliding_names(directory, out) @@ -654,8 +771,8 @@ def _declares_dims(directory: Path, attrs: Mapping) -> bool: """ if attrs.get("dims"): return True - table = _one_spelling(directory, ANNOTATION_STEM, (TABLE_SUFFIX,)) - return table is not None and _read_pragma(table)[0] is not None + table = _one_spelling(directory, ANNOTATION_STEM, TABLE_SUFFIXES) + return table is not None and _read_table_dims(table)[0] is not None def _merge_sets( @@ -870,24 +987,24 @@ def _load_set(directory: Path, attrs: Mapping, dims, **kwargs) -> AnnotationSet: _refuse_stray_tables( directory, (ANNOTATION_STEM, VERTEX_STEM), - f"which name no part of a set. A set states {ANNOTATION_STEM}" - f"{TABLE_SUFFIX} and, where it has vertices, {VERTEX_STEM}{TABLE_SUFFIX}.", + f"which name no part of a set. A set states {ANNOTATION_STEM} and, " + f"where it has vertices, {VERTEX_STEM}, each as a " + f"{' or a '.join(TABLE_SUFFIXES)} table.", ) - table = _one_spelling(directory, ANNOTATION_STEM, (TABLE_SUFFIX,)) + table = _one_spelling(directory, ANNOTATION_STEM, TABLE_SUFFIXES) if table is None: msg = ( - f"{quote_path(directory)} holds no {ANNOTATION_STEM}{TABLE_SUFFIX} " - f"and no {BLESSED_NAME}, so it states no annotations and carries " - "none." + f"{quote_path(directory)} holds no {ANNOTATION_STEM} table and no " + f"{BLESSED_NAME}, so it states no annotations and carries none." ) raise ParameterError(msg) - declared, skip = _read_pragma(table) + declared, skip = _read_table_dims(table) stated = _declared_dims(attrs, dims, directory, declared, table) frame = _read_set_table(table, stated, "no annotations", skip=skip) # Found as the annotations table is: a set spells each of its parts # once, and a `vertices.CSV` beside a `vertices.csv` is two spellings of # one part rather than a table nobody reads. - vertex_path = _one_spelling(directory, VERTEX_STEM, (TABLE_SUFFIX,)) + vertex_path = _one_spelling(directory, VERTEX_STEM, TABLE_SUFFIXES) vertices = None if vertex_path is not None: vertices = _read_set_table( @@ -895,7 +1012,7 @@ def _load_set(directory: Path, attrs: Mapping, dims, **kwargs) -> AnnotationSet: stated, "no vertices", ordered=True, - skip=_vertex_comments(vertex_path), + skip=_vertex_declaration(vertex_path), ) # The read dimensions rather than the given ones: they are the same # names, already a tuple, so a file spelling one dimension as a bare @@ -905,14 +1022,14 @@ def _load_set(directory: Path, attrs: Mapping, dims, **kwargs) -> AnnotationSet: def _load_file(path: Path, dims, **kwargs) -> AnnotationSet: """Load the set a bare table holds.""" - if path.suffix.lower() != TABLE_SUFFIX: + if path.suffix.casefold() not in TABLE_SUFFIXES: + named = " or ".join(TABLE_SUFFIXES) msg = ( f"{quote_path(path)} is not a table an annotation set is read from. " - f"A bare set is a {TABLE_SUFFIX} file; a set with vertices is a " - "directory." + f"A bare set is a {named} file; a set with vertices is a directory." ) raise ParameterError(msg) - declared, skip = _read_pragma(path) + declared, skip = _read_table_dims(path) # A bare table states no attributes of its own, so a caller may hand it # some -- and the dimensions they name are the ones its cells are read # in, since nothing can be read before that is known. @@ -924,21 +1041,26 @@ def _load_file(path: Path, dims, **kwargs) -> AnnotationSet: ) -def _vertex_comments(path: Path) -> int: +def _vertex_declaration(path: Path) -> int: """ Refuse a vertices table which declares dimensions, and return the lines above its header -- which, since it declares none, is none of them. - Vertices are read in the dimensions of the set they belong to. The lines - above a header are skipped only where a declaration is among them, so a - vertices table has no preamble to skip. + Vertices are read in the dimensions of the set they belong to, whether + they would declare them above a header or in a footer. The lines above a + header are skipped only where a declaration is among them, so a vertices + table has no preamble to skip. """ - declared, skip = _read_pragma(path) + declared, skip = _read_table_dims(path) if declared is not None: + where = ( + f"states {DIMS_KEY}" + if _is_parquet(path) + else f"declares {_DIMS_PRAGMA} above its header" + ) msg = ( - f"{quote_path(path)} declares {_DIMS_PRAGMA} above its header. " - "Vertices are read in the dimensions of the set they belong to, " - "which states them once." + f"{quote_path(path)} {where}. Vertices are read in the dimensions " + "of the set they belong to, which states them once." ) raise ParameterError(msg) return skip @@ -968,9 +1090,11 @@ def _declared_dims( stated = declared elif tuple(str(x) for x in iterate(stated)) != declared: spelled = "was given" if dims is not None else "is stated in its attributes" + named = table or source + where = f"in its {DIMS_KEY}" if _is_parquet(named) else "above its header" msg = ( - f"{quote_path(table or source)} declares the dimensions " - f"{', '.join(declared)} above its header, but " + f"{quote_path(named)} declares the dimensions " + f"{', '.join(declared)} {where}, but " f"{', '.join(str(x) for x in iterate(stated))} {spelled}. " "Dimensions may be spelled twice where the two agree." ) diff --git a/dascore/core/annotations.py b/dascore/core/annotations.py index 3e455f2a7..da0f65a40 100644 --- a/dascore/core/annotations.py +++ b/dascore/core/annotations.py @@ -54,7 +54,12 @@ from dascore.utils.intervals import normalize_value, value_kind from dascore.utils.mapping import FrozenDict from dascore.utils.misc import iterate, to_str, validate_acquisition_key -from dascore.utils.tables import parse_cell +from dascore.utils.tables import ( + parquet_table, + parse_cell, + write_parquet, + write_parquet_table, +) from dascore.utils.time import to_datetime64, to_timedelta64 # Columns any set may carry, whatever dimensions it declares. @@ -87,9 +92,20 @@ ATTRS_STEM = "attrs" ANNOTATION_STEM = "annotations" VERTEX_STEM = "vertices" -TABLE_SUFFIX = ".csv" +# The encodings a table takes, the suffix naming which one a file holds. +# CSV is the floor: it needs nothing beyond the standard library, so a set +# can always be written. Parquet is the same tables with their types kept, +# for a set too big to want text; it needs pyarrow. +TABLE_SUFFIXES = (".csv", ".parquet") +TABLE_SUFFIX = TABLE_SUFFIXES[0] OBJECT_SUFFIXES = (".json", ".yaml", ".yml") +# What a parquet table names its dimensions in, since it has no comment +# line to declare them in and its footer is the place a format states what +# its columns cannot. A JSON document, as GeoParquet's `geo` key holds one; +# namespaced, so a file may carry both without either reading the other's. +DIMS_KEY = "dascore:dims" + # What a range column is spelled with. _START, _END = "_start", "_end" @@ -766,6 +782,41 @@ def to_csv(self, path=None) -> str: path Where to write the text, or None to only return it. """ + self._refuse_bare_vertices() + return _write_table(self._df, path) + + def to_parquet(self, path) -> pathlib.Path: + """ + Write the annotations as one parquet file. + + The parquet spelling of + [to_csv](`dascore.core.annotations.AnnotationSet.to_csv`), for a + set too big to want text. It keeps what a CSV cannot: a column + parquet has a type for comes back as that type rather than as a + spelling to be guessed at, and the dimensions travel in the file's + own metadata rather than having to be stated again. A column with + no one type is written as JSON, which keeps the value of each cell + but not every python type it may have been held in -- a tuple comes + back as a list. + + Needs pyarrow, which CSV does not; a set of regions can always be + written as a table, whatever is installed. + + Parameters + ---------- + path + Where to write the file. + + Returns + ------- + The path written to, so a save reads straight back. + """ + self._refuse_bare_vertices() + write_parquet(self._df, path, {DIMS_KEY: json.dumps(list(self.dims))}) + return pathlib.Path(path) + + def _refuse_bare_vertices(self) -> None: + """Refuse to write a set of shapes as one table, which has one grain.""" if not self._vertices.empty: msg = ( "This set holds vertices, which a bare table has no row for. " @@ -773,9 +824,8 @@ def to_csv(self, path=None) -> str: "annotations." ) raise ParameterError(msg) - return _write_table(self._df, path) - def save(self, path) -> pathlib.Path: + def save(self, path, format: str = "csv") -> pathlib.Path: """ Write the set to a directory, creating it if needed. @@ -793,9 +843,15 @@ def save(self, path) -> pathlib.Path: belongs to, and what each of them states for itself travels in the attributes, so the flat spelling loses nothing. + The tables are CSV unless another encoding is asked for. Parquet + writes the same parts under the same names, with its own suffix, and + keeps a column's type rather than its spelling wherever it has one + for it; it needs pyarrow, where CSV needs nothing. + Writing states the whole directory, so a part this set does not - have is removed rather than left behind. A stale vertices table, or - the YAML the attributes used to be spelled in, would otherwise sit + have is removed rather than left behind. A stale vertices table, the + YAML the attributes used to be spelled in, or the CSV a set was + written as before it was written as parquet, would otherwise sit beside what was written and leave a directory which loaded before the save refusing to load after it. @@ -803,6 +859,8 @@ def save(self, path) -> pathlib.Path: ---------- path The directory to write into. + format + The encoding the tables are written in: ``csv`` or ``parquet``. Returns ------- @@ -815,33 +873,40 @@ def save(self, path) -> pathlib.Path: # Defaults are dropped from the document, so it says what the set # says; dims has no default, so it is always written, and the # attributes name their own model, which is what the file holds. + suffix = _table_suffix(format) document = self._attrs.model_dump(mode="json", exclude_defaults=True) - annotation_text = _write_table(self._df) - vertex_text = None if self._vertices.empty else _write_table(self._vertices) + spelled = {ANNOTATION_STEM: _spell_table(self._df, suffix, self.dims)} + if not self._vertices.empty: + spelled[VERTEX_STEM] = _spell_table(self._vertices, suffix) directory = pathlib.Path(path) directory.mkdir(parents=True, exist_ok=True) - vertex_table = directory / f"{VERTEX_STEM}{TABLE_SUFFIX}" attrs_file = directory / f"{ATTRS_STEM}{OBJECT_SUFFIXES[0]}" + writing = {attrs_file, *(directory / f"{x}{suffix}" for x in spelled)} # Only the spellings this format claims: a notes.txt or an # attrs.bak beside them participates in no convention and is not # this function's to delete. The suffix is matched without regard # to case, as the loader matches it. + claimed = { + ATTRS_STEM: OBJECT_SUFFIXES, + ANNOTATION_STEM: TABLE_SUFFIXES, + VERTEX_STEM: TABLE_SUFFIXES, + } superseded = [ x for x in directory.iterdir() - if x.stem == ATTRS_STEM - and x.suffix.casefold() in OBJECT_SUFFIXES - and x != attrs_file + if x.suffix.casefold() in claimed.get(x.stem, ()) and x not in writing ] - if vertex_text is None: - superseded.append(vertex_table) - for stale in superseded: - stale.unlink(missing_ok=True) + # Written before the superseded parts are cleared, not after: a + # write which fails partway -- a full disk, a permission changed + # under it -- then leaves the set it was replacing still in the + # directory, and a reader finds two spellings of one part and says + # so, rather than finding the set gone. with open(attrs_file, "w") as stream: json.dump(document, stream, indent=2) - _write_text(annotation_text, directory / f"{ANNOTATION_STEM}{TABLE_SUFFIX}") - if vertex_text is not None: - _write_text(vertex_text, vertex_table) + for stem, payload in spelled.items(): + _write_spelled(payload, directory / f"{stem}{suffix}") + for stale in superseded: + stale.unlink(missing_ok=True) return directory # --- what the set holds @@ -1586,6 +1651,40 @@ def _refuse_ambiguous_values(frame: pd.DataFrame) -> None: raise ParameterError(msg) +def _table_suffix(format: str) -> str: + """Return the suffix an encoding is named by, refusing an unknown one.""" + suffix = f".{str(format).lower().lstrip('.')}" + if suffix not in TABLE_SUFFIXES: + named = ", ".join(x.lstrip(".") for x in TABLE_SUFFIXES) + msg = f"{format!r} is not a table encoding; a set is written as {named}." + raise ParameterError(msg) + return suffix + + +def _spell_table(frame: pd.DataFrame, suffix: str, dims: Sequence[str] | None = None): + """ + Spell one table for the encoding its suffix names. + + Spelled before the directory is touched, so whichever encoding is + asked for, a table which cannot be written raises with the stored set + still whole. + """ + if suffix == TABLE_SUFFIX: + return _write_table(frame) + # A parquet file has no comment line to declare its dimensions in, so + # they go in the metadata its footer holds. + metadata = None if dims is None else {DIMS_KEY: json.dumps(list(dims))} + return parquet_table(frame, metadata) + + +def _write_spelled(payload, path) -> None: + """Write what `_spell_table` spelled, whichever encoding it is.""" + if isinstance(payload, str): + _write_text(payload, path) + else: + write_parquet_table(payload, path) + + def _write_table(frame: pd.DataFrame, path=None) -> str: """Return a frame as CSV text, optionally writing it to a path.""" _refuse_ambiguous_values(frame) diff --git a/dascore/utils/pd.py b/dascore/utils/pd.py index be9dd106f..19c5874c4 100644 --- a/dascore/utils/pd.py +++ b/dascore/utils/pd.py @@ -393,7 +393,13 @@ def _filter_contains(query_dict, df, bool_index): """Filter based on rows containing specified values.""" for key, val in query_dict.items(): _check_misdirected_range_query(key, val, df) - bool_index = np.logical_and(bool_index, df[key].isin(val)) + # An ellipsis names no value, so it matches nothing and is dropped + # before `isin` sees it: pandas' arrow-backed string columns refuse a + # value arrow has no type for, where its numpy-backed ones quietly + # never match. The range check above still sees it -- an open bound + # in a membership query is what that check exists to name. + wanted = [x for x in val if x is not ...] + bool_index = np.logical_and(bool_index, df[key].isin(wanted)) return bool_index diff --git a/dascore/utils/tables.py b/dascore/utils/tables.py index bbd141f3b..a05e94508 100644 --- a/dascore/utils/tables.py +++ b/dascore/utils/tables.py @@ -14,12 +14,21 @@ from __future__ import annotations import csv +import datetime +import json +from collections.abc import Mapping, Sized from pathlib import Path +import numpy as np import pandas as pd from dascore.exceptions import ParameterError +from dascore.utils.misc import optional_import, to_str from dascore.utils.paths import quote_path +from dascore.utils.time import to_datetime64 + +# The metadata key a parquet file names its document columns in. +DOCUMENT_KEY = "dascore:documents" def read_table(path: Path, what: str = "nothing", skip: int = 0) -> pd.DataFrame: @@ -131,6 +140,317 @@ def _check_widths(reader, header: list[str], path: Path, start: int = 2) -> None raise ParameterError(msg) +def write_parquet(frame: pd.DataFrame, path, metadata: Mapping | None = None) -> None: + """ + Write a dataframe as one parquet file, keeping the values it holds. + + Parquet stores types, so a column comes back as what it was written as + rather than as text a reader has to guess at. A column with no single + type -- one holding both text and booleans, or a model, or a nested + mapping -- has no parquet type of its own; each of its cells is written + as a JSON document instead, and the file names those columns in its + metadata so a reader gets the values back rather than their spelling. + + Parameters + ---------- + frame + The table to write. + path + Where to write it. + metadata + Key-value strings for the file's own metadata, which + [read_parquet](`dascore.utils.tables.read_parquet`) hands back. A + format states here what its table cannot say in a column. + + Examples + -------- + Needs pyarrow, which the doctest run does not have, so this one is not + executed; `tests/test_utils/test_tables.py` runs the same round trip. + + >>> import tempfile # doctest: +SKIP + >>> from pathlib import Path + >>> import pandas as pd + >>> from dascore.utils.tables import read_parquet, write_parquet + >>> frame = pd.DataFrame({"group": ["rail"], "value": [True]}) + >>> with tempfile.TemporaryDirectory() as folder: # doctest: +SKIP + ... path = Path(folder) / "coupling.parquet" + ... write_parquet(frame, path, {"dascore:dims": "distance"}) + ... out, stated = read_parquet(path) + >>> out.equals(frame), stated["dascore:dims"] # doctest: +SKIP + (True, 'distance') + """ + write_parquet_table(parquet_table(frame, metadata), path) + + +def parquet_table(frame: pd.DataFrame, metadata: Mapping | None = None): + """ + Return a dataframe as the parquet table it is written from. + + Spelled out before anything is written, so a caller storing several + tables at once can prepare them all before it touches the directory. + See [write_parquet](`dascore.utils.tables.write_parquet`), which is + this and the write together. + """ + arrow = optional_import("pyarrow", required_for="parquet tables") + spelled, documents = {}, [] + for label, name in zip(frame.columns, _named(frame.columns), strict=True): + series = frame[label] + if _one_type(series): + spelled[name] = series + continue + documents.append(name) + spelled[name] = series.map(_document) + stated = {str(k): str(v) for k, v in (metadata or {}).items()} + if reserved := sorted(x for x in stated if _is_reserved(x)): + msg = ( + f"The metadata key(s) {', '.join(reserved)} are not a caller's to " + f"state: {DOCUMENT_KEY} is what this writer names its own document " + "columns in, and pandas and ARROW: are what the file already holds " + "for the reader which wrote them." + ) + raise ParameterError(msg) + if documents: + stated[DOCUMENT_KEY] = json.dumps(documents) + table = arrow.Table.from_pandas( + pd.DataFrame(spelled, index=frame.index), preserve_index=False + ) + # Added to what pyarrow wrote rather than replacing it: the pandas key + # it puts there is what a pandas reader uses to rebuild an index and the + # dtypes it can, and dropping it would make this file say less than + # pyarrow wrote. + kept = {**(table.schema.metadata or {}), **stated} + return table.replace_schema_metadata(kept) + + +def write_parquet_table(table, path) -> None: + """Write a prepared parquet table, as `parquet_table` returns it.""" + parquet = optional_import("pyarrow.parquet", required_for="parquet tables") + parquet.write_table(table, path) + + +def read_parquet( + path, what: str = "nothing", empty: bool = False +) -> tuple[pd.DataFrame, dict[str, str]]: + """ + Read one parquet file, and whatever it states about itself. + + Returns the table and its metadata, with the columns + [write_parquet](`dascore.utils.tables.write_parquet`) stored as JSON + documents read back as the values they hold. + + Parameters + ---------- + path + The parquet file to read. + what + What a table with no columns fails to state, for that error message. + empty + Whether a table with no columns is allowed, for a format in which + that is how something holding nothing is written. + """ + parquet = optional_import("pyarrow.parquet", required_for="parquet tables") + try: + table = parquet.read_table(path) + stated = _stated_metadata(table.schema.metadata) + # Inside the block as well: pandas metadata another writer left is + # read here, and a malformed one raises where nothing else would + # name the file it came from. + frame = table.to_pandas() + except Exception as error: + # Any error pyarrow raises: it reports a truncated file, an + # unreadable one and something which is not parquet at all through + # several types of its own, and the caller's format names them all + # the same way. + msg = f"Could not read {quote_path(path)}: {error}." + raise ParameterError(msg) from error + if not len(frame.columns) and not empty: + msg = f"{quote_path(path)} has no columns, so it states {what}." + raise ParameterError(msg) + for name in _document_columns(stated.pop(DOCUMENT_KEY, "[]"), path): + if name not in frame.columns: + msg = ( + f"{quote_path(path)} names {name!r} as a column of documents, " + "and holds no such column." + ) + raise ParameterError(msg) + # Held as object: the cells are whatever their documents state, and + # letting pandas re-infer a type from them would hand back a column + # of a type the file never said it had. + read = [_read_document(x, name, path) for x in frame[name]] + frame[name] = pd.Series(read, index=frame.index, dtype=object) + return frame, stated + + +def read_parquet_metadata(path) -> dict[str, str]: + """ + Return what a parquet file states about itself, reading no rows. + + The footer alone, so a caller which needs what a file says before it + decides how to read the table -- which dimensions its columns are in, + say -- does not pay for the table to find out. + + Parameters + ---------- + path + The parquet file to read. + """ + parquet = optional_import("pyarrow.parquet", required_for="parquet tables") + try: + schema = parquet.read_schema(path) + except Exception as error: + msg = f"Could not read {quote_path(path)}: {error}." + raise ParameterError(msg) from error + return _stated_metadata(schema.metadata) + + +def _document_columns(stated: str, path) -> list[str]: + """Read the columns a file names as documents, refusing what it cannot mean.""" + try: + names = json.loads(stated) + except ValueError as error: + msg = ( + f"{quote_path(path)} states {DOCUMENT_KEY} as {stated!r}, which is " + f"not a JSON document: {error}." + ) + raise ParameterError(msg) from error + if not isinstance(names, list) or not all(isinstance(x, str) for x in names): + msg = ( + f"{quote_path(path)} states {DOCUMENT_KEY} as {stated!r}; it names " + "the columns which hold documents, so it is a list of names." + ) + raise ParameterError(msg) + return names + + +def _named(columns) -> list[str]: + """ + Return a table's column names as parquet holds them, which is as text. + + A frame may label a column with anything hashable and a CSV writes + whatever that prints as, but arrow takes names only as strings -- so + they are spelled here rather than left to fail deep inside a conversion. + Two labels which spell alike would name one column, which is the same + refusal a frame naming a column twice already gets. + """ + named = [str(x) for x in columns] + if len(set(named)) != len(named): + repeated = sorted({x for x in named if named.count(x) > 1}) + msg = ( + f"The column(s) {', '.join(repeated)} are named more than once once " + "their names are spelled as text, which is how parquet holds them; " + "one column states one thing." + ) + raise ParameterError(msg) + return named + + +def _is_reserved(key: str) -> bool: + """Whether a metadata key names something the file states for itself.""" + return key == DOCUMENT_KEY or key == "pandas" or key.startswith("ARROW:") + + +def _one_type(series: pd.Series) -> bool: + """Whether a column holds one type parquet has a column shape for.""" + if series.dtype != object: + return True + # Text is the one type an object column may still hold: pandas gives + # some string columns object dtype, and a whole column of text is + # exactly what parquet stores as a string column. + return all(isinstance(x, str) for x in series if _is_stated(x)) + + +def _is_stated(value) -> bool: + """Whether a cell states anything, whatever kind of thing it holds.""" + if value is None: + return False + if isinstance(value, Sized) and not isinstance(value, str): + # A container states itself; asking pandas whether it is null + # answers once per element, which has no truth value. + return True + return not pd.isna(value) + + +def _document(value): + """Spell one cell of a column parquet has no single type for.""" + if not _is_stated(value): + return None + return json.dumps(_documented(value), default=str) + + +def _documented(value): + """ + Return a value as the json types it is made of. + + A missing value nested inside one becomes null rather than the text of + whatever spelling of missing it was: `pd.NA` written as "" would + come back as a string a reader takes for a value. + + Numpy's scalars are the ones that matter: `np.int64` is not an `int` + and `np.bool_` is not a `bool`, so json falls back to spelling them as + text -- and a column of numbers written by numpy would come back as a + column of strings, which is exactly what this encoding exists not to + do. A time has no json type at all and is spelled as DASCore spells + every time, which reads back as one. + + A value none of this can spell -- a `Decimal`, a `bytes`, an object of + someone's own class -- is written as its text, which is what the CSV + encoding does with it too. That is the one place this loses a type, and + the reason a column parquet has a type for is never sent this way. + """ + if not _is_stated(value): + return None + # At nanoseconds, whatever resolution the value arrived in: a set holds + # its times at nanoseconds, and one spelling per instant is what makes + # two writes of one set the same file. + if isinstance(value, np.datetime64): + return to_str(np.datetime64(value, "ns")) + if isinstance(value, np.timedelta64): + return to_str(np.timedelta64(value, "ns")) + if hasattr(value, "model_dump"): + return value.model_dump(mode="json") + if isinstance(value, np.generic): + return _documented(value.item()) + if isinstance(value, datetime.datetime | datetime.date | pd.Timestamp): + return _documented(to_datetime64(value)) + if isinstance(value, Mapping): + return {str(k): _documented(v) for k, v in value.items()} + if isinstance(value, list | tuple | set | frozenset | np.ndarray): + return [_documented(x) for x in value] + return value + + +def _read_document(cell, name: str, path) -> object: + """Read one cell of a column the file names as documents.""" + if not isinstance(cell, str): + # Every spelling of a null becomes the one which was written: a + # column of documents holds values, and None is the value pandas + # hands back for a cell parquet stored as null. + return None if not _is_stated(cell) else cell + try: + return json.loads(cell) + except ValueError as error: + msg = ( + f"The column {name!r} of {quote_path(path)} holds " + f"{cell!r}, which is not a JSON document: {error}." + ) + raise ParameterError(msg) from error + + +def _stated_metadata(metadata) -> dict[str, str]: + """Return the key-value metadata a file states, as text.""" + out = {} + for key, value in (metadata or {}).items(): + # Decoded leniently, keys as well as values: this is another + # writer's metadata as often as it is ours, and a key which is not + # UTF-8 is something to report rather than something to hide. + name = key.decode(errors="replace") + # What pyarrow writes for itself, which is not the caller's to read. + if name == "pandas" or name.startswith("ARROW:"): + continue + out[name] = value.decode(errors="replace") + return out + + def row_cells(row) -> dict[str, str]: """ Return a row's stated cells, an empty one meaning unset. diff --git a/dascore/utils/time.py b/dascore/utils/time.py index 336a91b81..08d5b3091 100644 --- a/dascore/utils/time.py +++ b/dascore/utils/time.py @@ -164,8 +164,17 @@ def _float_to_datetime(ser: pd.Series) -> pd.Series: @to_datetime64.register(pd.arrays.StringArray) +@to_datetime64.register(pd.arrays.ArrowStringArray) def _string_array_to_datetime64(arr: pd.arrays.StringArray): - """Convert pandas StringArray to datetime64.""" + """ + Convert a pandas string array to datetime64. + + Both backings, since which one pandas gives text is not the caller's + choice: a `str` column is arrow-backed wherever pyarrow is installed + and numpy-backed where it is not. Neither class is a subclass of the + other -- they share only `BaseStringArray` -- so registering one does + not dispatch the other. + """ out = pd.to_datetime(arr, errors="coerce", format="mixed") return out.to_numpy(dtype="datetime64[ns]") @@ -313,8 +322,9 @@ def _series_to_timedelta64_series(ser: pd.Series) -> pd.Series: @to_timedelta64.register(pd.arrays.StringArray) +@to_timedelta64.register(pd.arrays.ArrowStringArray) def _string_array_to_timedelta64(arr: pd.arrays.StringArray): - """Convert pandas StringArray to timedelta64.""" + """Convert a pandas string array, of either backing, to timedelta64.""" out = pd.to_timedelta(arr, errors="coerce") return out.to_numpy(dtype="timedelta64[ns]") diff --git a/environment.yml b/environment.yml index 6150b359c..8a32ae726 100644 --- a/environment.yml +++ b/environment.yml @@ -15,6 +15,7 @@ dependencies: - matplotlib>=3.5 - scipy>=1.15.0 - findiff + - pyarrow - jupyter - nbformat - pint diff --git a/pyproject.toml b/pyproject.toml index f55fece7c..6147d4377 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,6 +66,9 @@ dependencies = [ extras = [ "xarray", + # An annotation set may be stored as parquet rather than CSV; CSV is + # the floor which needs nothing. + "pyarrow", "netCDF4", "h5netcdf", "findiff", diff --git a/tests/test_core/test_annotation_loader.py b/tests/test_core/test_annotation_loader.py index d47056f0c..40df30b49 100644 --- a/tests/test_core/test_annotation_loader.py +++ b/tests/test_core/test_annotation_loader.py @@ -15,10 +15,17 @@ except ImportError: yaml = None +try: + import pyarrow + import pyarrow.parquet +except ImportError: + pyarrow = None + import dascore as dc from dascore.core.annotation_loader import find_annotations -from dascore.core.annotations import Line, Moveout +from dascore.core.annotations import DIMS_KEY, Line, Moveout from dascore.exceptions import InvalidAnnotationError, ParameterError +from dascore.utils.tables import DOCUMENT_KEY, write_parquet DIMS = ("distance", "time") @@ -549,7 +556,7 @@ def test_no_annotations_table(self, tmp_path): """A directory without one states no annotations.""" directory = tmp_path / "picks" directory.mkdir() - with pytest.raises(InvalidAnnotationError, match=r"no annotations\.csv"): + with pytest.raises(InvalidAnnotationError, match="no annotations table"): dc.annotations(directory, dims=DIMS) def test_a_table_which_cannot_be_read(self, regions, tmp_path): @@ -805,7 +812,7 @@ def test_a_tree_which_holds_no_set_at_all(self, tmp_path): (half / "attrs.json").write_text('{"dims": ["time"]}') # The half-set is not what is wrong here: this directory states no # annotations of its own, and holds no set to make it a collection. - with pytest.raises(InvalidAnnotationError, match=r"holds no annotations\.csv"): + with pytest.raises(InvalidAnnotationError, match="holds no annotations table"): dc.annotations(root, dims=("time",)) def test_a_tree_of_collections(self, regions, tmp_path): @@ -1360,7 +1367,7 @@ def test_carried_twice(self, data, regions): """A directory states what it carries once.""" regions.save(data / ".annotations") (data / ".annotations.csv").write_text("group,time\nq,1\n") - with pytest.raises(InvalidAnnotationError, match="twice over"): + with pytest.raises(InvalidAnnotationError, match="more than once"): dc.annotations(data) def test_the_wrong_kind_of_thing(self, data): @@ -1398,3 +1405,255 @@ def test_the_data_directory_keeps_its_own_attrs(self, data, regions): (data / "attrs.json").write_text('{"object_type": "SomethingElse"}') regions.save(data / ".annotations") assert dc.annotations(data) == regions + + +def _forge(frame: pd.DataFrame, path, documents: str) -> None: + """Write a parquet file whose document footer this library did not write.""" + table = pyarrow.Table.from_pandas(frame, preserve_index=False) + kept = {**(table.schema.metadata or {}), DOCUMENT_KEY: documents} + pyarrow.parquet.write_table(table.replace_schema_metadata(kept), path) + + +@pytest.mark.skipif(pyarrow is None, reason="pyarrow is not installed") +class TestParquet: + """The same tables, with their types kept, for a set too big to want text.""" + + @pytest.fixture + def mixed(self) -> dc.AnnotationSet: + """A set whose columns hold what a CSV would have to spell as text.""" + frame = pd.DataFrame( + { + "id": ["r1", "r2"], + "group": ["noise", "quiet"], + "value": ["car", True], + "tags": [("road", "car"), None], + "time_start": [ + np.datetime64("2020-01-01T00:00:10"), + np.datetime64("2020-01-01T00:00:20"), + ], + "time_end": [ + np.datetime64("2020-01-01T00:00:12"), + np.datetime64("2020-01-01T00:00:22"), + ], + "score": [0.9, 0.2], + "checked": [True, False], + "meta": [{"a": 1}, None], + } + ) + return dc.AnnotationSet(frame, dims=DIMS, acquisition_key="NET.ARR.00.das") + + def test_a_bare_table(self, mixed, tmp_path): + """A set of regions is one file, and reads back as the set it was.""" + loaded = dc.annotations(mixed.to_parquet(tmp_path / "picks.parquet")) + assert loaded.to_dataframe().equals(mixed.to_dataframe()) + + def test_the_dimensions_travel_with_the_file(self, mixed, tmp_path): + """A parquet file states its dimensions where it can: its footer.""" + path = mixed.to_parquet(tmp_path / "picks.parquet") + assert dc.annotations(path).dims == DIMS + + def test_restating_the_dimensions(self, mixed, tmp_path): + """Agreement is allowed, disagreement is not, as with every spelling.""" + path = mixed.to_parquet(tmp_path / "picks.parquet") + assert dc.annotations(path, dims=DIMS).dims == DIMS + with pytest.raises(InvalidAnnotationError, match="where the two agree"): + dc.annotations(path, dims=("depth",)) + + def test_kinds_a_csv_would_lose(self, mixed, tmp_path): + """A column with no one type is written as documents, not as text.""" + loaded = dc.annotations(mixed.to_parquet(tmp_path / "picks.parquet")) + assert [type(x).__name__ for x in loaded.to_dataframe()["value"]] == [ + "str", + "bool", + ] + assert loaded[0].extra["meta"] == {"a": 1} + assert loaded[0].tags == ("road", "car") + + def test_text_stays_text(self, tmp_path): + """A typed format has a boolean, so a cell reading 'true' is the word.""" + frame = pd.DataFrame({"group": ["a"], "note": ["true"], "time": [1.0]}) + picks = dc.AnnotationSet(frame, dims=("time",)) + loaded = dc.annotations(picks.to_parquet(tmp_path / "picks.parquet")) + assert loaded[0].extra["note"] == "true" + + def test_a_value_a_csv_would_refuse(self, tmp_path): + """Text a table would read back as a boolean is safe where types are kept.""" + frame = pd.DataFrame({"group": ["a"], "value": ["true"], "time": [1.0]}) + picks = dc.AnnotationSet(frame, dims=("time",)) + with pytest.raises(ParameterError, match="a table would read back"): + picks.to_csv() + loaded = dc.annotations(picks.to_parquet(tmp_path / "picks.parquet")) + assert loaded[0].value == "true" + + def test_a_set_of_no_annotations(self, tmp_path): + """A set which states nothing writes a table which states nothing.""" + empty = dc.AnnotationSet(None, dims=("time",)) + assert dc.annotations(empty.save(tmp_path / "picks", format="parquet")) == empty + loaded = dc.annotations(empty.to_parquet(tmp_path / "picks.parquet")) + assert len(loaded) == 0 + + def test_an_empty_set_beside_a_full_one(self, picks, tmp_path): + """One set holding nothing does not take its collection down with it.""" + root = tmp_path / "sets" + picks.save(root / "phasenet", format="parquet") + dc.AnnotationSet(None, dims=("time",)).save(root / "empty", format="parquet") + assert len(dc.annotations(root)) == len(picks) + + def test_numbers_numpy_made(self, tmp_path): + """A value from numpy is the number it is, not the text str() gives.""" + frame = pd.DataFrame( + { + "group": ["g", "g", "h"], + "value": [np.int64(3), 5, "text"], + "time": [1.0, 2.0, 3.0], + } + ) + picks = dc.AnnotationSet(frame, dims=("time",)) + loaded = dc.annotations(picks.save(tmp_path / "picks", format="parquet")) + assert [x.value for x in loaded] == [3, 5, "text"] + + def test_a_time_inside_a_document(self, tmp_path): + """A time has no JSON type, so it is spelled as DASCore spells one.""" + frame = pd.DataFrame( + { + "group": ["a", "b"], + "meta": [{"when": np.datetime64("2020-01-01T00:00:01")}, "note"], + "time": [1.0, 2.0], + } + ) + picks = dc.AnnotationSet(frame, dims=("time",)) + loaded = dc.annotations(picks.save(tmp_path / "picks", format="parquet")) + assert loaded[0].extra["meta"] == {"when": "2020-01-01T00:00:01.000000000"} + + @pytest.mark.parametrize( + ("stated", "message"), + [(json.dumps("group"), "it is a list of names"), ("{oops", "not a JSON")], + ) + def test_a_footer_this_cannot_read(self, stated, message, tmp_path): + """A footer another writer left is read, or named where it cannot be.""" + path = tmp_path / "picks.parquet" + _forge(pd.DataFrame({"group": ["a"], "time": [1.0]}), path, stated) + with pytest.raises(InvalidAnnotationError, match=message): + dc.annotations(path, dims=("time",)) + + def test_a_directory(self, with_vertices, curve, tmp_path): + """Every part a set states is written under its own name.""" + directory = with_vertices.save(tmp_path / "picks", format="parquet") + assert sorted(x.name for x in directory.iterdir()) == [ + "annotations.parquet", + "attrs.json", + "vertices.parquet", + ] + loaded = dc.annotations(directory) + assert loaded == with_vertices + assert loaded[1].geometry.basis == curve + + def test_a_collection(self, regions, picks, tmp_path): + """A set is a set whichever encoding it is written in.""" + root = tmp_path / "sets" + regions.save(root / "hand", format="parquet") + picks.save(root / "phasenet") + assert sorted(dc.annotations(root).attrs.sets) == ["hand", "phasenet"] + + def test_carried_beside_data(self, regions, tmp_path): + """The hidden name takes the parquet spelling too.""" + directory = tmp_path / "data" + directory.mkdir() + regions.to_parquet(directory / ".annotations.parquet") + assert dc.annotations(directory).to_dataframe().equals(regions.to_dataframe()) + + def test_carried_whatever_case_it_names(self, regions, tmp_path): + """The carried name is matched as every other table name is.""" + directory = tmp_path / "data" + directory.mkdir() + regions.to_parquet(directory / ".annotations.PARQUET") + assert dc.annotations(directory).to_dataframe().equals(regions.to_dataframe()) + + def test_a_typed_column_which_cannot_be_a_dimension(self, tmp_path): + """Stating a type is not stating one a coordinate can be.""" + path = tmp_path / "picks.parquet" + write_parquet(pd.DataFrame({"group": ["a"], "time": [True]}), path) + with pytest.raises(InvalidAnnotationError, match="where it states numbers"): + dc.annotations(path, dims=("time",)) + + def test_a_typed_vertex_order_which_is_not_a_number(self, with_vertices, tmp_path): + """A vertex states its place in the order as a number, typed or not.""" + directory = with_vertices.save(tmp_path / "picks", format="parquet") + vertices = with_vertices.to_vertices() + stamps = np.array(["2020-01-01", "2020-01-02"], dtype="datetime64[ns]") + vertices["seq"] = list(stamps) * (len(vertices) // 2) + list( + stamps[: len(vertices) % 2] + ) + write_parquet(vertices, directory / "vertices.parquet") + with pytest.raises(InvalidAnnotationError, match="where it states a number"): + dc.annotations(directory) + + def test_the_other_encoding_is_superseded(self, regions, tmp_path): + """A set written twice states itself once, not once per encoding.""" + directory = regions.save(tmp_path / "picks") + assert (directory / "annotations.csv").exists() + regions.save(directory, format="parquet") + assert not (directory / "annotations.csv").exists() + assert dc.annotations(directory) == regions + regions.save(directory) + assert not (directory / "annotations.parquet").exists() + + def test_both_encodings_at_once(self, regions, tmp_path): + """A directory holding both says two things; neither is chosen.""" + directory = regions.save(tmp_path / "picks") + regions.to_parquet(directory / "annotations.parquet") + with pytest.raises(InvalidAnnotationError, match="each of its parts once"): + dc.annotations(directory) + + def test_a_bare_table_refuses_vertices(self, with_vertices, tmp_path): + """One file states one grain, whatever its encoding.""" + with pytest.raises(ParameterError, match="a bare table has no row for"): + with_vertices.to_parquet(tmp_path / "picks.parquet") + + def test_an_unknown_encoding(self, regions, tmp_path): + """A set is written in an encoding it has, and says which it has.""" + with pytest.raises(ParameterError, match="not a table encoding"): + regions.save(tmp_path / "picks", format="feather") + + def test_vertices_declare_nothing(self, with_vertices, tmp_path): + """Vertices are read in the dimensions of the set they belong to.""" + directory = with_vertices.save(tmp_path / "picks", format="parquet") + frame = with_vertices.to_vertices() + write_parquet(frame, directory / "vertices.parquet", {DIMS_KEY: '["time"]'}) + with pytest.raises(InvalidAnnotationError, match="states them once"): + dc.annotations(directory) + + def test_dimensions_which_are_not_a_document(self, regions, tmp_path): + """What the footer states is read, and named where it is not readable.""" + path = tmp_path / "picks.parquet" + write_parquet(regions.to_dataframe(), path, {DIMS_KEY: "time, distance"}) + with pytest.raises(InvalidAnnotationError, match="not a JSON document"): + dc.annotations(path) + + def test_dimensions_which_name_none(self, regions, tmp_path): + """A file which declares its dimensions names them.""" + path = tmp_path / "picks.parquet" + write_parquet(regions.to_dataframe(), path, {DIMS_KEY: "[]"}) + with pytest.raises(InvalidAnnotationError, match="names none"): + dc.annotations(path) + + def test_a_stray_parquet_table(self, regions, tmp_path): + """A near-miss on the convention is a near-miss in either encoding.""" + directory = regions.save(tmp_path / "picks") + regions.to_parquet(directory / "annotation.parquet") + with pytest.raises(InvalidAnnotationError, match=r"annotation\.parquet"): + dc.annotations(directory) + + def test_a_file_which_is_not_parquet(self, tmp_path): + """Whatever pyarrow makes of it, the error names the file.""" + path = tmp_path / "picks.parquet" + path.write_text("group,time\na,1\n") + with pytest.raises(InvalidAnnotationError, match="Could not read"): + dc.annotations(path, dims=("time",)) + + def test_a_table_which_is_neither(self, tmp_path): + """A bare set is a table, and the message names the encodings it takes.""" + path = tmp_path / "picks.txt" + path.write_text("group\na\n") + with pytest.raises(InvalidAnnotationError, match=r"\.csv or \.parquet"): + dc.annotations(path, dims=("time",)) diff --git a/tests/test_utils/test_pd.py b/tests/test_utils/test_pd.py index 1c868d341..d6795dd94 100644 --- a/tests/test_utils/test_pd.py +++ b/tests/test_utils/test_pd.py @@ -22,6 +22,11 @@ ) from dascore.utils.time import to_datetime64, to_timedelta64 +try: + import pyarrow +except ImportError: + pyarrow = None + @pytest.fixture() def random_df_from_patch(random_patch): @@ -313,6 +318,13 @@ def test_ellipsis_kept_for_non_interval_column(self, example_df_2): out = filter_df(example_df_2, first_name=("Jason", ...)) assert np.all(out == example_df_2["first_name"].isin(["Jason"])) + @pytest.mark.skipif(pyarrow is None, reason="pyarrow is not installed") + def test_ellipsis_with_an_arrow_backed_column(self, example_df_2): + """An arrow-backed column refuses a value arrow has no type for.""" + df = example_df_2.astype({"first_name": "string[pyarrow]"}) + out = filter_df(df, first_name=("Jason", ...)) + assert np.all(out == df["first_name"].isin(["Jason"])) + def test_open_bound_ignored_for_unknown_column(self, example_df_2): """ Unknown columns are forwarded to patch level select, so an open bound diff --git a/tests/test_utils/test_tables.py b/tests/test_utils/test_tables.py index 2f89136e9..33b1a6c24 100644 --- a/tests/test_utils/test_tables.py +++ b/tests/test_utils/test_tables.py @@ -4,18 +4,32 @@ import csv +import numpy as np import pandas as pd import pytest +from dascore.core.annotations import Line from dascore.exceptions import ParameterError from dascore.utils.tables import ( + DOCUMENT_KEY, ordered_rows, + parquet_table, parse_cell, + read_parquet, + read_parquet_metadata, read_table, require_columns, require_stated, row_cells, + write_parquet, ) +from dascore.utils.time import to_datetime64, to_timedelta64 + +try: + import pyarrow + import pyarrow.parquet +except ImportError: + pyarrow = None def _write(path, text: str, name: str = "table.csv", encoding: str = "utf-8"): @@ -209,3 +223,183 @@ def test_float_keeps_its_point(self): def test_text(self): """Anything else is the string it was written as.""" assert parse_cell("car") == "car" + + +def _forge(frame: pd.DataFrame, path, documents: str) -> None: + """Write a parquet file whose document footer this library did not write.""" + table = pyarrow.Table.from_pandas(frame, preserve_index=False) + kept = {**(table.schema.metadata or {}), DOCUMENT_KEY: documents} + pyarrow.parquet.write_table(table.replace_schema_metadata(kept), path) + + +@pytest.mark.skipif(pyarrow is None, reason="pyarrow is not installed") +class TestParquet: + """Parquet keeps what a column holds, and what the file says about itself.""" + + def test_types_survive(self, tmp_path): + """A column comes back as what it was written as.""" + frame = pd.DataFrame({"a": [1.5], "b": [True], "c": ["text"]}) + path = tmp_path / "table.parquet" + write_parquet(frame, path) + out, _ = read_parquet(path) + assert out.equals(frame) + + def test_a_column_of_no_one_type(self, tmp_path): + """A column parquet has no shape for is written as documents.""" + frame = pd.DataFrame({"value": ["car", True, 3]}) + path = tmp_path / "table.parquet" + write_parquet(frame, path) + out, _ = read_parquet(path) + assert list(out["value"]) == ["car", True, 3] + + def test_a_nested_cell(self, tmp_path): + """A mapping is a document, and reads back as the mapping it was.""" + path = tmp_path / "table.parquet" + write_parquet(pd.DataFrame({"meta": [{"a": [1, 2]}, None]}), path) + out, _ = read_parquet(path) + assert list(out["meta"]) == [{"a": [1, 2]}, None] + + def test_a_column_of_text_stays_a_column_of_text(self, tmp_path): + """Text is a type parquet has, so it is not written as documents.""" + frame = pd.DataFrame({"note": pd.Series(["a", None], dtype=object)}) + assert frame["note"].dtype == object # the branch this is about + table = parquet_table(frame) + assert DOCUMENT_KEY.encode() not in (table.schema.metadata or {}) + # Whichever width of string arrow picks, it is a string column. + assert "string" in str(table.schema.field("note").type) + path = tmp_path / "table.parquet" + write_parquet(frame, path) + out, _ = read_parquet(path) + assert out["note"][0] == "a" and pd.isna(out["note"][1]) + + def test_a_cell_holding_a_list(self, tmp_path): + """A container states itself; asking pandas answers once per element.""" + path = tmp_path / "table.parquet" + write_parquet(pd.DataFrame({"tags": [["a", "b"], None]}), path) + out, _ = read_parquet(path) + assert list(out["tags"]) == [["a", "b"], None] + + def test_numbers_numpy_made(self, tmp_path): + """A numpy scalar is the number it holds, not the text str() gives.""" + path = tmp_path / "table.parquet" + frame = pd.DataFrame({"value": [np.int64(3), "text", np.bool_(True)]}) + write_parquet(frame, path) + out, _ = read_parquet(path) + assert list(out["value"]) == [3, "text", True] + + def test_a_time_of_any_spelling(self, tmp_path): + """A time has no JSON type, so every spelling of one is written alike.""" + path = tmp_path / "table.parquet" + stamp = pd.Timestamp("2020-01-01T00:00:01") + # Three spellings of one instant, which arrive at different + # resolutions: a Timestamp reads back at microseconds, the others at + # nanoseconds, and the file holds one of them. + frame = pd.DataFrame( + {"when": [stamp, stamp.to_pydatetime(), stamp.to_datetime64(), "text"]} + ) + write_parquet(frame, path) + out, _ = read_parquet(path) + written = list(out["when"]) + assert written[3] == "text" + assert len(set(written[:3])) == 1 + assert to_datetime64(written[0]) == np.datetime64("2020-01-01T00:00:01") + + def test_a_duration_and_a_model(self, tmp_path): + """Neither has a JSON type; a model dumps itself, a duration is text.""" + path = tmp_path / "table.parquet" + line = Line(start={"distance": 0.0}, end={"distance": 10.0}) + frame = pd.DataFrame({"cell": [np.timedelta64(5, "s"), line, "text"]}) + write_parquet(frame, path) + out, _ = read_parquet(path) + held = list(out["cell"]) + assert to_timedelta64(held[0]) == np.timedelta64(5, "s") + assert Line(**held[1]) == line + + def test_a_missing_value_inside_a_document(self, tmp_path): + """Missing is missing, not the text of whichever spelling it arrived in.""" + path = tmp_path / "table.parquet" + write_parquet(pd.DataFrame({"meta": [{"score": pd.NA, "n": 1}, "text"]}), path) + out, _ = read_parquet(path) + assert list(out["meta"]) == [{"score": None, "n": 1}, "text"] + + def test_a_document_column_keeps_its_own_type(self, tmp_path): + """A column of documents holds what its cells hold, not what pandas infers.""" + path = tmp_path / "table.parquet" + write_parquet(pd.DataFrame({"code": pd.Series([1, 2], dtype=object)}), path) + out, _ = read_parquet(path) + assert out["code"].dtype == object + assert list(out["code"]) == [1, 2] + + def test_metadata_round_trips(self, tmp_path): + """What a format states about its table comes back to it.""" + path = tmp_path / "table.parquet" + write_parquet(pd.DataFrame({"a": [1]}), path, {"dascore:dims": '["time"]'}) + _, stated = read_parquet(path) + assert stated == {"dascore:dims": '["time"]'} + + def test_metadata_without_the_table(self, tmp_path): + """The footer alone, for a caller which only needs what it says.""" + path = tmp_path / "table.parquet" + write_parquet(pd.DataFrame({"a": [1]}), path, {"dascore:dims": '["time"]'}) + assert read_parquet_metadata(path) == {"dascore:dims": '["time"]'} + + def test_no_columns_refused(self, tmp_path): + """An empty table states nothing, as an empty CSV does.""" + path = tmp_path / "table.parquet" + write_parquet(pd.DataFrame(), path) + with pytest.raises(ParameterError, match="states no track"): + read_parquet(path, what="no track") + + def test_a_file_which_is_not_parquet(self, tmp_path): + """Whatever pyarrow raises, the caller sees its own error.""" + path = _write(tmp_path, "a,b\n1,2\n", name="table.parquet") + with pytest.raises(ParameterError, match="Could not read"): + read_parquet(path) + with pytest.raises(ParameterError, match="Could not read"): + read_parquet_metadata(path) + + @pytest.mark.parametrize("key", [DOCUMENT_KEY, "pandas", "ARROW:schema"]) + def test_a_key_the_file_states_for_itself(self, key, tmp_path): + """A caller states neither this writer's keys nor the format's own.""" + path = tmp_path / "table.parquet" + with pytest.raises(ParameterError, match="not a caller's to state"): + write_parquet(pd.DataFrame({"a": [1]}), path, {key: '["a"]'}) + + def test_a_column_named_something_other_than_text(self, tmp_path): + """Arrow names a column with text, so a label is spelled as text.""" + path = tmp_path / "table.parquet" + write_parquet(pd.DataFrame({"distance": [1.0], 7: ["x"]}), path) + out, _ = read_parquet(path) + assert list(out.columns) == ["distance", "7"] + + def test_two_labels_which_spell_alike(self, tmp_path): + """One column states one thing, whatever the two labels were.""" + path = tmp_path / "table.parquet" + frame = pd.DataFrame({7: ["a"], "7": ["b"]}) + with pytest.raises(ParameterError, match="named more than once"): + write_parquet(frame, path) + + def test_a_file_whose_pandas_metadata_is_not_json(self, tmp_path): + """Another writer's metadata is read here, so a bad one is named here.""" + path = tmp_path / "table.parquet" + table = pyarrow.Table.from_pandas( + pd.DataFrame({"a": [1]}), preserve_index=False + ) + kept = {**(table.schema.metadata or {}), b"pandas": b"not json"} + pyarrow.parquet.write_table(table.replace_schema_metadata(kept), path) + with pytest.raises(ParameterError, match="Could not read"): + read_parquet(path) + + def test_a_document_column_which_is_not_there(self, tmp_path): + """A file naming a column it does not hold says so.""" + path = tmp_path / "table.parquet" + _forge(pd.DataFrame({"a": [1]}), path, '["b"]') + with pytest.raises(ParameterError, match="holds no such column"): + read_parquet(path) + + def test_a_document_which_does_not_parse(self, tmp_path): + """A cell a file names as a document is read as one, or named.""" + path = tmp_path / "table.parquet" + _forge(pd.DataFrame({"a": ["{oops"]}), path, '["a"]') + with pytest.raises(ParameterError, match="not a JSON document"): + read_parquet(path) diff --git a/tests/test_utils/test_time.py b/tests/test_utils/test_time.py index 0f6c6d145..8abeec9a9 100644 --- a/tests/test_utils/test_time.py +++ b/tests/test_utils/test_time.py @@ -13,6 +13,11 @@ import dascore as dc from dascore.compat import random_state from dascore.exceptions import TimeError, UnitError + +try: + import pyarrow +except ImportError: + pyarrow = None from dascore.utils.time import ( is_datetime64, is_timedelta64, @@ -238,6 +243,16 @@ def test_pandas_string_array(self): ) assert np.all(out == expected) + @pytest.mark.skipif(pyarrow is None, reason="pyarrow is not installed") + def test_arrow_backed_string_array(self): + """Which backing pandas gives text is not the caller's choice.""" + arr = pd.array(self.date_strs, dtype="string[pyarrow]") + out = to_datetime64(arr) + expected = np.array([dc.to_datetime64(x) for x in self.date_strs]).astype( + "datetime64[ns]" + ) + assert np.all(out == expected) + class TestToTimeDelta64: """Tests for creating timedeltas.""" @@ -364,6 +379,13 @@ def test_pandas_string_array(self): assert np.all(out[:2] == expected[:2]) assert pd.isnull(out[2]) + @pytest.mark.skipif(pyarrow is None, reason="pyarrow is not installed") + def test_arrow_backed_string_array(self): + """A string column is arrow-backed wherever pyarrow is installed.""" + arr = pd.array(["1s", "2s"], dtype="string[pyarrow]") + out = to_timedelta64(arr) + assert np.all(out == np.array([1, 2]).astype("timedelta64[s]")) + def test_unsupported_type(self): """Ensure unsupported types raise.""" with pytest.raises(NotImplementedError):