Chunk spools by data size and by unit-bearing lengths

dascore/core/spool.py

@@ -216,20 +216,35 @@ def chunk(
         kwargs
             kwargs are used to specify the dimension along which to chunk, eg:
             `time=10` chunks along the time axis in 10 second increments.
+            The value may also be a quantity: one of the coordinate's own
+            units (`time=10 * s`) or a data size (`time=25 * megabytes`),
+            which chunks so each patch's data array is about that large.
+            `overlap` accepts the same forms.
 
         Examples
         --------
         >>> import dascore as dc
-        >>> from dascore.units import s
+        >>> from dascore.units import s, megabytes
         >>>
         >>> spool = dc.get_example_spool("random_das")
         >>> # get spools with time duration of 10 seconds
         >>> time_chunked = spool.chunk(time=10, overlap=1)
+        >>> # the same, with the units stated explicitly
+        >>> unit_chunked = spool.chunk(time=10 * s)
+        >>> # get patches whose data arrays are at most ~1 MB
+        >>> size_chunked = spool.chunk(time=1 * megabytes)
         >>> # merge along time axis
         >>> time_merged = spool.chunk(time=...)
 
         Notes
         -----
+        A data size measures the patch's data array only; coordinates and
+        attrs are extra, as are any copies a later processing step makes,
+        so the patch as a whole is somewhat larger. The sample count is
+        rounded down, so the data never exceeds the requested size, and a
+        merge of patches with different dtypes is sized against the dtype
+        they upcast to.
+
         [`Spool.concatenate`](`dascore.BaseSpool.concatenate`) performs a
         similar operation but disregards the coordinate values.
 
@@ -818,7 +833,7 @@ def chunk(
     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
+        from dascore.utils.chunk_plan import ChunkPlan, _combined_dtype
 
         if len(kwargs) != 1:
             msg = (
@@ -848,6 +863,11 @@ def concatenate(self, check_behavior: WARN_LEVELS = "warn", **kwargs) -> Self:
             )
             member_frames.append(members)
             first = group_rows.iloc[0].to_dict()
+            if "_dtype" in group_rows.columns:
+                # concatenation upcasts like a merge does, so the group's
+                # dtype is what the members combine to, not the first row's
+                combined = _combined_dtype(group_rows["_dtype"])
+                first["_dtype"] = "" if combined is None else str(combined)
             if has_envelope:
                 first[f"{dim}_min"] = group_rows[f"{dim}_min"].min()
                 first[f"{dim}_max"] = group_rows[f"{dim}_max"].max()

dascore/io/index/backend.py

@@ -483,6 +483,7 @@ def write_sources(self, records: list[SourceRecord]) -> None:
                             patch.n_dims,
                             patch.dims,
                             patch.shape,
+                            patch.dtype,
                             patch.sample_count_total,
                             patch.time_min,
                             patch.time_max,

dascore/io/index/catalog.py

@@ -38,7 +38,7 @@
     InvalidSpoolQueryError,
     Query,
 )
-from dascore.io.index.schema import SPOOL_HIDDEN_COLUMNS
+from dascore.io.index.schema import SPOOL_HIDDEN_COLUMNS, SPOOL_PRIVATE_RENAMES
 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
@@ -922,7 +922,7 @@ def to_df(self) -> pd.DataFrame:
                     "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"}
+                columns=dict(SPOOL_PRIVATE_RENAMES)
             )
             # SQL identifies overlapping source patches. Expose the selected
             # envelopes, matching spool.get_contents() and the exact trim

dascore/io/index/indexer.py

@@ -31,7 +31,7 @@
     hive_path_attrs,
     summaries_to_records,
 )
