[Fix] add filter to Polars read parquet + skip BufferPool eviction of unknown objects

stratum/optimizer/physical/_source_execs.py

@@ -10,6 +10,8 @@
 """
 from __future__ import annotations
 
+import json
+
 import numpy as np
 import pandas as pd
 import polars as pl
@@ -63,17 +65,70 @@ def rechunk_pl_frame(df, rows_per_chunk=128_000):
     "comment": "comment_prefix",
     "quotechar": "quote_char",
     "low_memory": "low_memory",
-    "storage_options": "storage_options",
 }
 
 #: pandas parquet option -> polars parquet option. polars takes no positional
 #: options at all, hence the ``read_args`` guard in ``_translate_read_options``.
 _POLARS_PARQUET_RENAMES = {
     "columns": "columns",
-    "storage_options": "storage_options",
 }
 
 
+#: fsspec (s3fs) credential key -> polars object-store key, same value.
+_OBJECT_STORE_KEYS = {
+    "key": "aws_access_key_id",
+    "secret": "aws_secret_access_key",
+    "endpoint_url": "aws_endpoint_url",
+}
+
+
+def _translate_storage_options(value):
+    """pandas ``storage_options`` (fsspec) -> polars ``storage_options`` (object store).
+
+    pandas hands these to fsspec (s3fs, gcsfs); polars reads cloud paths with
+    its own object store, whose keys differ. An untranslated fsspec key is not an
+    error there: it is ignored, the request goes out unauthenticated, and the
+    read fails later with a 401. The common credential spellings are mapped;
+    every other key is passed on unchanged, so options already written in
+    object-store spelling keep working.
+    """
+    if value is None:
+        return "storage_options", None
+    translated = {}
+    for name, option in value.items():
+        if name == "token":  # gcsfs
+            if isinstance(option, dict):
+                translated["service_account_key"] = json.dumps(option)
+            elif option == "anon":
+                translated["skip_signature"] = "true"
+            elif option not in (None, "google_default", "cloud"):
+                translated["service_account"] = str(option)
+            # else: application default credentials, the object store's default
+        elif name == "anon":  # s3fs
+            if option:
+                translated["skip_signature"] = "true"
+        else:
+            translated[_OBJECT_STORE_KEYS.get(name, name)] = option
+    return "storage_options", translated
+
+
+def _polars_cloud_path(file_path):
+    """Spell a cloud dataset directory the way polars lists it.
+
+    pyarrow treats ``gs://bucket/table`` as a dataset directory when no object
+    has that exact name; polars only lists a prefix that ends in ``/`` and
+    otherwise asks for the single object, which does not exist. A cloud path
+    without a file extension or glob is taken to be such a directory.
+    """
+    if (not isinstance(file_path, str) or "://" not in file_path
+            or file_path.startswith("file://") or file_path.endswith("/")
+            or any(ch in file_path for ch in "*?[")):
+        return file_path
+    if "." in file_path.rsplit("/", 1)[-1]:
+        return file_path
+    return file_path + "/"
+
+
 def _translate_header(value):
     """pandas ``header`` -> polars ``has_header``.
 
@@ -106,11 +161,58 @@ def _translate_encoding(value):
         f"'utf8-lossy' only).")
 
 
+#: pyarrow ``filters`` operator -> polars predicate on (column expr, value).
+_PYARROW_FILTER_OPS = {
+    "=": lambda col, v: col == v,
+    "==": lambda col, v: col == v,
+    "!=": lambda col, v: col != v,
+    "<": lambda col, v: col < v,
+    ">": lambda col, v: col > v,
+    "<=": lambda col, v: col <= v,
+    ">=": lambda col, v: col >= v,
+    "in": lambda col, v: col.is_in(list(v)),
+    "not in": lambda col, v: ~col.is_in(list(v)),
+}
+
+
+def _translate_filters(value):
+    """pandas ``filters`` (pyarrow DNF) -> one polars predicate expression.
+
+    pyarrow accepts a list of ``(column, op, value)`` tuples, ANDed, or a list of
+    such lists, ORed -- told apart the way pyarrow does, by whether the first
+    element's first item is a column name. The predicate filters rows just like
+    pyarrow does; :class:`PolarsReadParquet` applies it to a ``scan_parquet``,
+    so hive partitions and row groups it excludes are never read. A
+    ``pyarrow.compute.Expression`` filter has no polars spelling and is rejected.
+    """
+    if value is None:
+        return "filters", None
+    if not isinstance(value, (list, tuple)) or not value:
+        raise ValueError(
+            f"filters={value!r} has no polars equivalent: only pyarrow's "
+            f"list-of-tuples (DNF) form translates.")
+    disjuncts = [value] if isinstance(value[0][0], str) else value
+    predicate = None
+    for conjunct in disjuncts:
+        term = None
+        for clause in conjunct:
+            if len(clause) != 3 or clause[1] not in _PYARROW_FILTER_OPS:
+                raise ValueError(
+                    f"filter {clause!r} has no polars equivalent: expected "
+                    f"(column, op, value) with op in {sorted(_PYARROW_FILTER_OPS)}.")
+            column, op, operand = clause
+            expr = _PYARROW_FILTER_OPS[op](pl.col(column), operand)
+            term = expr if term is None else term & expr
+        predicate = term if predicate is None else predicate | term
+    return "filters", predicate
+
+
 #: pandas option -> callable(value) -> (polars option, polars value), for the
 #: options that need the value rewritten and not just the name.
 _POLARS_CSV_VALUE_TRANSLATORS = {
     "header": _translate_header,
     "encoding": _translate_encoding,
+    "storage_options": _translate_storage_options,
 }
 
 
