Fix/streaming bulk reads

README.md

@@ -237,6 +237,14 @@ never served as a cache hit. Eviction is LRU by fetch time against a 10 GiB
 budget (`RWPYTOOLS_CACHE_MAX_BYTES`). The cache is thread-safe within a
 process; it is not coordinated across processes.
 
+Reading a cached export is memory-bounded: `start`, `end` and `symbols` are
+pushed into a batch-by-batch scan of feather and parquet exports, so a narrow
+window over a multi-GB export holds only the matching rows, never the whole
+file. Iterating a large export one window at a time (e.g. a year-by-year
+backfill of `liquid_universe_prices`) therefore stays flat, and freed Arrow
+memory is returned to the OS after every read so long-running workers do not
+ratchet up. The per-pair FX exports (Feather V1) are small and read whole.
+
 ```python
 client.cache.total_bytes()
 client.cache.entries()

src/rwpytools/formats.py

@@ -42,14 +42,18 @@
 --------------
 
 :func:`read_frame` narrows by date range and symbol while reading. For the
-Arrow formats the predicate runs on the Arrow table *before* pandas
-conversion, so ``symbols=["BTCUSDT"]`` against a 20M-row export only ever
-materialises BTCUSDT's rows. :func:`max_timestamp` finds an export's
-watermark by reading just its date column.
+Arrow formats the predicate is pushed into a batch-by-batch dataset scan,
+so ``start="2024-01-01", symbols=["BTCUSDT"]`` against a 20M-row export
+never holds more than the matching rows plus one record batch — the whole
+export is never materialised. :func:`max_timestamp` finds an export's
+watermark by reading just its date column. After each read, Arrow's
+allocator is asked to return freed memory to the OS
+(:func:`release_arrow_memory`), so long-running processes do not ratchet up.
 """
 
 from __future__ import annotations
 
+import contextlib
 import datetime as _dt
 import enum
 import logging
@@ -61,6 +65,7 @@
 import pandas as pd
 import pyarrow as pa
 import pyarrow.compute as pc
+import pyarrow.dataset as ds
 import pyarrow.feather
 import pyarrow.parquet
 
@@ -344,59 +349,118 @@ def _arrow_bound(column_type: Any, day: _dt.date) -> Any:
     return None
 
 
-def _filter_table(
-    table: Any,
+def _filter_expression(
+    schema: Any,
     *,
     date_column: str | None,
     start: _dt.date | None,
     end: _dt.date | None,
     symbol_column: str | None,
     symbols: Sequence[str] | None,
 ) -> tuple[Any, bool]:
-    """Apply what filters Arrow can evaluate natively.
+    """Translate the date / symbol filters into an Arrow dataset expression.
 
-    Returns ``(table, fully_filtered)``; when ``fully_filtered`` is False
-    the caller finishes the job in pandas (e.g. dictionary-encoded symbols
-    or a date column stored in an unusual type).
+    Returns ``(expression, complete)``. ``expression`` is ``None`` when no
+    filter applies to this schema. ``complete`` is False when some filter
+    could not be expressed (e.g. a dictionary-encoded symbol column or a
+    date column in an unusual type); the caller then finishes the job in
+    pandas, so the result is always fully filtered.
     """
 
     complete = True
-    mask: Any = None
+    expression: Any = None
 
     def _and(clause: Any) -> None:
-        nonlocal mask
-        mask = clause if mask is None else pc.and_(mask, clause)
+        nonlocal expression
+        expression = clause if expression is None else (expression & clause)
 
-    if symbols and symbol_column and symbol_column in table.column_names:
-        column = table[symbol_column]
-        try:
-            value_set = pa.array(list(symbols)).cast(column.type)
-            _and(pc.is_in(column, value_set=value_set))
-        except (pa.ArrowInvalid, pa.ArrowNotImplementedError, pa.ArrowTypeError):
+    names = set(schema.names)
+    if symbols and symbol_column and symbol_column in names:
+        column_type = schema.field(symbol_column).type
+        if pa.types.is_dictionary(column_type):
             complete = False
-
-    if (start is not None or end is not None) and date_column and date_column in table.column_names:
-        column = table[date_column]
-        try:
-            if start is not None:
-                low = _arrow_bound(column.type, start)
-                if low is None:
-                    raise pa.ArrowNotImplementedError(str(column.type))
-                _and(pc.greater_equal(column, low))
-            if end is not None:
-                high = _arrow_bound(column.type, end + _dt.timedelta(days=1))
-                if high is None:
-                    raise pa.ArrowNotImplementedError(str(column.type))
-                _and(pc.less(column, high))
-        except (pa.ArrowInvalid, pa.ArrowNotImplementedError, pa.ArrowTypeError):
+        else:
+            try:
+                value_set = pa.array(list(symbols)).cast(column_type)
+                _and(pc.field(symbol_column).isin(value_set))
+            except (pa.ArrowInvalid, pa.ArrowNotImplementedError, pa.ArrowTypeError):
+                complete = False
+
+    if (start is not None or end is not None) and date_column and date_column in names:
+        column_type = schema.field(date_column).type
+        low = _arrow_bound(column_type, start) if start is not None else None
+        high = _arrow_bound(column_type, end + _dt.timedelta(days=1)) if end is not None else None
+        if (start is not None and low is None) or (end is not None and high is None):
             complete = False
+        else:
+            if low is not None:
+                _and(pc.field(date_column) >= low)
+            if high is not None:
+                _and(pc.field(date_column) < high)
+
+    return expression, complete
+
+
+def _scan_filtered(
+    path: Path,
+    fmt: DataFormat,
+    *,
+    date_column: str | None,
+    start: _dt.date | None,
+    end: _dt.date | None,
+    symbol_column: str | None,
+    symbols: Sequence[str] | None,
+    memory_pool: Any = None,
+) -> tuple[Any, bool] | None:
+    """Scan an Arrow-format export with the filter pushed into the scan.
+
+    The file is read record batch by record batch and each batch is
+    filtered as it is decoded, so peak memory is the *matching* rows plus
+    one batch — never the whole export. That is what keeps a narrow
+    ``start`` / ``end`` cheap on a multi-GB file, on the first read and on
+    every different window after it.
+
+    Returns ``(table, complete)``, or ``None`` when the file cannot be
+    scanned as a dataset (Feather V1, which the per-pair FX exports use, is
+    not an Arrow IPC file) and the caller must fall back to a full read.
+    """
 
-    if mask is not None:
-        # Null comparisons yield null; treat them as "not selected".
-        table = table.filter(pc.fill_null(mask, False))
+    try:
+        dataset = ds.dataset(str(path), format="ipc" if fmt is DataFormat.FEATHER else "parquet")
+    except (pa.ArrowInvalid, pa.ArrowNotImplementedError, OSError):
+        return None
+    expression, complete = _filter_expression(
+        dataset.schema,
+        date_column=date_column,
+        start=start,
+        end=end,
+        symbol_column=symbol_column,
+        symbols=symbols,
+    )
+    try:
+        table = dataset.to_table(filter=expression, memory_pool=memory_pool)
+    except (pa.ArrowInvalid, pa.ArrowNotImplementedError, pa.ArrowTypeError):
+        # An expression the scanner rejects for this file: read it all and
+        # let pandas filter instead.
+        table = dataset.to_table(memory_pool=memory_pool)
+        complete = False
     return table, complete
 
 
+def release_arrow_memory() -> None:
+    """Return Arrow's freed-but-cached memory to the operating system.
+
+    Arrow's allocator (jemalloc / mimalloc) keeps freed memory for reuse, so
+    after reading a large export a process's RSS stays near its peak even
+    once every table is gone — and a long-running worker that reads export
+    after export ratchets upward until it hits its memory limit. Called
+    after every read of a cached export.
+    """
+
+    with contextlib.suppress(AttributeError, NotImplementedError):
+        pa.default_memory_pool().release_unused()
+
+
 def _read_arrow_table(path: Path, fmt: DataFormat, columns: Sequence[str] | None = None) -> Any:
     cols = list(columns) if columns is not None else None
     if fmt is DataFormat.FEATHER:
@@ -419,25 +483,38 @@ def read_frame(
 
     With no filters this is :func:`read_auto`. ``end`` is inclusive of the
     whole day, so an intraday column keeps every bar on ``end``.
+
+    Feather (Arrow IPC) and parquet exports are scanned with the filter
+    pushed down (see :func:`_scan_filtered`), so memory scales with the
+    rows returned rather than with the size of the export.
     """
 
     fmt = detect_format(object_name)
     filtering = bool(symbols) or start is not None or end is not None
-    if not filtering:
-        return read(path, fmt, max_bytes=max_bytes)
-    _check_memory_guard(path, max_bytes=max_bytes)
-    kwargs: dict[str, Any] = {
-        "date_column": date_column,
-        "start": start,
-        "end": end,
-        "symbol_column": symbol_column,
-        "symbols": symbols,
-    }
-    if fmt in (DataFormat.FEATHER, DataFormat.PARQUET):
-        table, complete = _filter_table(_read_arrow_table(path, fmt), **kwargs)
-        frame = _arrow_to_df(table)
-        return frame if complete else filter_frame(frame, **kwargs)
-    return filter_frame(_BY_FORMAT[fmt].reader(path), **kwargs)
+    try:
+        if not filtering:
+            return read(path, fmt, max_bytes=max_bytes)
+        _check_memory_guard(path, max_bytes=max_bytes)
+        kwargs: dict[str, Any] = {
+            "date_column": date_column,
+            "start": start,
+            "end": end,
+            "symbol_column": symbol_column,
+            "symbols": symbols,
+        }
+        if fmt in (DataFormat.FEATHER, DataFormat.PARQUET):
+            scanned = _scan_filtered(path, fmt, **kwargs)
+            if scanned is None:  # Feather V1: no batch index, read it whole
+                table = _read_arrow_table(path, fmt)
+                complete = False
+            else:
+                table, complete = scanned
+            frame = _arrow_to_df(table)
+            del table
+            return frame if complete else filter_frame(frame, **kwargs)
+        return filter_frame(_BY_FORMAT[fmt].reader(path), **kwargs)
+    finally:
+        release_arrow_memory()
 
 
 def max_timestamp(path: Path, *, object_name: str, column: str) -> pd.Timestamp | None:
@@ -455,6 +532,8 @@ def max_timestamp(path: Path, *, object_name: str, column: str) -> pd.Timestamp
         except (KeyError, pa.ArrowInvalid):
             return None
         series = _arrow_to_df(table)[column]
+        del table
+        release_arrow_memory()
     else:
         compression: CsvCompression = "gzip" if fmt is DataFormat.CSV_GZ else None
         try:
@@ -535,5 +614,6 @@ def read_csv_chunks(
     "read_frame",
     "read_head",
     "register_dtype_hints",
+    "release_arrow_memory",
     "to_timestamps",
 ]

tests/test_session.py

@@ -690,3 +690,115 @@ def test_snapshot_export_is_served_from_cache_without_replanning_a_download(
     assert not live.called
     assert again.route.bulk_action is not None and again.route.bulk_action.value == "cache"
     assert "no date watermark" in again.route.reason
+
+
+# ---- results are never shared between calls ------------------------------------------
+#
+# Each get() must return data read fresh for that call. Any in-memory reuse of
+# a parsed export (e.g. an LRU of DataFrames) has to keep these passing: the
+# seam refresh writes live values into the bulk half in place, and callers
+# are free to mutate the frame they are handed.
+
+
+def _mock_provisional_prices(respx_mock: respx.MockRouter) -> None:
+    table = pa.table(
+        {
+            "ticker": ["HBM", "HBM"],
+            "date": pa.array([dt.date(2026, 9, 21), dt.date(2026, 9, 22)], pa.date32()),
+            "close": [27.0, 27.5],
+            "volume": [4_000_000.0, 3_176_000.0],
+        }
+    )
+    buf = io.BytesIO()
+    pyarrow.feather.write_feather(table, buf)
+    _mock_file(
+        respx_mock,
+        namespace="equities",
+        dataset="statarb",
+        schema="prices",
+        bucket="equity_factors_research_pod",
+        obj="statarb/prices.feather",
+        payload=buf.getvalue(),
+        fmt="feather",
+    )
+    respx_mock.get("/v1/equities/statarb/prices").respond(
+        json=_page([{"ticker": "HBM", "date": "2026-09-22", "close": 27.74, "volume": 5_836_000.0}])
+    )
+
+
+@respx.mock(base_url=FAKE_API, assert_all_called=False)
+def test_a_bulk_only_read_is_never_changed_by_an_earlier_live_top_up(
+    respx_mock: respx.MockRouter, client: Client
+) -> None:
+    _mock_provisional_prices(respx_mock)
+    topped_up = client.get("statarb_prices")
+    assert topped_up.seam_updates == 1  # live values were applied to that result...
+
+    # ...but the export itself, read again for the same window, is untouched.
+    bulk = client.get("statarb_prices", end="2026-09-22", source="bulk").df
+    assert bulk["volume"].tolist() == [4_000_000.0, 3_176_000.0]
+    assert bulk["close"].tolist() == [27.0, 27.5]
+
+
+@respx.mock(base_url=FAKE_API, assert_all_called=False)
+def test_mutating_a_result_never_changes_the_next_one(
+    respx_mock: respx.MockRouter, client: Client
+) -> None:
+    _mock_provisional_prices(respx_mock)
+    first = client.get("statarb_prices", source="bulk")
+    first.df.loc[0, "close"] = -1.0
+    first.df["extra"] = 1
+    first.df.sort_values("close", ascending=False, inplace=True)
+
+    second = client.get("statarb_prices", source="bulk")
+    assert second.df["close"].tolist() == [27.0, 27.5]
+    assert "extra" not in second.df.columns
+    assert second.df is not first.df
+
+
+@respx.mock(base_url=FAKE_API, assert_all_called=False)
+def test_windowed_gets_scan_the_export_instead_of_loading_it(
+    respx_mock: respx.MockRouter, client: Client, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """End to end: a year-by-year backfill through Client.get() never reads a
+    whole non-templated export into memory (only the one-column watermark
+    scan touches every row)."""
+
+    import rwpytools.formats as formats_module
+
+    days = pd.bdate_range("2015-01-01", "2018-12-31")
+    table = pa.table(
+        {
+            "ticker": ["AAPL"] * len(days) + ["MSFT"] * len(days),
+            "date": pa.array(list(days.date) * 2, pa.date32()),
+            "close": list(range(len(days) * 2)),
+        }
+    )
+    buf = io.BytesIO()
+    pyarrow.feather.write_feather(table, buf, chunksize=100)
+    _mock_file(
+        respx_mock,
+        namespace="equities",
+        dataset="equity.factors",
+        schema="prices",
+        bucket="equity_factors_research_pod",
+        obj="equity_factors/prices.feather",
+        payload=buf.getvalue(),
+        fmt="feather",
+    )
+    real = formats_module._read_arrow_table
+
+    def _only_single_columns(path: Any, fmt: Any, columns: Any = None) -> Any:
+        assert columns is not None and len(columns) == 1, "a whole export was loaded"
+        return real(path, fmt, columns)
+
+    monkeypatch.setattr(formats_module, "_read_arrow_table", _only_single_columns)
+
+    rows = 0
+    for year in range(2015, 2019):
+        result = client.get(
+            "liquid_universe_prices", start=f"{year}-01-01", end=f"{year}-12-31", source="bulk"
+        )
+        assert set(pd.to_datetime(result.df["date"]).dt.year) == {year}
+        rows += len(result)
+    assert rows == len(days) * 2