boti-data 1.3.7__tar.gz → 1.3.9__tar.gz

{boti_data-1.3.7 → boti_data-1.3.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: boti-data
-Version: 1.3.7
+Version: 1.3.9
 Summary: Data infrastructure for the Boti ecosystem
 Author: Luis Valverde
 Maintainer-email: Luis Valverde <lvalverdeb@gmail.com>
@@ -884,6 +884,129 @@ with DataHelper.session(cluster_factory=LocalCluster) as client:
 
 ---
 
+## Sizing connections and fetch concurrency
+
+More database connections do not mean more throughput. Past a modest number they
+mean *less*, because connections compete for the same cores, disks and locks. Both
+the [PostgreSQL wiki](https://wiki.postgresql.org/wiki/Number_Of_Database_Connections)
+and the [HikariCP pool-sizing guide](https://github.com/brettwooldridge/HikariCP/wiki/About-Pool-Sizing)
+give the same starting formula:
+
+```
+connections = (core_count * 2) + effective_spindle_count
+```
+
+Count physical cores, not hyperthreads, and take `effective_spindle_count` as zero
+when the working set is cached in RAM. For a 4-core server with one disk that is
+about 10 connections in total — not per client, not per worker, **in total**.
+
+The characteristic failure of getting this wrong is counter-intuitive: the database
+sits nearly idle, at a few percent CPU, while clients time out. The connections are
+all queued rather than working. Raising the pool size makes it worse.
+
+`boti-data` exposes three separate numbers that multiply into that total. Size the
+product, not the individual knobs.
+
+### 1. Client-side pool: `pool_size + max_overflow`
+
+`SqlDatabaseConfig` defaults to `pool_size=5` and `max_overflow=10`, so a single
+engine bursts to **15** connections, not 5. `max_overflow` is three times
+`pool_size`, so the steady-state number is a third of the real ceiling.
+
+```python
+config = SqlDatabaseConfig(
+    connection_url="postgresql+psycopg://...",
+    pool_size=5,       # steady state
+    max_overflow=10,   # burst headroom — the ceiling is 5 + 10 = 15
+)
+```
+
+Engines are cached per distinct configuration, so two differently-configured helpers
+against the same database hold two pools: 30 connections, not 15.
+
+### 2. Distributed fetch concurrency: `n_workers × max_concurrent_fetches`
+
+This is the one that surprises people. `max_concurrent_fetches` (default `4`, capped
+at `16`) bounds concurrent partition fetches **per worker process**, because the
+semaphore enforcing it lives in process memory. It is not a cluster-wide limit.
+
+```
+20 Dask workers × max_concurrent_fetches=4 = up to 80 concurrent connections
+```
+
+To work backwards from a budget:
+
+```python
+max_concurrent_fetches = max(1, connection_budget // n_workers)
+```
+
+With the 10-connection budget above and 5 workers, that is `2` — not the default `4`.
+
+### 3. Worker connections are unpooled by design
+
+Worker engines always use `NullPool`. Connection pools cannot be pickled and would
+bind to the wrong event loop, so each partition fetch opens a fresh connection and
+closes it. That is correct for distribution, but it means every fetch pays a full TCP
+handshake, TLS negotiation and authentication — on a remote database, tens of
+milliseconds each, before a single row is read. A 200-partition load pays it 200 times.
+
+The fix is not to pool inside the workers, but to put a pooler in front of the
+database, so that warm server-side connections are reused while workers keep opening
+cheap local ones. **The pooling mode matters**, because it interacts with `query_only`:
+
+| Pooler mode | Connection reuse | `query_only=True` |
+|---|---|---|
+| Session (e.g. PgBouncer `session`) | Yes | **Supported** — the read-only `SET` lasts for the session |
+| Transaction (e.g. PgBouncer `transaction`) | Higher | **Not supported** — see below |
+
+`query_only=True` issues its read-only statement on the SQLAlchemy `connect` event, so
+it applies to a physical connection. Under transaction-mode pooling a physical
+connection can be handed to a different logical client between transactions, and that
+session state is not guaranteed to survive the handoff. The failure is silent: a load
+that believes it is read-only quietly stops being so.
+
+With transaction-mode pooling, set `query_only=False` and enforce read-only access
+with database grants instead, where it cannot be lost in a handoff.
+
+### Measuring it
+
+Set `diagnostics=True` on a partitioned load and the completion line reports how long
+fetches spent queueing for a slot:
+
+```
+Partitioned SQL load completed strategy=offset partitions=4 rows=40 result_partitions=1
+  elapsed=0.04s prepare_stmt=0.000s load_plan=0.031s
+  gate_waits=0/4 gate_wait_total=0.000s gate_wait_max_seen=0.000s
+```
+
+`gate_waits=0/4` means none of the four fetches had to wait. A rising
+`gate_wait_total` alongside a database that is not CPU-bound is the signature of the
+failure above — concurrency queued rather than working. The other stage timers cannot
+show this, because they are all measured *inside* the gate.
+
+The fields are omitted entirely when no fetch ran in the logging process. A lazy load
+returns a Dask graph before anything is fetched, and in a distributed run the fetches
+happen in other processes; reporting `0.000s` there would say "no contention" when the
+truth is "not measured here".
+
+For lazy loads, read the counters after `.compute()` instead:
+
+```python
+from boti_data import fetch_gate_stats, reset_fetch_gate_stats
+
+reset_fetch_gate_stats()
+result = helper.dask.load().compute()
+
+for gate_key, stats in fetch_gate_stats().items():
+    print(gate_key, stats.waited_fetches, "/", stats.fetches, stats.total_wait_seconds)
+```
+
+`fetch_gate_stats()` reports the **calling process only**. On a distributed cluster each
+worker accumulates its own counters, so run it on the workers (e.g. via
+`client.run(fetch_gate_stats)`) to see the whole picture.
+
+---
+
 ## Choosing between distributed and non-distributed
 
 Use the following as a guide:

{boti_data-1.3.7 → boti_data-1.3.9}/README.md

@@ -845,6 +845,129 @@ with DataHelper.session(cluster_factory=LocalCluster) as client:
 
 ---
 
+## Sizing connections and fetch concurrency
+
+More database connections do not mean more throughput. Past a modest number they
+mean *less*, because connections compete for the same cores, disks and locks. Both
+the [PostgreSQL wiki](https://wiki.postgresql.org/wiki/Number_Of_Database_Connections)
+and the [HikariCP pool-sizing guide](https://github.com/brettwooldridge/HikariCP/wiki/About-Pool-Sizing)
+give the same starting formula:
+
+```
+connections = (core_count * 2) + effective_spindle_count
+```
+
+Count physical cores, not hyperthreads, and take `effective_spindle_count` as zero
+when the working set is cached in RAM. For a 4-core server with one disk that is
+about 10 connections in total — not per client, not per worker, **in total**.
+
+The characteristic failure of getting this wrong is counter-intuitive: the database
+sits nearly idle, at a few percent CPU, while clients time out. The connections are
+all queued rather than working. Raising the pool size makes it worse.
+
+`boti-data` exposes three separate numbers that multiply into that total. Size the
+product, not the individual knobs.
+
+### 1. Client-side pool: `pool_size + max_overflow`
+
+`SqlDatabaseConfig` defaults to `pool_size=5` and `max_overflow=10`, so a single
+engine bursts to **15** connections, not 5. `max_overflow` is three times
+`pool_size`, so the steady-state number is a third of the real ceiling.
+
+```python
+config = SqlDatabaseConfig(
+    connection_url="postgresql+psycopg://...",
+    pool_size=5,       # steady state
+    max_overflow=10,   # burst headroom — the ceiling is 5 + 10 = 15
+)
+```
+
+Engines are cached per distinct configuration, so two differently-configured helpers
+against the same database hold two pools: 30 connections, not 15.
+
+### 2. Distributed fetch concurrency: `n_workers × max_concurrent_fetches`
+
+This is the one that surprises people. `max_concurrent_fetches` (default `4`, capped
+at `16`) bounds concurrent partition fetches **per worker process**, because the
+semaphore enforcing it lives in process memory. It is not a cluster-wide limit.
+
+```
+20 Dask workers × max_concurrent_fetches=4 = up to 80 concurrent connections
+```
+
+To work backwards from a budget:
+
+```python
+max_concurrent_fetches = max(1, connection_budget // n_workers)
+```
+
+With the 10-connection budget above and 5 workers, that is `2` — not the default `4`.
+
+### 3. Worker connections are unpooled by design
+
+Worker engines always use `NullPool`. Connection pools cannot be pickled and would
+bind to the wrong event loop, so each partition fetch opens a fresh connection and
+closes it. That is correct for distribution, but it means every fetch pays a full TCP
+handshake, TLS negotiation and authentication — on a remote database, tens of
+milliseconds each, before a single row is read. A 200-partition load pays it 200 times.
+
+The fix is not to pool inside the workers, but to put a pooler in front of the
+database, so that warm server-side connections are reused while workers keep opening
+cheap local ones. **The pooling mode matters**, because it interacts with `query_only`:
+
+| Pooler mode | Connection reuse | `query_only=True` |
+|---|---|---|
+| Session (e.g. PgBouncer `session`) | Yes | **Supported** — the read-only `SET` lasts for the session |
+| Transaction (e.g. PgBouncer `transaction`) | Higher | **Not supported** — see below |
+
+`query_only=True` issues its read-only statement on the SQLAlchemy `connect` event, so
+it applies to a physical connection. Under transaction-mode pooling a physical
+connection can be handed to a different logical client between transactions, and that
+session state is not guaranteed to survive the handoff. The failure is silent: a load
+that believes it is read-only quietly stops being so.
+
+With transaction-mode pooling, set `query_only=False` and enforce read-only access
+with database grants instead, where it cannot be lost in a handoff.
+
+### Measuring it
+
+Set `diagnostics=True` on a partitioned load and the completion line reports how long
+fetches spent queueing for a slot:
+
+```
+Partitioned SQL load completed strategy=offset partitions=4 rows=40 result_partitions=1
+  elapsed=0.04s prepare_stmt=0.000s load_plan=0.031s
+  gate_waits=0/4 gate_wait_total=0.000s gate_wait_max_seen=0.000s
+```
+
+`gate_waits=0/4` means none of the four fetches had to wait. A rising
+`gate_wait_total` alongside a database that is not CPU-bound is the signature of the
+failure above — concurrency queued rather than working. The other stage timers cannot
+show this, because they are all measured *inside* the gate.
+
+The fields are omitted entirely when no fetch ran in the logging process. A lazy load
+returns a Dask graph before anything is fetched, and in a distributed run the fetches
+happen in other processes; reporting `0.000s` there would say "no contention" when the
+truth is "not measured here".
+
+For lazy loads, read the counters after `.compute()` instead:
+
+```python
+from boti_data import fetch_gate_stats, reset_fetch_gate_stats
+
+reset_fetch_gate_stats()
+result = helper.dask.load().compute()
+
+for gate_key, stats in fetch_gate_stats().items():
+    print(gate_key, stats.waited_fetches, "/", stats.fetches, stats.total_wait_seconds)
+```
+
+`fetch_gate_stats()` reports the **calling process only**. On a distributed cluster each
+worker accumulates its own counters, so run it on the workers (e.g. via
+`client.run(fetch_gate_stats)`) to see the whole picture.
+
+---
+
 ## Choosing between distributed and non-distributed
 
 Use the following as a guide:

{boti_data-1.3.7 → boti_data-1.3.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "boti-data"
-version = "1.3.7"
+version = "1.3.9"
 description = "Data infrastructure for the Boti ecosystem"
 readme = "README.md"
 requires-python = ">=3.13"

{boti_data-1.3.7 → boti_data-1.3.9}/src/boti_data/__init__.py

@@ -60,6 +60,9 @@ __all__ = [
     "SqlPartitionSpec",
     "SqlPartitionedLoadRequest",
     "SqlPartitionedLoader",
+    "GateWaitStats",
+    "fetch_gate_stats",
+    "reset_fetch_gate_stats",
     "SqlModelRegistry",
     "SqlRepository",
     "SqlUnitOfWork",
@@ -119,6 +122,9 @@ _LAZY = {
     # dask+pandas — the one heavy import inside the otherwise-clean db package.
     "SqlPartitionedLoader": ("boti_data.db.partitioned_loader", "SqlPartitionedLoader"),
     "SqlPartitionedLoadRequest": ("boti_data.db.partitioned_loader", "SqlPartitionedLoadRequest"),
+    "GateWaitStats": ("boti_data.db.partitioned_execution", "GateWaitStats"),
+    "fetch_gate_stats": ("boti_data.db.partitioned_execution", "fetch_gate_stats"),
+    "reset_fetch_gate_stats": ("boti_data.db.partitioned_execution", "reset_fetch_gate_stats"),
     # dask/pandas/polars/pyarrow-touching — the dataframe-shaped, bulk-oriented side.
     "ConnectionCatalog": ("boti_data.connection_catalog", "ConnectionCatalog"),
     "ParquetDataConfig": ("boti_data.parquet", "ParquetDataConfig"),

{boti_data-1.3.7 → boti_data-1.3.9}/src/boti_data/db/partitioned_execution.py

@@ -10,6 +10,13 @@ import dask.dataframe as dd
 import pandas as pd
 from sqlalchemy.sql import Select
 
+from boti_data.db.partitioned_fetch_gate import (
+    GateWaitStats,
+    _timed_fetch_gate,
+    fetch_gate_stats,
+    format_gate_wait_suffix,
+    reset_fetch_gate_stats,
+)
 from boti_data.db.partitioned_row_materialization import (
     align_and_coerce_partition,
     arrow_align_and_coerce_partition,
@@ -27,8 +34,16 @@ from boti_data.db.sql_engine import (
     _get_worker_engine_identity,
 )
 
-_FETCH_GATES: dict[tuple[str, int], threading.BoundedSemaphore] = {}
-_FETCH_GATES_LOCK = threading.Lock()
+# The fetch gate and its contention counters live in partitioned_fetch_gate.py
+# (long-file headroom) and are re-exported here so existing import paths — and
+# the tests and callers that use them — keep working unchanged.
+__all__ = [
+    "GateWaitStats",
+    "SqlPartitionExecutor",
+    "fetch_gate_stats",
+    "format_gate_wait_suffix",
+    "reset_fetch_gate_stats",
+]
 
 # Thread-local caches: avoid recreating engine objects and event loops on every partition fetch.
 # The engine itself is cheap state (dialect, URL, event listeners); caching it per worker
@@ -76,16 +91,6 @@ def _get_cached_worker_async_engine(config: WorkerSqlConfig) -> Any:
     return engine
 
 
-def _get_fetch_gate(gate_key: str, max_concurrent_fetches: int) -> threading.BoundedSemaphore:
-    cache_key = (gate_key, max_concurrent_fetches)
-    with _FETCH_GATES_LOCK:
-        gate = _FETCH_GATES.get(cache_key)
-        if gate is None:
-            gate = threading.BoundedSemaphore(max_concurrent_fetches)
-            _FETCH_GATES[cache_key] = gate
-        return gate
-
-
 class SqlPartitionExecutor:
     """Execute planned partitions and assemble lazy Dask outputs."""
 
@@ -216,8 +221,7 @@ class SqlPartitionExecutor:
         meta_dtypes: dict[str, str],
         use_arrow: bool = True,
     ) -> pd.DataFrame:
-        gate = _get_fetch_gate(gate_key, max_concurrent_fetches)
-        with gate:
+        with _timed_fetch_gate(gate_key, max_concurrent_fetches):
             engine = _get_cached_worker_sync_engine(config)
             with engine.connect() as conn:
                 if not use_arrow:
@@ -253,8 +257,7 @@ class SqlPartitionExecutor:
         NullPool means the engine carries no loop-bound connection pool state, so
         both the loop and the engine are safe to reuse across calls in the same thread.
         """
-        gate = _get_fetch_gate(gate_key, max_concurrent_fetches)
-        with gate:
+        with _timed_fetch_gate(gate_key, max_concurrent_fetches):
 
             async def _fetch() -> pd.DataFrame:
                 engine = _get_cached_worker_async_engine(config)

boti_data-1.3.9/src/boti_data/db/partitioned_fetch_gate.py

@@ -0,0 +1,138 @@
+"""Partition-fetch concurrency gate and its contention accounting.
+
+Split out of partitioned_execution.py purely for long-file headroom. Every name
+here is re-exported from partitioned_execution.py so existing import paths keep
+working unchanged.
+
+The gate bounds how many partition fetches run at once, and the counters record
+how long fetches spent queueing for a slot. That second half exists because the
+gate is otherwise invisible: the ``diagnostics=True`` stage timers (``db_fetch``,
+``prepare_stmt``, ``load_plan``) are all measured *inside* the gate, so under
+saturation total elapsed time rises while every stage timer still looks healthy.
+An idle database plus timing-out callers is the classic symptom, and without
+these counters it is indistinguishable from a database that is simply slow.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from time import perf_counter
+
+_FETCH_GATES: dict[tuple[str, int], threading.BoundedSemaphore] = {}
+_FETCH_GATES_LOCK = threading.Lock()
+
+
+def _get_fetch_gate(gate_key: str, max_concurrent_fetches: int) -> threading.BoundedSemaphore:
+    cache_key = (gate_key, max_concurrent_fetches)
+    with _FETCH_GATES_LOCK:
+        gate = _FETCH_GATES.get(cache_key)
+        if gate is None:
+            gate = threading.BoundedSemaphore(max_concurrent_fetches)
+            _FETCH_GATES[cache_key] = gate
+        return gate
+
+
+@dataclass(frozen=True)
+class GateWaitStats:
+    """Time spent waiting for a partition-fetch slot, accumulated per process.
+
+    ``max_wait_seconds`` is cumulative for the life of the process and cannot be
+    differenced into a per-load figure the way the other three can (a running
+    maximum carries no record of when it was set). :meth:`since` therefore
+    carries it through unchanged; it reads as "worst wait seen in this process
+    so far", not "worst wait during this load".
+    """
+
+    fetches: int = 0
+    waited_fetches: int = 0
+    total_wait_seconds: float = 0.0
+    max_wait_seconds: float = 0.0
+
+    def since(self, baseline: GateWaitStats | None) -> GateWaitStats:
+        """Return the stats accumulated since ``baseline`` was captured."""
+        if baseline is None:
+            return self
+        return GateWaitStats(
+            fetches=self.fetches - baseline.fetches,
+            waited_fetches=self.waited_fetches - baseline.waited_fetches,
+            total_wait_seconds=self.total_wait_seconds - baseline.total_wait_seconds,
+            max_wait_seconds=self.max_wait_seconds,
+        )
+
+
+# An uncontended BoundedSemaphore.acquire() returns in ~1us; anything at or above
+# a millisecond means the fetch genuinely queued behind another one.
+_GATE_WAIT_EPSILON_SECONDS = 0.001
+
+_GATE_WAIT_STATS: dict[str, GateWaitStats] = {}
+_GATE_WAIT_LOCK = threading.Lock()
+
+
+def _record_gate_wait(gate_key: str, waited_seconds: float) -> None:
+    with _GATE_WAIT_LOCK:
+        current = _GATE_WAIT_STATS.get(gate_key, GateWaitStats())
+        _GATE_WAIT_STATS[gate_key] = GateWaitStats(
+            fetches=current.fetches + 1,
+            waited_fetches=current.waited_fetches
+            + (1 if waited_seconds >= _GATE_WAIT_EPSILON_SECONDS else 0),
+            total_wait_seconds=current.total_wait_seconds + waited_seconds,
+            max_wait_seconds=max(current.max_wait_seconds, waited_seconds),
+        )
+
+
+def fetch_gate_stats() -> dict[str, GateWaitStats]:
+    """Snapshot partition-fetch gate contention for **this process**, by gate key.
+
+    Scope matters for interpreting the numbers. Partition fetches are dispatched
+    as Dask tasks, so in a distributed run each worker process accumulates its
+    own counters and this function only ever sees the calling process's share.
+    A lazy load returns before any fetch has run at all, so call this after
+    ``.compute()``, not after the load call returns.
+    """
+    with _GATE_WAIT_LOCK:
+        return dict(_GATE_WAIT_STATS)
+
+
+def reset_fetch_gate_stats() -> None:
+    """Clear accumulated gate-wait counters for this process."""
+    with _GATE_WAIT_LOCK:
+        _GATE_WAIT_STATS.clear()
+
+
+def format_gate_wait_suffix(gate_key: str, baseline: GateWaitStats | None) -> str:
+    """Format gate contention for a diagnostics line, or ``""`` if nothing ran here.
+
+    Shared by the sync loader and the async gateway so both diagnostics lines
+    report contention identically.
+
+    Returning empty is deliberate rather than reporting zeros. A lazy load
+    returns a Dask graph before any partition has been fetched, and in a
+    distributed run the fetches then happen inside other worker processes
+    entirely — ``gate_wait_total=0.000s`` there would read as "no contention"
+    when the truth is "not measured here". Callers wanting figures for a lazy
+    load should read :func:`fetch_gate_stats` after ``.compute()``.
+    """
+    delta = fetch_gate_stats().get(gate_key, GateWaitStats()).since(baseline)
+    if delta.fetches <= 0:
+        return ""
+    return (
+        f" gate_waits={delta.waited_fetches}/{delta.fetches}"
+        f" gate_wait_total={delta.total_wait_seconds:.3f}s"
+        f" gate_wait_max_seen={delta.max_wait_seconds:.3f}s"
+    )
+
+
+@contextmanager
+def _timed_fetch_gate(gate_key: str, max_concurrent_fetches: int) -> Iterator[None]:
+    """Hold a fetch slot, recording how long it took to get one."""
+    gate = _get_fetch_gate(gate_key, max_concurrent_fetches)
+    started = perf_counter()
+    gate.acquire()
+    _record_gate_wait(gate_key, perf_counter() - started)
+    try:
+        yield
+    finally:
+        gate.release()

{boti_data-1.3.7 → boti_data-1.3.9}/src/boti_data/db/partitioned_loader.py

@@ -11,7 +11,11 @@ from boti.core.lifecycle_pickle import PicklableLifecycleCoreMixin
 from boti.core.logger import Logger
 from sqlalchemy.engine import url as sqlalchemy_url
 
-from boti_data.db.partitioned_execution import SqlPartitionExecutor
+from boti_data.db.partitioned_execution import (
+    SqlPartitionExecutor,
+    fetch_gate_stats,
+    format_gate_wait_suffix,
+)
 from boti_data.db.partitioned_planner import SqlPartitionPlanner
 from boti_data.db.partitioned_types import (
     SqlPartitionedLoadRequest,
@@ -200,6 +204,7 @@ class SqlPartitionedLoader(PicklableLifecycleCoreMixin, LifecycleCore):
         t_plan_fetch = perf_counter()
         prepared_stmt = self._planner.prepare_statement(request)
         t_load = perf_counter()
+        gate_baseline = fetch_gate_stats().get(self._gate_key) if request.diagnostics else None
         result = executor.load_plan(
             plan,
             as_pandas=request.as_pandas,
@@ -217,6 +222,7 @@ class SqlPartitionedLoader(PicklableLifecycleCoreMixin, LifecycleCore):
                 f"elapsed={t_end - started:.2f}s "
                 f"prepare_stmt={t_load - t_plan_fetch:.3f}s "
                 f"load_plan={t_end - t_load:.3f}s"
+                f"{format_gate_wait_suffix(self._gate_key, gate_baseline)}"
             )
         return result
 

{boti_data-1.3.7 → boti_data-1.3.9}/src/boti_data/enrichment/protocol.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 from typing import Any, Protocol, runtime_checkable
 
+import dask.dataframe as dd
 import pandas as pd
 
 
@@ -17,13 +18,15 @@ import pandas as pd
 class DataFrameTransformer(Protocol):
     """Generic DataFrame transform contract.
 
-    Implementations receive a ``pd.DataFrame`` and arbitrary keyword
-    arguments (used by higher layers to pass context) and must return
-    a ``pd.DataFrame``.
+    Implementations receive a ``pd.DataFrame`` or ``dd.DataFrame`` and
+    arbitrary keyword arguments (used by higher layers to pass context)
+    and must return a frame of the same kind they were given.
 
     This protocol is intentionally minimal — ``**kwargs`` keeps it
     flexible so that etl-core can pass a ``TransformContext`` without
     boti-data knowing about it.
     """
 
-    async def transform(self, df: pd.DataFrame, **kwargs: Any) -> pd.DataFrame: ...
+    async def transform(
+        self, df: pd.DataFrame | dd.DataFrame, **kwargs: Any
+    ) -> pd.DataFrame | dd.DataFrame: ...

{boti_data-1.3.7 → boti_data-1.3.9}/src/boti_data/enrichment/transformers.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 from typing import Any
 
+import dask.dataframe as dd
 import pandas as pd
 from boti.core.security import has_dunder_identifier
 
@@ -46,16 +47,30 @@ class TypeCaster:
         ``"int64"``, ``"float64"``, ``"bool"``, ``"string"``, etc.
 
     Missing columns are silently skipped.
+
+    Dask frames
+    -----------
+    ``dd.DataFrame.astype()`` has no ``errors=`` kwarg, and a dask
+    column's dtype must be uniform across every partition — so there is
+    no dask equivalent of pandas' per-column "cast what you can, leave
+    the rest" behaviour. Against a ``dd.DataFrame``, a cast that would
+    have been silently skipped under ``errors="ignore"`` instead raises
+    immediately. Only list columns you're confident are castable when
+    the input is a dask frame.
     """
 
     def __init__(self, casts: dict[str, str]) -> None:
         self._casts = dict(casts)
 
-    async def transform(self, df: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
+    async def transform(
+        self, df: pd.DataFrame | dd.DataFrame, **kwargs: Any
+    ) -> pd.DataFrame | dd.DataFrame:
         applicable = {col: dt for col, dt in self._casts.items() if col in df.columns}
-        if applicable:
-            df = df.astype(applicable, errors="ignore")
-        return df
+        if not applicable:
+            return df
+        if isinstance(df, dd.DataFrame):
+            return df.astype(applicable)
+        return df.astype(applicable, errors="ignore")
 
 
 class RowFilter:
@@ -79,7 +94,9 @@ class RowFilter:
     def __init__(self, predicates: list[str]) -> None:
         self._predicates = list(predicates)
 
-    async def transform(self, df: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
+    async def transform(
+        self, df: pd.DataFrame | dd.DataFrame, **kwargs: Any
+    ) -> pd.DataFrame | dd.DataFrame:
         for expr in self._predicates:
             _reject_dunder_expr(expr)
             mask = df.eval(expr)  # nosec CWE-94 -- dunder-chain rejected above; see class docstring's Security section
@@ -113,7 +130,9 @@ class DerivedColumn:
     def __init__(self, columns: dict[str, str]) -> None:
         self._columns = dict(columns)
 
-    async def transform(self, df: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
+    async def transform(
+        self, df: pd.DataFrame | dd.DataFrame, **kwargs: Any
+    ) -> pd.DataFrame | dd.DataFrame:
         for name, expr in self._columns.items():
             _reject_dunder_expr(expr)
             df[name] = df.eval(expr)  # nosec CWE-94 -- dunder-chain rejected above; see class docstring's Security section
@@ -141,7 +160,9 @@ class Deduplicator:
         self._subset = subset
         self._keep = keep
 
-    async def transform(self, df: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
+    async def transform(
+        self, df: pd.DataFrame | dd.DataFrame, **kwargs: Any
+    ) -> pd.DataFrame | dd.DataFrame:
         before = len(df)
         df = df.drop_duplicates(subset=self._subset, keep=self._keep).reset_index(
             drop=True,