pyadata-cli 0.5.0__py3-none-any.whl

adata/core/info.py

@@ -0,0 +1,276 @@
+from __future__ import annotations
+
+from typing import Optional, Tuple, Dict, Any, Union
+
+import numpy as np
+
+from adata.elements import spec
+from adata.elements.read import dataframe_columns, element_len, resolve_index
+from adata.storage import is_dataset, is_group, is_hdf5_dataset
+
+
+def _decode_attr(value: Any) -> Any:
+    if isinstance(value, bytes):
+        return value.decode("utf-8")
+    return value
+
+
+def _base_result() -> Dict[str, Any]:
+    return {
+        "type": "unknown",
+        "export_as": None,
+        "encoding": None,
+        "shape": None,
+        "dtype": None,
+        "details": "",
+        "version": None,
+        "ordered": None,
+    }
+
+
+#: Declared encoding-type -> (reported type, suggested export format).
+_ENCODING_TYPES = {
+    spec.ANNDATA: ("anndata", None),
+    spec.RAW: ("raw", None),
+    spec.DICT: ("dict", "json"),
+    spec.DATAFRAME: ("dataframe", "csv"),
+    spec.ARRAY: ("array", "npy"),
+    spec.STRING_ARRAY: ("string-array", "csv"),
+    spec.CATEGORICAL: ("categorical", "csv"),
+    spec.CSR_MATRIX: ("sparse-matrix", "mtx"),
+    spec.CSC_MATRIX: ("sparse-matrix", "mtx"),
+    spec.NULLABLE_INTEGER: ("nullable-array", "csv"),
+    spec.NULLABLE_BOOLEAN: ("nullable-array", "csv"),
+    spec.NULLABLE_STRING_ARRAY: ("nullable-array", "csv"),
+    spec.NUMERIC_SCALAR: ("scalar", "json"),
+    spec.STRING: ("scalar", "json"),
+    spec.NULL: ("null", "json"),
+    spec.AWKWARD_ARRAY: ("awkward-array", "json"),
+}
+
+
+def get_entry_type(entry: Any) -> Dict[str, Any]:
+    """Describe an element: its type, shape, and how it can be exported.
+
+    The declared `encoding-type` is authoritative and is consulted first;
+    structural inspection is only a fallback, for files written by anndata
+    0.7.x and earlier which carry no encoding attributes. Inferring first
+    would let a group merely *containing* a member called `obs_names` be
+    misreported as a dataframe.
+    """
+    result = _base_result()
+
+    enc, enc_ver = spec.encoding_of(entry)
+    result["encoding"] = enc
+    result["version"] = enc_ver
+
+    if is_dataset(entry):
+        result["shape"] = entry.shape
+        result["dtype"] = str(entry.dtype)
+
+    if enc in _ENCODING_TYPES:
+        result["type"], result["export_as"] = _ENCODING_TYPES[enc]
+        _describe(entry, result, enc)
+        return result
+
+    return _infer_untagged(entry, result)
+
+
+def _describe(entry: Any, result: Dict[str, Any], enc: str) -> None:
+    """Fill in the human-readable detail line for a tagged element."""
+    if enc in spec.SPARSE_TYPES:
+        shape = entry.attrs.get("shape", None)
+        shape_str = f"{int(shape[0])}x{int(shape[1])}" if shape is not None else "?"
+        kind = enc.replace("_matrix", "").upper()
+        result["details"] = f"Sparse {kind} matrix {shape_str}"
+        return
+
+    if enc == spec.CATEGORICAL:
+        result["ordered"] = bool(spec.decode_attr(entry.attrs.get("ordered", False)))
+        n_codes = _safe_len(entry.get("codes"))
+        n_cats = _safe_len(entry.get("categories"))
+        order = ", ordered" if result["ordered"] else ""
+        result["details"] = f"Categorical [{n_codes} values, {n_cats} categories{order}]"
+        return
+
+    if enc == spec.DATAFRAME:
+        cols = dataframe_columns(entry)
+        legacy = " (legacy __categories)" if "__categories" in entry else ""
+        result["details"] = f"DataFrame with {len(cols)} columns{legacy}"
+        return
+
+    if enc in spec.MASKED_TYPES:
+        n = _safe_len(entry.get("values"))
+        na = spec.decode_attr(entry.attrs.get("na-value", None))
+        na_str = f", na-value={na}" if na is not None else ""
+        result["details"] = f"Nullable array [{n} values] ({enc}{na_str})"
+        return
+
+    if enc == spec.STRING:
+        result["details"] = "String scalar"
+        return
+
+    if enc == spec.NUMERIC_SCALAR:
+        result["details"] = f"Numeric scalar ({entry.dtype})"
+        return
+
+    if enc == spec.NULL:
+        result["details"] = "Null value"
+        return
+
+    if enc == spec.AWKWARD_ARRAY:
+        result["details"] = f"Awkward array (length={entry.attrs.get('length', '?')})"
+        return
+
+    if enc == spec.STRING_ARRAY:
+        n = entry.shape[0] if getattr(entry, "shape", None) else "?"
+        result["details"] = f"String array [{n}]"
+        return
+
+    if enc in spec.MAPPING_TYPES:
+        result["details"] = f"Group with {len(list(entry.keys()))} keys"
+        return
+
+    if enc == spec.ARRAY:
+        result["details"] = _array_details(entry)
+
+
+def _safe_len(obj: Any) -> Any:
+    if obj is None:
+        return "?"
+    shape = getattr(obj, "shape", None)
+    return shape[0] if shape else "?"
+
+
+def _array_details(entry: Any) -> str:
+    if entry.shape == ():
+        return f"Scalar value ({entry.dtype})"
+    if entry.ndim == 1:
+        return f"1D array [{entry.shape[0]}] ({entry.dtype})"
+    if entry.ndim == 2:
+        return f"Dense matrix {entry.shape[0]}x{entry.shape[1]} ({entry.dtype})"
+    return f"{entry.ndim}D array {entry.shape} ({entry.dtype})"
+
+
+def _infer_untagged(entry: Any, result: Dict[str, Any]) -> Dict[str, Any]:
+    """Classify an element that carries no encoding attributes (anndata <= 0.7)."""
+    if is_dataset(entry):
+        if "categories" in entry.attrs:
+            result["type"] = "categorical"
+            result["export_as"] = "csv"
+            result["version"] = result["version"] or "0.1.0"
+            n_cats = "?"
+            if is_hdf5_dataset(entry):
+                try:
+                    n_cats = entry.file[entry.attrs["categories"]].shape[0]
+                except Exception:
+                    n_cats = "?"
+            result["details"] = (
+                f"Legacy categorical [{entry.shape[0]} values, {n_cats} categories]"
+            )
+            return result
+
+        if entry.shape == ():
+            result["type"] = "scalar"
+            result["export_as"] = "json"
+            result["details"] = f"Scalar value ({entry.dtype})"
+            return result
+
+        result["type"] = "dense-matrix" if entry.ndim == 2 else "array"
+        result["export_as"] = "npy"
+        result["details"] = _array_details(entry)
+        return result
+
+    if is_group(entry):
+        if "codes" in entry and "categories" in entry:
+            result["type"] = "categorical"
+            result["export_as"] = "csv"
+            result["details"] = (
+                f"Categorical [{_safe_len(entry.get('codes'))} values, "
+                f"{_safe_len(entry.get('categories'))} categories]"
+            )
+            return result
+
+        if "values" in entry and "mask" in entry:
+            result["type"] = "nullable-array"
+            result["export_as"] = "csv"
+            result["details"] = f"Nullable array [{_safe_len(entry.get('values'))} values]"
+            return result
+
+        if "_index" in entry.attrs or "obs_names" in entry or "var_names" in entry:
+            result["type"] = "dataframe"
+            result["export_as"] = "csv"
+            result["version"] = result["version"] or "0.1.0"
+            legacy = " (legacy v0.1.0)" if "__categories" in entry else ""
+            result["details"] = (
+                f"DataFrame with {len(dataframe_columns(entry))} columns{legacy}"
+            )
+            return result
+
+        result["type"] = "dict"
+        result["export_as"] = "json"
+        result["details"] = f"Group with {len(list(entry.keys()))} keys"
+        return result
+
+    return result
+
+
+def format_type_info(info: Dict[str, Any]) -> str:
+    type_colors = {
+        "anndata": "cyan",
+        "raw": "cyan",
+        "nullable-array": "blue",
+        "string-array": "green",
+        "null": "dim",
+        "awkward-array": "magenta",
+        "dataframe": "green",
+        "sparse-matrix": "magenta",
+        "dense-matrix": "blue",
+        "array": "blue",
+        "dict": "yellow",
+        "categorical": "green",
+        "scalar": "white",
+        "unknown": "red",
+    }
+
+    color = type_colors.get(info["type"], "white")
+    return f"[{color}]<{info['type']}>[/]"
+
+
+def axis_len(file: Any, axis: str) -> int:
+    """Number of rows along `axis` ("obs" or "var").
+
+    Resolves the index through :func:`resolve_index`, so this works whether the
+    index is a plain dataset or -- as anndata >= 0.11 writes it -- a
+    `nullable-string-array` group of `values` and `mask`.
+    """
+    if axis not in file:
+        raise KeyError(f"'{axis}' not found in the file.")
+
+    group = file[axis]
+    if not is_group(group):
+        raise TypeError(f"'{axis}' is not a group.")
+
+    if axis not in ("obs", "var"):
+        raise ValueError(f"Invalid axis '{axis}'. Must be 'obs' or 'var'.")
+
+    index, index_name = resolve_index(group, axis)
+    try:
+        return element_len(index)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(
+            f"Cannot determine length of '{axis}': index {index_name!r} "
+            f"has no usable length ({exc})."
+        ) from exc
+
+
+def get_axis_group(file: Any, axis: str) -> Tuple[Any, int, str]:
+    """Return the ``(group, length, index_name)`` triple for an axis."""
+    if axis not in ("obs", "var"):
+        raise ValueError("axis must be 'obs' or 'var'.")
+    if axis not in file:
+        raise KeyError(f"'{axis}' not found in the file.")
+
+    group = file[axis]
+    _, index_name = resolve_index(group, axis)
+    return group, axis_len(file, axis), index_name

