diff --git a/tafra/__init__.py b/tafra/__init__.py index b90c531..161c434 100644 --- a/tafra/__init__.py +++ b/tafra/__init__.py @@ -14,7 +14,8 @@ """ from importlib.metadata import version as _version -__version__ = _version('tafra') + +__version__ = _version("tafra") from .base import Tafra, object_formatter from .group import Union, GroupBy, Transform, IterateBy, InnerJoin, LeftJoin, CrossJoin diff --git a/tafra/base.py b/tafra/base.py index 04f445f..e1a2af5 100644 --- a/tafra/base.py +++ b/tafra/base.py @@ -12,9 +12,10 @@ ----- Created on April 25, 2020 """ + from __future__ import annotations -__all__ = ['Tafra'] +__all__ = ["Tafra"] from pathlib import Path import re @@ -29,9 +30,20 @@ import numpy as np from .protocol import Series, DataFrame, Cursor # just for mypy... -from typing import (Any, Callable, Mapping, Sequence, - Sized, Iterable, Iterator, KeysView, ValuesView, ItemsView, - IO, TYPE_CHECKING) +from typing import ( + Any, + Callable, + Mapping, + Sequence, + Sized, + Iterable, + Iterator, + KeysView, + ValuesView, + ItemsView, + IO, + TYPE_CHECKING, +) from typing_extensions import Concatenate, Literal, ParamSpec from typing import cast from io import TextIOWrapper @@ -40,42 +52,37 @@ from .csvreader import CSVReader -P = ParamSpec('P') +P = ParamSpec("P") # default object formats object_formatter = ObjectFormatter() -object_formatter['Decimal'] = lambda x: x.astype(float) +object_formatter["Decimal"] = lambda x: x.astype(float) NAMEDTUPLE_TYPE: dict[str, type[Any]] = { - 'int': int, - 'float': float, - 'bool': bool, - 'str': str, - 'date': date, - 'datetime': datetime, - 'object': str, + "int": int, + "float": float, + "bool": bool, + "str": str, + "date": date, + "datetime": datetime, + "object": str, } RECORD_TYPE: dict[str, Callable[[Any], Any]] = { - 'int': int, - 'float': float, - 'bool': bool, - 'str': str, - 'date': lambda x: x.isoformat(), - 'datetime': lambda x: x.isoformat(), - 'object': str, + "int": int, + "float": float, + "bool": bool, + "str": str, + "date": lambda x: x.isoformat(), + "datetime": lambda x: x.isoformat(), + "object": str, } Scalar = str | int | float | bool -_Mapping = ( - Mapping[str, Any] - | Mapping[int, Any] - | Mapping[float, Any] - | Mapping[bool, Any] -) +_Mapping = Mapping[str, Any] | Mapping[int, Any] | Mapping[float, Any] | Mapping[bool, Any] _Element = tuple[str | int | float | np.ndarray[Any, Any], Any] | list[Any] | _Mapping InitVar = ( tuple[str, Any] @@ -124,6 +131,7 @@ class Tafra: methods. """ + data: dc.InitVar[InitVar] dtypes: dc.InitVar[InitVar | None] = None validate: dc.InitVar[bool] = True @@ -133,6 +141,7 @@ class Tafra: _dtypes: dict[str, str] = dc.field(init=False) if TYPE_CHECKING: + def __init__( self, data: InitVar, @@ -169,12 +178,12 @@ def __post_init__(self, data: InitVar, dtypes: InitVar | None, validate: bool) - rows = n_rows if self.check_rows and rows != n_rows: - raise ValueError('`Tafra` must have consistent row counts.') + raise ValueError("`Tafra` must have consistent row counts.") elif rows < n_rows: # pragma: no cover rows = n_rows if rows is None: - raise ValueError('No data provided in constructor statement.') + raise ValueError("No data provided in constructor statement.") self.update_dtypes_inplace(self._dtypes, _from_init=True) # must coalesce all dtypes immediately, other functions assume a @@ -214,8 +223,10 @@ def _check_initvar(self, values: InitVar) -> dict[str, Any]: _values = cast(dict[Any, Any], values) # type: ignore[unreachable] if not isinstance(_values, dict): - raise TypeError('Must contain `Dict`, `Mapping`, `Sequence`, Iterable, or Iterator, ' - f'got `{type(_values)}`') + raise TypeError( + "Must contain `Dict`, `Mapping`, `Sequence`, Iterable, or Iterator, " + f"got `{type(_values)}`" + ) # cast all keys to strings if they are not # must copy first as mutating the dict changes next(iterator) @@ -240,14 +251,17 @@ def _parse_sequence(self, values: Sequence[_Element]) -> dict[Any, Any]: # is the key an ndarray? turn it into a scalar if isinstance(head[0], np.ndarray) and len(np.atleast_1d(head[0])) == 1: # mypy doesn't get that we've checked the head of values as an ndarray - _values = {key.item(): value for key, value in - cast(Iterable[tuple[np.ndarray[Any, Any], Any]], values)} + _values = { + key.item(): value + for key, value in cast(Iterable[tuple[np.ndarray[Any, Any], Any]], values) + } else: _values = dict(cast(Iterable[tuple[Any, Any]], values)) else: - raise TypeError('Sequence must contain `Dict`, `Mapping`, or `Sequence`, ' - f'got `{type(head)}`') + raise TypeError( + f"Sequence must contain `Dict`, `Mapping`, or `Sequence`, got `{type(head)}`" + ) return _values @@ -267,17 +281,25 @@ def _parse_iterable(self, values: Iterable[_Element]) -> dict[Any, Any]: # is the key an ndarray? turn it into a scalar if isinstance(head[0], np.ndarray) and len(np.atleast_1d(head[0])) == 1: # mypy doesn't get that we've checked the head of values as an ndarray - _values = {key.item(): value for key, value in chain( - cast(Iterable[tuple[np.ndarray[Any, Any], Any]], [head]), - cast(Iterator[tuple[np.ndarray[Any, Any], Any]], iter_values))} + _values = { + key.item(): value + for key, value in chain( + cast(Iterable[tuple[np.ndarray[Any, Any], Any]], [head]), + cast(Iterator[tuple[np.ndarray[Any, Any], Any]], iter_values), + ) + } else: - _values = dict(chain( - cast(Iterable[tuple[Any, Any]], [head]), - cast(Iterator[tuple[Any, Any]], iter_values))) + _values = dict( + chain( + cast(Iterable[tuple[Any, Any]], [head]), + cast(Iterator[tuple[Any, Any]], iter_values), + ) + ) else: - raise TypeError('Iterable must contain `Dict`, `Mapping`, or `Sequence`, ' - f'got `{type(head)}`') + raise TypeError( + f"Iterable must contain `Dict`, `Mapping`, or `Sequence`, got `{type(head)}`" + ) return _values @@ -298,25 +320,31 @@ def _parse_iterator(self, values: Iterator[_Element]) -> dict[Any, Any]: # is the key an ndarray? turn it into a scalar if isinstance(head[0], np.ndarray) and len(np.atleast_1d(head[0])) == 1: # mypy doesn't get that we've checked the head of values as an ndarray - _values = {key.item(): value for key, value in chain( - cast(Iterable[tuple[np.ndarray[Any, Any], Any]], [head]), - cast(Iterator[tuple[np.ndarray[Any, Any], Any]], values))} + _values = { + key.item(): value + for key, value in chain( + cast(Iterable[tuple[np.ndarray[Any, Any], Any]], [head]), + cast(Iterator[tuple[np.ndarray[Any, Any], Any]], values), + ) + } else: - _values = dict(chain( - cast(Iterable[tuple[Any, Any]], [head]), - cast(Iterator[tuple[Any, Any]], values))) + _values = dict( + chain( + cast(Iterable[tuple[Any, Any]], [head]), + cast(Iterator[tuple[Any, Any]], values), + ) + ) else: - raise TypeError('Iterator must contain `Dict`, `Mapping`, or `Sequence`, ' - f'got `{type(head)}`') + raise TypeError( + f"Iterator must contain `Dict`, `Mapping`, or `Sequence`, got `{type(head)}`" + ) return _values def __getitem__( - self, - item: ( - str | int | slice | Sequence[str | int | bool] | np.ndarray[Any, Any] - )) -> Any: + self, item: (str | int | slice | Sequence[str | int | bool] | np.ndarray[Any, Any]) + ) -> Any: # return type is actually np.ndarray | 'Tafra' but mypy requires user to type check # in either case, what we return is a "slice" of the `Tafra` if isinstance(item, str): @@ -338,32 +366,30 @@ def __getitem__( return self._aindex(cast(Sequence[int | bool], item)) else: - raise TypeError(f'Type {type(item)} not supported.') + raise TypeError(f"Type {type(item)} not supported.") - def __setitem__(self, item: str, - value: np.ndarray[Any, Any] | Sequence[Any] | Any) -> None: + def __setitem__(self, item: str, value: np.ndarray[Any, Any] | Sequence[Any] | Any) -> None: self._ensure_valid(item, value, set_item=True) def __repr__(self) -> str: - if not hasattr(self, '_rows'): - return f'Tafra(data={self._data}, dtypes={self._dtypes}, rows=n/a)' - return f'Tafra(data={self._data}, dtypes={self._dtypes}, rows={self._rows})' + if not hasattr(self, "_rows"): + return f"Tafra(data={self._data}, dtypes={self._dtypes}, rows=n/a)" + return f"Tafra(data={self._data}, dtypes={self._dtypes}, rows={self._rows})" def __str__(self) -> str: return self.__repr__() def __len__(self) -> int: - assert self._data is not None, \ - 'Interal error: Cannot construct a Tafra with no data.' + assert self._data is not None, "Internal error: Cannot construct a Tafra with no data." return self._rows - def __iter__(self) -> Iterator['Tafra']: + def __iter__(self) -> Iterator["Tafra"]: return (self._iindex(i) for i in range(self._rows)) - def __rshift__(self, other: Callable[['Tafra'], 'Tafra']) -> 'Tafra': + def __rshift__(self, other: Callable[["Tafra"], "Tafra"]) -> "Tafra": return self.pipe(other) - def iterrows(self) -> Iterator['Tafra']: + def iterrows(self) -> Iterator["Tafra"]: """ Yield rows as `Tafra`. Use `itertuples()` for better performance. @@ -374,7 +400,7 @@ def iterrows(self) -> Iterator['Tafra']: """ yield from self.__iter__() - def itertuples(self, name: str | None = 'Tafra') -> Iterator[tuple[Any, ...]]: + def itertuples(self, name: str | None = "Tafra") -> Iterator[tuple[Any, ...]]: """ Yield rows as `NamedTuple`, or if `name` is `None`, yield rows as `tuple`. @@ -416,9 +442,9 @@ def _update_rows(self) -> None: iter_values = iter(self._data.values()) self._rows = len(next(iter_values)) if self.check_rows and not all(len(v) == self._rows for v in iter_values): - raise TypeError('Uneven length of data.') + raise TypeError("Uneven length of data.") - def _slice(self, _slice: slice) -> 'Tafra': + def _slice(self, _slice: slice) -> "Tafra": """ Use a `slice` to slice the `Tafra`. @@ -433,13 +459,12 @@ def _slice(self, _slice: slice) -> 'Tafra': The sliced `Tafra`. """ return Tafra( - {column: np.atleast_1d(value[_slice]) - for column, value in self._data.items()}, + {column: np.atleast_1d(value[_slice]) for column, value in self._data.items()}, self._dtypes, - validate=False + validate=False, ) - def _iindex(self, index: int) -> 'Tafra': + def _iindex(self, index: int) -> "Tafra": """ Use a `int` to slice the `Tafra`. @@ -453,13 +478,12 @@ def _iindex(self, index: int) -> 'Tafra': The sliced `Tafra`. """ return Tafra( - {column: value[[index]] - for column, value in self._data.items()}, + {column: value[[index]] for column, value in self._data.items()}, self._dtypes, - validate=False + validate=False, ) - def _aindex(self, index: Sequence[int | bool]) -> 'Tafra': + def _aindex(self, index: Sequence[int | bool]) -> "Tafra": """ Use numpy advanced indexing to slice the `Tafra`. @@ -473,13 +497,12 @@ def _aindex(self, index: Sequence[int | bool]) -> 'Tafra': The sliced `Tafra`. """ return Tafra( - {column: value[index] - for column, value in self._data.items()}, + {column: value[index] for column, value in self._data.items()}, self._dtypes, - validate=False + validate=False, ) - def _ndindex(self, index: np.ndarray[Any, Any]) -> 'Tafra': + def _ndindex(self, index: np.ndarray[Any, Any]) -> "Tafra": """ Use `numpy.ndarray` indexing to slice the `Tafra`. @@ -493,13 +516,12 @@ def _ndindex(self, index: np.ndarray[Any, Any]) -> 'Tafra': The sliced `Tafra`. """ if index.ndim != 1: - raise IndexError(f'Indexing np.ndarray must ndim == 1, got ndim == {index.ndim}') + raise IndexError(f"Indexing np.ndarray must ndim == 1, got ndim == {index.ndim}") return Tafra( - {column: value[index] - for column, value in self._data.items()}, + {column: value[index] for column, value in self._data.items()}, self._dtypes, - validate=False + validate=False, ) def _repr_pretty_(self, p: Any, cycle: bool) -> None: @@ -521,9 +543,9 @@ def _repr_pretty_(self, p: Any, cycle: bool) -> None: Calls p.text and returns. """ if cycle: - p.text('Tafra(...)') + p.text("Tafra(...)") else: - p.text(self._pretty_format(lambda s: ' ' + pprint.pformat(s, indent=1)[1:].strip())) + p.text(self._pretty_format(lambda s: " " + pprint.pformat(s, indent=1)[1:].strip())) def _repr_html_(self) -> str: """ @@ -545,18 +567,21 @@ def _pretty_format(self, formatter: Callable[[object], str]) -> str: string: str The formatted string for printing. """ - PATTERN = r'(, dtype=[a-z]+)(?=\))' + PATTERN = r"(, dtype=[a-z]+)(?=\))" - return '\n'.join([ - 'Tafra(data = {', - f'{re.sub(PATTERN, "", formatter(self._data))},', - 'dtypes = {', - f'{re.sub(PATTERN, "", formatter(self._dtypes))},', - f'rows = {self._rows})' - ]) + return "\n".join( + [ + "Tafra(data = {", + f"{re.sub(PATTERN, '', formatter(self._data))},", + "dtypes = {", + f"{re.sub(PATTERN, '', formatter(self._dtypes))},", + f"rows = {self._rows})", + ] + ) - def pformat(self, indent: int = 1, width: int = 80, depth: int | None = None, - compact: bool = False) -> str: + def pformat( + self, indent: int = 1, width: int = 80, depth: int | None = None, compact: bool = False + ) -> str: """ Format for pretty printing. Parameters are passed to `pprint.PrettyPrinter`. @@ -581,11 +606,14 @@ def pformat(self, indent: int = 1, width: int = 80, depth: int | None = None, A formatted string for pretty printing. """ return self._pretty_format( - lambda s: indent * ' ' + pprint.pformat( - s, indent, width, depth, compact=compact)[1:].strip()) + lambda s: ( + indent * " " + pprint.pformat(s, indent, width, depth, compact=compact)[1:].strip() + ) + ) - def pprint(self, indent: int = 1, width: int = 80, depth: int | None = None, - compact: bool = False) -> None: + def pprint( + self, indent: int = 1, width: int = 80, depth: int | None = None, compact: bool = False + ) -> None: """ Pretty print. Parameters are passed to `pprint.PrettyPrinter`. @@ -624,8 +652,9 @@ def _html_thead(columns: Iterable[Any]) -> str: HTML: str The HTML table head. """ - return '\n\n{th}\n\n' \ - .format(th='\n'.join(f'{c}' for c in columns)) + return "\n\n{th}\n\n".format( + th="\n".join(f"{c}" for c in columns) + ) @staticmethod def _html_tr(row: Iterable[Any]) -> str: @@ -642,8 +671,7 @@ def _html_tr(row: Iterable[Any]) -> str: HTML: str The HTML table row. """ - return '\n{td}\n' \ - .format(td='\n'.join(f'{td}' for td in row)) + return "\n{td}\n".format(td="\n".join(f"{td}" for td in row)) @staticmethod def _html_tbody(tr: Iterable[str]) -> str: @@ -660,8 +688,7 @@ def _html_tbody(tr: Iterable[str]) -> str: HTML: str The HTML table body. """ - return '\n{tr}\n' \ - .format(tr='\n'.join(tr)) + return "\n{tr}\n".format(tr="\n".join(tr)) @staticmethod def _html_table(thead: str, tbody: str) -> str: @@ -681,7 +708,7 @@ def _html_table(thead: str, tbody: str) -> str: HTML: str The HTML table. """ - return f'\n{thead}\n{tbody}\n
' + return f"\n{thead}\n{tbody}\n
" def to_html(self, n: int = 20) -> str: """ @@ -697,23 +724,28 @@ def to_html(self, n: int = 20) -> str: HTML: str The HTML table representation. """ - thead = self._html_thead(chain([''], self._data.keys())) + thead = self._html_thead(chain([""], self._data.keys())) tr = chain( - [self._html_tr(chain( - ['dtype'], - (self._dtypes[column] for column in self._data.keys()) - ))], - (self._html_tr(chain( - [i], - (v[i] for v in self._data.values()) - )) - for i in range(min(n, self._rows))) + [ + self._html_tr( + chain(["dtype"], (self._dtypes[column] for column in self._data.keys())) + ) + ], + ( + self._html_tr(chain([i], (v[i] for v in self._data.values()))) + for i in range(min(n, self._rows)) + ), ) tbody = self._html_tbody(tr) return self._html_table(thead, tbody) - def _ensure_valid(self, column: str, value: np.ndarray[Any, Any] | Sequence[Any] | Any, - check_rows: bool = True, set_item: bool = False) -> None: + def _ensure_valid( + self, + column: str, + value: np.ndarray[Any, Any] | Sequence[Any] | Any, + check_rows: bool = True, + set_item: bool = False, + ) -> None: """ Validate values as an `np.ndarray` of equal length to `rows` before assignment. Will attempt to create a `np.ndarray` if `value` is not one @@ -758,25 +790,26 @@ def _ensure_valid(self, column: str, value: np.ndarray[Any, Any] | Sequence[Any] elif not isinstance(value, Sized): # type: ignore[unreachable] value = np.full(rows, value) - assert isinstance(value, np.ndarray), \ - 'Internal error: `Tafra` only supports assigning `ndarray`.' + assert isinstance(value, np.ndarray), ( + "Internal error: `Tafra` only supports assigning `ndarray`." + ) if value.ndim > 1: sq_value = value.squeeze() if sq_value.ndim > 1: - raise ValueError('`ndarray` or `np.squeeze(ndarray)` must have ndim == 1.') + raise ValueError("`ndarray` or `np.squeeze(ndarray)` must have ndim == 1.") elif sq_value.ndim == 1: # if value was a single item, squeeze returns zero length item - warnings.warn('`np.squeeze(ndarray)` applied to set ndim == 1.') + warnings.warn("`np.squeeze(ndarray)` applied to set ndim == 1.") value = sq_value - assert value.ndim >= 1, \ - 'Interal error: `Tafra` only supports assigning ndim == 1.' + assert value.ndim >= 1, "Internal error: `Tafra` only supports assigning ndim == 1." if check_rows and len(value) != rows: raise ValueError( - '`Tafra` must have consistent row counts.\n' - f'This `Tafra` has {rows} rows. Assigned {_type} has {len(value)} rows.') + "`Tafra` must have consistent row counts.\n" + f"This `Tafra` has {rows} rows. Assigned {_type} has {len(value)} rows." + ) # special parsing of various object types parsed_value = object_formatter.parse_dtype(value) @@ -788,7 +821,7 @@ def _ensure_valid(self, column: str, value: np.ndarray[Any, Any] | Sequence[Any] self._data[column] = value self._dtypes[column] = self._format_dtype(value.dtype) - def parse_object_dtypes(self) -> 'Tafra': + def parse_object_dtypes(self) -> "Tafra": """ Parse the object dtypes using the `ObjectFormatter` instance. """ @@ -823,7 +856,7 @@ def _validate_columns(self, columns: Iterable[str]) -> None: """ for column in columns: if column not in self._data: - raise ValueError(f'Column {column} does not exist in `tafra`.') + raise ValueError(f"Column {column} does not exist in `tafra`.") def _validate_dtypes(self, dtypes: dict[str, Any]) -> dict[str, str]: """ @@ -862,10 +895,10 @@ def _format_dtype(dtype: Any) -> str: """ _dtype = dtype if isinstance(dtype, np.dtype) else np.dtype(dtype) name = _dtype.type.__name__ - if 'str' in name: - return 'str' + if "str" in name: + return "str" - return name.replace('_', '') + return name.replace("_", "") @staticmethod def _reduce_dtype(dtype: Any) -> str: @@ -883,7 +916,7 @@ def _reduce_dtype(dtype: Any) -> str: The parsed dtype. """ name = np.dtype(dtype).type.__name__ - m = re.search(r'([a-z]+)', name) + m = re.search(r"([a-z]+)", name) if m: return m.group(1) @@ -891,8 +924,13 @@ def _reduce_dtype(dtype: Any) -> str: return name # pragma: no cover @classmethod - def from_records(cls, records: Iterable[Iterable[Any]], columns: Iterable[str], - dtypes: Iterable[Any] | None = None, **kwargs: Any) -> 'Tafra': + def from_records( + cls, + records: Iterable[Iterable[Any]], + columns: Iterable[str], + dtypes: Iterable[Any] | None = None, + **kwargs: Any, + ) -> "Tafra": """ Construct a `Tafra` from an Iterator of records, e.g. from a SQL query. The records should be a nested Iterable, but can also be fed a cursor method such as @@ -920,11 +958,11 @@ def from_records(cls, records: Iterable[Iterable[Any]], columns: Iterable[str], return Tafra( {column: value for column, value in zip(columns, zip(*records))}, {column: value for column, value in zip(columns, dtypes)}, - **kwargs + **kwargs, ) @classmethod - def from_series(cls, s: Series, dtype: str | None = None, **kwargs: Any) -> 'Tafra': + def from_series(cls, s: Series, dtype: str | None = None, **kwargs: Any) -> "Tafra": """ Construct a `Tafra` from a `pandas.Series`. If `dtype` is not given, take from `pandas.Series.dtype`. @@ -946,15 +984,12 @@ def from_series(cls, s: Series, dtype: str | None = None, **kwargs: Any) -> 'Taf dtype = s.dtype dtypes = {s.name: cls._format_dtype(dtype)} - return cls( - {s.name: s.values.astype(dtypes[s.name])}, - dtypes, - **kwargs - ) + return cls({s.name: s.values.astype(dtypes[s.name])}, dtypes, **kwargs) @classmethod - def from_dataframe(cls, df: DataFrame, dtypes: dict[str, Any] | None = None, - **kwargs: Any) -> 'Tafra': + def from_dataframe( + cls, df: DataFrame, dtypes: dict[str, Any] | None = None, **kwargs: Any + ) -> "Tafra": """ Construct a `Tafra` from a `pandas.DataFrame`. If `dtypes` are not given, take from `pandas.DataFrame.dtypes`. @@ -979,11 +1014,11 @@ def from_dataframe(cls, df: DataFrame, dtypes: dict[str, Any] | None = None, return cls( {c: df[c].values.astype(dtypes[c]) for c in df.columns}, {c: dtypes[c] for c in df.columns}, - **kwargs + **kwargs, ) @classmethod - def read_sql(cls, query: str, cur: Cursor) -> 'Tafra': + def read_sql(cls, query: str, cur: Cursor) -> "Tafra": """ Execute a SQL SELECT statement using a `pyodbc.Cursor` and return a Tuple of column names and an Iterator of records. @@ -1012,7 +1047,7 @@ def read_sql(cls, query: str, cur: Cursor) -> 'Tafra': return Tafra.from_records(chain([head], cur.fetchall()), columns, dtypes) @classmethod - def read_sql_chunks(cls, query: str, cur: Cursor, chunksize: int = 100) -> Iterator['Tafra']: + def read_sql_chunks(cls, query: str, cur: Cursor, chunksize: int = 100) -> Iterator["Tafra"]: """ Execute a SQL SELECT statement using a `pyodbc.Cursor` and return a Tuple of column names and an Iterator of records. @@ -1047,10 +1082,14 @@ def chunks(iterable: Iterable[Any], chunksize: int = 1000) -> Iterator[Iterable[ yield Tafra.from_records(chunk, columns, dtypes) @classmethod - def read_csv(cls, csv_file: str | Path | TextIOWrapper | IO[str], guess_rows: int = 5, - missing: str | None = '', dtypes: dict[str, Any] | None = None, - **csvkw: dict[str, Any] - ) -> 'Tafra': + def read_csv( + cls, + csv_file: str | Path | TextIOWrapper | IO[str], + guess_rows: int = 5, + missing: str | None = "", + dtypes: dict[str, Any] | None = None, + **csvkw: dict[str, Any], + ) -> "Tafra": """ Read a CSV file with a header row, infer the types of each column, and return a Tafra containing the file's contents. @@ -1075,13 +1114,13 @@ def read_csv(cls, csv_file: str | Path | TextIOWrapper | IO[str], guess_rows: in tafra: Tafra The constructed `Tafra`. """ - reader = CSVReader(cast(str | Path | TextIOWrapper, csv_file), - guess_rows, missing, **csvkw) + reader = CSVReader(cast(str | Path | TextIOWrapper, csv_file), guess_rows, missing, **csvkw) return Tafra(reader.read(), dtypes=dtypes) @classmethod - def as_tafra(cls, maybe_tafra: Tafra | DataFrame | Series | dict[str, Any] | Any - ) -> Tafra | None: + def as_tafra( + cls, maybe_tafra: Tafra | DataFrame | Series | dict[str, Any] | Any + ) -> Tafra | None: """ Returns the unmodified `tafra`` if already a `Tafra`, else construct a `Tafra` from known types or subtypes of `DataFrame` or `dict`. @@ -1106,19 +1145,19 @@ def as_tafra(cls, maybe_tafra: Tafra | DataFrame | Series | dict[str, Any] | Any elif isinstance(maybe_tafra, Series): # pragma: no cover return cls.from_series(maybe_tafra) - elif type(maybe_tafra).__name__ == 'Series': # pragma: no cover + elif type(maybe_tafra).__name__ == "Series": # pragma: no cover return cls.from_series(cast(Series, maybe_tafra)) elif isinstance(maybe_tafra, DataFrame): # pragma: no cover return cls.from_dataframe(maybe_tafra) - elif type(maybe_tafra).__name__ == 'DataFrame': # pragma: no cover + elif type(maybe_tafra).__name__ == "DataFrame": # pragma: no cover return cls.from_dataframe(cast(DataFrame, maybe_tafra)) elif isinstance(maybe_tafra, dict): return cls(maybe_tafra) - raise TypeError(f'Unknown type `{type(maybe_tafra)}` for conversion to `Tafra`') + raise TypeError(f"Unknown type `{type(maybe_tafra)}` for conversion to `Tafra`") @property def columns(self) -> tuple[str, ...]: @@ -1134,7 +1173,7 @@ def columns(self) -> tuple[str, ...]: @columns.setter def columns(self, value: Any) -> None: - raise ValueError('Assignment to `columns` is forbidden.') + raise ValueError("Assignment to `columns` is forbidden.") @property def rows(self) -> int: @@ -1151,7 +1190,7 @@ def rows(self) -> int: @rows.setter def rows(self, value: Any) -> None: - raise ValueError('Assignment to `rows` is forbidden.') + raise ValueError("Assignment to `rows` is forbidden.") @property # type: ignore def data(self) -> dict[str, np.ndarray[Any, Any]]: @@ -1167,7 +1206,7 @@ def data(self) -> dict[str, np.ndarray[Any, Any]]: @data.setter def data(self, value: Any) -> None: - raise ValueError('Assignment to `data` is forbidden.') + raise ValueError("Assignment to `data` is forbidden.") @property # type: ignore def dtypes(self) -> dict[str, str]: @@ -1183,7 +1222,7 @@ def dtypes(self) -> dict[str, str]: @dtypes.setter def dtypes(self, value: Any) -> None: - raise ValueError('Assignment to `dtypes` is forbidden.') + raise ValueError("Assignment to `dtypes` is forbidden.") @property def size(self) -> int: @@ -1199,7 +1238,7 @@ def size(self) -> int: @size.setter def size(self, value: Any) -> None: - raise ValueError('Assignment to `size` is forbidden.') + raise ValueError("Assignment to `size` is forbidden.") @property def ndim(self) -> int: @@ -1215,7 +1254,7 @@ def ndim(self) -> int: @ndim.setter def ndim(self, value: Any) -> None: - raise ValueError('Assignment to `ndim` is forbidden.') + raise ValueError("Assignment to `ndim` is forbidden.") @property def shape(self) -> tuple[int, int]: @@ -1231,7 +1270,7 @@ def shape(self) -> tuple[int, int]: @shape.setter def shape(self, value: Any) -> None: - raise ValueError('Assignment to `shape` is forbidden.') + raise ValueError("Assignment to `shape` is forbidden.") def row_map(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Iterator[Any]: """ @@ -1256,8 +1295,9 @@ def row_map(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Iterator """ return (fn(tf, *args, **kwargs) for tf in self.__iter__()) - def tuple_map(self, fn: Callable[..., Any], *args: Any, - name: str | None = 'Tafra', **kwargs: Any) -> Iterator[Any]: + def tuple_map( + self, fn: Callable[..., Any], *args: Any, name: str | None = "Tafra", **kwargs: Any + ) -> Iterator[Any]: """ Map a function over rows. This is faster than `row_map()`. To apply to specific columns, use `select()` first. @@ -1288,8 +1328,7 @@ def tuple_map(self, fn: Callable[..., Any], *args: Any, An iterator to map the function. """ if name is None: - return (fn(row, *args, **kwargs) - for row in zip(*self._data.values())) + return (fn(row, *args, **kwargs) for row in zip(*self._data.values())) return (fn(tf, *args, **kwargs) for tf in self.itertuples(name)) def col_map(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Iterator[Any]: @@ -1316,8 +1355,9 @@ def col_map(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Iterator return (fn(value, *args, **kwargs) for column, value in self.itercols()) - def key_map(self, fn: Callable[..., Any], - *args: Any, **kwargs: Any) -> Iterator[tuple[str, Any]]: + def key_map( + self, fn: Callable[..., Any], *args: Any, **kwargs: Any + ) -> Iterator[tuple[str, Any]]: """ Map a function over columns like `col_map()`, but return `tuple` of the key with the function result. To apply to specific columns, use `select()` @@ -1341,11 +1381,12 @@ def key_map(self, fn: Callable[..., Any], """ return ((column, fn(value, *args, **kwargs)) for column, value in self.itercols()) - def pipe(self, fn: Callable[Concatenate['Tafra', P], 'Tafra'], - *args: Any, **kwargs: Any) -> 'Tafra': + def pipe( + self, fn: Callable[Concatenate["Tafra", P], "Tafra"], *args: Any, **kwargs: Any + ) -> "Tafra": """ Apply a function to the `Tafra` and return the resulting `Tafra`. Primarily - used to build a tranformer pipeline. + used to build a transformer pipeline. Parameters ---------- @@ -1365,7 +1406,7 @@ def pipe(self, fn: Callable[Concatenate['Tafra', P], 'Tafra'], """ return fn(self, *args, **kwargs) - def select(self, columns: Iterable[str]) -> 'Tafra': + def select(self, columns: Iterable[str]) -> "Tafra": """ Use column names to slice the `Tafra` columns analogous to SQL SELECT. This does not copy the data. Call `copy()` to obtain a copy of the sliced @@ -1388,10 +1429,10 @@ def select(self, columns: Iterable[str]) -> 'Tafra': return Tafra( {column: self._data[column] for column in columns}, {column: self._dtypes[column] for column in columns}, - validate=False + validate=False, ) - def head(self, n: int = 5) -> 'Tafra': + def head(self, n: int = 5) -> "Tafra": """ Display the head of the `Tafra`. @@ -1406,7 +1447,7 @@ def head(self, n: int = 5) -> 'Tafra': """ return self._slice(slice(n)) - def tail(self, n: int = 5) -> 'Tafra': + def tail(self, n: int = 5) -> "Tafra": """ Return the last `n` rows. @@ -1422,8 +1463,7 @@ def tail(self, n: int = 5) -> 'Tafra': """ return self._slice(slice(-n, None)) - def sort(self, columns: str | Iterable[str], - reverse: bool = False) -> 'Tafra': + def sort(self, columns: str | Iterable[str], reverse: bool = False) -> "Tafra": """ Return a new `Tafra` sorted by the given columns. @@ -1447,7 +1487,7 @@ def sort(self, columns: str | Iterable[str], return result._slice(slice(None, None, -1)) return result - def sample(self, n: int, seed: int | None = None) -> 'Tafra': + def sample(self, n: int, seed: int | None = None) -> "Tafra": """ Return a random sample of `n` rows. @@ -1465,12 +1505,12 @@ def sample(self, n: int, seed: int | None = None) -> 'Tafra': The sampled `Tafra`. """ if n > self._rows: - raise ValueError(f'Cannot sample {n} rows from {self._rows} rows.') + raise ValueError(f"Cannot sample {n} rows from {self._rows} rows.") rng = np.random.default_rng(seed) idx = rng.choice(self._rows, size=n, replace=False) return self._ndindex(idx) - def drop_duplicates(self, columns: Iterable[str] | None = None) -> 'Tafra': + def drop_duplicates(self, columns: Iterable[str] | None = None) -> "Tafra": """ Remove duplicate rows, keeping the first occurrence. @@ -1492,7 +1532,7 @@ def drop_duplicates(self, columns: Iterable[str] | None = None) -> 'Tafra': else: col_arrays = [self._data[c] for c in cols] # check for StringDType — can't use structured array - if any(c.dtype.kind == 'T' for c in col_arrays): + if any(c.dtype.kind == "T" for c in col_arrays): seen: dict[Any, int] = {} idx_list: list[int] = [] for i, k in enumerate(zip(*col_arrays)): @@ -1501,16 +1541,16 @@ def drop_duplicates(self, columns: Iterable[str] | None = None) -> 'Tafra': idx_list.append(i) idx = np.array(idx_list, dtype=np.intp) else: - dt = np.dtype([(f'f{i}', c.dtype) for i, c in enumerate(col_arrays)]) + dt = np.dtype([(f"f{i}", c.dtype) for i, c in enumerate(col_arrays)]) key = np.empty(self._rows, dtype=dt) for i, c in enumerate(col_arrays): - key[f'f{i}'] = c + key[f"f{i}"] = c _, idx = np.unique(key, return_index=True) idx.sort() return self._ndindex(idx) - def value_counts(self, column: str) -> 'Tafra': + def value_counts(self, column: str) -> "Tafra": """ Count occurrences of each unique value in a column. @@ -1529,12 +1569,9 @@ def value_counts(self, column: str) -> 'Tafra': data = self._data[column] unique, counts = np.unique(data, return_counts=True) order = np.argsort(-counts) - return Tafra( - {column: unique[order], 'count': counts[order]}, - validate=False - ) + return Tafra({column: unique[order], "count": counts[order]}, validate=False) - def describe(self) -> 'Tafra': + def describe(self) -> "Tafra": """ Summary statistics for numeric columns: count, mean, std, min, 25%, 50%, 75%, max. @@ -1545,26 +1582,28 @@ def describe(self) -> 'Tafra': A `Tafra` with a `'stat'` column and one column per numeric column in the original. """ - stats = ['count', 'mean', 'std', 'min', '25%', '50%', '75%', 'max'] + stats = ["count", "mean", "std", "min", "25%", "50%", "75%", "max"] result: dict[str, np.ndarray[Any, Any]] = { - 'stat': np.array(stats, dtype=np.dtypes.StringDType(na_object=None)) # type: ignore[call-arg] + "stat": np.array(stats, dtype=np.dtypes.StringDType(na_object=None)) # type: ignore[call-arg] } for col, val in self._data.items(): - if val.dtype.kind in ('i', 'u', 'f'): + if val.dtype.kind in ("i", "u", "f"): fval = val.astype(float) - result[col] = np.array([ - float(len(fval)), - np.mean(fval), - np.std(fval), - np.min(fval), - np.percentile(fval, 25), - np.percentile(fval, 50), - np.percentile(fval, 75), - np.max(fval), - ]) + result[col] = np.array( + [ + float(len(fval)), + np.mean(fval), + np.std(fval), + np.min(fval), + np.percentile(fval, 25), + np.percentile(fval, 50), + np.percentile(fval, 75), + np.max(fval), + ] + ) return Tafra(result, validate=False) - def shift(self, n: int = 1) -> 'Tafra': + def shift(self, n: int = 1) -> "Tafra": """ Shift all columns by `n` rows. Positive shifts forward (lag), negative shifts backward (lead). Vacated rows are filled with @@ -1591,21 +1630,21 @@ def shift(self, n: int = 1) -> 'Tafra': result: dict[str, np.ndarray[Any, Any]] = {} for col, val in self._data.items(): - if val.dtype.kind in ('i', 'u', 'f'): + if val.dtype.kind in ("i", "u", "f"): shifted = np.empty(self._rows, dtype=float) shifted[:] = np.nan if n > 0: - shifted[n:] = val[:self._rows - n] + shifted[n:] = val[: self._rows - n] else: - shifted[:self._rows + n] = val[-n:] + shifted[: self._rows + n] = val[-n:] result[col] = shifted else: shifted_obj = np.empty(self._rows, dtype=object) shifted_obj[:] = None if n > 0: - shifted_obj[n:] = val[:self._rows - n] + shifted_obj[n:] = val[: self._rows - n] else: - shifted_obj[:self._rows + n] = val[-n:] + shifted_obj[: self._rows + n] = val[-n:] result[col] = shifted_obj return Tafra(result, validate=False) @@ -1664,7 +1703,7 @@ def get(self, key: str, default: Any = None) -> Any: """ return self._data.get(key, default) - def update(self, other: 'Tafra') -> 'Tafra': + def update(self, other: "Tafra") -> "Tafra": """ Update the data and dtypes of this `Tafra` with another `Tafra`. Length of rows must match, while data of different `dtype` will overwrite. @@ -1682,7 +1721,7 @@ def update(self, other: 'Tafra') -> 'Tafra': tafra.update_inplace(other) return tafra - def update_inplace(self, other: 'Tafra') -> None: + def update_inplace(self, other: "Tafra") -> None: """ Inplace version. @@ -1707,8 +1746,9 @@ def update_inplace(self, other: 'Tafra') -> None: for column, value in other._data.items(): if len(value) != rows: raise ValueError( - 'Other `Tafra` must have consistent row count. ' - f'This `Tafra` has {rows} rows, other `Tafra` has {len(value)} rows.') + "Other `Tafra` must have consistent row count. " + f"This `Tafra` has {rows} rows, other `Tafra` has {len(value)} rows." + ) self._data[column] = value self.update_dtypes_inplace(other._dtypes) @@ -1727,7 +1767,7 @@ def _coalesce_dtypes(self) -> None: if column not in self._dtypes: self._dtypes[column] = self._format_dtype(self._data[column].dtype) - def update_dtypes(self, dtypes: dict[str, Any]) -> 'Tafra': + def update_dtypes(self, dtypes: dict[str, Any]) -> "Tafra": """ Apply new dtypes. @@ -1745,8 +1785,7 @@ def update_dtypes(self, dtypes: dict[str, Any]) -> 'Tafra': tafra.update_dtypes_inplace(dtypes) return tafra - def update_dtypes_inplace(self, dtypes: dict[str, Any], - _from_init: bool = False) -> None: + def update_dtypes_inplace(self, dtypes: dict[str, Any], _from_init: bool = False) -> None: """ Inplace version. @@ -1769,7 +1808,7 @@ def update_dtypes_inplace(self, dtypes: dict[str, Any], self._validate_columns([column]) if isinstance(dtype, np.dtype): raw_dtypes[column] = dtype - elif isinstance(dtype, str) and dtype == 'str' and not _from_init: + elif isinstance(dtype, str) and dtype == "str" and not _from_init: # 'str' label → StringDType with na_object=None so the # column can hold None values. Skip during __post_init__ # to preserve the original dtype. @@ -1790,22 +1829,26 @@ def update_dtypes_inplace(self, dtypes: dict[str, Any], # and the source is already a string type. This happens when # __post_init__ round-trips through formatted labels ('str'). # Explicit StringDType() or specific 'Tafra': + def rename(self, renames: dict[str, str]) -> "Tafra": """ Rename columns in the `Tafra` from a `dict`. @@ -1847,7 +1890,7 @@ def rename_inplace(self, renames: dict[str, str]) -> None: self._dtypes[new] = self._dtypes.pop(cur) return None - def delete(self, columns: Iterable[str]) -> 'Tafra': + def delete(self, columns: Iterable[str]) -> "Tafra": """ Remove a column from `data` and `dtypes`. @@ -1866,11 +1909,9 @@ def delete(self, columns: Iterable[str]) -> 'Tafra': self._validate_columns(columns) return Tafra( - {column: value.copy() for column, value in self._data.items() - if column not in columns}, - {column: value for column, value in self._dtypes.items() - if column not in columns}, - validate=False + {column: value.copy() for column, value in self._data.items() if column not in columns}, + {column: value for column, value in self._dtypes.items() if column not in columns}, + validate=False, ) def delete_inplace(self, columns: Iterable[str]) -> None: @@ -1897,7 +1938,7 @@ def delete_inplace(self, columns: Iterable[str]) -> None: _ = self._data.pop(column, None) _ = self._dtypes.pop(column, None) - def copy(self, order: 'Literal["K", "A", "C", "F"]' = 'C') -> 'Tafra': + def copy(self, order: 'Literal["K", "A", "C", "F"]' = "C") -> "Tafra": """ Create a copy of a `Tafra`. @@ -1914,15 +1955,16 @@ def copy(self, order: 'Literal["K", "A", "C", "F"]' = 'C') -> 'Tafra': A copied `Tafra`. """ return Tafra( - {column: value.copy(order=order) - for column, value in self._data.items()}, + {column: value.copy(order=order) for column, value in self._data.items()}, self._dtypes.copy(), - validate=False + validate=False, ) - def coalesce(self, column: str, fills: Iterable[ - Iterable[None | str | int | float | bool | np.ndarray[Any, Any]] - ]) -> np.ndarray[Any, Any]: + def coalesce( + self, + column: str, + fills: Iterable[Iterable[None | str | int | float | bool | np.ndarray[Any, Any]]], + ) -> np.ndarray[Any, Any]: """ Fill `None` values from `fills`. Analogous to `SQL COALESCE` or `pandas.fillna()`. @@ -1964,9 +2006,11 @@ def coalesce(self, column: str, fills: Iterable[ return value - def coalesce_inplace(self, column: str, fills: Iterable[ - Iterable[None | str | int | float | bool | np.ndarray[Any, Any]] - ]) -> None: + def coalesce_inplace( + self, + column: str, + fills: Iterable[Iterable[None | str | int | float | bool | np.ndarray[Any, Any]]], + ) -> None: """ In-place version. @@ -1988,8 +2032,7 @@ def coalesce_inplace(self, column: str, fills: Iterable[ self._data[column] = self.coalesce(column, fills) self.update_dtypes_inplace({column: self._data[column].dtype}) - def _cast_record(self, dtype: str, data: np.ndarray[Any, Any], - cast_null: bool) -> float | None: + def _cast_record(self, dtype: str, data: np.ndarray[Any, Any], cast_null: bool) -> float | None: """ Casts needed to generate records for database insert. @@ -2014,12 +2057,13 @@ def _cast_record(self, dtype: str, data: np.ndarray[Any, Any], """ _dtype = self._reduce_dtype(dtype) value: Any = RECORD_TYPE[_dtype](data.item()) - if cast_null and _dtype == 'float' and np.isnan(data.item()): + if cast_null and _dtype == "float" and np.isnan(data.item()): return None return value - def to_records(self, columns: Iterable[str] | None = None, - cast_null: bool = True) -> Iterator[tuple[Any, ...]]: + def to_records( + self, columns: Iterable[str] | None = None, cast_null: bool = True + ) -> Iterator[tuple[Any, ...]]: """ Return a `Iterator` of `Tuple`, each being a record (i.e. row) and allowing heterogeneous typing. Useful for e.g. sending records back to a @@ -2044,16 +2088,19 @@ def to_records(self, columns: Iterable[str] | None = None, columns = [columns] self._validate_columns(columns) - return (tuple( - None if len(self._data[c]) <= row else self._cast_record( - self._dtypes[c], self._data[c][[row]], - cast_null + return ( + tuple( + None + if len(self._data[c]) <= row + else self._cast_record(self._dtypes[c], self._data[c][[row]], cast_null) + for c in columns ) - for c in columns) - for row in range(self._rows)) + for row in range(self._rows) + ) - def to_list(self, columns: Iterable[str] | None = None, - inner: bool = False) -> list[np.ndarray[Any, Any]] | list[list[Any]]: + def to_list( + self, columns: Iterable[str] | None = None, inner: bool = False + ) -> list[np.ndarray[Any, Any]] | list[list[Any]]: """ Return a list of homogeneously typed columns (as `numpy.ndarray`). If a generator is needed, use `to_records()`. If `inner == True` each column @@ -2082,9 +2129,12 @@ def to_list(self, columns: Iterable[str] | None = None, return [list(self._data[c]) for c in columns] return [self._data[c] for c in columns] - def to_tuple(self, columns: Iterable[str] | None = None, name: str | None = 'Tafra', - inner: bool = False, - ) -> tuple[np.ndarray[Any, Any]] | tuple[tuple[Any, ...]]: + def to_tuple( + self, + columns: Iterable[str] | None = None, + name: str | None = "Tafra", + inner: bool = False, + ) -> tuple[np.ndarray[Any, Any]] | tuple[tuple[Any, ...]]: """ Return a `NamedTuple` or `Tuple`. If a generator is needed, use `to_records()`. If `inner == True` each column will be cast from @@ -2163,7 +2213,7 @@ def to_pandas(self, columns: Iterable[str] | None = None) -> DataFrame: try: import pandas as pd # type: ignore except ImportError as e: # pragma: no cover - raise ImportError('`pandas` does not appear to be installed.') + raise ImportError("`pandas` does not appear to be installed.") if columns is None: columns = self.columns @@ -2172,13 +2222,13 @@ def to_pandas(self, columns: Iterable[str] | None = None) -> DataFrame: columns = [columns] self._validate_columns(columns) - return pd.DataFrame({ - column: pd.Series(value) for column, value in self._data.items() - if column in columns - }) + return pd.DataFrame( + {column: pd.Series(value) for column, value in self._data.items() if column in columns} + ) - def to_csv(self, filename: str | Path | TextIOWrapper | IO[str], - columns: Iterable[str] | None = None) -> None: + def to_csv( + self, filename: str | Path | TextIOWrapper | IO[str], columns: Iterable[str] | None = None + ) -> None: """ Write the `Tafra` to a CSV. @@ -2198,29 +2248,30 @@ def to_csv(self, filename: str | Path | TextIOWrapper | IO[str], self._validate_columns(columns) if isinstance(filename, (str, Path)): - f = open(filename, 'w', newline='') + f = open(filename, "w", newline="") should_close = True elif isinstance(filename, TextIOWrapper): - if 'w' not in filename.mode: - raise ValueError(f'file must be opened in write mode, got {filename.mode}') + if "w" not in filename.mode: + raise ValueError(f"file must be opened in write mode, got {filename.mode}") f = filename should_close = False - f.reconfigure(newline='') + f.reconfigure(newline="") else: raise TypeError( - f'`filename` must be `str`, `Path`, or `TextIOWrapper`, got `{type(filename)}`') + f"`filename` must be `str`, `Path`, or `TextIOWrapper`, got `{type(filename)}`" + ) - writer = csv.writer(f, delimiter=',', quotechar='"') + writer = csv.writer(f, delimiter=",", quotechar='"') writer.writerow((column for column in self._data.keys() if column in columns)) writer.writerows(self.to_records(columns)) if should_close: f.close() - def union(self, other: 'Tafra') -> 'Tafra': + def union(self, other: "Tafra") -> "Tafra": """ Helper function to implement `tafra.group.Union.apply()`. @@ -2239,7 +2290,7 @@ def union(self, other: 'Tafra') -> 'Tafra': """ return Union().apply(self, other) - def union_inplace(self, other: 'Tafra') -> None: + def union_inplace(self, other: "Tafra") -> None: """ Inplace version. @@ -2260,9 +2311,12 @@ def union_inplace(self, other: 'Tafra') -> None: """ Union().apply_inplace(self, other) - def group_by(self, columns: Iterable[str], aggregation: InitAggregation | None = None, - iter_fn: Mapping[str, Callable[[np.ndarray[Any, Any]], Any]] | None = None, - ) -> 'Tafra': + def group_by( + self, + columns: Iterable[str], + aggregation: InitAggregation | None = None, + iter_fn: Mapping[str, Callable[[np.ndarray[Any, Any]], Any]] | None = None, + ) -> "Tafra": """ Helper function to implement `tafra.group.GroupBy.apply()`. @@ -2295,9 +2349,12 @@ def group_by(self, columns: Iterable[str], aggregation: InitAggregation | None = iter_fn = {} return GroupBy(columns, aggregation, iter_fn).apply(self) - def transform(self, columns: Iterable[str], aggregation: InitAggregation | None = None, - iter_fn: dict[str, Callable[[np.ndarray[Any, Any]], Any]] | None = None, - ) -> 'Tafra': + def transform( + self, + columns: Iterable[str], + aggregation: InitAggregation | None = None, + iter_fn: dict[str, Callable[[np.ndarray[Any, Any]], Any]] | None = None, + ) -> "Tafra": """ Helper function to implement `tafra.group.Transform.apply()`. @@ -2330,7 +2387,7 @@ def transform(self, columns: Iterable[str], aggregation: InitAggregation | None iter_fn = {} return Transform(columns, aggregation, iter_fn).apply(self) - def iterate_by(self, columns: Iterable[str]) -> Iterator['GroupDescription']: + def iterate_by(self, columns: Iterable[str]) -> Iterator["GroupDescription"]: """ Helper function to implement `tafra.group.IterateBy.apply()`. @@ -2352,8 +2409,12 @@ def iterate_by(self, columns: Iterable[str]) -> Iterator['GroupDescription']: """ yield from IterateBy(columns).apply(self) - def inner_join(self, right: 'Tafra', on: Iterable[tuple[str, str, str]], - select: Iterable[str] | None = None) -> 'Tafra': + def inner_join( + self, + right: "Tafra", + on: Iterable[tuple[str, str, str]], + select: Iterable[str] | None = None, + ) -> "Tafra": """ Helper function to implement `tafra.group.InnerJoin.apply()`. @@ -2391,8 +2452,12 @@ def inner_join(self, right: 'Tafra', on: Iterable[tuple[str, str, str]], select = [] return InnerJoin(on, select).apply(self, right) - def left_join(self, right: 'Tafra', on: Iterable[tuple[str, str, str]], - select: Iterable[str] | None = None) -> 'Tafra': + def left_join( + self, + right: "Tafra", + on: Iterable[tuple[str, str, str]], + select: Iterable[str] | None = None, + ) -> "Tafra": """ Helper function to implement `tafra.group.LeftJoin.apply()`. @@ -2430,15 +2495,13 @@ def left_join(self, right: 'Tafra', on: Iterable[tuple[str, str, str]], select = [] return LeftJoin(on, select).apply(self, right) - def cross_join(self, right: 'Tafra', - select: Iterable[str] | None = None) -> 'Tafra': + def cross_join(self, right: "Tafra", select: Iterable[str] | None = None) -> "Tafra": """ Helper function to implement `tafra.group.CrossJoin.apply()`. A cross join. - Analogy to SQL CROSS JOIN, or `pandas.merge(..., how='outer') using temporary - columns of static value to intersect all rows`. + Analogy to SQL CROSS JOIN, or `pandas.merge(left, right, how='cross')`. Parameters ---------- @@ -2459,7 +2522,7 @@ def cross_join(self, right: 'Tafra', select = [] return CrossJoin([], select).apply(self, right) - def chunks(self, n: int, sort_by: Iterable[str] | None = None) -> list['Tafra']: + def chunks(self, n: int, sort_by: Iterable[str] | None = None) -> list["Tafra"]: """ Split into `n` roughly equal-sized `Tafra` chunks. @@ -2477,7 +2540,7 @@ def chunks(self, n: int, sort_by: Iterable[str] | None = None) -> list['Tafra']: The chunked `Tafra` instances. """ if n < 1: - raise ValueError('n must be >= 1') + raise ValueError("n must be >= 1") source = self._sorted(sort_by) if sort_by else self splits = np.array_split(np.arange(source._rows), n) @@ -2485,12 +2548,13 @@ def chunks(self, n: int, sort_by: Iterable[str] | None = None) -> list['Tafra']: Tafra( {col: val[idx] for col, val in source._data.items()}, source._dtypes.copy(), - validate=False + validate=False, ) - for idx in splits if len(idx) > 0 + for idx in splits + if len(idx) > 0 ] - def chunk_rows(self, size: int, sort_by: Iterable[str] | None = None) -> list['Tafra']: + def chunk_rows(self, size: int, sort_by: Iterable[str] | None = None) -> list["Tafra"]: """ Split into chunks of at most `size` rows each. @@ -2508,15 +2572,14 @@ def chunk_rows(self, size: int, sort_by: Iterable[str] | None = None) -> list['T The chunked `Tafra` instances. """ if size < 1: - raise ValueError('size must be >= 1') + raise ValueError("size must be >= 1") n = max(1, (self._rows + size - 1) // size) return self.chunks(n, sort_by=sort_by) def partition( - self, columns: Iterable[str], - sort_by: Iterable[str] | None = None - ) -> list[tuple[tuple[Any, ...], 'Tafra']]: + self, columns: Iterable[str], sort_by: Iterable[str] | None = None + ) -> list[tuple[tuple[Any, ...], "Tafra"]]: """ Split by unique values in `columns`, preserving group integrity. @@ -2536,13 +2599,13 @@ def partition( from .group import GroupSet unique, group_indices = GroupSet._build_group_indices(self, columns) - result: list[tuple[tuple[Any, ...], 'Tafra']] = [] + result: list[tuple[tuple[Any, ...], "Tafra"]] = [] for key, rows in zip(unique, group_indices): sub = Tafra( {col: val[rows] for col, val in self._data.items()}, self._dtypes.copy(), - validate=False + validate=False, ) if sort_by: sub = sub._sorted(sort_by) @@ -2550,7 +2613,7 @@ def partition( return result - def _sorted(self, sort_by: Iterable[str]) -> 'Tafra': + def _sorted(self, sort_by: Iterable[str]) -> "Tafra": """Return a new Tafra sorted by the given columns.""" cols = list(sort_by) self._validate_columns(cols) @@ -2562,11 +2625,11 @@ def _sorted(self, sort_by: Iterable[str]) -> 'Tafra': return Tafra( {col: val[order] for col, val in self._data.items()}, self._dtypes.copy(), - validate=False + validate=False, ) @classmethod - def concat(cls, tafras: Iterable['Tafra']) -> 'Tafra': + def concat(cls, tafras: Iterable["Tafra"]) -> "Tafra": """ Concatenate multiple `Tafra` instances row-wise. @@ -2582,20 +2645,19 @@ def concat(cls, tafras: Iterable['Tafra']) -> 'Tafra': """ tafra_list = list(tafras) if not tafra_list: - raise ValueError('No tafras to concatenate.') + raise ValueError("No tafras to concatenate.") columns = list(tafra_list[0]._data.keys()) col_set = set(columns) for i, t in enumerate(tafra_list[1:], 1): if set(t._data.keys()) != col_set: raise ValueError( - f'Tafra at index {i} has columns {list(t._data.keys())}, ' - f'expected {columns}.') + f"Tafra at index {i} has columns {list(t._data.keys())}, expected {columns}." + ) return cls( - {col: np.concatenate([t._data[col] for t in tafra_list]) - for col in columns}, + {col: np.concatenate([t._data[col] for t in tafra_list]) for col in columns}, tafra_list[0]._dtypes.copy(), - validate=False + validate=False, ) @@ -2603,16 +2665,17 @@ def to_field_name(maybe_text: str | int | float) -> str: # pragma: no cover text = str(maybe_text) # Remove invalid characters - mid_text = re.sub('[^0-9a-zA-Z]', '', text) + mid_text = re.sub("[^0-9a-zA-Z]", "", text) # Remove leading characters until we find a letter - final_text = re.sub('^[^a-zA-Z]+', '', mid_text) + final_text = re.sub("^[^a-zA-Z]+", "", mid_text) - if final_text == '': - final_text = 'field_' + mid_text + if final_text == "": + final_text = "field_" + mid_text return final_text + def _in_notebook() -> bool: # pragma: no cover """ Checks if running in a Jupyter Notebook. @@ -2623,12 +2686,23 @@ def _in_notebook() -> bool: # pragma: no cover """ try: from IPython import get_ipython - if 'IPKernelApp' in get_ipython().config: # type: ignore[no-untyped-call,unused-ignore] + + if "IPKernelApp" in get_ipython().config: # type: ignore[no-untyped-call,unused-ignore] return True except Exception as e: pass return False + # Import here to resolve circular dependency -from .group import (GroupBy, Transform, IterateBy, InnerJoin, LeftJoin, CrossJoin, Union, - InitAggregation, GroupDescription) +from .group import ( + GroupBy, + Transform, + IterateBy, + InnerJoin, + LeftJoin, + CrossJoin, + Union, + InitAggregation, + GroupDescription, +) diff --git a/tafra/formatter.py b/tafra/formatter.py index f2967c5..89c8047 100644 --- a/tafra/formatter.py +++ b/tafra/formatter.py @@ -12,6 +12,7 @@ ----- Created on April 25, 2020 """ + from __future__ import annotations from typing import Any, Callable, Iterator, MutableMapping @@ -19,8 +20,10 @@ import numpy as np -class ObjectFormatter(dict[str, Callable[[np.ndarray[Any, Any]], np.ndarray[Any, Any]]], - MutableMapping[str, Callable[[np.ndarray[Any, Any]], np.ndarray[Any, Any]]]): +class ObjectFormatter( + dict[str, Callable[[np.ndarray[Any, Any]], np.ndarray[Any, Any]]], + MutableMapping[str, Callable[[np.ndarray[Any, Any]], np.ndarray[Any, Any]]], +): """ A dictionary that contains mappings for formatting objects. Some numpy objects should be cast to other types, e.g. the `decimal.Decimal` type cannot @@ -33,11 +36,14 @@ class ObjectFormatter(dict[str, Callable[[np.ndarray[Any, Any]], np.ndarray[Any, looked up from the first element of the `np.ndarray`, i.e. `type(array[0]).__name__`. """ + test_array = np.arange(4) - def __setitem__(self, dtype: str, - value: Callable[[np.ndarray[Any, Any]], np.ndarray[Any, Any]], - ) -> None: + def __setitem__( + self, + dtype: str, + value: Callable[[np.ndarray[Any, Any]], np.ndarray[Any, Any]], + ) -> None: """ Set the dtype formatter. """ @@ -45,13 +51,13 @@ def __setitem__(self, dtype: str, result = value(self.test_array) except Exception as e: raise ValueError( - 'Must provide a function that takes an `np.ndarray` and returns ' - 'an np.ndarray.') from e + "Must provide a function that takes an `np.ndarray` and returns an np.ndarray." + ) from e if not isinstance(result, np.ndarray): raise ValueError( - 'Must provide a function that takes an `np.ndarray` and returns ' - 'an np.ndarray.') + "Must provide a function that takes an `np.ndarray` and returns an np.ndarray." + ) dict.__setitem__(self, dtype, value) @@ -72,8 +78,8 @@ def __repr__(self) -> str: def __str__(self) -> str: if self.__len__() < 1: - return r'{}' - return '{' + '\n'.join(f'{c}: {v}' for c, v in self.items()) + '}' + return r"{}" + return "{" + "\n".join(f"{c}: {v}" for c, v in self.items()) + "}" def __iter__(self) -> Iterator[Any]: yield from dict.__iter__(self) @@ -98,7 +104,7 @@ def parse_dtype(self, value: np.ndarray[Any, Any]) -> np.ndarray[Any, Any] | Non value, modified: Tuple(np.ndarray, bool) The `np.ndarray` and whether it was modified or not. """ - if value.dtype.kind != 'O': + if value.dtype.kind != "O": return None type_name = type(value[0]).__name__ @@ -107,7 +113,7 @@ def parse_dtype(self, value: np.ndarray[Any, Any]) -> np.ndarray[Any, Any] | Non return value # convert object arrays of strings to StringDType - if type_name == 'str': + if type_name == "str": return value.astype(np.dtypes.StringDType(na_object=None)) # type: ignore[call-arg] return None diff --git a/tafra/group.py b/tafra/group.py index db30356..207f7c5 100644 --- a/tafra/group.py +++ b/tafra/group.py @@ -61,6 +61,9 @@ # 'U' (fixed-width unicode result_array. _VECTORIZED_AGGS: dict[int, Callable[..., np.ndarray[Any, Any]]] = {} @@ -472,6 +475,20 @@ def _encode_columns( encoded.append(codes) codebooks.append(uniq) else: + # Shift signed integer columns to non-negative for positional encoding. + if c.dtype.kind == "i" and len(c) > 0: + c64 = c.astype(np.int64, copy=False) + c_min = int(c64.min()) + if c_min < 0: + c_max = int(c64.max()) + if c_max - c_min > np.iinfo(np.int64).max: + raise ValueError( + f"Column value range ({c_min} to {c_max}) exceeds " + f"int64 after shift — cannot encode for groupby." + ) + encoded.append(c64 - c_min) + codebooks.append(None) + continue encoded.append(c) codebooks.append(None) return encoded, codebooks @@ -488,7 +505,14 @@ def _encode_columns_paired( left_enc = [] right_enc = [] for lc, rc in zip(left_cols, right_cols): - if lc.dtype.kind in ("T", "U", "S", "O"): + # Verify both sides have compatible kinds before encoding + l_kind, r_kind = lc.dtype.kind, rc.dtype.kind + if (l_kind in _STRING_KINDS) != (r_kind in _STRING_KINDS): + raise TypeError( + f"Cannot encode columns with incompatible dtype kinds: " + f"left={lc.dtype} (kind={l_kind}), right={rc.dtype} (kind={r_kind})" + ) + if l_kind in _STRING_KINDS: combined = np.concatenate([lc, rc]) if _HAS_ACCEL and len(combined) >= 50_000: codes, _ = _c_encode_strings(combined.astype(object)) @@ -497,6 +521,20 @@ def _encode_columns_paired( left_enc.append(codes[: len(lc)]) right_enc.append(codes[len(lc) :]) else: + # For signed integer columns, shift to non-negative using + # COMBINED min so both sides share the same baseline. + if l_kind == "i" and len(lc) > 0 and len(rc) > 0: + combined_min = min(int(lc.min()), int(rc.min())) + if combined_min < 0: + combined_max = max(int(lc.max()), int(rc.max())) + if combined_max - combined_min > np.iinfo(np.int64).max: + raise ValueError( + f"Column value range ({combined_min} to {combined_max}) " + f"exceeds int64 after shift — cannot encode for join." + ) + left_enc.append(lc.astype(np.int64, copy=False) - combined_min) + right_enc.append(rc.astype(np.int64, copy=False) - combined_min) + continue left_enc.append(lc) right_enc.append(rc) return left_enc, right_enc @@ -515,7 +553,10 @@ def _build_composite_key(encoded: list[np.ndarray[Any, Any]]) -> np.ndarray[Any, if len(encoded[0]) == 0: return np.array([], dtype=np.int64) - # compute cardinality of each column + # Precondition: columns are already non-negative integer codes. + # Callers _encode_columns (GroupBy/Transform) and _encode_columns_paired + # (joins) ensure this by encoding strings via np.unique and shifting + # negative integers to non-negative range. cards = [int(c.max()) + 1 for c in encoded] # check for overflow: product of all cardinalities must fit int64 @@ -530,14 +571,14 @@ def _build_composite_key(encoded: list[np.ndarray[Any, Any]]) -> np.ndarray[Any, if not overflow: if _HAS_ACCEL: return _c_composite_key( - tuple(c.astype(np.int64) for c in encoded), + tuple(c.astype(np.int64, copy=False) for c in encoded), tuple(cards), ) # Python fallback: flat integer key via positional encoding key = np.zeros(len(encoded[0]), dtype=np.int64) multiplier = 1 for c, card in zip(reversed(encoded), reversed(cards)): - key += c.astype(np.int64) * multiplier + key += c.astype(np.int64, copy=False) * multiplier multiplier *= card return key else: @@ -938,6 +979,11 @@ class Join(GroupSet): on: Iterable[tuple[str, str, str]] select: Iterable[str] + def __post_init__(self) -> None: + # Materialize iterables to prevent generator exhaustion on reuse. + self.on = list(self.on) + self.select = list(self.select) + def _validate_dtypes( self, l_table: "Tafra", r_table: "Tafra" ) -> tuple[dict[str, np.dtype[Any]], dict[str, np.dtype[Any]]]: @@ -1066,10 +1112,26 @@ def _sort_merge_indices( def _resolve_join_cols( self, left_t: "Tafra", right_t: "Tafra" ) -> tuple[list[str], dict[str, str]]: - """Compute deduplicated output columns and dtypes.""" + """Compute deduplicated output columns and dtypes. + + Validates that all select columns exist in at least one table. + """ + # Use a frozenset for O(1) membership testing. + # select is already materialized to a list in __post_init__. + select_set = frozenset(self.select) + + # Validate select columns exist + if select_set: + valid_cols = set(left_t._data.keys()) | set(right_t._data.keys()) + invalid = select_set - valid_cols + if invalid: + raise KeyError( + f"Column(s) {tuple(sorted(invalid))} in `select` not found in either table." + ) + seen_cols: dict[str, None] = {} for c in chain(left_t._data.keys(), right_t._data.keys()): - if not self.select or c in self.select: + if not select_set or c in select_set: seen_cols[c] = None join_cols = list(seen_cols.keys()) dtypes: dict[str, str] = { @@ -1651,10 +1713,9 @@ def apply(self, left_t: "Tafra", right_t: "Tafra") -> "Tafra": @dc.dataclass class CrossJoin(Join): """ - A cross join. + A cross join (Cartesian product). - Analogy to SQL CROSS JOIN, or `pandas.merge(..., how='outer') - using temporary columns of static value to intersect all rows`. + Analogy to SQL CROSS JOIN, or `pandas.merge(left, right, how='cross')`. Parameters ---------- diff --git a/tafra/protocol.py b/tafra/protocol.py index ff8b529..baed32a 100644 --- a/tafra/protocol.py +++ b/tafra/protocol.py @@ -12,6 +12,7 @@ ----- Created on April 25, 2020 """ + from __future__ import annotations import numpy as np @@ -31,6 +32,7 @@ class DataFrame(Protocol): """ A fake class to satisfy typing of a ``pandas.DataFrame`` without a dependency. """ + _data: dict[str, Series] columns: list[str] dtypes: list[str] @@ -41,11 +43,13 @@ def __getitem__(self, column: str) -> Series: def __setitem__(self, column: str, value: np.ndarray[Any, Any]) -> None: raise NotImplementedError + @runtime_checkable class Cursor(Protocol): """ A fake class to satisfy typing of a ``pyodbc.Cursor`` without a dependency. """ + description: tuple[tuple[str, type[Any], int | None, int, int, int, bool]] def __iter__(self) -> Iterator[tuple[Any, ...]]: diff --git a/test/bench_tafra.py b/test/bench_tafra.py index 82f49b9..e8f6467 100644 --- a/test/bench_tafra.py +++ b/test/bench_tafra.py @@ -5,6 +5,7 @@ python test/bench_tafra.py pytest test/bench_tafra.py -v (if pytest-benchmark is installed) """ + from __future__ import annotations import time @@ -24,34 +25,40 @@ def timer(label: str): def make_tafra(n_rows: int, n_groups: int) -> Tafra: rng = np.random.default_rng(42) - return Tafra({ - 'group_a': rng.integers(0, n_groups, size=n_rows).astype('int64'), - 'group_b': rng.choice(['x', 'y', 'z'], size=n_rows), - 'value1': rng.standard_normal(n_rows), - 'value2': rng.standard_normal(n_rows), - 'value3': rng.integers(0, 1000, size=n_rows).astype('float64'), - }) + return Tafra( + { + "group_a": rng.integers(0, n_groups, size=n_rows).astype("int64"), + "group_b": rng.choice(["x", "y", "z"], size=n_rows), + "value1": rng.standard_normal(n_rows), + "value2": rng.standard_normal(n_rows), + "value3": rng.integers(0, 1000, size=n_rows).astype("float64"), + } + ) def make_join_tafras(left_rows: int, right_rows: int) -> tuple: rng = np.random.default_rng(42) - left = Tafra({ - 'key': rng.integers(0, 100, size=left_rows).astype('int64'), - 'left_val': rng.standard_normal(left_rows), - }) - right = Tafra({ - 'key': rng.integers(0, 100, size=right_rows).astype('int64'), - 'right_val': rng.standard_normal(right_rows), - }) + left = Tafra( + { + "key": rng.integers(0, 100, size=left_rows).astype("int64"), + "left_val": rng.standard_normal(left_rows), + } + ) + right = Tafra( + { + "key": rng.integers(0, 100, size=right_rows).astype("int64"), + "right_val": rng.standard_normal(right_rows), + } + ) return left, right def bench_construction(n_rows: int, n_iter: int = 100) -> None: rng = np.random.default_rng(42) data = { - 'a': rng.standard_normal(n_rows), - 'b': rng.standard_normal(n_rows), - 'c': rng.integers(0, 100, size=n_rows).astype('int64'), + "a": rng.standard_normal(n_rows), + "b": rng.standard_normal(n_rows), + "c": rng.integers(0, 100, size=n_rows).astype("int64"), } with timer(f"Tafra() validated, {n_rows} rows x {n_iter} iters"): @@ -68,13 +75,13 @@ def bench_getitem(n_rows: int, n_iter: int = 10000) -> None: with timer(f"__getitem__ str, {n_iter} iters"): for _ in range(n_iter): - t['value1'] + t["value1"] with timer(f"__getitem__ list[str], {n_iter} iters"): for _ in range(n_iter): - t[['value1', 'value2']] + t[["value1", "value2"]] - mask = t['group_a'] < 25 + mask = t["group_a"] < 25 with timer(f"__getitem__ bool mask, {n_iter} iters"): for _ in range(n_iter): t[mask] @@ -86,30 +93,30 @@ def bench_setitem(n_rows: int, n_iter: int = 1000) -> None: with timer(f"__setitem__ existing col, {n_iter} iters"): for _ in range(n_iter): - t['value1'] = new_values + t["value1"] = new_values with timer(f"__setitem__ new col, {n_iter} iters"): for i in range(n_iter): - t[f'new_{i}'] = new_values + t[f"new_{i}"] = new_values def bench_groupby(n_rows: int, n_groups: int) -> None: t = make_tafra(n_rows, n_groups) gb = GroupBy( - group_by_cols=['group_a'], - aggregation={'mean_v1': (np.mean, 'value1'), 'sum_v2': (np.sum, 'value2')}, + group_by_cols=["group_a"], + aggregation={"mean_v1": (np.mean, "value1"), "sum_v2": (np.sum, "value2")}, iter_fn={}, ) with timer(f"GroupBy 1 col, {n_groups} groups, {n_rows} rows"): gb.apply(t) gb2 = GroupBy( - group_by_cols=['group_a', 'group_b'], - aggregation={'mean_v1': (np.mean, 'value1')}, + group_by_cols=["group_a", "group_b"], + aggregation={"mean_v1": (np.mean, "value1")}, iter_fn={}, ) - unique_combos = len(set(zip(t['group_a'], t['group_b']))) + unique_combos = len(set(zip(t["group_a"], t["group_b"]))) with timer(f"GroupBy 2 cols, ~{unique_combos} groups, {n_rows} rows"): gb2.apply(t) @@ -118,8 +125,8 @@ def bench_transform(n_rows: int, n_groups: int) -> None: t = make_tafra(n_rows, n_groups) tr = Transform( - group_by_cols=['group_a'], - aggregation={'mean_v1': (np.mean, 'value1'), 'sum_v2': (np.sum, 'value2')}, + group_by_cols=["group_a"], + aggregation={"mean_v1": (np.mean, "value1"), "sum_v2": (np.sum, "value2")}, iter_fn={}, ) with timer(f"Transform 1 col, {n_groups} groups, {n_rows} rows"): @@ -129,7 +136,7 @@ def bench_transform(n_rows: int, n_groups: int) -> None: def bench_iterateby(n_rows: int, n_groups: int) -> None: t = make_tafra(n_rows, n_groups) - ib = IterateBy(group_by_cols=['group_a']) + ib = IterateBy(group_by_cols=["group_a"]) with timer(f"IterateBy 1 col, {n_groups} groups, {n_rows} rows"): list(ib.apply(t)) @@ -138,7 +145,7 @@ def bench_innerjoin(left_rows: int, right_rows: int) -> None: left, right = make_join_tafras(left_rows, right_rows) ij = InnerJoin( - on=[('key', 'key', '==')], + on=[("key", "key", "==")], select=[], ) with timer(f"InnerJoin {left_rows}x{right_rows}"): @@ -149,7 +156,7 @@ def bench_leftjoin(left_rows: int, right_rows: int) -> None: left, right = make_join_tafras(left_rows, right_rows) lj = LeftJoin( - on=[('key', 'key', '==')], + on=[("key", "key", "==")], select=[], ) with timer(f"LeftJoin {left_rows}x{right_rows}"): @@ -160,10 +167,10 @@ def bench_build_group_indices(n_rows: int, n_groups: int) -> None: t = make_tafra(n_rows, n_groups) with timer(f"_build_group_indices 1 col, {n_rows} rows"): - GroupBy._build_group_indices(t, ['group_a']) + GroupBy._build_group_indices(t, ["group_a"]) with timer(f"_build_group_indices 2 cols, {n_rows} rows"): - GroupBy._build_group_indices(t, ['group_a', 'group_b']) + GroupBy._build_group_indices(t, ["group_a", "group_b"]) def bench_validate_columns(n_rows: int, n_iter: int = 10000) -> None: @@ -181,11 +188,11 @@ def bench_ensure_valid(n_rows: int, n_iter: int = 1000) -> None: with timer(f"_ensure_valid ndarray, {n_iter} iters"): for _ in range(n_iter): - t._ensure_valid('value1', arr) + t._ensure_valid("value1", arr) with timer(f"_ensure_valid scalar, {n_iter} iters"): for _ in range(n_iter): - t._ensure_valid('value1', 1.0) + t._ensure_valid("value1", 1.0) def bench_copy(n_rows: int, n_iter: int = 100) -> None: @@ -201,11 +208,11 @@ def bench_select(n_rows: int, n_iter: int = 1000) -> None: with timer(f"select() 2 cols, {n_iter} iters"): for _ in range(n_iter): - t.select(['value1', 'value2']) + t.select(["value1", "value2"]) def bench_format_dtype(n_iter: int = 10000) -> None: - dt = np.dtype('float64') + dt = np.dtype("float64") with timer(f"_format_dtype, {n_iter} iters"): for _ in range(n_iter): @@ -214,6 +221,7 @@ def bench_format_dtype(n_iter: int = 10000) -> None: def bench_parse_dtype(n_iter: int = 10000) -> None: from tafra import object_formatter + arr = np.ones(100) with timer(f"parse_dtype (non-object), {n_iter} iters"): @@ -221,63 +229,72 @@ def bench_parse_dtype(n_iter: int = 10000) -> None: object_formatter.parse_dtype(arr) -def _heavy_worker(sub: 'Tafra') -> 'Tafra': +def _heavy_worker(sub: "Tafra") -> "Tafra": """Simulate expensive per-group work (~20-50ms): model fit + forecast.""" - v = sub['value1'] + v = sub["value1"] t = np.linspace(0, 1, 2000) result = np.zeros_like(t) for _ in range(100): result += np.sin(np.mean(v) * t) * np.exp(-np.std(v) * t) result += np.cos(np.var(v) * t) * np.log1p(np.abs(t * np.mean(v))) - return Tafra({ - 'group_a': sub['group_a'][:1], - 'forecast': np.array([np.sum(result)]), - }, validate=False) + return Tafra( + { + "group_a": sub["group_a"][:1], + "forecast": np.array([np.sum(result)]), + }, + validate=False, + ) -def _light_worker(sub: 'Tafra') -> 'Tafra': +def _light_worker(sub: "Tafra") -> "Tafra": """Light per-group work: simple aggregations.""" - return Tafra({ - 'group_a': sub['group_a'][:1], - 'mean_v1': np.array([np.mean(sub['value1'])]), - 'std_v1': np.array([np.std(sub['value1'])]), - }, validate=False) + return Tafra( + { + "group_a": sub["group_a"][:1], + "mean_v1": np.array([np.mean(sub["value1"])]), + "std_v1": np.array([np.std(sub["value1"])]), + }, + validate=False, + ) def _warmup(x: int) -> int: return x -def bench_partition_vs_groupby(n_rows: int, n_groups: int, - pool: 'ProcessPoolExecutor', - n_workers: int) -> None: +def bench_partition_vs_groupby( + n_rows: int, n_groups: int, pool: "ProcessPoolExecutor", n_workers: int +) -> None: t = make_tafra(n_rows, n_groups) print(f" [{n_groups} groups, {n_rows} rows, {n_workers} workers]") with timer(" group_by (light agg)"): - t.group_by(['group_a'], { - 'mean_v1': (np.mean, 'value1'), - 'std_v1': (np.std, 'value1'), - }) + t.group_by( + ["group_a"], + { + "mean_v1": (np.mean, "value1"), + "std_v1": (np.std, "value1"), + }, + ) with timer(" partition+serial (light)"): - parts = t.partition(['group_a']) + parts = t.partition(["group_a"]) results = [_light_worker(sub) for _, sub in parts] Tafra.concat(results) with timer(" partition+serial (heavy)"): - parts = t.partition(['group_a']) + parts = t.partition(["group_a"]) results = [_heavy_worker(sub) for _, sub in parts] Tafra.concat(results) with timer(f" partition+{n_workers} workers (heavy)"): - parts = t.partition(['group_a']) + parts = t.partition(["group_a"]) results = list(pool.map(_heavy_worker, [sub for _, sub in parts])) Tafra.concat(results) -if __name__ == '__main__': +if __name__ == "__main__": print("=" * 60) print("Tafra Performance Benchmarks") print("=" * 60) @@ -319,6 +336,7 @@ def bench_partition_vs_groupby(n_rows: int, n_groups: int, print("\n--- Partition vs GroupBy ---") from concurrent.futures import ProcessPoolExecutor import os + n_workers = min(os.cpu_count() or 4, 8) pool = ProcessPoolExecutor(max_workers=n_workers) list(pool.map(_warmup, range(n_workers))) # warm up workers diff --git a/test/bench_vs_pandas_vs_polars.py b/test/bench_vs_pandas_vs_polars.py index 6d0f7c1..08de31e 100644 --- a/test/bench_vs_pandas_vs_polars.py +++ b/test/bench_vs_pandas_vs_polars.py @@ -3,6 +3,7 @@ Run with: python test/bench_vs_pandas_vs_polars.py """ + from __future__ import annotations import time @@ -15,12 +16,14 @@ try: import polars as pl + HAS_POLARS = True except ImportError: HAS_POLARS = False try: from tafra._accel import groupby_sum as _c_test + HAS_ACCEL = True except ImportError: HAS_ACCEL = False @@ -80,41 +83,40 @@ def fmt(val, best=False): def print_row(label, t_accel, t_pure, t_pandas, t_polars): """Print one benchmark row, marking the best.""" - vals = {'accel': t_accel, 'pure': t_pure, 'pandas': t_pandas, 'polars': t_polars} + vals = {"accel": t_accel, "pure": t_pure, "pandas": t_pandas, "polars": t_polars} real_vals = {k: v for k, v in vals.items() if v is not None} best_key = min(real_vals, key=real_vals.get) if real_vals else None - print(f" {label:<35s}" - f" {fmt(t_accel, best_key == 'accel')}" - f" {fmt(t_pure, best_key == 'pure')}" - f" {fmt(t_pandas, best_key == 'pandas')}" - f" {fmt(t_polars, best_key == 'polars')}") + print( + f" {label:<35s}" + f" {fmt(t_accel, best_key == 'accel')}" + f" {fmt(t_pure, best_key == 'pure')}" + f" {fmt(t_pandas, best_key == 'pandas')}" + f" {fmt(t_polars, best_key == 'polars')}" + ) def print_header(): accel_label = "Tafra+C" if HAS_ACCEL else "" - print(f" {'Benchmark':<35s}" - f" {accel_label:>9s}" - f" {'Tafra':>9s}" - f" {'pandas':>9s}" - f" {'polars':>9s}") - print(f" {'-'*35} {'-'*9} {'-'*9} {'-'*9} {'-'*9}") + print(f" {'Benchmark':<35s} {accel_label:>9s} {'Tafra':>9s} {'pandas':>9s} {'polars':>9s}") + print(f" {'-' * 35} {'-' * 9} {'-' * 9} {'-' * 9} {'-' * 9}") # ============================================================ # Benchmarks # ============================================================ + def bench_construction(): print("\n--- Construction (100k rows, 5 cols) ---") print_header() rng = np.random.default_rng(42) data = { - 'a': rng.standard_normal(100_000), - 'b': rng.standard_normal(100_000), - 'c': rng.integers(0, 100, size=100_000).astype('int64'), - 'd': rng.standard_normal(100_000), - 'e': rng.standard_normal(100_000), + "a": rng.standard_normal(100_000), + "b": rng.standard_normal(100_000), + "c": rng.integers(0, 100, size=100_000).astype("int64"), + "d": rng.standard_normal(100_000), + "e": rng.standard_normal(100_000), } # construction doesn't use accel, same for both @@ -129,16 +131,16 @@ def bench_column_access(): print_header() rng = np.random.default_rng(42) n = 100_000 - tf = Tafra({'x': rng.standard_normal(n), 'y': rng.standard_normal(n)}) + tf = Tafra({"x": rng.standard_normal(n), "y": rng.standard_normal(n)}) df = pd.DataFrame(tf.data) n_iter = 10_000 - t_tafra = median_of(lambda: [tf['x'] for _ in range(n_iter)]) - t_pandas = median_of(lambda: [df['x'] for _ in range(n_iter)]) + t_tafra = median_of(lambda: [tf["x"] for _ in range(n_iter)]) + t_pandas = median_of(lambda: [df["x"] for _ in range(n_iter)]) t_polars = None if HAS_POLARS: plf = pl.DataFrame(tf.data) - t_polars = median_of(lambda: [plf['x'] for _ in range(n_iter)]) + t_polars = median_of(lambda: [plf["x"] for _ in range(n_iter)]) print_row("tf['x'] x 10k", t_tafra, t_tafra, t_pandas, t_polars) @@ -155,11 +157,13 @@ def row_fn(a: float, b: float, c: float) -> float: (100_000, "100k rows"), (1_000_000, "1M rows"), ]: - tf = Tafra({ - 'a': rng.standard_normal(n_rows), - 'b': rng.standard_normal(n_rows), - 'c': rng.standard_normal(n_rows), - }) + tf = Tafra( + { + "a": rng.standard_normal(n_rows), + "b": rng.standard_normal(n_rows), + "c": rng.standard_normal(n_rows), + } + ) df = pd.DataFrame(tf.data) def tuple_mapper(r: tuple) -> tuple: @@ -167,17 +171,18 @@ def tuple_mapper(r: tuple) -> tuple: n_rep = max(1, min(7, 50_000 // n_rows)) t_tafra = median_of(lambda: list(tf.tuple_map(tuple_mapper, name=None)), n=n_rep) - t_pandas = median_of( - lambda: [row_fn(r.a, r.b, r.c) for r in df.itertuples()], n=n_rep) + t_pandas = median_of(lambda: [row_fn(r.a, r.b, r.c) for r in df.itertuples()], n=n_rep) t_polars = None if HAS_POLARS: plf = pl.DataFrame(tf.data) - t_polars = median_of(lambda: plf.with_columns( - pl.struct(['a', 'b', 'c']).map_elements( - lambda s: row_fn(s['a'], s['b'], s['c']), - return_dtype=pl.Float64 - ).alias('result') - ), n=n_rep) + t_polars = median_of( + lambda: plf.with_columns( + pl.struct(["a", "b", "c"]) + .map_elements(lambda s: row_fn(s["a"], s["b"], s["c"]), return_dtype=pl.Float64) + .alias("result") + ), + n=n_rep, + ) print_row(label, t_tafra, t_tafra, t_pandas, t_polars) @@ -192,24 +197,33 @@ def bench_vectorized_expr(): (100_000, "100k rows"), (1_000_000, "1M rows"), ]: - tf = Tafra({ - 'a': rng.standard_normal(n_rows), - 'b': rng.standard_normal(n_rows), - 'c': rng.standard_normal(n_rows), - }) + tf = Tafra( + { + "a": rng.standard_normal(n_rows), + "b": rng.standard_normal(n_rows), + "c": rng.standard_normal(n_rows), + } + ) df = pd.DataFrame(tf.data) # tafra/numpy: direct array ops - t_tafra = median_of(lambda: np.sqrt(tf['a']**2 + tf['b']**2) + np.log1p(np.abs(tf['c']))) + t_tafra = median_of( + lambda: np.sqrt(tf["a"] ** 2 + tf["b"] ** 2) + np.log1p(np.abs(tf["c"])) + ) # pandas: same numpy ops on .values, or pandas API - t_pandas = median_of(lambda: np.sqrt(df['a']**2 + df['b']**2) + np.log1p(np.abs(df['c']))) + t_pandas = median_of( + lambda: np.sqrt(df["a"] ** 2 + df["b"] ** 2) + np.log1p(np.abs(df["c"])) + ) t_polars = None if HAS_POLARS: plf = pl.DataFrame(tf.data) - t_polars = median_of(lambda: plf.with_columns( - ((pl.col('a')**2 + pl.col('b')**2).sqrt() - + (pl.col('c').abs() + 1).log()).alias('result') - )) + t_polars = median_of( + lambda: plf.with_columns( + ( + (pl.col("a") ** 2 + pl.col("b") ** 2).sqrt() + (pl.col("c").abs() + 1).log() + ).alias("result") + ) + ) print_row(label, t_tafra, t_tafra, t_pandas, t_polars) @@ -226,42 +240,53 @@ def bench_groupby(): (1_000_000, 100, "1M rows, 100 grp"), (1_000_000, 10000, "1M rows, 10k grp"), ]: - tf = Tafra({ - 'group': rng.integers(0, n_groups, size=n_rows).astype('int64'), - 'value': rng.standard_normal(n_rows), - }) + tf = Tafra( + { + "group": rng.integers(0, n_groups, size=n_rows).astype("int64"), + "value": rng.standard_normal(n_rows), + } + ) df = pd.DataFrame(tf.data) - t_accel, t_pure = _tafra_time(lambda: tf.group_by( - ['group'], {'mean_val': (np.mean, 'value'), 'sum_val': (np.sum, 'value')})) + t_accel, t_pure = _tafra_time( + lambda: tf.group_by( + ["group"], {"mean_val": (np.mean, "value"), "sum_val": (np.sum, "value")} + ) + ) t_pandas = median_of( - lambda: df.groupby('group')['value'].agg(['mean', 'sum']).reset_index()) + lambda: df.groupby("group")["value"].agg(["mean", "sum"]).reset_index() + ) t_polars = None if HAS_POLARS: plf = pl.DataFrame(tf.data) - t_polars = median_of(lambda: plf.group_by('group').agg( - pl.col('value').mean().alias('mean_val'), - pl.col('value').sum().alias('sum_val'), - )) + t_polars = median_of( + lambda: plf.group_by("group").agg( + pl.col("value").mean().alias("mean_val"), + pl.col("value").sum().alias("sum_val"), + ) + ) print_row(label, t_accel, t_pure, t_pandas, t_polars) # multi-column for n_rows, label in [(100_000, "100k, 2 col, ~300 grp"), (1_000_000, "1M, 2 col, ~300 grp")]: - tf = Tafra({ - 'g1': rng.integers(0, 100, size=n_rows).astype('int64'), - 'g2': rng.choice(['x', 'y', 'z'], size=n_rows), - 'value': rng.standard_normal(n_rows), - }) + tf = Tafra( + { + "g1": rng.integers(0, 100, size=n_rows).astype("int64"), + "g2": rng.choice(["x", "y", "z"], size=n_rows), + "value": rng.standard_normal(n_rows), + } + ) df = pd.DataFrame(tf.data) - t_accel, t_pure = _tafra_time(lambda: tf.group_by( - ['g1', 'g2'], {'mean_val': (np.mean, 'value')})) - t_pandas = median_of( - lambda: df.groupby(['g1', 'g2'])['value'].agg('mean').reset_index()) + t_accel, t_pure = _tafra_time( + lambda: tf.group_by(["g1", "g2"], {"mean_val": (np.mean, "value")}) + ) + t_pandas = median_of(lambda: df.groupby(["g1", "g2"])["value"].agg("mean").reset_index()) t_polars = None if HAS_POLARS: plf = pl.DataFrame(tf.data) - t_polars = median_of(lambda: plf.group_by('g1', 'g2').agg( - pl.col('value').mean().alias('mean_val'))) + t_polars = median_of( + lambda: plf.group_by("g1", "g2").agg(pl.col("value").mean().alias("mean_val")) + ) print_row(label, t_accel, t_pure, t_pandas, t_polars) @@ -275,21 +300,26 @@ def bench_transform(): (100_000, 100, "100k rows, 100 grp"), (1_000_000, 1000, "1M rows, 1k grp"), ]: - tf = Tafra({ - 'group': rng.integers(0, n_groups, size=n_rows).astype('int64'), - 'value': rng.standard_normal(n_rows), - }) + tf = Tafra( + { + "group": rng.integers(0, n_groups, size=n_rows).astype("int64"), + "value": rng.standard_normal(n_rows), + } + ) df = pd.DataFrame(tf.data) - t_accel, t_pure = _tafra_time(lambda: tf.transform( - ['group'], {'mean_val': (np.mean, 'value')})) + t_accel, t_pure = _tafra_time( + lambda: tf.transform(["group"], {"mean_val": (np.mean, "value")}) + ) t_pandas = median_of( - lambda: df.assign(mean_val=df.groupby('group')['value'].transform('mean'))) + lambda: df.assign(mean_val=df.groupby("group")["value"].transform("mean")) + ) t_polars = None if HAS_POLARS: plf = pl.DataFrame(tf.data) - t_polars = median_of(lambda: plf.with_columns( - pl.col('value').mean().over('group').alias('mean_val'))) + t_polars = median_of( + lambda: plf.with_columns(pl.col("value").mean().over("group").alias("mean_val")) + ) print_row(label, t_accel, t_pure, t_pandas, t_polars) @@ -298,53 +328,61 @@ def bench_joins(): print_header() rng = np.random.default_rng(42) - for n, label in [(1000, "1k x 1k"), (5000, "5k x 5k"), (10000, "10k x 10k"), - (50000, "50k x 50k")]: - left_tf = Tafra({ - 'key': rng.integers(0, 100, size=n).astype('int64'), - 'left_val': rng.standard_normal(n), - }) - right_tf = Tafra({ - 'key': rng.integers(0, 100, size=n).astype('int64'), - 'right_val': rng.standard_normal(n), - }) + for n, label in [ + (1000, "1k x 1k"), + (5000, "5k x 5k"), + (10000, "10k x 10k"), + (50000, "50k x 50k"), + ]: + left_tf = Tafra( + { + "key": rng.integers(0, 100, size=n).astype("int64"), + "left_val": rng.standard_normal(n), + } + ) + right_tf = Tafra( + { + "key": rng.integers(0, 100, size=n).astype("int64"), + "right_val": rng.standard_normal(n), + } + ) left_df = pd.DataFrame(left_tf.data) right_df = pd.DataFrame(right_tf.data) - t_accel, t_pure = _tafra_time(lambda: left_tf.inner_join( - right_tf, [('key', 'key', '==')])) - t_pandas = median_of(lambda: pd.merge( - left_df, right_df, on='key', how='inner')) + t_accel, t_pure = _tafra_time(lambda: left_tf.inner_join(right_tf, [("key", "key", "==")])) + t_pandas = median_of(lambda: pd.merge(left_df, right_df, on="key", how="inner")) t_polars = None if HAS_POLARS: left_pl = pl.DataFrame(left_tf.data) right_pl = pl.DataFrame(right_tf.data) - t_polars = median_of(lambda: left_pl.join(right_pl, on='key', how='inner')) + t_polars = median_of(lambda: left_pl.join(right_pl, on="key", how="inner")) print_row(label, t_accel, t_pure, t_pandas, t_polars) print("\n--- Left Join (equi) ---") print_header() for n, label in [(1000, "1k x 1k"), (5000, "5k x 5k"), (50000, "50k x 50k")]: - left_tf = Tafra({ - 'key': rng.integers(0, 100, size=n).astype('int64'), - 'left_val': rng.standard_normal(n), - }) - right_tf = Tafra({ - 'key': rng.integers(0, 100, size=n).astype('int64'), - 'right_val': rng.standard_normal(n), - }) + left_tf = Tafra( + { + "key": rng.integers(0, 100, size=n).astype("int64"), + "left_val": rng.standard_normal(n), + } + ) + right_tf = Tafra( + { + "key": rng.integers(0, 100, size=n).astype("int64"), + "right_val": rng.standard_normal(n), + } + ) left_df = pd.DataFrame(left_tf.data) right_df = pd.DataFrame(right_tf.data) - t_accel, t_pure = _tafra_time(lambda: left_tf.left_join( - right_tf, [('key', 'key', '==')])) - t_pandas = median_of(lambda: pd.merge( - left_df, right_df, on='key', how='left')) + t_accel, t_pure = _tafra_time(lambda: left_tf.left_join(right_tf, [("key", "key", "==")])) + t_pandas = median_of(lambda: pd.merge(left_df, right_df, on="key", how="left")) t_polars = None if HAS_POLARS: left_pl = pl.DataFrame(left_tf.data) right_pl = pl.DataFrame(right_tf.data) - t_polars = median_of(lambda: left_pl.join(right_pl, on='key', how='left')) + t_polars = median_of(lambda: left_pl.join(right_pl, on="key", how="left")) print_row(label, t_accel, t_pure, t_pandas, t_polars) @@ -375,29 +413,32 @@ def numba_fn(a: np.ndarray, b: np.ndarray, c: np.ndarray) -> np.ndarray: (100_000, "100k rows"), (1_000_000, "1M rows"), ]: - tf = Tafra({ - 'a': rng.standard_normal(n_rows), - 'b': rng.standard_normal(n_rows), - 'c': rng.standard_normal(n_rows), - }) + tf = Tafra( + { + "a": rng.standard_normal(n_rows), + "b": rng.standard_normal(n_rows), + "c": rng.standard_normal(n_rows), + } + ) df = pd.DataFrame(tf.data) # tafra: direct ndarray access - t_tafra = median_of(lambda: numba_fn(tf['a'], tf['b'], tf['c'])) + t_tafra = median_of(lambda: numba_fn(tf["a"], tf["b"], tf["c"])) # pandas: .values to get ndarray - t_pandas = median_of(lambda: numba_fn(df['a'].values, df['b'].values, df['c'].values)) + t_pandas = median_of(lambda: numba_fn(df["a"].values, df["b"].values, df["c"].values)) t_polars = None if HAS_POLARS: plf = pl.DataFrame(tf.data) - t_polars = median_of(lambda: numba_fn( - plf['a'].to_numpy(), plf['b'].to_numpy(), plf['c'].to_numpy())) + t_polars = median_of( + lambda: numba_fn(plf["a"].to_numpy(), plf["b"].to_numpy(), plf["c"].to_numpy()) + ) print_row(label, t_tafra, t_tafra, t_pandas, t_polars) -if __name__ == '__main__': +if __name__ == "__main__": print("=" * 80) libs = f"pandas {pd.__version__}, numpy {np.__version__}" if HAS_POLARS: diff --git a/test/test_tafra.py b/test/test_tafra.py index 017e95f..c186ee8 100644 --- a/test/test_tafra.py +++ b/test/test_tafra.py @@ -2774,7 +2774,7 @@ def test_inner_join_select_columns(self) -> None: assert "other" not in t.columns def test_inner_join_select_nonexistent(self) -> None: - """Nonexistent column in select is ignored.""" + """Nonexistent column in select raises KeyError.""" left = Tafra( { "k": np.array([1, 2, 3]), @@ -2787,13 +2787,12 @@ def test_inner_join_select_nonexistent(self) -> None: "rv": np.array([200, 300]), } ) - t = left.inner_join( - right, - [("k", "k", "==")], - select=["k", "lv", "rv", "nonexistent"], - ) - assert "nonexistent" not in t.columns - assert len(t) == 2 + with pytest.raises(KeyError, match="nonexistent"): + left.inner_join( + right, + [("k", "k", "==")], + select=["k", "lv", "rv", "nonexistent"], + ) def test_inner_join_duplicate_column_left_wins(self) -> None: """Shared column name, left wins.""" @@ -3999,3 +3998,213 @@ def test_integers_as_objects(self) -> None: codes, nu = encode_strings(arr) assert nu == 3 np.testing.assert_array_equal(codes, [0, 1, 0, 2, 1]) + + +class TestJoinSelectValidation: + """Tests for select parameter validation and generator handling.""" + + def test_generator_select_inner_join(self) -> None: + """Generator passed as select should not exhaust on reuse.""" + from tafra.group import InnerJoin + + left = Tafra({"k": np.array([1, 2]), "v": np.array([10, 20])}) + right = Tafra({"k": np.array([1, 2]), "w": np.array([100, 200])}) + + def gen_select() -> Iterator[str]: + yield "k" + yield "v" + + j = InnerJoin(on=[("k", "k", "==")], select=gen_select()) + r1 = j.apply(left, right) + assert "k" in r1.columns + assert "v" in r1.columns + # Second call should also work (generator was materialized in __post_init__) + r2 = j.apply(left, right) + assert "k" in r2.columns + assert "v" in r2.columns + + def test_generator_on_inner_join(self) -> None: + """Generator passed as on should not exhaust on reuse.""" + from tafra.group import InnerJoin + + left = Tafra({"k": np.array([1, 2]), "v": np.array([10, 20])}) + right = Tafra({"k": np.array([1, 2]), "w": np.array([100, 200])}) + + def gen_on() -> Iterator[tuple[str, str, str]]: + yield ("k", "k", "==") + + j = InnerJoin(on=gen_on(), select=[]) + r1 = j.apply(left, right) + assert r1.rows == 2 + r2 = j.apply(left, right) + assert r2.rows == 2 + + def test_invalid_select_inner_join(self) -> None: + """InnerJoin should raise KeyError for nonexistent select columns.""" + left = Tafra({"k": np.array([1, 2]), "v": np.array([10, 20])}) + right = Tafra({"k": np.array([1, 2]), "w": np.array([100, 200])}) + with pytest.raises(KeyError, match="nonexistent"): + left.inner_join(right, on=[("k", "k", "==")], select=["nonexistent"]) + + def test_invalid_select_left_join(self) -> None: + """LeftJoin should raise KeyError for nonexistent select columns.""" + left = Tafra({"k": np.array([1, 2]), "v": np.array([10, 20])}) + right = Tafra({"k": np.array([1, 2]), "w": np.array([100, 200])}) + with pytest.raises(KeyError, match="nonexistent"): + left.left_join(right, on=[("k", "k", "==")], select=["nonexistent"]) + + def test_invalid_select_cross_join(self) -> None: + """CrossJoin should raise KeyError for nonexistent select columns.""" + left = Tafra({"k": np.array([1, 2]), "v": np.array([10, 20])}) + right = Tafra({"j": np.array([3, 4]), "w": np.array([100, 200])}) + with pytest.raises(KeyError, match="nonexistent"): + left.cross_join(right, select=["nonexistent"]) + + +class TestCoverageGaps: + """Tests targeting specific uncovered code paths.""" + + def test_left_join_empty_right_datetime_null_fill(self) -> None: + """Empty right table with datetime column uses NaT fill.""" + left = Tafra({"k": np.array([1, 2])}) + right = Tafra( + { + "k": np.array([], dtype=np.int64), + "dt": np.array([], dtype="datetime64[ns]"), + } + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + t = left.left_join(right, on=[("k", "k", "==")]) + assert t.rows == 2 + assert np.isnat(t["dt"][0]) + + def test_left_join_empty_right_timedelta_null_fill(self) -> None: + """Empty right table with timedelta column uses NaT fill.""" + left = Tafra({"k": np.array([1, 2])}) + right = Tafra( + { + "k": np.array([], dtype=np.int64), + "td": np.array([], dtype="timedelta64[ns]"), + } + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + t = left.left_join(right, on=[("k", "k", "==")]) + assert t.rows == 2 + assert np.isnat(t["td"][0]) + + def test_left_join_empty_right_int_object_warning(self) -> None: + """Empty right table with int column emits warning and casts to object.""" + left = Tafra({"k": np.array([1, 2])}) + right = Tafra( + { + "k": np.array([], dtype=np.int64), + "count": np.array([], dtype=np.int64), + } + ) + with pytest.warns(UserWarning, match="cast to object"): + t = left.left_join(right, on=[("k", "k", "==")]) + assert t["count"].dtype == object + + def test_left_join_empty_right_string_null_fill(self) -> None: + """Empty right table with StringDType column fills None.""" + left = Tafra({"k": np.array([1, 2])}) + right = Tafra( + { + "k": np.array([], dtype=np.int64), + "name": np.array([], dtype=np.dtypes.StringDType()), + } + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + t = left.left_join(right, on=[("k", "k", "==")]) + assert t.rows == 2 + assert t["name"][0] is None + + def test_left_join_nonequi_null_keys(self) -> None: + """Non-equi left join with NaN keys: null rows get unmatched.""" + left = Tafra( + { + "k": np.array([1.0, np.nan, 3.0]), + "lv": np.array([10, 20, 30]), + } + ) + right = Tafra( + { + "k": np.array([0.5, 2.5]), + "rv": np.array([100, 200]), + } + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + t = left.left_join(right, on=[("k", "k", ">=")]) + # NaN row should be unmatched (null key) + nan_rows = [i for i in range(t.rows) if t["lv"][i] == 20] + for i in nan_rows: + assert t["rv"][i] is None + + def test_inner_join_nonequi_null_keys(self) -> None: + """Non-equi inner join with NaN keys: null rows excluded.""" + left = Tafra( + { + "k": np.array([1.0, np.nan, 3.0]), + "lv": np.array([10, 20, 30]), + } + ) + right = Tafra( + { + "k": np.array([0.5, 2.5]), + "rv": np.array([100, 200]), + } + ) + t = left.inner_join(right, on=[("k", "k", ">=")]) + # NaN row (lv=20) should NOT appear in results + assert 20 not in list(t["lv"]) + + def test_composite_key_overflow(self) -> None: + """Composite key overflow raises ValueError.""" + from tafra.group import GroupSet + + # Two columns with huge cardinality to trigger overflow + big = np.array([0, 2**31], dtype=np.int64) + with pytest.raises(ValueError, match="overflow"): + GroupSet._build_composite_key([big, big]) + + def test_negative_int_groupby_multi_column(self) -> None: + """GroupBy on multiple columns with negative integer keys.""" + t = Tafra( + { + "g1": np.array([-2, -2, -1, -1, 0, 0], dtype=np.int64), + "g2": np.array([1, 2, 1, 2, 1, 2], dtype=np.int64), + "v": np.array([10.0, 20.0, 30.0, 40.0, 50.0, 60.0]), + } + ) + r = t.group_by(["g1", "g2"], {"v": np.sum}) + assert r.rows == 6 + # Each combination is unique — sums should equal original values + for i in range(r.rows): + g1, g2, v = r["g1"][i], r["g2"][i], r["v"][i] + expected = t["v"][(t["g1"] == g1) & (t["g2"] == g2)].sum() + assert v == expected, f"Mismatch for g1={g1}, g2={g2}: {v} != {expected}" + + def test_negative_int_join_multi_column(self) -> None: + """Multi-column join with negative integer keys on both sides.""" + left = Tafra( + { + "k1": np.array([-5, -3, 0, 2], dtype=np.int64), + "k2": np.array([1, 2, 3, 4], dtype=np.int64), + "lv": np.array([10, 20, 30, 40]), + } + ) + right = Tafra( + { + "k1": np.array([-3, 0, 2, 5], dtype=np.int64), + "k2": np.array([2, 3, 4, 1], dtype=np.int64), + "rv": np.array([100, 200, 300, 400]), + } + ) + r = left.inner_join(right, on=[("k1", "k1", "=="), ("k2", "k2", "==")]) + assert r.rows == 3 + assert list(r["lv"]) == [20, 30, 40] + assert list(r["rv"]) == [100, 200, 300]