-from dascore.io.index.schema import SPOOL_HIDDEN_COLUMNS
+from dascore.io.index.schema import SPOOL_HIDDEN_COLUMNS, SPOOL_PRIVATE_RENAMES
 from dascore.utils.misc import _iter_filesystem
 from dascore.utils.paths import (
     coerce_to_local_path,
@@ -459,7 +459,7 @@ def get_contents(self, _attrs=None, _coords=None, **kwargs) -> pd.DataFrame:
         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"})
+        return df.rename(columns=dict(SPOOL_PRIVATE_RENAMES))
 
     __call__ = get_contents
 

dascore/io/index/ingest.py

@@ -127,6 +127,7 @@ class PatchRecord:
     source_patch_id: str
     dims: str
     shape: str
+    dtype: str
     n_dims: int
     sample_count_total: int | None
     time_min: int | None
@@ -395,6 +396,7 @@ def patch_record(summary: PatchSummary) -> PatchRecord:
         source_patch_id=normalize_source_patch_id(summary.source_patch_id),
         dims=",".join(summary.dims),
         shape=",".join(str(x) for x in shape),
+        dtype=str(summary.dtype or ""),
         n_dims=len(summary.dims),
         sample_count_total=int(np.prod(shape)) if shape else None,
         time_min=time_min,
@@ -496,7 +498,7 @@ def summaries_to_records(
 _PATCH_ROW_FIELDS = tuple(
     f.name
     for f in fields(PatchRecord)
-    if f.name not in ("source_patch_id", "dims", "shape", "attrs", "coords")
+    if f.name not in ("source_patch_id", "dims", "shape", "dtype", "attrs", "coords")
 )
 # Backends with no boolean type hand these columns back as 0/1, which
 # def_key would hash differently than the bool a scan produced.
@@ -609,6 +611,7 @@ def assemble_source_records(
                     source_patch_id=normalize_source_patch_id(patch.source_patch_id),
                     dims=_py_scalar(patch.dims) or "",
                     shape=_py_scalar(patch.shape) or "",
+                    dtype=_py_scalar(patch.dtype) or "",
                     attrs=typed,
                     coords=tuple(coords),
                     **{f: _py_scalar(getattr(patch, f)) for f in _PATCH_ROW_FIELDS},

dascore/io/index/planned.py

@@ -62,6 +62,11 @@ def _num(value) -> float | None:
     return float(value)
 
 
+def _dtype_str(value) -> str:
+    """Convert a stored element dtype to its string, "" when unknown."""
+    return "" if value is None or pd.isnull(value) else str(value)
+
+
 def _coord_record_from_row(
     row: Mapping, name: str, dims: tuple[str, ...] | None = None
 ) -> CoordRecord | None:
@@ -268,6 +273,11 @@ def _output_records(
             source_patch_id=str(output_id),
             dims=dims,
             shape="",
+            # the plan carries the element dtype privately so a chained
+            # chunk can still size patches by their memory footprint.
+            # NaN is truthy, so `or ""` alone would store the string
+            # "nan" and poison every later np.dtype() of this column.
+            dtype=_dtype_str(row.get("_dtype")),
             n_dims=len(dim_names),
             sample_count_total=None,
             time_min=_ns(row.get("time_min")),

dascore/io/index/schema.py

@@ -25,7 +25,7 @@
 from typing import NamedTuple, get_args, get_type_hints
 
 # Version of the index schema, independent of dascore's version.
-INDEX_VERSION = 5
+INDEX_VERSION = 6
 # Identity string so any tool can sanity-check what it opened.
 WHAT_IS_THIS = "dascore_spool_index"
 
@@ -100,6 +100,7 @@ class PatchRow(NamedTuple):
     n_dims: int
     dims: str
     shape: str
+    dtype: str  # the data array's dtype, eg "float64"
     sample_count_total: int | None
     time_min: int | None  # epoch ns; NULL for relative-time patches
     time_max: int | None
@@ -265,6 +266,7 @@ def _columns(row_type: type[NamedTuple]) -> MappingProxyType[str, str]:
         "n_dims",
         "dims",
         "shape",
+        "dtype",
         "sample_count_total",
         "coord_def_id",
         "def_key",
@@ -297,6 +299,13 @@ def _columns(row_type: type[NamedTuple]) -> MappingProxyType[str, str]:
 # non-private columns.
 SPOOL_HIDDEN_COLUMNS = ("n_dims", "sample_count_total", "shape")
 
+# Structural columns the spool relation carries *privately*. The leading
+# underscore is load-bearing, not cosmetic: chunk's merge-compatibility
+# grouping and conflict policing both compare all non-private columns, so
+# a public `dtype` would raise CoordMergeError on every merge of patches
+# with differing element types.
+SPOOL_PRIVATE_RENAMES = MappingProxyType({"patch_id": "_patch_id", "dtype": "_dtype"})
+
 # 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,

dascore/units.py

@@ -469,6 +469,67 @@ def is_percent(value: Any) -> bool:
     return isinstance(value, Quantity) and value.units == get_unit("percent")
 
 
+def is_data_size(value: Any) -> bool:
+    """
+    Return True if value is a quantity of information (bytes, bits, MB, ...).
+
+    Pint treats information as dimensionless, so a compatibility check
+    against bytes also passes for percents and bare dimensionless
+    quantities. The base unit is the only reliable discriminator.
+
+    Parameters
+    ----------
+    value
+        Any value of any type to test if it is a data size quantity.
+
+    Examples
+    --------
+    >>> import dascore as dc
+    >>> from dascore.units import is_data_size
+    >>>
+    >>> assert is_data_size(25 * dc.units.megabytes)
+    >>> assert is_data_size(dc.get_quantity("1 MiB"))
+    >>>
+    >>> # Percents, strain and plain numbers are not sizes.
+    >>> assert not is_data_size(dc.get_quantity("50%"))
+    >>> assert not is_data_size(25)
+    """
+    return isinstance(value, Quantity) and value.to_base_units().units == get_unit(
+        "bit"
+    )
+
+
+def get_byte_count(value: Quantity) -> float:
+    """
+    Return the number of bytes a data size quantity represents.
+
+    Parameters
+    ----------
+    value
+        A quantity of information (eg 25 * dc.units.megabytes).
+
+    Notes
+    -----
+    This is the only correct way to get a byte count from a quantity.
+    In particular [`to_float`](`dascore.utils.time.to_float`) is not:
+    pint converts a dimensionless quantity to its base units, which for
+    information is *bits*, so anything routing a size through a plain
+    `float()` conversion is eight times too large.
+
+    Examples
+    --------
+    >>> import dascore as dc
+    >>> from dascore.units import get_byte_count
+    >>>
+    >>> assert get_byte_count(25 * dc.units.megabytes) == 25_000_000
+    >>> assert get_byte_count(dc.get_quantity("1 MiB")) == 1_048_576
+    """
+    if not is_data_size(value):
+        msg = f"Expected a data size quantity (eg '25 MB'), got {value!r}."
+        raise UnitError(msg)
+    return value.to("byte").magnitude
+
+
 def maybe_convert_percent_to_fraction(obj):
     """
     Iterate an object and convert any percentages to fractions.