adata/core/query.py

@@ -0,0 +1,268 @@
+"""A small predicate language for selecting obs/var rows.
+
+Deliberately not SQL. The point is to express the common filters -- one or two
+comparisons on annotation columns -- without pulling a query engine into a tool
+whose selling point is streaming with light dependencies. Anything more
+involved is better served by exporting to CSV and using duckdb.
+
+Grammar::
+
+    expr    := or_expr
+    or_expr := and_expr ( "or" and_expr )*
+    and_expr:= term ( "and" term )*
+    term    := "not" term | "(" expr ")" | comparison
+    comparison := IDENT OP VALUE
+    OP      := == | != | < | <= | > | >= | in | not in
+
+Values are bare words, quoted strings, or comma-separated lists for `in`.
+Comparison against a column is done on its string form for equality and
+membership, and numerically for the ordering operators.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Sequence
+
+import numpy as np
+
+_TOKEN = re.compile(
+    r"""
+    \s*(?:
+        (?P<lparen>\()
+      | (?P<rparen>\))
+      | (?P<op>==|!=|<=|>=|<|>)
+      | (?P<word>"[^"]*"|'[^']*'|[^\s()]+)
+    )
+    """,
+    re.VERBOSE,
+)
+
+_KEYWORDS = {"and", "or", "not", "in"}
+
+
+class QueryError(ValueError):
+    """Raised when a query cannot be parsed or refers to a missing column."""
+
+
+@dataclass
+class Token:
+    kind: str
+    text: str
+
+
+def tokenize(expr: str) -> List[Token]:
+    tokens: List[Token] = []
+    pos = 0
+    while pos < len(expr):
+        match = _TOKEN.match(expr, pos)
+        if match is None:
+            if expr[pos:].strip() == "":
+                break
+            raise QueryError(f"Cannot parse query at {expr[pos:]!r}")
+        pos = match.end()
+        for kind in ("lparen", "rparen", "op", "word"):
+            text = match.group(kind)
+            if text is None:
+                continue
+            if kind == "word" and text.lower() in _KEYWORDS:
+                tokens.append(Token(text.lower(), text.lower()))
+            else:
+                tokens.append(Token(kind, text))
+            break
+    return tokens
+
+
+def _unquote(text: str) -> str:
+    if len(text) >= 2 and text[0] == text[-1] and text[0] in "\"'":
+        return text[1:-1]
+    return text
+
+
+#: A predicate maps a column reader to a boolean mask over the chunk.
+Predicate = Callable[[Dict[str, List[str]]], np.ndarray]
+
+
+class _Parser:
+    def __init__(self, tokens: Sequence[Token]) -> None:
+        self.tokens = list(tokens)
+        self.pos = 0
+
+    def peek(self) -> Optional[Token]:
+        return self.tokens[self.pos] if self.pos < len(self.tokens) else None
+
+    def take(self) -> Token:
+        token = self.peek()
+        if token is None:
+            raise QueryError("Unexpected end of query.")
+        self.pos += 1
+        return token
+
+    def parse(self) -> Predicate:
+        node = self.parse_or()
+        if self.peek() is not None:
+            raise QueryError(f"Unexpected {self.peek().text!r} in query.")
+        return node
+
+    def parse_or(self) -> Predicate:
+        left = self.parse_and()
+        while self.peek() is not None and self.peek().kind == "or":
+            self.take()
+            right = self.parse_and()
+            left = _combine(left, right, np.logical_or)
+        return left
+
+    def parse_and(self) -> Predicate:
+        left = self.parse_term()
+        while self.peek() is not None and self.peek().kind == "and":
+            self.take()
+            right = self.parse_term()
+            left = _combine(left, right, np.logical_and)
+        return left
+
+    def parse_term(self) -> Predicate:
+        token = self.peek()
+        if token is None:
+            raise QueryError("Unexpected end of query.")
+
+        if token.kind == "not":
+            self.take()
+            inner = self.parse_term()
+            return lambda cols: np.logical_not(inner(cols))
+
+        if token.kind == "lparen":
+            self.take()
+            inner = self.parse_or()
+            closing = self.take()
+            if closing.kind != "rparen":
+                raise QueryError("Unbalanced parentheses in query.")
+            return inner
+
+        return self.parse_comparison()
+
+    def parse_comparison(self) -> Predicate:
+        name_token = self.take()
+        if name_token.kind != "word":
+            raise QueryError(f"Expected a column name, got {name_token.text!r}.")
+        column = _unquote(name_token.text)
+
+        op_token = self.take()
+        if op_token.kind == "not":
+            following = self.take()
+            if following.kind != "in":
+                raise QueryError("Expected 'in' after 'not'.")
+            return _membership(column, self._take_list(), negate=True)
+        if op_token.kind == "in":
+            return _membership(column, self._take_list(), negate=False)
+        if op_token.kind != "op":
+            raise QueryError(
+                f"Expected a comparison operator after {column!r}, "
+                f"got {op_token.text!r}."
+            )
+
+        value_token = self.take()
+        if value_token.kind != "word":
+            raise QueryError(f"Expected a value, got {value_token.text!r}.")
+        return _comparison(column, op_token.text, _unquote(value_token.text))
+
+    def _take_list(self) -> List[str]:
+        token = self.take()
+        if token.kind == "lparen":
+            items: List[str] = []
+            while True:
+                nxt = self.take()
+                if nxt.kind == "rparen":
+                    break
+                items.extend(
+                    v.strip() for v in _unquote(nxt.text).split(",") if v.strip()
+                )
+            return items
+        if token.kind != "word":
+            raise QueryError(f"Expected a value list, got {token.text!r}.")
+        return [v.strip() for v in _unquote(token.text).split(",") if v.strip()]
+
+
+def _combine(left: Predicate, right: Predicate, op: Any) -> Predicate:
+    return lambda cols: op(left(cols), right(cols))
+
+
+def _column(cols: Dict[str, List[str]], name: str) -> np.ndarray:
+    if name not in cols:
+        raise QueryError(
+            f"Column {name!r} not found. Available: {', '.join(sorted(cols))}"
+        )
+    return np.asarray(cols[name], dtype=str)
+
+
+def _as_float(values: np.ndarray, column: str) -> np.ndarray:
+    """Parse a string column as floats, treating unparseable entries as NaN."""
+    out = np.full(len(values), np.nan, dtype=np.float64)
+    for i, v in enumerate(values):
+        try:
+            out[i] = float(v)
+        except (TypeError, ValueError):
+            continue
+    return out
+
+
+def _comparison(column: str, op: str, value: str) -> Predicate:
+    def run(cols: Dict[str, List[str]]) -> np.ndarray:
+        values = _column(cols, column)
+        if op == "==":
+            return values == value
+        if op == "!=":
+            return values != value
+
+        try:
+            threshold = float(value)
+        except ValueError as exc:
+            raise QueryError(
+                f"Operator {op!r} needs a number, but got {value!r}."
+            ) from exc
+
+        numeric = _as_float(values, column)
+        with np.errstate(invalid="ignore"):
+            if op == "<":
+                return numeric < threshold
+            if op == "<=":
+                return numeric <= threshold
+            if op == ">":
+                return numeric > threshold
+            return numeric >= threshold
+
+    return run
+
+
+def _membership(column: str, options: Sequence[str], negate: bool) -> Predicate:
+    wanted = set(options)
+
+    def run(cols: Dict[str, List[str]]) -> np.ndarray:
+        values = _column(cols, column)
+        mask = np.isin(values, list(wanted))
+        return np.logical_not(mask) if negate else mask
+
+    return run
+
+
+def compile_query(expr: str) -> Predicate:
+    """Compile a query string into a predicate over a chunk of columns."""
+    tokens = tokenize(expr)
+    if not tokens:
+        raise QueryError("Query is empty.")
+    return _Parser(tokens).parse()
+
+
+def referenced_columns(expr: str) -> List[str]:
+    """Names that look like column references, so only those need reading."""
+    tokens = tokenize(expr)
+    names: List[str] = []
+    for i, token in enumerate(tokens):
+        if token.kind != "word":
+            continue
+        following = tokens[i + 1] if i + 1 < len(tokens) else None
+        if following is not None and following.kind in ("op", "in", "not"):
+            name = _unquote(token.text)
+            if name not in names:
+                names.append(name)
+    return names

adata/core/read.py

@@ -0,0 +1,28 @@
+"""Column reading. The implementation now lives in :mod:`adata.elements.read`.
+
+Kept as a stable import path for existing callers and tests.
+"""
+
+from __future__ import annotations
+
+from adata.elements.read import (
+    col_chunk_as_strings,
+    decode_str_array,
+    element_len,
+    read_categorical_column,
+    read_masked_column,
+    read_str_all,
+    read_str_chunk,
+    resolve_index,
+)
+
+__all__ = [
+    "col_chunk_as_strings",
+    "decode_str_array",
+    "element_len",
+    "read_categorical_column",
+    "read_masked_column",
+    "read_str_all",
+    "read_str_chunk",
+    "resolve_index",
+]

adata/core/select.py

@@ -0,0 +1,122 @@
+"""Selecting rows of an axis by column value, without loading the store.
+
+Both `subset --query` and `split --by` need the same thing: walk an axis's
+annotation columns in chunks and decide which rows to keep. Only the columns a
+query actually mentions are read.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+from adata.core.info import get_axis_group
+from adata.core.query import QueryError, compile_query, referenced_columns
+from adata.elements.read import col_chunk_as_strings, dataframe_columns
+
+
+def _available_columns(group: Any, index_name: str) -> List[str]:
+    cols = dataframe_columns(group, index_name)
+    return cols + [index_name]
+
+
+def _read_chunk(
+    group: Any,
+    columns: List[str], start: int, end: int, cache: Dict
+) -> Dict[str, List[str]]:
+    return {
+        name: col_chunk_as_strings(group, name, start, end, cache)
+        for name in columns
+    }
+
+
+def select_indices(
+    root: Any,
+    axis: str,
+    expr: str,
+    *,
+    chunk_size: int = 100_000,
+) -> np.ndarray:
+    """Row indices of `axis` matching the query `expr`.
+
+    Only the columns the query references are read, so filtering on one
+    annotation does not touch the rest of the frame.
+    """
+    group, n_rows, index_name = get_axis_group(root, axis)
+    predicate = compile_query(expr)
+
+    wanted = referenced_columns(expr)
+    available = _available_columns(group, index_name)
+    missing = [c for c in wanted if c not in available]
+    if missing:
+        raise QueryError(
+            f"Column(s) {', '.join(missing)} not found in '{axis}'. "
+            f"Available: {', '.join(sorted(available))}"
+        )
+
+    cache: Dict = {}
+    kept: List[np.ndarray] = []
+    for start in range(0, n_rows, chunk_size):
+        end = min(start + chunk_size, n_rows)
+        mask = np.asarray(
+            predicate(_read_chunk(group, wanted, start, end, cache)), dtype=bool
+        )
+        if mask.any():
+            kept.append(np.nonzero(mask)[0] + start)
+
+    if not kept:
+        return np.empty(0, dtype=np.int64)
+    return np.concatenate(kept).astype(np.int64)
+
+
+def group_indices(
+    root: Any,
+    axis: str,
+    column: str,
+    *,
+    chunk_size: int = 100_000,
+) -> Tuple[Dict[str, np.ndarray], List[str]]:
+    """Group `axis` row indices by the value of `column`.
+
+    Returns the groups and the distinct labels in order of first appearance,
+    so output ordering follows the data rather than the hash of the labels.
+    """
+    group, n_rows, index_name = get_axis_group(root, axis)
+
+    available = _available_columns(group, index_name)
+    if column not in available:
+        raise KeyError(
+            f"Column {column!r} not found in '{axis}'. "
+            f"Available: {', '.join(sorted(available))}"
+        )
+
+    cache: Dict = {}
+    buckets: Dict[str, List[np.ndarray]] = {}
+    order: List[str] = []
+
+    for start in range(0, n_rows, chunk_size):
+        end = min(start + chunk_size, n_rows)
+        values = np.asarray(
+            col_chunk_as_strings(group, column, start, end, cache), dtype=str
+        )
+        for label in dict.fromkeys(values.tolist()):
+            if label not in buckets:
+                buckets[label] = []
+                order.append(label)
+            buckets[label].append(np.nonzero(values == label)[0] + start)
+
+    return (
+        {k: np.concatenate(v).astype(np.int64) for k, v in buckets.items()},
+        order,
+    )
+
+
+def read_names(root: Any, axis: str, indices: Optional[np.ndarray] = None) -> List[str]:
+    """Read an axis's index labels, optionally only at `indices`."""
+    group, n_rows, index_name = get_axis_group(root, axis)
+    cache: Dict = {}
+    names = col_chunk_as_strings(group, index_name, 0, n_rows, cache)
+    if indices is None:
+        return names
+    return [names[i] for i in indices]