lazynwb 1.0.0.dev2__tar.gz → 1.0.0.dev3__tar.gz

{lazynwb-1.0.0.dev2/src/lazynwb.egg-info → lazynwb-1.0.0.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lazynwb
-Version: 1.0.0.dev2
+Version: 1.0.0.dev3
 Summary: An attempt to speed-up access to large NWB (Neurodata Without Borders) files stored in the cloud.
 Author-email: Ben Hardcastle <ben.hardcastle@alleninstitue.org>
 License: MIT

{lazynwb-1.0.0.dev2 → lazynwb-1.0.0.dev3}/pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "aiosqlite>=0.20.0",
     "tomli>=2.0.0; python_version < '3.11'",
 ]
-version = "1.0.0.dev2"
+version = "1.0.0.dev3"
 classifiers = [
     "Development Status :: 3 - Alpha",          # https://pypi.org/classifiers/
     "Programming Language :: Python :: 3",

{lazynwb-1.0.0.dev2 → lazynwb-1.0.0.dev3}/src/lazynwb/lazyframe.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+import functools
+import json
 import logging
+import operator
 from collections.abc import Iterable, Iterator, Mapping
 
 import polars as pl
@@ -14,6 +17,8 @@ import lazynwb.types_
 
 logger = logging.getLogger(__name__)
 
+_POLARS_DYNAMIC_PREDICATE_TOKEN = "dynamic_pred"
+
 
 def scan_nwb(
     source: lazynwb.types_.PathLike | Iterable[lazynwb.types_.PathLike],
@@ -133,6 +138,17 @@ def scan_nwb(
         else:
             logger.debug(f"Batch size set to {batch_size} rows per batch")
 
+        predicate, dynamic_predicate_count = _remove_polars_dynamic_predicates(
+            predicate
+        )
+        if dynamic_predicate_count:
+            logger.debug(
+                "Removed %d Polars dynamic predicate conjunct(s) from %r scan "
+                "predicate; Python IO plugins do not receive the dynamic TopK state",
+                dynamic_predicate_count,
+                table_path,
+            )
+
         if predicate is not None:
             # - if we have a predicate, we'll fetch the minimal df, apply predicate, then fetch remaining columns in with_columns
             initial_columns = predicate.meta.root_names()
@@ -282,6 +298,136 @@ def scan_nwb(
     )
 
 
+def _remove_polars_dynamic_predicates(
+    predicate: pl.Expr | None,
+) -> tuple[pl.Expr | None, int]:
+    if predicate is None or not _predicate_contains_polars_dynamic_predicate(
+        predicate
+    ):
+        return predicate, 0
+
+    conjuncts = _split_conjunctive_predicate(predicate)
+    retained_conjuncts: list[pl.Expr] = []
+    dropped_conjunct_count = 0
+    for conjunct in conjuncts:
+        if not _predicate_contains_polars_dynamic_predicate(conjunct):
+            retained_conjuncts.append(conjunct)
+            continue
+        if _is_standalone_polars_dynamic_predicate(conjunct):
+            dropped_conjunct_count += 1
+            continue
+        logger.debug(
+            "Cannot isolate Polars dynamic predicate from pushed predicate %s; "
+            "leaving predicate unchanged",
+            conjunct,
+        )
+        return predicate, 0
+
+    if dropped_conjunct_count == 0:
+        return predicate, 0
+    if not retained_conjuncts:
+        return None, dropped_conjunct_count
+    return functools.reduce(operator.and_, retained_conjuncts), dropped_conjunct_count
+
+
+def _split_conjunctive_predicate(predicate: pl.Expr) -> list[pl.Expr]:
+    if not _is_binary_and_predicate(predicate):
+        return [predicate]
+
+    conjuncts: list[pl.Expr] = []
+    try:
+        children = predicate.meta.pop()
+    except BaseException as exc:
+        if not _is_polars_panic_exception(exc):
+            raise
+        logger.debug(
+            "Polars panicked while splitting pushed predicate %s; leaving it intact",
+            predicate,
+        )
+        return [predicate]
+
+    for child in children:
+        conjuncts.extend(_split_conjunctive_predicate(child))
+    return conjuncts
+
+
+def _is_binary_and_predicate(predicate: pl.Expr) -> bool:
+    payload = _predicate_json_payload(predicate)
+    if not isinstance(payload, dict):
+        return False
+    binary_expr = payload.get("BinaryExpr")
+    return isinstance(binary_expr, dict) and binary_expr.get("op") == "And"
+
+
+def _is_standalone_polars_dynamic_predicate(predicate: pl.Expr) -> bool:
+    payload = _predicate_json_payload(predicate)
+    if not isinstance(payload, dict):
+        return False
+    display_expr = payload.get("Display")
+    if not isinstance(display_expr, dict):
+        return False
+    fmt_str = display_expr.get("fmt_str")
+    return isinstance(fmt_str, str) and fmt_str.startswith(
+        _POLARS_DYNAMIC_PREDICATE_TOKEN
+    )
+
+
+def _predicate_contains_polars_dynamic_predicate(predicate: pl.Expr) -> bool:
+    predicate_text = str(predicate)
+    if _POLARS_DYNAMIC_PREDICATE_TOKEN not in predicate_text:
+        return False
+
+    payload = _predicate_json_payload(predicate)
+    if payload is None:
+        return True
+    return _json_payload_contains_polars_dynamic_predicate(payload)
+
+
+def _predicate_json_payload(predicate: pl.Expr) -> object | None:
+    try:
+        serialized = predicate.meta.serialize(format="json")
+    except BaseException as exc:
+        if not _is_polars_panic_exception(exc):
+            raise
+        logger.debug(
+            "Polars panicked while serializing pushed predicate %s to JSON",
+            predicate,
+        )
+        return None
+
+    if not isinstance(serialized, str):
+        return None
+    try:
+        return json.loads(serialized)
+    except json.JSONDecodeError:
+        logger.debug("Could not decode pushed predicate JSON: %s", serialized)
+        return None
+
+
+def _json_payload_contains_polars_dynamic_predicate(payload: object) -> bool:
+    if isinstance(payload, dict):
+        display_expr = payload.get("Display")
+        if isinstance(display_expr, dict):
+            fmt_str = display_expr.get("fmt_str")
+            if isinstance(fmt_str, str) and fmt_str.startswith(
+                _POLARS_DYNAMIC_PREDICATE_TOKEN
+            ):
+                return True
+        return any(
+            _json_payload_contains_polars_dynamic_predicate(value)
+            for value in payload.values()
+        )
+    if isinstance(payload, list):
+        return any(
+            _json_payload_contains_polars_dynamic_predicate(value) for value in payload
+        )
+    return False
+
+
+def _is_polars_panic_exception(exc: BaseException) -> bool:
+    return type(exc).__name__ == "PanicException"
+
+
 def _get_limited_path_to_row_indices(
     source: Iterable[lazynwb.types_.PathLike],
     table_path: str,

{lazynwb-1.0.0.dev2 → lazynwb-1.0.0.dev3/src/lazynwb.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lazynwb
-Version: 1.0.0.dev2
+Version: 1.0.0.dev3
 Summary: An attempt to speed-up access to large NWB (Neurodata Without Borders) files stored in the cloud.
 Author-email: Ben Hardcastle <ben.hardcastle@alleninstitue.org>
 License: MIT

{lazynwb-1.0.0.dev2 → lazynwb-1.0.0.dev3}/tests/test_lazyframe.py

@@ -4,16 +4,16 @@ import tempfile
 import uuid
 from datetime import datetime, timezone
 
+import h5py
 import polars as pl
 import pytest
-
-import h5py
 from pynwb import NWBHDF5IO, NWBFile
 from pynwb.epoch import TimeIntervals
 
 import lazynwb
 import lazynwb.tables
 
+
 def test_polars_dtype_inference(local_hdf5_path):
     schema = lazynwb.tables.get_table_schema(
         local_hdf5_path,
@@ -160,6 +160,41 @@ def test_scan_nwb_predicate_pushdown(local_hdf5_path):
     ).all(), "Predicate pushdown filter on internal column was not applied correctly"
     assert len(filtered_internal_df) == len(lf.collect().filter(internal_expr)), "Filtered DataFrame length does not match length when collecting and filtering separately"
 
+
+def test_scan_nwb_sort_head_ignores_polars_dynamic_topk_predicate(
+    local_hdf5_path: pathlib.Path,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Polars 1.39+ may push optimizer-only dynamic TopK predicates into IO plugins."""
+    caplog.set_level(logging.DEBUG, logger="lazynwb.lazyframe")
+
+    lf = lazynwb.scan_nwb(
+        source=local_hdf5_path,
+        table_path="/units",
+        disable_progress=True,
+    )
+    query = (
+        lf.filter(pl.col("structure") == "VISp")
+        .sort(lazynwb.TABLE_INDEX_COLUMN_NAME, descending=True)
+        .head(1)
+        .select(lazynwb.TABLE_INDEX_COLUMN_NAME, "spike_times")
+    )
+    plan = query.explain()
+
+    result = query.collect()
+    expected = (
+        lf.collect()
+        .filter(pl.col("structure") == "VISp")
+        .sort(lazynwb.TABLE_INDEX_COLUMN_NAME, descending=True)
+        .head(1)
+        .select(lazynwb.TABLE_INDEX_COLUMN_NAME, "spike_times")
+    )
+
+    assert result.equals(expected)
+    if "dynamic_predicate" in plan:
+        assert "Removed 1 Polars dynamic predicate conjunct" in caplog.text
+
+
 def test_scan_nwb_raises_on_missing(local_hdf5_path):
     """Test that scan_nwb raises an error when the table is not found and raise_on_missing=True."""
     # Try with a non-existent table path