@@ -235,7 +337,8 @@ class PolarsReadParquet(ReadParquet):
     is_abstract = False
 
     _RENAMES = _POLARS_PARQUET_RENAMES
-    _VALUE_TRANSLATORS = {}
+    _VALUE_TRANSLATORS = {"filters": _translate_filters,
+                          "storage_options": _translate_storage_options}
 
     def on_impl_selected(self, ctx) -> None:
         _check_read_options(self.read_args, self.read_kwargs, self._RENAMES,
@@ -246,7 +349,17 @@ def process(self, mode: str, inputs: list):
         read_kwargs = _translate_read_options(
             read_args, read_kwargs, self._RENAMES, self._VALUE_TRANSLATORS,
             type(self).__name__)
-        return pl.read_parquet(file_path, **read_kwargs)
+        file_path = _polars_cloud_path(file_path)
+        predicate = read_kwargs.pop("filters", None)
+        if predicate is None:
+            return pl.read_parquet(file_path, **read_kwargs)
+        # Filter before projecting: like pyarrow, a filter may name a column
+        # (e.g. the hive partition key) that `columns` leaves out.
+        columns = read_kwargs.pop("columns", None)
+        frame = pl.scan_parquet(file_path, **read_kwargs).filter(predicate)
+        if columns is not None:
+            frame = frame.select(columns)
+        return frame.collect()
 
 
 class NumpyLoad(FileReadOp):

stratum/runtime/_buffer_pool.py

@@ -12,6 +12,7 @@
 from stratum._config import FLAGS
 from stratum.runtime._object_serialization import (
     SpilledObject,
+    UnsupportedSpillType,
     delete_object,
     deserialize_object,
     serialize_object,
@@ -87,6 +88,7 @@ class _Entry:
     data: Any = None
     handle: SpilledObject | None = None
     pin_count: int = 0
+    spillable: bool = True  # False once serialization refused the data
 
     @property
     def is_spilled(self) -> bool:
@@ -259,14 +261,22 @@ def check_for_eviction(self) -> bool:
             if self.memory_usage <= self.memory_budget:
                 break
             entry = self.live_variable_map[key]
-            if entry.pin_count > 0 or entry.is_spilled:
+            if entry.pin_count > 0 or entry.is_spilled or not entry.spillable:
+                continue
+            try:
+                self._evict(key, entry)
+            except UnsupportedSpillType as exc:
+                # Admitted but not spillable (e.g. a CV splitter object): keep it
+                # resident and try the next LRU entry instead of failing the plan.
+                entry.spillable = False
+                logger.debug(f"Keeping {key} in memory, it cannot be spilled: {exc}")
                 continue
-            self._evict(key, entry)
             evicted = True
         if self.memory_usage > self.memory_budget:
             logger.warning(
                 f"Budget {prettify_bytes(self.memory_budget)} still exceeded after eviction "
-                f"(usage {prettify_bytes(self.memory_usage)}); all remaining entries are pinned."
+                f"(usage {prettify_bytes(self.memory_usage)}); all remaining entries are "
+                f"pinned or cannot be spilled."
             )
         return evicted
 

stratum/runtime/_object_serialization.py

@@ -41,6 +41,14 @@
 _PICKLABLE_PRIMITIVE = (str, int, float, bool, bytes)
 
 
+class UnsupportedSpillType(ValueError):
+    """The object (or something inside it) has no spill format.
+
+    The buffer pool admits such objects (it sizes them as UNKNOWN_SIZE), so it
+    catches this and keeps the entry in memory rather than failing the plan.
+    """
+
+
 def _native_format(obj: Any) -> str | None:
     """Return the portable file format for a structured leaf, else None."""
     for typ, fmt in _STRUCTURED_FORMAT:
@@ -113,9 +121,16 @@ def build(node: Any) -> Any:
             return {k: build(v) for k, v in node.items()}
         if isinstance(node, _PICKLABLE_PRIMITIVE) or node is None:
             return node  # primitive: inline in the skeleton
-        raise ValueError(f"Unsupported type for serialization: {type(node)}")
-
-    skeleton = build(obj)
+        raise UnsupportedSpillType(f"Unsupported type for serialization: {type(node)}")
+
+    try:
+        skeleton = build(obj)
+    except UnsupportedSpillType:
+        # Leaves written before the unsupported node was reached would be
+        # orphaned: no handle will ever point at them.
+        for ref in leaves:
+            Path(ref.path).unlink(missing_ok=True)
+        raise
     # The skeleton is pickled, so it is Python-only. That is a deliberate
     # trade-off: the heavy data (frames/arrays) lives in portable native files
     # above, which is what another language would actually consume. If a

tests/optimizer/physical/test_source_execs.py

@@ -8,6 +8,8 @@
 backend, and re-executing either is stable, which is only true if no ``process``
 consults anything outside the op.
 """
+import shutil
+import tempfile
 import unittest
 
 import numpy as np
@@ -23,7 +25,8 @@
 from stratum.optimizer.physical._source_execs import (
     InMemoryFrame, NumpyLoad, PandasInMemoryFrame, PandasReadCSV,
     PandasReadParquet, PolarsInMemoryFrame, PolarsReadCSV, PolarsReadParquet,
-    ReadCSV, ReadParquet, lower_data_source)
+    ReadCSV, ReadParquet, _polars_cloud_path, _translate_storage_options,
+    lower_data_source)
 from stratum.runtime._buffer_pool import BufferPool
 from tests._helpers import csv_file, npy_file, parquet_file
 
@@ -308,6 +311,48 @@ def test_untranslatable_value_is_rejected(self):
             with self.assertRaises(ValueError):
                 op.process("fit_transform", [])
 
+    def test_parquet_filters_select_the_plan_backend(self):
+        # pyarrow `filters` used to fail the whole polars plan at selection time.
+        with parquet_file(self.df) as path:
+            data = st.as_data_op(path).skb.apply_func(
+                pd.read_parquet, filters=[("x", ">", 1)])
+            ops, *_ = optimize(data, polars_conf(dataframe_ops=True))
+        self.assertIsInstance(ops[-1], PolarsReadParquet)
+
+    def test_untranslatable_parquet_filter_is_rejected(self):
+        op = PolarsReadParquet(file_path="x.parquet",
+                               read_kwargs={"filters": [("x", "like", "a%")]})
+        with self.assertRaises(ValueError):
+            op.process("fit_transform", [])
+
+    def test_fsspec_storage_options_become_object_store_keys(self):
+        # Passed verbatim, polars ignored gcsfs' `token` and read the bucket
+        # unauthenticated (401).
+        cases = [
+            ({"token": "/keys/sa.json"}, {"service_account": "/keys/sa.json"}),
+            ({"token": {"type": "service_account"}},
+             {"service_account_key": '{"type": "service_account"}'}),
+            ({"token": "anon"}, {"skip_signature": "true"}),
+            ({"token": "google_default"}, {}),
+            ({"key": "k", "secret": "s", "anon": False},
+             {"aws_access_key_id": "k", "aws_secret_access_key": "s"}),
+            ({"aws_region": "eu-west-1"}, {"aws_region": "eu-west-1"}),
+        ]
+        for pandas_options, polars_options in cases:
+            with self.subTest(pandas_options):
+                self.assertEqual(("storage_options", polars_options),
+                                 _translate_storage_options(pandas_options))
+
+    def test_cloud_dataset_directory_gets_trailing_slash(self):
+        for path, expected in [
+                ("gs://bucket/lake/table", "gs://bucket/lake/table/"),
+                ("s3://bucket/table/", "s3://bucket/table/"),
+                ("gs://bucket/table/part-0.parquet", "gs://bucket/table/part-0.parquet"),
+                ("gs://bucket/table/**/*.parquet", "gs://bucket/table/**/*.parquet"),
+                ("/local/table", "/local/table")]:
+            with self.subTest(path):
+                self.assertEqual(expected, _polars_cloud_path(path))
+
     def test_positional_read_options_are_rejected(self):
         # polars takes only the path positionally, so a pandas positional
         # option (pd.read_parquet's `engine`) would bind to something else.
@@ -316,5 +361,46 @@ def test_positional_read_options_are_rejected(self):
             op.process("fit_transform", [])
 
 
+class TestPolarsParquetFilters(unittest.TestCase):
+    """``PolarsReadParquet`` honours pyarrow ``filters`` exactly like
+    ``pd.read_parquet``, on a hive-partitioned dataset."""
+
+    def setUp(self):
+        self.root = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, self.root, ignore_errors=True)
+        pd.DataFrame({"year": [2019, 2020, 2021, 2022] * 5, "x": range(20),
+                      "kind": list("abcde") * 4}).to_parquet(
+            self.root, partition_cols=["year"])
+
+    def assert_matches_pandas(self, **options):
+        expected = pd.read_parquet(self.root, **options)
+        result = PolarsReadParquet(file_path=self.root, read_kwargs=options).process(
+            "fit_transform", [])
+        columns = options.get("columns") or ["x", "kind"]
+        self.assertEqual(columns, [c for c in result.columns if c in columns])
+        key = ["x"] if "x" in columns else columns
+        self.assertEqual(
+            expected[columns].sort_values(key).values.tolist(),
+            [list(row) for row in result.select(columns).sort(key).rows()])
+
+    def test_conjunction(self):
+        self.assert_matches_pandas(filters=[("year", ">=", 2020), ("x", "<", 15)])
+
+    def test_disjunction_of_conjunctions(self):
+        self.assert_matches_pandas(filters=[[("year", "==", 2019)],
+                                            [("year", "=", 2022), ("kind", "!=", "c")]])
+
+    def test_membership(self):
+        self.assert_matches_pandas(filters=[("kind", "in", ["a", "b"]),
+                                            ("year", "not in", {2021})])
+
+    def test_filter_on_column_left_out_of_projection(self):
+        self.assert_matches_pandas(columns=["kind", "x"],
+                                   filters=[("year", "<=", 2020)])
+
+    def test_none_filters_read_everything(self):
+        self.assert_matches_pandas(filters=None)
+
+
 if __name__ == "__main__":
     unittest.main()

tests/runtime/test_buffer_pool.py

@@ -254,6 +254,26 @@ def test_pinned_entries_not_evicted(self):
         # b is next LRU among unpinned, so b should have been spilled instead.
         self.assertTrue(self.pool.live_variable_map[b].is_spilled)
 
+    def test_unspillable_entries_stay_resident(self):
+        # The pool admits objects the serializer has no format for (a CV
+        # splitter, say); evicting one used to crash the plan. It stays in
+        # memory and the next LRU entry is spilled instead.
+        class Opaque:
+            pass
+
+        opaque, holder, a, b, c = (_make_op(n) for n in ("opaque", "holder", "a", "b", "c"))
+        self.pool.put(opaque, Opaque())
+        self.pool.put(holder, [_arr(400), Opaque()])  # 3_200 B, spill fails mid-way
+        self.pool.put(a, _arr(500))
+        self.pool.put(b, _arr(500))  # 11_201 B, over budget
+        self.assertFalse(self.pool.live_variable_map[opaque].is_spilled)
+        self.assertFalse(self.pool.live_variable_map[holder].is_spilled)
+        self.assertTrue(self.pool.live_variable_map[a].is_spilled)
+        # Only a's spill is on disk: the holder's half-written leaf is cleaned up.
+        self.assertEqual(1, len(list(self.spill_root.iterdir())))
+        self.pool.put(c, _arr(500))
+        self.assertTrue(self.pool.live_variable_map[b].is_spilled)
+
     def test_pin_refcount(self):
         a = _make_op("a")
         self.pool.put(a, "data")

tests/runtime/test_object_serialization.py

@@ -152,6 +152,9 @@ class Foo:
 
         with self.assertRaises(ValueError):
             serialize_object([pd.DataFrame({"a": [1]}), Foo()], self._stem())
+        # The frame's leaf was written before Foo was reached; nothing may be
+        # left behind for a spill that never happened.
+        self.assertEqual(list(self.root.iterdir()), [])
 
     def test_delete_removes_skeleton_and_leaves(self):
         h = serialize_object((np.arange(3), np.arange(4)), self._stem())