From bd9a0c14f221bfb46068ef4fb42c03da4bd5b861 Mon Sep 17 00:00:00 2001 From: Elias Strauss Date: Tue, 6 Oct 2026 00:31:56 +0200 Subject: [PATCH 1/2] [Fix] Read pandas-spelled parquet sources with polars under greedy PolarsReadParquet only forwarded `columns` and `storage_options`, so a pd.read_parquet source that works on pandas broke under implementation_selector="greedy": - `filters=` (pyarrow DNF) failed the whole plan at selection time. It is now translated to one polars predicate and applied to a scan_parquet before projecting `columns`, so pruned hive partitions are never read. - fsspec `storage_options` (gcsfs `token`, s3fs `key`/`secret`/`anon`) were ignored by polars' object store, so cloud reads went out unauthenticated (401). They are now mapped to object-store keys; other keys pass through. - A cloud dataset directory without a trailing slash was looked up as a single object (404). Such paths now get the slash polars needs to list them. --- stratum/optimizer/physical/_source_execs.py | 121 +++++++++++++++++- tests/optimizer/physical/test_source_execs.py | 88 ++++++++++++- 2 files changed, 204 insertions(+), 5 deletions(-) diff --git a/stratum/optimizer/physical/_source_execs.py b/stratum/optimizer/physical/_source_execs.py index f37244aa..af1e606f 100644 --- a/stratum/optimizer/physical/_source_execs.py +++ b/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): diff --git a/tests/optimizer/physical/test_source_execs.py b/tests/optimizer/physical/test_source_execs.py index 35e192bb..083757f1 100644 --- a/tests/optimizer/physical/test_source_execs.py +++ b/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() From 334f8126ec43ebfa2d1e176e13f45b984b30dc19 Mon Sep 17 00:00:00 2001 From: Elias Strauss Date: Tue, 6 Oct 2026 00:31:57 +0200 Subject: [PATCH 2/2] [Fix] Keep unspillable buffer pool entries resident instead of crashing The pool admits objects it cannot size (counted as UNKNOWN_SIZE), but the serializer refuses anything that is not a frame, array, primitive or container of those. Over budget, eviction picked such an entry -- in practice the splitter from mark_as_X(cv=...) -- and the ValueError failed the whole plan. Serialization now raises UnsupportedSpillType and removes leaf files it already wrote for a container it then gave up on. Eviction catches it, marks the entry unspillable, keeps it in memory, and moves on to the next LRU candidate. --- stratum/runtime/_buffer_pool.py | 16 +++++++++++++--- stratum/runtime/_object_serialization.py | 21 ++++++++++++++++++--- tests/runtime/test_buffer_pool.py | 20 ++++++++++++++++++++ tests/runtime/test_object_serialization.py | 3 +++ 4 files changed, 54 insertions(+), 6 deletions(-) diff --git a/stratum/runtime/_buffer_pool.py b/stratum/runtime/_buffer_pool.py index e790243a..25a11b3d 100644 --- a/stratum/runtime/_buffer_pool.py +++ b/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 diff --git a/stratum/runtime/_object_serialization.py b/stratum/runtime/_object_serialization.py index e3e22ae2..0886e307 100644 --- a/stratum/runtime/_object_serialization.py +++ b/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 diff --git a/tests/runtime/test_buffer_pool.py b/tests/runtime/test_buffer_pool.py index c17fabcb..0656bee8 100644 --- a/tests/runtime/test_buffer_pool.py +++ b/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") diff --git a/tests/runtime/test_object_serialization.py b/tests/runtime/test_object_serialization.py index fa90f837..a12422b0 100644 --- a/tests/runtime/test_object_serialization.py +++ b/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())