diff --git a/README.md b/README.md index 70ddf66..9d73c0d 100644 --- a/README.md +++ b/README.md @@ -90,8 +90,11 @@ For a guided walkthrough, start with the - **Caching** — `cache` keeps an in-memory, column- and partition-level cache for iterative work; `cache_parquet` materializes date-partitioned Parquet on local disk or S3, fetching only the partitions a query needs. -- **Distributed execution** — `execute_on_ray` splits a `LazyFrame` by calendar period - and runs the partitions across an existing Ray cluster. +- **Distributed execution** — `execute_on_ray` splits a `LazyFrame` across an existing Ray + cluster, one task per partition. Build partitions with `by_time` (calendar windows), + `by_value` (discrete keys), `by_range` (numeric buckets), `by_key` (enumerated keys), or an + explicit `ReadPartition` list. (For multi-stage distributed pipelines, + [Polars Cloud](https://docs.cloud.pola.rs/polars-cloud/) is the more strategic option.) - **Ergonomics** — `iter_rows` for memory-efficient row iteration, `debug` to inspect what Polars pushes into a source, and `disable_optimizations` to compare against plain Polars. diff --git a/docs/wiki/API-Reference.md b/docs/wiki/API-Reference.md index 9db6620..8e431db 100644 --- a/docs/wiki/API-Reference.md +++ b/docs/wiki/API-Reference.md @@ -130,13 +130,34 @@ common sub-plan elimination. ### `execute_on_ray` ```python -lf.piot.execute_on_ray(*, date_column, time_unit, return_as="arrow", - remote_options=None, max_concurrency=100) +lf.piot.execute_on_ray(partitions, *, return_as="arrow", + remote_options=None, max_concurrency=100, + preserve_partition_order=None) ``` -Split the LazyFrame into calendar periods and execute each on an already-initialised Ray -cluster. `time_unit` is `"daily"`, `"monthly"`, or `"yearly"`. Requires `ray.init()` to -have been called and a bounded predicate on `date_column`. +Distribute the LazyFrame across an already-initialised Ray cluster, running one task per +partition. `partitions` is either a partitioner or an explicit iterable of +`ReadPartition(predicate, key)` (a `RayPartition` additionally carries per-task +`remote_options`). Build partitions with: + +- `by_time(column, every)` — calendar windows derived from the pushed-down date range + (`every` is `"1mo"`/`"2w"`/`"5d"`/`"1q"`/`"1y"` or an integer number of days). +- `by_value(column, values=None)` — one task per discrete value; derived from the pushed-down + `IN` filter when `values` is omitted. +- `by_range(column, every)` — fixed-width numeric buckets over the pushed-down range. +- `by_key(partitions, by, *, partition_remote_options=None)` — equality on caller-enumerated + keys; `by` is a column name, list, selector, or a `pl.Expr` (e.g. `pl.col("id").hash() % N`). + `partition_remote_options` sets per-partition Ray options from a struct column or `{key: dict}`. +- `discrete_partitions` / `cartesian_partitions` — explicit `col.is_in(...)` member lists and + `date_window × bucket` products. + +A partitioner requires a bounded predicate on its column. Requires `ray.init()` to have been +called. As a legacy shortcut, `execute_on_ray(date_column=..., time_unit="daily"|"monthly"|"yearly")` +is equivalent to `partitions=by_time(date_column, ...)`. + +Chaining multiple `execute_on_ray` calls relies on predicate pushdown surviving intervening +operations — partition once at the outermost boundary. For multi-stage distributed pipelines, +prefer [Polars Cloud](https://docs.cloud.pola.rs/polars-cloud/). ### `sink_delta` diff --git a/polars_io_tools/io_sources/__init__.py b/polars_io_tools/io_sources/__init__.py index a69e4e3..f1e101d 100644 --- a/polars_io_tools/io_sources/__init__.py +++ b/polars_io_tools/io_sources/__init__.py @@ -20,6 +20,7 @@ from .lazy_narwhals_reader import * from .lazy_probe import probe, probe as _lazy_probe from .lazy_sql_reader import * +from .partitions import * from .pushdown_combine import * from .pushdown_pivot import * from .pushdown_unpivot import * @@ -39,6 +40,7 @@ from .lazy_clickhouse_writer import sink_clickhouse # noqa: TC004 from .lazy_data_generator import scan_synthetic_panel, scan_synthetic_regression from .lazy_iter_rows import iter_rows # noqa: TC004 + from .partitions import KeyPartitions, ReadPartition, by_key, by_range, by_time, by_value from .pushdown_combine import FilterSpec, pushdown_combine from .pushdown_pivot import pushdown_pivot from .pushdown_unpivot import pushdown_unpivot @@ -50,7 +52,9 @@ # we also can't import it *inside* the execute_on_ray method of the # PIOTOperations class, because thit needs to be defined at the module level # for the `functools.wraps` decorator to work. That's why we use a stub here. - from .lazy_ray import execute_on_ray as _execute_on_ray_proto + from .lazy_ray import ( + execute_on_ray as _execute_on_ray_proto, + ) else: def _execute_on_ray_proto(*_a, **_kw): ... diff --git a/polars_io_tools/io_sources/lazy_ray.py b/polars_io_tools/io_sources/lazy_ray.py index 93e6a2e..1811287 100644 --- a/polars_io_tools/io_sources/lazy_ray.py +++ b/polars_io_tools/io_sources/lazy_ray.py @@ -1,15 +1,28 @@ +import contextlib import io -from collections.abc import Iterator -from datetime import datetime, timedelta -from typing import Literal +import math +from collections import deque +from collections.abc import Callable, Iterable, Iterator, Mapping +from dataclasses import dataclass, field +from typing import Any, Literal import cloudpickle import polars as pl -import portion +import polars.selectors as cs from tqdm import tqdm -from .range_visitor import convert_expr_to_datetime_range -from .util import _extend_interval +from .partitions import ( + KeyPartitions, + Partitioner, + ReadPartition, + as_partition_list, + by_range, + by_time, + by_value, + cartesian_partitions, + discrete_partitions, + retained_columns, +) try: import ray @@ -19,85 +32,45 @@ from .util import register_io_source_with_is_pure -# As we did in the `lazy_parquet_cache` module, we expose the function to users, in case -# they want to use it directly or though Polars' `.pipe()` syntax; however, the canonical -# useage is to call the function as a method on the `LazyFrame`'s `piot` namespace. -__all__ = ("execute_on_ray",) +# As we did in the `lazy_parquet_cache` module, we expose the functions to users, in case +# they want to use them directly or though Polars' `.pipe()` syntax; however, the canonical +# useage is to call the functions as methods on the `LazyFrame`'s `piot` namespace. +__all__ = ( + "RayPartition", + "ReadPartition", + "by_range", + "by_time", + "by_value", + "cartesian_partitions", + "discrete_partitions", + "execute_on_ray", +) + + +@dataclass(frozen=True) +class RayPartition(ReadPartition): + """A :class:`ReadPartition` carrying per-task Ray options. + + The shared ``predicate`` / ``key`` come from :class:`ReadPartition`; ``remote_options`` are + Ray-specific overrides shallow-merged over the calling function's uniform base options. + """ + remote_options: Mapping[str, Any] | None = field(default=None) -def _partition_specs( - date_min: datetime, - date_max: datetime, - unit: Literal["daily", "monthly", "yearly"], -) -> list[tuple[datetime, datetime]]: - """ - Returns a list of (start, end) tuples that cover the - range. Each `end` is exclusive, i.e. [start, end). - """ - if date_min is None or date_max is None: - return [] - - specs: list[tuple[datetime, datetime]] = [] - - if unit == "daily": - cur = date_min.date() - end = date_max.date() - while cur <= end: - nxt = cur + timedelta(days=1) - specs.append((datetime.combine(cur, datetime.min.time()), datetime.combine(nxt, datetime.min.time()))) - cur = nxt - - elif unit == "monthly": - cur = datetime(date_min.year, date_min.month, 1) # noqa: DTZ001 -- naive-UTC by design; bounds normalized via ValidatedDatetime - end = datetime(date_max.year, date_max.month, 1) # noqa: DTZ001 -- naive-UTC by design; bounds normalized via ValidatedDatetime - while cur <= end: - nxt = datetime(cur.year + 1, 1, 1) if cur.month == 12 else datetime(cur.year, cur.month + 1, 1) # noqa: DTZ001 -- naive-UTC by design; bounds normalized via ValidatedDatetime - specs.append((cur, nxt)) - cur = nxt - - elif unit == "yearly": - for yr in range(date_min.year, date_max.year + 1): - start = datetime(yr, 1, 1) # noqa: DTZ001 -- naive-UTC by design; bounds normalized via ValidatedDatetime - end = datetime(yr + 1, 1, 1) # noqa: DTZ001 -- naive-UTC by design; bounds normalized via ValidatedDatetime - specs.append((start, end)) - - return specs - - -def _trim_partition_specs( - specs: list[tuple[datetime, datetime]], - date_interval: "portion.Interval", - column_type: pl.DataType, -) -> list[tuple[datetime, datetime]]: - """ - Intersect each [start, end) partition with the user's temporal interval - and return only those with non-empty overlap. - Both the start and end of each partition are tightened to the - intersection. Since ``_execute_partition`` uses ``col < end`` (open - upper), any closed upper bound is converted to open via - ``_extend_interval``. - """ - trimmed: list[tuple[datetime, datetime]] = [] - for start, end in specs: - intersection = portion.closedopen(start, end) & date_interval - if not intersection.empty: - extended = _extend_interval(intersection, column_type) - trimmed.append((extended.lower, extended.upper)) - return trimmed +_TIME_UNIT_TO_INTERVAL = {"daily": "1d", "monthly": "1mo", "yearly": "1y"} @ray.remote def _execute_partition( - plan_bytes_or_ref: bytes | ray.ObjectRef, - date_col: str, - start: datetime, - end: datetime, + plan_bytes_or_ref: "bytes | ray.ObjectRef", + predicate_bytes: bytes, return_as: Literal["arrow", "ipc", "parquet"] = "arrow", ) -> bytes: """ - Deserialise plan, apply `[start, end)` filter, collect, return IPC bytes. - Works wheter or not Ray has already dereferenced `plan_bytes_or_ref`. + Deserialise plan and partition predicate, apply the filter, collect, and return the + result in the requested format. Works whether or not Ray has already dereferenced + ``plan_bytes_or_ref``. """ # Resolve the reference if we have it; don't worry # otherwise, since we've got the raw bytes @@ -107,8 +80,8 @@ def _execute_partition( plan_bytes = plan_bytes_or_ref lf = cloudpickle.loads(plan_bytes) + predicate = cloudpickle.loads(predicate_bytes) - predicate = (pl.col(date_col) >= start) & (pl.col(date_col) < end) out = lf.filter(predicate).collect() if return_as == "arrow": @@ -121,60 +94,56 @@ def _execute_partition( buf = io.BytesIO() out.write_ipc(buf) return buf.getvalue() + else: + raise ValueError(f"Unsupported return format: {return_as}") -def execute_on_ray( - self: pl.LazyFrame, - *, - date_column: str, - time_unit: Literal["daily", "monthly", "yearly"], - return_as: Literal["arrow", "ipc", "parquet"] = "arrow", - remote_options: dict | None = None, - max_concurrency: int | None = 100, - description: str | None = None, -) -> pl.LazyFrame: - """ - Execute a Polars LazyFrame on an *already initialised* Ray cluster, - distributing the work by calendar periods. +def _decode_blob(blob: bytes, return_as: str) -> pl.DataFrame: + if return_as == "arrow": + return pl.DataFrame(blob) + elif return_as == "parquet": + return pl.read_parquet(io.BytesIO(blob)) + elif return_as == "ipc": + return pl.read_ipc(io.BytesIO(blob)) + else: + raise ValueError(f"Unsupported return format: {return_as}") - The function returns **another** LazyFrame whose scan node is a - custom I/O source. No computation happens immediately; evaluation - is triggered only when the user calls `.collect()`. - Args: - date_column (str): Name of the datetime column that defines the partitioning axis. - The column must be of type `pl.Datetime` or `pl.Date`. - time_unit ({"daily", "monthly", "yearly"}): Granularity of the split. - return_as ({"arrow", "ipc", "parquet"}, default "arrow"): The format in which the Ray worker returns the data. - - "arrow" returns zero-copy Arrow buffers, - - "ipc" returns Arrow IPC buffers, - - "parquet" returns Parquet buffers. - remote_options (Optional[dict]): A dictionary of options for each Ray task. Please see the Ray - documentation for details: https://docs.ray.io/en/latest/_modules/ray/remote_function.html#RemoteFunction.options - You may wish to specify keys such as ``num_cpus``, ``num_gpus``, or - pass environment variables like POLARS_MAX_THREADS and POLARS_ENGINE_AFFINITY - in the runtime environment. If `None`, the task launches with standard defaults. - max_concurrency (Optional[int], default 100): The maximum number of concurrent tasks to run. We follow this pattern from the Ray documentation: - https://docs.ray.io/en/latest/ray-core/patterns/limit-pending-tasks.html. You may also wish to experiment - with using resource hints to manage concurrency, although we do not recommend doing so; please see the - following page for this alternative pattern: https://docs.ray.io/en/latest/ray-core/patterns/limit-running-tasks.html - description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). +# Shared core - Returns: - pl.LazyFrame: A new LazyFrame whose execution plan includes information about - running execution on a Ray cluster when the user calls `.collect()` +_VALID_RETURN_AS = ("arrow", "ipc", "parquet") - **NOTE:** This function requires a bounded predicate on the chosen `date_column` - """ - if not ray.is_initialized(): - raise RuntimeError("Ray is not initialised. Please call `ray.init()` before using `execute_on_ray`.") +def _output_schema(full_schema: pl.Schema, requested_cols: list[str] | None) -> pl.Schema: + if requested_cols is None: + return full_schema + return pl.Schema({c: full_schema[c] for c in requested_cols}) - remote_options = remote_options or {} - if not isinstance(remote_options, dict): - raise TypeError("`remote_options` must be a dict or None.") - original_lf = self +def _run_on_ray( + original_lf: pl.LazyFrame, + make_specs: Callable[["pl.Expr | None"], list[RayPartition]], + *, + predicate_required: bool, + return_as: Literal["arrow", "ipc", "parquet"], + base_remote_options: dict, + max_concurrency: int | None, + preserve_partition_order: bool, + description: str | None = None, +) -> pl.LazyFrame: + """ + Register a custom IO source that distributes ``original_lf`` across Ray by fanning out + one task per partition produced by ``make_specs``. + + ``make_specs`` receives the pushed-down predicate (or ``None``) and returns the list of + :class:`RayPartition` to execute. The core owns everything partition-scheme agnostic: + projection/predicate pushdown, predicate-column retention, task fan-out with bounded + concurrency, result reassembly, ``n_rows``/``batch_size`` handling, and cancellation. + """ + if return_as not in _VALID_RETURN_AS: + raise ValueError(f"`return_as` must be one of {_VALID_RETURN_AS}, got {return_as!r}.") + if max_concurrency is not None and max_concurrency <= 0: + raise ValueError(f"`max_concurrency` must be a positive integer or None, got {max_concurrency!r}.") def source_generator( with_columns: list[str] | None, @@ -182,132 +151,541 @@ def source_generator( n_rows: int | None, batch_size: int | None, ) -> Iterator[pl.DataFrame]: - lf = original_lf + full_schema = original_lf.collect_schema() + requested_cols = list(with_columns) if with_columns is not None else None - requested_cols = with_columns[:] if with_columns is not None else None + if predicate_required and predicate is None: + raise ValueError("`execute_on_ray` requires a bounded temporal predicate on the chosen `date_column`.") - # push-downs from the optimizer - if with_columns is not None: - if date_column not in with_columns: - with_columns = with_columns + [date_column] - lf = lf.select(with_columns) + specs = list(make_specs(predicate)) + if not specs: + yield pl.DataFrame(schema=_output_schema(full_schema, requested_cols)) + return + + # Columns required to evaluate every predicate that runs against the source (each + # partition predicate on the worker, plus the pushed-down predicate) must survive + # projection pushdown even if the user did not request them, and are dropped again from + # yielded frames. + predicates = [sp.predicate for sp in specs] + if predicate is not None: + predicates.append(predicate) + retained = retained_columns(predicates, requested_cols) + + lf = original_lf + added_cols: list[str] = [] + if requested_cols is not None: + if retained is None: + # A predicate's dependencies could not be resolved -- retain every column and drop + # the non-requested ones from yielded frames. + added_cols = [c for c in full_schema.names() if c not in requested_cols] + else: + added_cols = [c for c in retained if c not in requested_cols] + lf = lf.select(retained) if predicate is not None: lf = lf.filter(predicate) - # derive temporal bounds from the pushed-down predicate - if predicate is None: - raise ValueError(f"`execute_on_ray` requires a bounded temporal predicate on column '{date_column}'.") - date_interval = convert_expr_to_datetime_range(predicate, date_column, get_enclosure=False) + plan_bytes = cloudpickle.dumps(lf) - if date_interval.empty: - yield pl.DataFrame(schema=self.collect_schema()) - return + def make_task(sp: RayPartition) -> ray.ObjectRef: + opts = dict(base_remote_options) + if sp.remote_options: + opts.update(sp.remote_options) + pred_bytes = cloudpickle.dumps(sp.predicate) + return _execute_partition.options(**opts).remote(plan_bytes, pred_bytes, return_as) + + def fetch(ref: ray.ObjectRef, idx: int, sp: RayPartition) -> pl.DataFrame: + try: + blob = ray.get(ref) + except Exception as e: + err_msg = ( + f"Ray worker failed while executing partition {idx} (key={sp.key!r}) of lazy frame.\n" + f"Polars plan for this lazy frame:\n{lf.explain()}" + f"\n\nWhile running the above, received error: {e.__class__.__name__}:{e}" + ) + raise RuntimeError(err_msg) from e + return _decode_blob(blob, return_as) + + spec_iter = iter(enumerate(specs)) + effective_concurrency = max_concurrency if max_concurrency is not None else len(specs) + window = min(effective_concurrency, len(specs)) - if (date_interval.lower is portion.inf) or (date_interval.upper is -portion.inf): - raise ValueError("Un-bounded temporal predicate - unable to split the work.") + pbar = tqdm(total=len(specs), desc="execute_on_ray") - _min = date_interval.lower - _max = date_interval.upper + rows_yielded = 0 + any_yielded = False + stop = False + + def prepare(df: pl.DataFrame) -> tuple[list[pl.DataFrame], bool]: + """Drop helper columns, honour ``n_rows``, and slice by ``batch_size``. + + Returns the frames to yield and whether the global ``n_rows`` limit is now met. + """ + nonlocal rows_yielded + if added_cols: + df = df.drop(added_cols) + stop = False + if n_rows is not None: + remaining = n_rows - rows_yielded + if remaining <= 0: + return [], True + if len(df) >= remaining: + df = df.head(remaining) + stop = True + rows_yielded += len(df) + if batch_size is not None and len(df) > batch_size: + return [df.slice(i, batch_size) for i in range(0, len(df), batch_size)], stop + return [df], stop + + if preserve_partition_order: + # Ordered fan-out. Keep a spec-order deque of at most ``window`` in-flight tasks and + # block only on the head -- the next partition to yield. Later tasks run concurrently; + # their results stay in Ray's object store behind their refs (spillable) until we + # reach them, rather than being decoded and buffered on the driver. Outstanding + # results are bounded to ``window`` refs plus the one partition being materialised. + inflight: deque[tuple[ray.ObjectRef, int, RayPartition]] = deque() + try: + for _ in range(window): + idx, sp = next(spec_iter) + inflight.append((make_task(sp), idx, sp)) + while inflight: + ref, idx, sp = inflight.popleft() + df = fetch(ref, idx, sp) + del ref # release the head result promptly once decoded + frames, stop = prepare(df) + pbar.update() + # Refill the freed slot *before* yielding: a generator suspends at ``yield``, + # so submitting the replacement afterwards would idle a worker until the + # consumer pulls again. + if not stop: + with contextlib.suppress(StopIteration): + nxt_idx, nxt_sp = next(spec_iter) + inflight.append((make_task(nxt_sp), nxt_idx, nxt_sp)) + for fr in frames: + any_yielded = True + yield fr + if stop: + break + finally: + for ref, _, _ in inflight: + with contextlib.suppress(Exception): + ray.cancel(ref, force=False) + pbar.close() + else: + # Completion-order fan-out. Yield each partition as soon as it finishes. A ref is + # owned by ``pending`` until popped on completion, so nothing pins finished results. + pending: dict[ray.ObjectRef, tuple[int, RayPartition]] = {} + try: + for _ in range(window): + idx, sp = next(spec_iter) + pending[make_task(sp)] = (idx, sp) + while pending and not stop: + ready_refs, _ = ray.wait(list(pending), num_returns=1) + for ref in ready_refs: + idx, sp = pending.pop(ref) + # keep the pipeline full + with contextlib.suppress(StopIteration): + nxt_idx, nxt_sp = next(spec_iter) + pending[make_task(nxt_sp)] = (nxt_idx, nxt_sp) + pbar.update() + df = fetch(ref, idx, sp) + frames, stop = prepare(df) + for fr in frames: + any_yielded = True + yield fr + if stop: + break + finally: + # Cancel any still-pending tasks (n_rows satisfied early, or a failure occurred). + for ref in pending: + with contextlib.suppress(Exception): + ray.cancel(ref, force=False) + pbar.close() + + if not any_yielded: + yield pl.DataFrame(schema=_output_schema(full_schema, requested_cols)) + + return register_io_source_with_is_pure(source_generator, schema=lambda: original_lf.collect_schema(), explain_detail=description) + + +# Shared preconditions / helpers + + +def _require_ray() -> None: + if not ray.is_initialized(): + raise RuntimeError("Ray is not initialised. Please call `ray.init()` before using `execute_on_ray`.") - specs = _partition_specs(_min, _max, time_unit) - date_col_type = self.collect_schema()[date_column] - specs = _trim_partition_specs(specs, date_interval, date_col_type) +def _normalize_remote_options(remote_options: dict | None) -> dict: + remote_options = remote_options or {} + if not isinstance(remote_options, dict): + raise TypeError("`remote_options` must be a dict or None.") + return remote_options + + +def _resolve_ro(remote_options, key): + """Resolve per-partition remote options from a callable or a {key: dict} mapping.""" + if remote_options is None: + return None + if callable(remote_options): + return remote_options(key) + if isinstance(remote_options, Mapping): + return remote_options.get(key) + raise TypeError("per-partition `remote_options` must be a callable or a mapping.") + + +def _canonical_key(value): + """Canonicalize a key for duplicate detection so that distinct NaN objects (which are + unequal in a Python set but match under polars equality) collapse to one value.""" + if isinstance(value, tuple): + return tuple(_canonical_key(v) for v in value) + if isinstance(value, float) and math.isnan(value): # NaN + return "__nan__" + return value + + +def _prune_by_key(specs: list[RayPartition], key_cols: list[str], predicate: pl.Expr | None, schema: pl.Schema) -> list[RayPartition]: + """Drop equality partitions that cannot match the pushed-down ``predicate``. + + Sound and exact: each partition corresponds to a concrete key value on ``key_cols``, so a + row in that partition has those exact values. When ``predicate`` references *only* + ``key_cols``, evaluating it -- with polars' real semantics -- on a one-row-per-partition + frame of the key values is equivalent to evaluating it on any real row of the partition, + provided both preconditions below hold. We then keep only partitions whose key row + satisfies it. + + Preconditions for exactness (otherwise fail open, keeping every partition): + - **Non-float key dtype.** Float ``==`` is not observational (``-0.0``/``+0.0``, NaN), so a + matched row need not be observationally identical to the key. Every other dtype has + observational equality, making the key a faithful representative. + - **No Python UDF in the predicate.** A pushed row-wise native predicate is a + deterministic function of the row's values; a stateful/non-deterministic UDF is not, so + one evaluation on the key row would not predict the worker's evaluation. + + Pruning only decides whether to launch a task; the partition predicate is always applied + on the worker, so a retained partition never changes its result. + """ + if predicate is None or not specs: + return specs + try: + pred_cols = set(predicate.meta.root_names()) + except Exception: # noqa: BLE001 -- cannot introspect, fail open + return specs + if not pred_cols or not pred_cols.issubset(key_cols): + return specs + if any(c not in schema or schema[c].is_float() for c in key_cols): + return specs + if _has_python_udf(predicate): + return specs + try: + columns = {} + for i, c in enumerate(key_cols): + values = [(sp.key[i] if isinstance(sp.key, tuple) else sp.key) for sp in specs] + columns[c] = pl.Series(c, values, dtype=schema[c]) + keep_mask = pl.DataFrame(columns).select(predicate.alias("__keep__"))["__keep__"].to_list() + except Exception: # noqa: BLE001 -- evaluation failed, fail open + return specs + return [sp for sp, keep in zip(specs, keep_mask) if keep is True] + + +def _has_python_udf(expr: pl.Expr) -> bool: + """Whether ``expr`` contains a Python UDF (``map_elements`` / ``map_batches``), which can + be non-deterministic or stateful. Returns True conservatively if it cannot be verified.""" + try: + return "AnonymousFunction" in expr.meta.serialize(format="json") + except Exception: # noqa: BLE001 -- cannot verify, treat as unsafe + return True - if not specs: - yield pl.DataFrame(schema=self.collect_schema()) - return - plan_bytes = cloudpickle.dumps(lf) +def execute_on_ray( + self: pl.LazyFrame, + partitions: "Partitioner | Iterable[ReadPartition] | None" = None, + *, + date_column: str | None = None, + time_unit: Literal["daily", "monthly", "yearly"] | None = None, + return_as: Literal["arrow", "ipc", "parquet"] = "arrow", + remote_options: dict | None = None, + max_concurrency: int | None = 100, + preserve_partition_order: bool | None = None, + description: str | None = None, +) -> pl.LazyFrame: + """ + Execute a Polars LazyFrame on an *already initialised* Ray cluster, distributing the work + across one Ray task per :class:`ReadPartition`. - # create iterator of (idx, start, end) triples - spec_iter = iter(enumerate(specs)) + The function returns **another** LazyFrame whose scan node is a custom I/O source. No + computation happens immediately; evaluation is triggered only when the user calls + ``.collect()``. - future_to_idx: dict[ray.ObjectRef, int] = {} - pending_futures: set[ray.ObjectRef] = set() + Args: + partitions (Partitioner | Iterable[ReadPartition] | None): How to split the work. Pass a + partitioner from :mod:`polars_io_tools.io_sources.partitions` (``by_time``, ``by_value``, + ``by_range``) to derive slices from the filter pushed down at scan time, or an explicit + iterable of :class:`ReadPartition` / :class:`RayPartition` for hand-built slices. A + partitioner requires a bounded pushed-down predicate on its column. + date_column, time_unit: Legacy calendar shortcut, equivalent to + ``partitions=by_time(date_column, {daily: "1d", monthly: "1mo", yearly: "1y"}[time_unit])``. + Mutually exclusive with ``partitions``. + return_as ({"arrow", "ipc", "parquet"}, default "arrow"): The format in which the Ray + worker returns the data. + remote_options (Optional[dict]): Uniform Ray ``.options()`` for each task, overridden per + partition by any :class:`RayPartition.remote_options`. + max_concurrency (Optional[int], default 100): The maximum number of concurrent tasks. + preserve_partition_order (bool | None): If True, yield partitions in spec order; if False, + yield in completion order. Ordered mode blocks on the next partition in order while + later tasks run ahead, keeping outstanding results bounded by ``max_concurrency`` + (their unmaterialised results live in Ray's spillable object store). Defaults to + completion order for ``partitions``, and to spec order for the legacy calendar shortcut. + description: Optional free-form description of this source instance, attached to its + OpenTelemetry span (``explain_detail``). - # helper to launch one task - def submit(idx: int, span: tuple[datetime, datetime]): - f = _execute_partition.options(**remote_options).remote(plan_bytes, date_column, span[0], span[1], return_as) - future_to_idx[f] = idx - pending_futures.add(f) + Returns: + pl.LazyFrame: A new LazyFrame whose execution runs on a Ray cluster at ``.collect()``. - # prime up to `max_concurrency` tasks - effective_concurrency = max_concurrency if max_concurrency is not None else len(specs) - for _ in range(min(effective_concurrency, len(specs))): - idx, span = next(spec_iter) - submit(idx, span) + **WARNING:** Chaining multiple ``execute_on_ray`` calls (nesting one distributed source + inside another's plan) can have unintended consequences -- it relies on predicate pushdown + surviving intervening operations and can silently re-execute the upstream once per downstream + partition (N x N). Partition once at the outermost boundary. For multi-stage / cross-shuffle + distributed pipelines prefer Polars Cloud (https://docs.cloud.pola.rs/polars-cloud/). + """ + _require_ray() + base_remote_options = _normalize_remote_options(remote_options) + original_lf = self - pbar = tqdm(total=len(specs), desc="execute_on_ray") + # Legacy calendar shortcut: translate `date_column`/`time_unit` into a `by_time` partitioner + # and run it through the single partition path below (kept chronological, as it was before). + if date_column is not None or time_unit is not None: + if partitions is not None: + raise ValueError("Pass either `partitions` or the legacy `date_column`/`time_unit`, not both.") + if date_column is None or time_unit is None: + raise ValueError("The legacy calendar shortcut needs both `date_column` and `time_unit`.") + try: + every = _TIME_UNIT_TO_INTERVAL[time_unit] + except KeyError: + raise ValueError(f"time_unit must be one of {sorted(_TIME_UNIT_TO_INTERVAL)}, got {time_unit!r}.") from None + partitions = by_time(date_column, every) + if preserve_partition_order is None: + preserve_partition_order = True + + if partitions is None: + raise ValueError("Provide `partitions` (a partitioner or ReadPartition list), or the legacy `date_column`/`time_unit`.") + + if isinstance(partitions, KeyPartitions): + return _execute_on_ray_by( + self, + partitions.partitions, + partitions.by, + remote_options=remote_options, + partition_remote_options=partitions.partition_remote_options, + return_as=return_as, + max_concurrency=max_concurrency, + preserve_partition_order=bool(preserve_partition_order), + description=description, + ) + + # Materialise an explicit iterable once (a partitioner is re-usable and re-built per collect, + # but a bare generator would be exhausted after the first collection). + if not isinstance(partitions, Partitioner): + partitions = list(partitions) + + def make_specs(predicate: pl.Expr | None) -> list[RayPartition]: + resolved = as_partition_list(partitions, predicate) + if resolved is None: + raise ValueError("Could not derive partitions from the pushed-down predicate (no bounded range on the partition column).") + return [sp if isinstance(sp, RayPartition) else RayPartition(predicate=sp.predicate, key=sp.key) for sp in resolved] + + return _run_on_ray( + original_lf, + make_specs, + # A partitioner that needs the pushed predicate reports that by returning None from + # build(), which make_specs turns into a clear error -- so the source never needs to + # pre-require a predicate (an explicit list or by_value([...]) needs none). + predicate_required=False, + return_as=return_as, + base_remote_options=base_remote_options, + max_concurrency=max_concurrency, + preserve_partition_order=bool(preserve_partition_order), + description=description, + ) + + +def _partitions_to_frame(partitions, by) -> tuple[pl.DataFrame, list[str] | None, "pl.Expr | None"]: + """Normalize ``partitions`` into a DataFrame and resolve the key columns / expression. + + Returns ``(pdf, key_cols, key_expr)``. Exactly one of ``key_cols`` / ``key_expr`` is set. + When ``by`` is an expression, the single key column in ``pdf`` is named ``__key__``. + """ + # A polars selector is an `Expr` subclass, so it must be detected *before* the general + # expression branch and resolved through the frame's schema like column selectors. + is_selector = cs.is_selector(by) + + if isinstance(by, pl.Expr) and not is_selector: + if isinstance(partitions, pl.DataFrame): + if partitions.width != 1: + raise ValueError("When `by` is an expression, `partitions` must be a single column / Series / sequence.") + values = partitions.to_series() + elif isinstance(partitions, pl.Series): + values = partitions + elif isinstance(partitions, pl.LazyFrame): + values = partitions.collect().to_series() + else: + values = pl.Series("__key__", list(partitions)) + return values.rename("__key__").to_frame(), None, by + + # by is column name(s) or a selector + if isinstance(partitions, pl.LazyFrame): + pdf = partitions.collect() + elif isinstance(partitions, pl.DataFrame): + pdf = partitions + elif isinstance(partitions, pl.Series): + # For a scalar string `by`, align the series name with the key column. + pdf = (partitions.rename(by) if isinstance(by, str) else partitions).to_frame() + elif isinstance(by, str): + pdf = pl.Series(by, list(partitions)).to_frame() + else: + raise TypeError("`partitions` must be a DataFrame / LazyFrame when `by` selects multiple columns.") - ready_parts: dict[int, pl.DataFrame] = {} - next_to_yield = 0 - rows_yielded = 0 # needed for global `n_rows` limit - - while pending_futures or ready_parts: - if pending_futures: - ready_refs, _ = ray.wait(list(pending_futures), num_returns=1) - for ref in ready_refs: - pending_futures.remove(ref) - idx = future_to_idx[ref] - try: - blob = ray.get(ref) - except Exception as e: - err_msg = ( - f"Ray worker failed while executing partition {idx} of lazy frame.\nPolars plan for this lazy frame:\n{lf.explain()}" - ) - err_msg += f"\n\nWhile running the above, received error: {e.__class__.__name__}:{e}" - raise RuntimeError(err_msg) from e - - # immediately submit a new task if specs remain - try: - new_idx, new_span = next(spec_iter) - submit(new_idx, new_span) - except StopIteration: - pass + if isinstance(by, str): + key_cols = [by] + elif isinstance(by, (list, tuple)) and all(isinstance(c, str) for c in by): + key_cols = list(by) + else: + # treat as a selector + key_cols = pdf.select(by).columns - pbar.update() + return pdf, key_cols, None - blob = ray.get(ref) - df: pl.DataFrame - if return_as == "arrow": - result = pl.DataFrame(blob) - df = result - elif return_as == "parquet": - df = pl.read_parquet(io.BytesIO(blob)) - elif return_as == "ipc": - df = pl.read_ipc(io.BytesIO(blob)) - else: - raise ValueError(f"Unsupported return format: {return_as}") - ready_parts[idx] = df +def _execute_on_ray_by( + self: pl.LazyFrame, + partitions, + by, + *, + remote_options: dict | None = None, + partition_remote_options=None, + return_as: Literal["arrow", "ipc", "parquet"] = "arrow", + max_concurrency: int | None = 100, + preserve_partition_order: bool = False, + description: str | None = None, +) -> pl.LazyFrame: + """ + Distribute ``self`` across Ray by equality on caller-enumerated partition keys. - while next_to_yield in ready_parts: - df = ready_parts.pop(next_to_yield) - next_to_yield += 1 + ``partitions`` is a small frame with **one row per partition**; each row becomes a task + whose predicate is ``AND(col_i == row[col_i])`` over the ``by`` columns (a null key value + becomes ``col.is_null()``). ``by`` may be a column name, a list of names, a polars + selector, or a ``pl.Expr`` (e.g. ``pl.col("id").hash() % N``) evaluated against the + execution frame -- so partitioning needs no upstream ``with_columns`` and no schema + pollution. - # We need to drop the date column if the user did not originally request it. - if requested_cols is not None and date_column not in requested_cols: - df = df.drop(date_column) + For **grouped / member-list** partitioning (one shard = many ids), use ``by_value`` or + :func:`discrete_partitions` (``id.is_in(members)``) with :func:`execute_on_ray` rather than + this scalar-equality convenience. - if n_rows is not None: - remaining = n_rows - rows_yielded - if remaining <= 0: - return # limit already met - if len(df) > remaining: - df = df.head(remaining) - rows_yielded += len(df) + Args: + partitions: A ``DataFrame`` / ``LazyFrame`` (one row per partition), or a + ``Series`` / sequence of scalar keys (with a scalar ``by``). + by: Column name(s), a selector, or a ``pl.Expr`` identifying the partition key. + remote_options (Optional[dict]): Uniform base Ray options. + partition_remote_options: Per-partition Ray options as either the name of a struct + column in ``partitions``, or a ``{key: dict}`` mapping keyed by the partition key. + return_as ({"arrow", "ipc", "parquet"}, default "arrow"): Worker return format. + max_concurrency (Optional[int], default 100): Maximum concurrent tasks. + preserve_partition_order (bool, default False): Yield in spec order (blocks on the next + partition while later tasks run ahead; outstanding results bounded by + ``max_concurrency``) vs completion order (no order guarantee). - if batch_size is not None and len(df) > batch_size: - for i in range(0, len(df), batch_size): - yield df.slice(i, batch_size) - else: - yield df + Returns: + pl.LazyFrame: A new LazyFrame whose execution runs on a Ray cluster at ``.collect()``. - if next_to_yield == 0: # nothing ever yielded - yield pl.DataFrame(schema=self.collect_schema()) + Note: hash-bucket keys (a computed ``hash % N``) cannot be pruned at the source, so each + task scans the universe and filters after -- this bounds working-set memory, not scan + I/O. Discrete keys on a real, physically-prunable column can prune at the source. - pbar.close() + **WARNING:** Chaining multiple ``execute_on_ray*`` calls can have unintended + consequences. Partition once at the outermost boundary; for multi-stage distributed + pipelines prefer Polars Cloud (https://docs.cloud.pola.rs/polars-cloud/). + """ + _require_ray() + base_remote_options = _normalize_remote_options(remote_options) + original_lf = self - return register_io_source_with_is_pure(source_generator, schema=lambda: self.collect_schema(), explain_detail=description) + pdf, key_cols, key_expr = _partitions_to_frame(partitions, by) + + option_col = partition_remote_options if isinstance(partition_remote_options, str) else None + option_map = partition_remote_options if isinstance(partition_remote_options, Mapping) else None + if option_col is not None and option_col not in pdf.columns: + raise ValueError(f"`partition_remote_options` column {option_col!r} not found in `partitions`.") + # The options column must not be treated as a key column. + if option_col is not None and key_cols is not None: + key_cols = [c for c in key_cols if c != option_col] + if not key_cols: + raise ValueError("`by` must select at least one key column distinct from the options column.") + + # Reject duplicate partition keys up front (cheap, deterministic). NaN keys are + # canonicalized so that repeated NaNs (which compare unequal in a Python set but match + # under polars equality) are rejected. + if key_expr is not None: + _keys = pdf["__key__"].to_list() + elif len(key_cols) > 1: + _keys = [tuple(r) for r in pdf.select(key_cols).iter_rows()] + else: + _keys = pdf[key_cols[0]].to_list() + _seen: set = set() + for _k in _keys: + _canon = _canonical_key(_k) + if _canon in _seen: + raise ValueError(f"Duplicate partition key {_k!r} in `partitions`.") + _seen.add(_canon) + + # Resolve the output dtype of the key expression once, for a typed literal. + key_expr_dtype = None + if key_expr is not None: + key_expr_dtype = original_lf.select(key_expr.alias("__key__")).collect_schema()["__key__"] + + def make_specs(predicate: pl.Expr | None) -> list[RayPartition]: + exec_schema = original_lf.collect_schema() + specs: list[RayPartition] = [] + for row in pdf.iter_rows(named=True): + if key_expr is not None: + val = row["__key__"] + pred = key_expr.is_null() if val is None else (key_expr == pl.lit(val, dtype=key_expr_dtype)) + key = val + else: + conds = [] + key_vals = [] + for c in key_cols: + v = row[c] + key_vals.append(v) + if v is None: + conds.append(pl.col(c).is_null()) + else: + conds.append(pl.col(c) == pl.lit(v, dtype=exec_schema[c])) + pred = conds[0] + for extra in conds[1:]: + pred = pred & extra + key = tuple(key_vals) if len(key_cols) > 1 else key_vals[0] + + if option_col is not None: + ro = row[option_col] + else: + ro = _resolve_ro(option_map, key) + + specs.append(RayPartition(predicate=pred, key=key, remote_options=ro)) + # Exact key-value pruning (only when the pushed-down predicate touches key columns + # only). Pruning drops only partitions that cannot match; retained partitions are + # unchanged. + if key_cols is not None: + specs = _prune_by_key(specs, key_cols, predicate, exec_schema) + return specs + + return _run_on_ray( + original_lf, + make_specs, + predicate_required=False, + return_as=return_as, + base_remote_options=base_remote_options, + max_concurrency=max_concurrency, + preserve_partition_order=preserve_partition_order, + description=description, + ) diff --git a/polars_io_tools/io_sources/partitions.py b/polars_io_tools/io_sources/partitions.py new file mode 100644 index 0000000..920d046 --- /dev/null +++ b/polars_io_tools/io_sources/partitions.py @@ -0,0 +1,368 @@ +"""Backend-neutral read-partition specs shared by the SQL and distributed readers. + +A *read partition* is a slice of the input domain, defined by a polars predicate, that can be +read or executed independently. This differs from :class:`polars.io.partition.PartitionBy`, +which is write-side (it routes materialised rows to output files); these describe how to +*split a read* so a database scan or a distributed executor can process slices in parallel. + +Partitions come from two places: + +* **Eagerly** -- the caller enumerates them (an explicit value list, a hand-built + ``list[ReadPartition]``). +* **Derived** -- the split depends on the predicate Polars pushes down at scan time (calendar + windows from a date range, distinct values from an ``IN`` list). A :class:`Partitioner` + builds the concrete list from that predicate. + +Both collapse to the same currency -- a list of :class:`ReadPartition` -- so one vocabulary +drives both ``scan_db(..., partitions=...)`` and the Ray executor. A consumer that cannot honour +a partition's predicate (a database cannot run an arbitrary UDF, for instance) is expected to +reject it rather than silently drop the slice. +""" + +from __future__ import annotations + +import datetime +import itertools +import re +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any, Protocol, runtime_checkable + +import polars as pl +import portion + +from .range_visitor import convert_expr_to_datetime_range, convert_expr_to_range +from .set_visitor import convert_expr_to_valid_values + +__all__ = ( + "KeyPartitions", + "Partitioner", + "ReadPartition", + "by_key", + "by_range", + "by_time", + "by_value", + "cartesian_partitions", + "discrete_partitions", +) + + +@dataclass(frozen=True) +class ReadPartition: + """A single read partition: a predicate defining a disjoint slice of the input, plus a label. + + Args: + predicate: The filter that defines this slice. Consumers apply it as an ordinary polars + filter (distributed executor) or translate it to a SQL ``WHERE`` (database reader). + key: An optional label used for ordering, error messages, and per-partition option + lookup. For a multi-column key this is typically a tuple. + """ + + predicate: pl.Expr + key: Any = None + + +@runtime_checkable +class Partitioner(Protocol): + """Builds a list of :class:`ReadPartition` from the predicate pushed down at scan time. + + ``on`` names the column(s) the partitions filter on, so a consumer can keep them through + projection pushdown. ``build`` returns ``None`` when no partitioning is possible (e.g. the + predicate carries no bounded range on ``on``), signalling the consumer to fall back to a + single unpartitioned read. + """ + + on: str + + def build(self, predicate: pl.Expr | None) -> list[ReadPartition] | None: ... + + +_WORD_TO_INTERVAL = {"day": "1d", "week": "1w", "month": "1mo", "quarter": "1q", "year": "1y"} +_INTERVAL_RE = re.compile(r"\d+(?:d|w|mo|q|y)") +_MAX_WINDOWS = 100_000 + + +def _interval_str(every: str | int) -> str: + """Normalise a bucket size to a validated Polars interval string (e.g. ``"1mo"``, ``"5d"``).""" + if isinstance(every, bool): + raise TypeError("every must be an int or interval string, not bool") + if isinstance(every, int): + if every <= 0: + raise ValueError(f"every (days) must be positive, got {every}") + return f"{every}d" + text = _WORD_TO_INTERVAL.get(str(every).strip().lower(), str(every).strip().lower()) + if text[:1].isalpha(): # bare unit, e.g. "mo" -> "1mo" + text = "1" + text + if not _INTERVAL_RE.fullmatch(text): + raise ValueError(f"Could not parse every={every!r}; use e.g. '1mo', '2w', '5d', '1q', '1y' or an int of days") + return text + + +def _date_windows(interval: portion.Interval, every: str | int) -> list[tuple[Any, Any]] | None: + """Split a date ``interval`` into contiguous half-open ``[lo, hi)`` windows. + + Boundaries come from Polars' calendar-aware temporal functions (``dt.truncate`` + + ``date_range``), so month/quarter/year arithmetic and the interval-string grammar are not + reimplemented here. A disjoint interval is handled per atomic component (windows are generated + only over the ranges the predicate actually selects, not across the gaps between them), and the + ``_MAX_WINDOWS`` cap applies to the emitted windows rather than the enclosing span. + """ + if interval.empty: + return None + iv = _interval_str(every) + windows: dict[tuple[Any, Any], None] = {} # dict = ordered, de-duplicated + for atomic in interval: + lower, upper = atomic.lower, atomic.upper + if lower == -portion.inf or upper == portion.inf: + return None + start = pl.Series([lower]).dt.truncate(iv).item() + end = pl.Series([upper]).dt.offset_by(iv).item() # one step past ``upper`` closes the final window + is_date = isinstance(start, datetime.date) and not isinstance(start, datetime.datetime) + make_range = pl.date_range if is_date else pl.datetime_range + bounds = make_range(start, end, interval=iv, closed="both", eager=True) + inclusive_upper = atomic.right == portion.CLOSED + for lo, hi in itertools.pairwise(bounds): + if (lo <= upper) if inclusive_upper else (lo < upper): + windows[(lo, hi)] = None + if len(windows) > _MAX_WINDOWS: + return None + return list(windows) or None + + +@dataclass(frozen=True) +class _ByTime: + on: str + every: str | int + + def build(self, predicate: pl.Expr | None) -> list[ReadPartition] | None: + if predicate is None: + return None + interval = convert_expr_to_datetime_range(predicate, self.on, get_enclosure=False) + if interval.empty: + return [] # contradictory predicate selects nothing + windows = _date_windows(interval, self.every) + if windows is None: + return None # unbounded range -- cannot partition + col = pl.col(self.on) + return [ReadPartition(predicate=(col >= lo) & (col < hi), key=(lo, hi)) for lo, hi in windows] + + +def by_time(column: str, every: str | int = "1mo") -> Partitioner: + """Partition on a date/datetime column by calendar windows derived from the pushed-down range. + + Note: the pushed-down bounds are normalised to timezone-naive values, so partitioning a + timezone-aware ``Datetime`` column is not currently supported. + + Args: + column: The date/datetime column to partition on. + every: Window size -- an interval string (``"1mo"``, ``"2w"``, ``"5d"``, ``"1q"``, ``"1y"``) + or an integer number of days. + """ + _interval_str(every) # validate eagerly + return _ByTime(on=column, every=every) + + +def _membership_predicate(target: pl.Expr, value: Any) -> pl.Expr: + """Build a predicate matching ``value`` on ``target``. + + A scalar becomes ``target == value``, or ``target.is_null()`` when ``value`` is ``None`` + (``target == None`` would match nothing). A list/tuple/set becomes ``target.is_in(members)``, + OR-ed with ``target.is_null()`` when the members include ``None``. + """ + if isinstance(value, (list, tuple, set, frozenset)): + members = list(value) + non_null = [m for m in members if m is not None] + pred = target.is_in(non_null) + if len(non_null) != len(members): + pred = pred | target.is_null() + return pred + if value is None: + return target.is_null() + return target == value + + +@dataclass(frozen=True) +class _ByValue: + on: str + values: tuple[Any, ...] | None + + def build(self, predicate: pl.Expr | None) -> list[ReadPartition] | None: + if self.values is not None: + return [ + ReadPartition( + predicate=_membership_predicate(pl.col(self.on), v), key=v if not isinstance(v, (list, tuple, set, frozenset)) else tuple(v) + ) + for v in self.values + ] + if predicate is None: + return None + allowed = convert_expr_to_valid_values(predicate, self.on) + if allowed is None: + return None # cannot determine the value set -- fall back + col = pl.col(self.on) + return [ReadPartition(predicate=_membership_predicate(col, v), key=v) for v in allowed] + + +def by_value(column: str, values: Iterable[Any] | None = None) -> Partitioner: + """Partition on a discrete column, one partition per value or per member group. + + Args: + column: The column to partition on. + values: The partition values. Each element is either a scalar (``col == v``) or a + collection (``col.is_in(v)``). When ``None``, the distinct values are read from the + ``IN`` / equality filter pushed down on ``column``; if none can be determined, the + read runs unpartitioned. + """ + materialized = None if values is None else tuple(values) + return _ByValue(on=column, values=materialized) + + +@dataclass(frozen=True) +class _ByRange: + on: str + every: float | int + + def build(self, predicate: pl.Expr | None) -> list[ReadPartition] | None: + if predicate is None: + return None + interval = convert_expr_to_range(predicate, self.on) + if interval.empty: + return [] # contradictory predicate selects nothing + col = pl.col(self.on) + parts: list[ReadPartition] = [] + for atomic in interval: # bucket each selected range; skip the gaps between disjoint ranges + if atomic.lower == -portion.inf or atomic.upper == portion.inf: + return None # unbounded range -- cannot partition + inclusive_upper = atomic.right == portion.CLOSED + lo = atomic.lower + while (lo <= atomic.upper) if inclusive_upper else (lo < atomic.upper): + hi = lo + self.every + parts.append(ReadPartition(predicate=(col >= lo) & (col < hi), key=(lo, hi))) + lo = hi + if len(parts) > _MAX_WINDOWS: + return None + return parts or None + + +def by_range(column: str, every: float) -> Partitioner: + """Partition on a numeric column into fixed-width ``[lo, lo+every)`` buckets. + + The numeric range is read from the bounded filter pushed down on ``column``; if the filter + leaves the range unbounded, the read runs unpartitioned. + + Args: + column: The numeric column to partition on. + every: Bucket width. + """ + if every <= 0: + raise ValueError(f"every must be positive, got {every}") + return _ByRange(on=column, every=every) + + +@dataclass(frozen=True) +class KeyPartitions: + """Enumerated equality partitioning by caller-supplied keys. + + A spec (not a :class:`Partitioner`): the concrete partitions depend on the target frame's + schema and the pushed-down predicate, so a backend that supports it resolves them at scan + time. ``partition_remote_options`` is honoured only by distributed executors. + + Args: + partitions: One key per partition -- a frame/series/sequence of keys. + by: The key column name(s), selector, or expression (e.g. ``pl.col("id").hash() % n``). + partition_remote_options: Optional per-partition executor options (distributed backends only). + """ + + partitions: Any + by: Any + partition_remote_options: Any = None + + +def by_key(partitions: Any, by: Any, *, partition_remote_options: Any = None) -> KeyPartitions: + """Partition by equality on caller-enumerated keys (one partition per key). + + Args: + partitions: A frame/series/sequence with one key per partition. + by: Column name(s), a selector, or a ``pl.Expr`` identifying the partition key. + partition_remote_options: Optional per-partition executor options (distributed backends only). + """ + return KeyPartitions(partitions=partitions, by=by, partition_remote_options=partition_remote_options) + + +def discrete_partitions(column: str, groups: Iterable[Sequence] | Mapping[Any, Sequence]) -> list[ReadPartition]: + """Build ``column.is_in(members)`` partitions -- one per group / member list. + + Args: + column: The column to partition on. + groups: Either an iterable of member lists (labelled by position) or a mapping + ``{label: members}``. + """ + items = groups.items() if isinstance(groups, Mapping) else enumerate(groups) + return [ReadPartition(predicate=_membership_predicate(pl.col(column), list(members)), key=label) for label, members in items] + + +def cartesian_partitions( + *, + date_windows: Iterable[tuple], + buckets: Iterable | Mapping, + date_column: str, + bucket: Any, +) -> list[ReadPartition]: + """Build a flat ``{date_windows} x {buckets}`` partition set in a single pass. + + Each cell's predicate is ``(date_column in window) & (bucket predicate)``, applied as one + filter so the date-window component still pushes down into inner sources. + + Args: + date_windows: Iterable of ``(lower, upper)`` half-open ``[lower, upper)`` windows. + buckets: Either an iterable of bucket values (labelled by position) or a mapping + ``{label: value}``. A list/tuple/set becomes ``target.is_in(value)``; a scalar becomes + ``target == value``; ``None`` matches nulls. + date_column: The datetime column for the window predicate. + bucket: The bucket target -- a column name (``str``) or a ``pl.Expr``. + """ + target = pl.col(bucket) if isinstance(bucket, str) else bucket + bucket_items = list(buckets.items()) if isinstance(buckets, Mapping) else list(enumerate(buckets)) + parts: list[ReadPartition] = [] + for lower, upper in date_windows: + date_pred = (pl.col(date_column) >= lower) & (pl.col(date_column) < upper) + for label, value in bucket_items: + parts.append(ReadPartition(predicate=date_pred & _membership_predicate(target, value), key=((lower, upper), label))) + return parts + + +def as_partition_list(partitions: Partitioner | Iterable[ReadPartition], predicate: pl.Expr | None) -> list[ReadPartition] | None: + """Resolve ``partitions`` to a concrete list. + + Returns ``None`` only when a partitioner cannot derive a bounded split (the caller may then + fall back to a single unpartitioned read). An explicit iterable always resolves to a list -- + an empty one stays empty (a known-empty partition set), never ``None``. + """ + if isinstance(partitions, Partitioner): + return partitions.build(predicate) + resolved = list(partitions) + if any(not isinstance(p, ReadPartition) for p in resolved): + raise TypeError("partitions must be a Partitioner (e.g. by_time/by_value/by_range) or an iterable of ReadPartition") + return resolved + + +def retained_columns(predicates: Iterable[pl.Expr], with_columns: list[str] | None) -> list[str] | None: + """Projection a consumer must request so every predicate is evaluable, or ``None`` to keep all. + + Partition (and pushed-down) predicates are applied after the read, so the columns they touch + must survive projection pushdown even when the caller did not request them. Returns ``None`` + (meaning "keep every column") when nothing was projected or when a predicate exposes no + concrete root names (e.g. a selector), so a projection can never drop a column a filter needs. + """ + if with_columns is None: + return None + required: set[str] = set() + unresolved = False + for pred in predicates: + roots = pred.meta.root_names() + if not roots: + unresolved = True + required.update(roots) + if unresolved: + return None + return list(dict.fromkeys([*with_columns, *(c for c in required if c not in with_columns)])) diff --git a/polars_io_tools/tests/io_sources/test_lazy_ray.py b/polars_io_tools/tests/io_sources/test_lazy_ray.py index 334da28..582c44c 100644 --- a/polars_io_tools/tests/io_sources/test_lazy_ray.py +++ b/polars_io_tools/tests/io_sources/test_lazy_ray.py @@ -1,5 +1,6 @@ import datetime import random +import sys import time import polars as pl @@ -11,6 +12,15 @@ import polars_io_tools as cpl import polars_io_tools.io_sources.lazy_ray # explicit import needed if using pytest-xdist +# Ray's task cancellation is unstable on Windows (its Windows support is beta): cancelling an +# in-flight task while the Polars source generator is being closed early (n_rows satisfied, or a +# failing partition) can segfault the interpreter. The behaviour under test is platform-agnostic +# and fully covered on Linux/macOS, so these early-termination cases are skipped on Windows. +_ray_cancel_flaky_on_windows = pytest.mark.skipif( + sys.platform == "win32", + reason="Ray task cancellation on early source termination can crash the interpreter on Windows (Ray Windows support is beta).", +) + @pytest.fixture(scope="session", autouse=True) def shared_ray_cluster(): @@ -299,3 +309,753 @@ def test_execute_on_ray_pickle_with_filter(self, shared_ray_cluster): expected = ray_lf.sort("date").collect() result = lf_unpickled.sort("date").collect() assert_frame_equal(expected, result) + + +# --------------------------------------------------------------------------- +# Generic partitioning: execute_on_ray(partitions=...) with by_key / explicit specs / builders +# --------------------------------------------------------------------------- + +from datetime import date, timedelta + +from polars_io_tools.io_sources.lazy_ray import ( + RayPartition, + cartesian_partitions, + discrete_partitions, +) +from polars_io_tools.io_sources.partitions import by_key, by_time, by_value +from polars_io_tools.io_sources.pushdown_combine import FilterSpec, pushdown_combine + +N_IDS = 6 + + +def generate_panel(n_ids: int = N_IDS, n_days: int = 10) -> pl.LazyFrame: + """A small (id x date) cross-section with a deterministic value column.""" + dates = pl.date_range(date(2023, 1, 1), date(2023, 1, n_days), interval="1d", eager=True) + ids = pl.DataFrame({"id": list(range(n_ids))}) + panel = ids.join(pl.DataFrame({"date": dates}), how="cross") + panel = panel.with_columns(val=(pl.col("id") * 100 + pl.col("date").dt.day())) + return panel.lazy() + + +def _count_tasks(monkeypatch): + calls = {"n": 0} + original_options = cpl.io_sources.lazy_ray._execute_partition.options + + def counting_options(**kw): + stub = original_options(**kw) + orig_remote = stub.remote + + def counting_remote(*args, **kwargs): + calls["n"] += 1 + return orig_remote(*args, **kwargs) + + stub.remote = counting_remote # type: ignore[attr-defined] + return stub + + monkeypatch.setattr(cpl.io_sources.lazy_ray._execute_partition, "options", counting_options) + return calls + + +def test_by_column_scalar_keys(monkeypatch): + lf = generate_panel() + calls = _count_tasks(monkeypatch) + + result = lf.piot.execute_on_ray(by_key(range(N_IDS), "id")).sort("id", "date").collect() + expected = lf.sort("id", "date").collect() + + assert_frame_equal(result, expected) + assert calls["n"] == N_IDS + + +def test_by_expr_hash_bucket(monkeypatch): + lf = generate_panel() + calls = _count_tasks(monkeypatch) + n_buckets = 3 + + result = lf.piot.execute_on_ray(by_key(range(n_buckets), pl.col("id").hash() % n_buckets)).sort("id", "date").collect() + expected = lf.sort("id", "date").collect() + + assert_frame_equal(result, expected) + assert calls["n"] == n_buckets + + +def test_discrete_partitions_is_in(): + lf = generate_panel() + specs = discrete_partitions("id", [[0, 1], [2, 3], [4, 5]]) + result = lf.piot.execute_on_ray(specs).sort("id", "date").collect() + expected = lf.sort("id", "date").collect() + assert_frame_equal(result, expected) + + +def test_cartesian_date_x_bucket(monkeypatch): + lf = generate_panel(n_days=10) + calls = _count_tasks(monkeypatch) + windows = [(date(2023, 1, 1), date(2023, 1, 6)), (date(2023, 1, 6), date(2023, 1, 11))] + specs = cartesian_partitions( + date_windows=windows, + buckets=[[0, 1, 2], [3, 4, 5]], + date_column="date", + bucket="id", + ) + result = lf.piot.execute_on_ray(specs).sort("id", "date").collect() + expected = lf.sort("id", "date").collect() + assert_frame_equal(result, expected) + assert calls["n"] == 4 # 2 windows x 2 buckets + + +def test_column_retention_under_projection(): + """The partition key column must be retained for the worker predicate, then dropped.""" + lf = generate_panel() + # `id` is the partition key but is NOT selected downstream. + result = lf.piot.execute_on_ray(by_key(range(N_IDS), "id")).select("val").sort("val").collect() + expected = lf.select("val").sort("val").collect() + assert result.columns == ["val"] + assert_frame_equal(result, expected) + + +def test_empty_partition_result_schema(): + lf = generate_panel() + # A key that matches nothing. + result = lf.piot.execute_on_ray(by_key([999], "id")).collect() + assert result.height == 0 + assert result.schema == lf.collect_schema() + + +def test_duplicate_key_rejected(): + lf = generate_panel() + with pytest.raises(ValueError, match="Duplicate partition key"): + lf.piot.execute_on_ray(by_key([1, 1], "id")).collect() + + +def test_completion_order_and_preserve_order(): + lf = generate_panel() + specs = discrete_partitions("id", [[i] for i in range(N_IDS)]) + + # completion order (default): correct rows regardless of arrival order + got = lf.piot.execute_on_ray(specs).sort("id", "date").collect() + expected = lf.sort("id", "date").collect() + assert_frame_equal(got, expected) + + # preserve order: id blocks appear in spec order + ordered = lf.piot.execute_on_ray(specs, preserve_partition_order=True).collect() + ids = ordered["id"].to_list() + assert ids == sorted(ids) + + +def test_per_partition_remote_options_struct_column(monkeypatch): + lf = generate_panel(n_ids=3) + seen = [] + original_options = cpl.io_sources.lazy_ray._execute_partition.options + + def capturing_options(**kw): + seen.append(kw) + return original_options(**kw) + + monkeypatch.setattr(cpl.io_sources.lazy_ray._execute_partition, "options", capturing_options) + + parts = pl.DataFrame({"id": [0, 1, 2], "ray": [{"num_cpus": 1}, {"num_cpus": 1}, {"num_cpus": 1}]}) + result = lf.piot.execute_on_ray(by_key(parts, "id", partition_remote_options="ray")).sort("id", "date").collect() + expected = lf.sort("id", "date").collect() + assert_frame_equal(result, expected) + assert all(kw.get("num_cpus") == 1 for kw in seen) + + +def test_per_partition_remote_options_mapping(monkeypatch): + lf = generate_panel(n_ids=3) + seen = [] + original_options = cpl.io_sources.lazy_ray._execute_partition.options + + def capturing_options(**kw): + seen.append(kw) + return original_options(**kw) + + monkeypatch.setattr(cpl.io_sources.lazy_ray._execute_partition, "options", capturing_options) + + result = lf.piot.execute_on_ray(by_key([0, 1, 2], "id", partition_remote_options={0: {"num_cpus": 1}})).sort("id", "date").collect() + expected = lf.sort("id", "date").collect() + assert_frame_equal(result, expected) + + +def test_validation_errors(): + lf = generate_panel() + with pytest.raises(ValueError, match="return_as"): + lf.piot.execute_on_ray(discrete_partitions("id", [[0]]), return_as="bogus") + with pytest.raises(ValueError, match="max_concurrency"): + lf.piot.execute_on_ray(discrete_partitions("id", [[0]]), max_concurrency=0) + + +def test_invariant1_lookback_self_padding(): + """Byte-identical single-frame vs distributed collect through a FilterSpec(lookback). + + Each date-window partition must self-pad its lookback: the partition predicate must + reach the inner source so it expands the read window, computes the rolling value, and + trims after combine. If pruning consumed the predicate, boundary rows would be wrong. + """ + dates = [date(2024, 1, i) for i in range(1, 13)] + vals = [float(i) for i in range(1, 13)] + data = pl.LazyFrame({"date": dates, "val": vals}) + + def combine_with_rolling(s): + return s["data"].sort("date").with_columns(pl.col("val").rolling_sum(window_size=3, min_samples=1).alias("rs")) + + lf = pushdown_combine( + sources={"data": (data, {"date": FilterSpec(lookback=timedelta(days=3))})}, + combine=combine_with_rolling, + ) + + single = lf.sort("date").collect() + + windows = [ + (date(2024, 1, 1), date(2024, 1, 5)), + (date(2024, 1, 5), date(2024, 1, 9)), + (date(2024, 1, 9), date(2024, 1, 13)), + ] + specs = [RayPartition((pl.col("date") >= lo) & (pl.col("date") < hi), key=i) for i, (lo, hi) in enumerate(windows)] + distributed = lf.piot.execute_on_ray(specs).sort("date").collect() + + assert_frame_equal(single, distributed) + + +# --------------------------------------------------------------------------- +# Regression tests for review findings (selectors, null/expr keys, series name, NaN dup) +# --------------------------------------------------------------------------- + +import polars.selectors as cs + + +def test_by_selector_multi_column(): + lf = generate_panel(n_ids=3, n_days=4) + lf2 = lf.with_columns(grp=(pl.col("id") % 2)) + parts = pl.DataFrame({"id": [0, 1, 2], "grp": [0, 1, 0]}) + result = lf2.piot.execute_on_ray(by_key(parts, cs.by_name("id", "grp"))).sort("id", "date").collect() + expected = lf2.sort("id", "date").collect() + assert_frame_equal(result, expected) + + +def test_selector_expanded_predicate_retains_columns(): + """A predicate whose root_names cannot be resolved must not lose columns under projection.""" + lf = pl.LazyFrame({"a": [1, -1, 1], "b": [1, 1, -1], "v": [10, 20, 30]}) + specs = [ + RayPartition(pl.all_horizontal(pl.col("a", "b") > 0), key="pos"), + RayPartition(~pl.all_horizontal(pl.col("a", "b") > 0), key="neg"), + ] + result = lf.piot.execute_on_ray(specs).select("v").sort("v").collect() + expected = lf.select("v").sort("v").collect() + assert result.columns == ["v"] + assert_frame_equal(result, expected) + + +def test_by_expr_null_key(): + lf = pl.LazyFrame({"id": [None, 1, 2], "v": [10, 20, 30]}, schema={"id": pl.Int64, "v": pl.Int64}) + result = lf.piot.execute_on_ray(by_key([None, 1, 2], pl.col("id"))).sort("v").collect() + expected = lf.sort("v").collect() + assert_frame_equal(result, expected) + + +def test_by_series_name_mismatch(): + lf = generate_panel(n_ids=3, n_days=4) + # Series named differently from `by`; should be aligned to the key column. + result = lf.piot.execute_on_ray(by_key(pl.Series("whatever", [0, 1, 2]), "id")).sort("id", "date").collect() + expected = lf.sort("id", "date").collect() + assert_frame_equal(result, expected) + + +def test_duplicate_nan_keys_rejected(): + lf = pl.LazyFrame({"x": [1.0, float("nan")], "v": [1, 2]}) + nan = float("nan") + with pytest.raises(ValueError, match="Duplicate partition key"): + lf.piot.execute_on_ray(by_key([nan, nan], "x")).collect() + + +# --------------------------------------------------------------------------- +# Predicate-pushdown key pruning (fail-open) +# --------------------------------------------------------------------------- + + +def test_pruning_equality_key(monkeypatch): + """A downstream `id == 3` should launch only the matching partition.""" + lf = generate_panel() + calls = _count_tasks(monkeypatch) + result = lf.piot.execute_on_ray(by_key(range(N_IDS), "id")).filter(pl.col("id") == 3).sort("date").collect() + expected = lf.filter(pl.col("id") == 3).sort("date").collect() + assert_frame_equal(result, expected) + assert calls["n"] == 1 + + +def test_partitions_primitive_no_pruning_correct_under_filter(monkeypatch): + """execute_on_ray does not prune caller-supplied explicit specs; results stay correct + under a downstream filter and every partition runs.""" + lf = generate_panel() + calls = _count_tasks(monkeypatch) + specs = discrete_partitions("id", [[0, 1], [2, 3], [4, 5]]) + result = lf.piot.execute_on_ray(specs).filter(pl.col("id").is_in([0, 5])).sort("id", "date").collect() + expected = lf.filter(pl.col("id").is_in([0, 5])).sort("id", "date").collect() + assert_frame_equal(result, expected) + assert calls["n"] == 3 # caller controls the spec set; no pruning + + +def test_cartesian_no_pruning_correct_under_filter(monkeypatch): + lf = generate_panel(n_days=10) + calls = _count_tasks(monkeypatch) + windows = [(date(2023, 1, 1), date(2023, 1, 6)), (date(2023, 1, 6), date(2023, 1, 11))] + specs = cartesian_partitions(date_windows=windows, buckets=[[0, 1, 2], [3, 4, 5]], date_column="date", bucket="id") + result = ( + lf.piot.execute_on_ray(specs).filter(pl.col("date") < date(2023, 1, 6)).filter(pl.col("id").is_in([0, 1, 2])).sort("id", "date").collect() + ) + expected = lf.filter(pl.col("date") < date(2023, 1, 6)).filter(pl.col("id").is_in([0, 1, 2])).sort("id", "date").collect() + assert_frame_equal(result, expected) + assert calls["n"] == 4 # all cells run; caller controls the spec set + + +def test_pruning_fail_open_unrelated_predicate(monkeypatch): + """A downstream predicate on a non-key column must NOT prune any partition.""" + lf = generate_panel() + calls = _count_tasks(monkeypatch) + result = lf.piot.execute_on_ray(by_key(range(N_IDS), "id")).filter(pl.col("val") > 500).sort("id", "date").collect() + expected = lf.filter(pl.col("val") > 500).sort("id", "date").collect() + assert_frame_equal(result, expected) + assert calls["n"] == N_IDS # nothing pruned + + +def test_pruning_never_drops_matching(monkeypatch): + """Pruning must not drop a partition that can still contribute rows (correctness).""" + lf = generate_panel() + result = lf.piot.execute_on_ray(by_key(range(N_IDS), "id")).filter(pl.col("id") >= 2).sort("id", "date").collect() + expected = lf.filter(pl.col("id") >= 2).sort("id", "date").collect() + assert_frame_equal(result, expected) + + +def test_pruning_no_false_prune_nan_in_is_in(): + """NaN member in an is_in must not be dropped into a false prune.""" + lf = pl.DataFrame({"x": [2.0, float("nan")], "v": [1, 2]}).lazy() + specs = discrete_partitions("x", [[2.0, float("nan")]]) + result = lf.piot.execute_on_ray(specs).filter(pl.col("x").is_in([1.0, float("nan")])).collect() + expected = lf.filter(pl.col("x").is_in([1.0, float("nan")])).collect() + assert_frame_equal(result.sort("v"), expected.sort("v")) + assert result.height == 1 # the NaN row must survive + + +def test_pruning_no_false_prune_float_coercion(): + """Int/float coercion must not cause a false prune.""" + v = 2**53 + 1 + lf = pl.DataFrame({"id": [v], "w": [7]}, schema={"id": pl.Int64, "w": pl.Int64}).lazy() + result = lf.piot.execute_on_ray(by_key([v], "id")).filter(pl.col("id") == float(2**53)).collect() + expected = lf.filter(pl.col("id") == float(2**53)).collect() + assert_frame_equal(result, expected) + assert result.height == expected.height + + +def test_by_expr_cast_no_false_prune(): + """A cast expression key must not false-prune (expression keys are never pruned).""" + lf = pl.DataFrame({"id": [2]}, schema={"id": pl.Int64}).lazy() + result = lf.piot.execute_on_ray(by_key([True], pl.col("id").cast(pl.Boolean))).filter(pl.col("id") == 2).collect() + expected = lf.filter(pl.col("id") == 2).collect() + assert_frame_equal(result, expected) + assert result.height == 1 + + +def test_partitions_multicolumn_null_no_false_prune(): + """Arbitrary specs with null-matching predicates must not be pruned.""" + schema = pl.Schema({"c": pl.String, "d": pl.Int64}) + lf = pl.DataFrame({"c": [None], "d": [1]}, schema=schema).lazy() + specs = [RayPartition(pl.col("c").is_in(["b", None], nulls_equal=True) & (pl.col("d") == 1), key="p")] + downstream = pl.col("c").is_in(["a", None], nulls_equal=True) & (pl.col("d") == 1) + result = lf.piot.execute_on_ray(specs).filter(downstream).collect() + expected = lf.filter(downstream).collect() + assert_frame_equal(result, expected) + assert result.height == 1 # the null row survives + + +def test_pruning_by_nan_key_exact(): + """Exact key evaluation keeps a NaN partition matched by the downstream predicate.""" + lf = pl.DataFrame({"x": [1.0, float("nan")], "v": [1, 2]}).lazy() + result = lf.piot.execute_on_ray(by_key([1.0, float("nan")], "x")).filter(pl.col("x").is_in([float("nan")])).collect() + expected = lf.filter(pl.col("x").is_in([float("nan")])).collect() + assert_frame_equal(result.sort("v"), expected.sort("v")) + + +def test_pruning_float_key_no_false_prune(): + """Float keys are not observational under == (-0.0 vs +0.0); must fail open.""" + lf = pl.DataFrame({"x": [-0.0], "v": [1]}).lazy() + predicate = (pl.lit(1.0) / pl.col("x")) < 0 + result = lf.piot.execute_on_ray(by_key([0.0], "x")).filter(predicate).collect() + expected = lf.filter(predicate).collect() + assert_frame_equal(result, expected) + assert result.height == 1 + + +def test_pruning_python_udf_no_false_prune(): + """A Python UDF predicate must not be pruned via single-key evaluation.""" + + class EverySecond: + def __init__(self): + self.n = 0 + + def __call__(self, _): + self.n += 1 + return self.n % 2 == 0 + + lf = pl.DataFrame({"id": [1, 1], "v": [10, 11]}, schema={"id": pl.Int64, "v": pl.Int64}).lazy() + predicate = pl.col("id").map_elements(EverySecond(), return_dtype=pl.Boolean) + result = lf.piot.execute_on_ray(by_key([1], "id")).filter(predicate).collect() + expected = lf.filter(predicate).collect() + assert_frame_equal(result.sort("v"), expected.sort("v")) + + +def test_cartesian_none_bucket_matches_nulls(): + """A None bucket value must match null rows via is_null(), not `== None`.""" + lf = pl.DataFrame( + { + "date": [date(2026, 1, 1), date(2026, 1, 1)], + "desk": [None, "rates"], + "value": [100, 200], + }, + schema={"date": pl.Date, "desk": pl.String, "value": pl.Int64}, + ).lazy() + specs = cartesian_partitions( + date_windows=[(date(2026, 1, 1), date(2026, 1, 2))], + buckets={"missing-desk": None, "rates": "rates"}, + date_column="date", + bucket="desk", + ) + result = lf.piot.execute_on_ray(specs).sort("value").collect() + expected = lf.sort("value").collect() + assert_frame_equal(result, expected) + assert result.height == 2 # both the null-desk and the rates row + + +def test_discrete_partitions_none_member_matches_nulls(): + """A None member in a discrete group must include null rows.""" + lf = pl.DataFrame({"id": [None, 1, 2], "v": [10, 20, 30]}, schema={"id": pl.Int64, "v": pl.Int64}).lazy() + specs = discrete_partitions("id", [[1, None], [2]]) + result = lf.piot.execute_on_ray(specs).sort("v").collect() + expected = lf.sort("v").collect() + assert_frame_equal(result, expected) + assert result.height == 3 + + +def test_selector_downstream_predicate_retains_columns(): + """A pushed-down selector-expanded predicate (empty root_names) must see all columns.""" + lf = pl.DataFrame({"a": [1, None], "b": [None, 2]}, schema={"a": pl.Int64, "b": pl.Int64}).lazy() + specs = discrete_partitions("a", [[1], [2]]) # non-selector partition preds + selector_predicate = pl.all_horizontal(pl.all().is_not_null()) + # Downstream selector predicate + projection to only "a": must evaluate over full schema. + result = lf.piot.execute_on_ray(specs).filter(selector_predicate).select("a").collect() + expected = lf.filter(selector_predicate).select("a").collect() + assert_frame_equal(result.sort("a"), expected.sort("a")) + assert result.height == 0 # no row has both a and b non-null + + +# --- Unified `partitions=` entry point (shared ReadPartition vocabulary) -------------------- + + +def _category_lazyframe(start, end): + dates = pl.datetime_range(start, end, interval="1d", eager=True) + return pl.LazyFrame( + { + "date": dates, + "category": [("A", "B", "C")[i % 3] for i in range(len(dates))], + "val": range(len(dates)), + } + ) + + +def test_unified_by_time_matches_calendar_and_single_frame(): + lf = generate_sample_lazyframe((s := datetime.datetime(2023, 1, 1)), (e := datetime.datetime(2023, 6, 30))) + pipeline = lf.filter(pl.col("quantity") > 10).with_columns(pl.col("price") * 2) + flt = pl.col("date").is_between(s, e) + + expected = pipeline.filter(flt).sort("date").collect() + unified = pipeline.piot.execute_on_ray(cpl.by_time("date", "1mo")).filter(flt).sort("date").collect() + legacy = pipeline.piot.execute_on_ray(date_column="date", time_unit="monthly").filter(flt).sort("date").collect() + + assert_frame_equal(unified, expected) + assert_frame_equal(unified, legacy) + + +def test_unified_by_value_derived_from_predicate(): + lf = _category_lazyframe(datetime.datetime(2023, 1, 1), datetime.datetime(2023, 3, 31)) + flt = pl.col("category").is_in(["A", "B"]) + + expected = lf.filter(flt).sort("date").collect() + got = lf.piot.execute_on_ray(cpl.by_value("category")).filter(flt).sort("date").collect() + + assert_frame_equal(got, expected) + + +def test_unified_explicit_read_partition_list(): + lf = generate_sample_lazyframe(datetime.datetime(2023, 1, 1), datetime.datetime(2023, 3, 31)) + mid = datetime.datetime(2023, 2, 15) + parts = [ + cpl.ReadPartition(pl.col("date") < mid, key="lo"), + cpl.ReadPartition(pl.col("date") >= mid, key="hi"), + ] + + expected = lf.sort("date").collect() + got = lf.piot.execute_on_ray(parts).sort("date").collect() + + assert_frame_equal(got, expected) + + +def test_unified_partitions_and_legacy_are_mutually_exclusive(): + lf = generate_sample_lazyframe(datetime.datetime(2023, 1, 1), datetime.datetime(2023, 3, 31)) + with pytest.raises(ValueError, match="either"): + lf.piot.execute_on_ray(cpl.by_time("date", "1mo"), date_column="date", time_unit="monthly") + + +def test_unified_by_value_explicit_needs_no_predicate(): + # An explicit value list is fully specified: it must run without any pushed-down predicate. + lf = _category_lazyframe(datetime.datetime(2023, 1, 1), datetime.datetime(2023, 2, 28)) + expected = lf.filter(pl.col("category").is_in(["A", "B"])).sort("date").collect() + got = lf.piot.execute_on_ray(by_value("category", ["A", "B"])).sort("date").collect() + assert_frame_equal(got, expected) + + +def test_unified_by_time_empty_range_returns_empty(): + # A contradictory date filter yields no partitions -> an empty frame, not an error. + lf = generate_sample_lazyframe(datetime.datetime(2023, 1, 1), datetime.datetime(2023, 12, 31)) + flt = (pl.col("date") >= datetime.datetime(2023, 6, 1)) & (pl.col("date") < datetime.datetime(2023, 1, 1)) + got = lf.piot.execute_on_ray(by_time("date", "1mo")).filter(flt).collect() + assert got.height == 0 + + +def test_explicit_generator_partitions_are_reusable(): + # A generator of ReadPartitions must be materialised once, so re-collecting the same + # LazyFrame does not silently yield an empty frame the second time. + lf = generate_sample_lazyframe(datetime.datetime(2023, 1, 1), datetime.datetime(2023, 3, 31)) + mid = datetime.datetime(2023, 2, 15) + + def gen(): + yield cpl.ReadPartition(pl.col("date") < mid, key="lo") + yield cpl.ReadPartition(pl.col("date") >= mid, key="hi") + + ray_lf = lf.piot.execute_on_ray(gen()) + first = ray_lf.sort("date").collect() + second = ray_lf.sort("date").collect() + assert first.height == lf.collect().height + assert_frame_equal(first, second) + + +# --------------------------------------------------------------------------- +# Bounded, ordered fan-out (deque window). Tasks are always submitted in spec +# order, so the k-th submission corresponds to spec index k; the helpers below +# exploit that to force deterministic completion / failure schedules. They patch +# ``.options`` (not ``.remote``) because the executor submits via +# ``_execute_partition.options(...).remote(...)``. +# --------------------------------------------------------------------------- + + +@ray.remote(num_cpus=0) +def _sleep_then(delay, boxed_ref): + time.sleep(delay) + return ray.get(boxed_ref[0]) + + +@ray.remote(num_cpus=0) +def _raise_or(fail, boxed_ref): + if fail: + raise RuntimeError("worker boom") + return ray.get(boxed_ref[0]) + + +def _patch_worker(monkeypatch, wrap): + """Route each real task's ObjectRef through ``wrap(k, submit_real)``. + + ``k`` is the 0-based submission index (== spec index); ``submit_real()`` submits the genuine + partition task and returns its ObjectRef. Refs are boxed in a list so Ray does not eagerly + resolve them as task arguments. + """ + original_options = cpl.io_sources.lazy_ray._execute_partition.options + counter = {"n": 0} + + def counting_options(**kw): + stub = original_options(**kw) + orig_remote = stub.remote + + def wrapped_remote(*args, **kwargs): + k = counter["n"] + counter["n"] += 1 + return wrap(k, lambda: orig_remote(*args, **kwargs)) + + stub.remote = wrapped_remote # type: ignore[attr-defined] + return stub + + monkeypatch.setattr(cpl.io_sources.lazy_ray._execute_partition, "options", counting_options) + + +def _patch_delays(monkeypatch, delays): + """Make the k-th submitted task complete after ``delays[k]`` extra seconds.""" + + def wrap(k, submit_real): + delay = delays[k] if k < len(delays) else 0.0 + return _sleep_then.remote(delay, [submit_real()]) + + _patch_worker(monkeypatch, wrap) + + +def _patch_failures(monkeypatch, fail_indices): + """Make tasks at the given spec indices raise inside the worker.""" + + def wrap(k, submit_real): + return _raise_or.remote(k in fail_indices, [submit_real()]) + + _patch_worker(monkeypatch, wrap) + + +@pytest.mark.parametrize( + "delays", + [ + [0.20, 0.15, 0.10, 0.05, 0.02, 0.0], # reverse: head is the straggler + [0.20, 0.0, 0.0, 0.0, 0.0, 0.0], # head straggler only + [0.0, 0.0, 0.20, 0.0, 0.0, 0.0], # middle straggler + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], # all equal + ], + ids=["reverse", "head", "middle", "equal"], +) +def test_ordered_preserves_order_under_reorder(monkeypatch, delays): + lf = generate_panel() + specs = discrete_partitions("id", [[i] for i in range(N_IDS)]) + _patch_delays(monkeypatch, delays) + + ordered = lf.piot.execute_on_ray(specs, preserve_partition_order=True, remote_options={"num_cpus": 0}).collect() + ids = ordered["id"].to_list() + assert ids == sorted(ids) + assert_frame_equal(ordered.sort("id", "date"), lf.sort("id", "date").collect()) + + +@pytest.mark.parametrize("max_concurrency", [1, 2, 100]) +def test_ordered_correct_for_window_sizes(monkeypatch, max_concurrency): + lf = generate_panel() + specs = discrete_partitions("id", [[i] for i in range(N_IDS)]) + # Reverse completion order stresses the window at every size. + _patch_delays(monkeypatch, [0.15, 0.12, 0.09, 0.06, 0.03, 0.0]) + + ordered = lf.piot.execute_on_ray(specs, preserve_partition_order=True, max_concurrency=max_concurrency, remote_options={"num_cpus": 0}).collect() + assert ordered["id"].to_list() == sorted(ordered["id"].to_list()) + assert_frame_equal(ordered.sort("id", "date"), lf.sort("id", "date").collect()) + + +def test_ordered_single_partition(): + lf = generate_panel(n_ids=1) + got = lf.piot.execute_on_ray(discrete_partitions("id", [[0]]), preserve_partition_order=True).sort("date").collect() + assert_frame_equal(got, lf.sort("date").collect()) + + +def test_ordered_window_never_exceeds_max_concurrency(monkeypatch): + """The deque must never keep more than ``max_concurrency`` tasks running at once.""" + max_conc = 3 + + @ray.remote + class _Gauge: + def __init__(self): + self.active = 0 + self.peak = 0 + + def enter(self): + self.active += 1 + self.peak = max(self.peak, self.active) + return self.active + + def leave(self): + self.active -= 1 + + def peak_value(self): + return self.peak + + gauge = _Gauge.remote() + + @ray.remote(num_cpus=0) + def fake_part(*_a, **_k): + active = ray.get(gauge.enter.remote()) + assert active <= max_conc + time.sleep(0.05) + ray.get(gauge.leave.remote()) + return pl.DataFrame({"id": [0], "date": [datetime.date(2023, 1, 1)], "val": [0]}).to_arrow() + + original_options = cpl.io_sources.lazy_ray._execute_partition.options + + def counting_options(**kw): + stub = original_options(**kw) + stub.remote = fake_part.remote # type: ignore[attr-defined] + return stub + + monkeypatch.setattr(cpl.io_sources.lazy_ray._execute_partition, "options", counting_options) + + specs = discrete_partitions("id", [[i] for i in range(12)]) + got = ( + generate_panel(n_ids=12) + .piot.execute_on_ray(specs, preserve_partition_order=True, max_concurrency=max_conc, remote_options={"num_cpus": 0}) + .collect() + ) + assert got.height == 12 # every partition was consumed + assert ray.get(gauge.peak_value.remote()) == max_conc + + +@_ray_cancel_flaky_on_windows +def test_ordered_n_rows_stops_submitting(monkeypatch): + """With n_rows satisfied by the first partition, only the initial window is submitted.""" + lf = generate_panel(n_ids=8) # each id block has 10 rows + specs = discrete_partitions("id", [[i] for i in range(8)]) + calls = {"n": 0} + original_options = cpl.io_sources.lazy_ray._execute_partition.options + + def counting_options(**kw): + stub = original_options(**kw) + orig_remote = stub.remote + + def counting_remote(*args, **kwargs): + calls["n"] += 1 + return orig_remote(*args, **kwargs) + + stub.remote = counting_remote # type: ignore[attr-defined] + return stub + + monkeypatch.setattr(cpl.io_sources.lazy_ray._execute_partition, "options", counting_options) + + got = lf.piot.execute_on_ray(specs, preserve_partition_order=True, max_concurrency=2).head(1).collect() + assert got.height == 1 + assert got["id"].to_list() == [0] + assert calls["n"] == 2, f"expected only the initial window submitted, got {calls['n']}" + + +@_ray_cancel_flaky_on_windows +def test_ordered_required_partition_failure_raises(monkeypatch): + lf = generate_panel() + specs = discrete_partitions("id", [[i] for i in range(N_IDS)]) + _patch_failures(monkeypatch, {0}) # the head is required -> must surface + with pytest.raises((RuntimeError, pl.exceptions.ComputeError), match="partition 0"): + lf.piot.execute_on_ray(specs, preserve_partition_order=True, remote_options={"num_cpus": 0}).collect() + + +@_ray_cancel_flaky_on_windows +def test_ordered_failure_beyond_satisfied_n_rows_is_ignored(monkeypatch): + """A later partition that fails is never observed once an ordered prefix satisfies n_rows.""" + lf = generate_panel(n_ids=6) + specs = discrete_partitions("id", [[i] for i in range(6)]) + # All submitted at once (max_concurrency defaults high); the last partition fails on the + # worker, but we stop after partition 0 and drop the rest without ever calling ray.get. + _patch_failures(monkeypatch, {5}) + got = lf.piot.execute_on_ray(specs, preserve_partition_order=True, remote_options={"num_cpus": 0}).head(1).collect() + assert got.height == 1 + assert got["id"].to_list() == [0] + + +@_ray_cancel_flaky_on_windows +def test_ordered_exact_n_rows_does_not_fetch_next(monkeypatch): + """When the first partition supplies *exactly* n_rows, the next partition must not be + fetched -- a failure in it would otherwise surface incorrectly.""" + lf = generate_panel(n_ids=2) # each id block has exactly 10 rows + specs = discrete_partitions("id", [[0], [1]]) + _patch_failures(monkeypatch, {1}) # partition 1 fails if we ever fetch it + got = lf.piot.execute_on_ray(specs, preserve_partition_order=True, remote_options={"num_cpus": 0}).head(10).collect() + assert got.height == 10 + assert got["id"].to_list() == [0] * 10 + + +@_ray_cancel_flaky_on_windows +def test_completion_order_failure_raises(monkeypatch): + lf = generate_panel() + specs = discrete_partitions("id", [[i] for i in range(N_IDS)]) + _patch_failures(monkeypatch, {3}) + with pytest.raises((RuntimeError, pl.exceptions.ComputeError), match="partition 3"): + lf.piot.execute_on_ray(specs, preserve_partition_order=False, remote_options={"num_cpus": 0}).collect() diff --git a/polars_io_tools/tests/io_sources/test_range_visitor.py b/polars_io_tools/tests/io_sources/test_range_visitor.py index 0fd9869..b925cc3 100644 --- a/polars_io_tools/tests/io_sources/test_range_visitor.py +++ b/polars_io_tools/tests/io_sources/test_range_visitor.py @@ -1584,85 +1584,51 @@ def test_mixed_bounds_compound(self): _test_expression_with_df(result, index_col, test_dates, [2, 3, 5, 6]) -class TestPartitionPruningLogic: - """ - Regression tests for the closed="left" bug: verify that partition specs - generated by _partition_specs are correctly trimmed by intersecting - with the extracted temporal interval. +class TestDateWindows: + """Calendar window generation (``partitions._date_windows``): calendar flooring and the + exclusive-vs-inclusive upper-bound handling that the ``closed="left"`` regression guards.""" - These tests exercise the exact code path used by execute_on_ray's - source_generator without needing Ray or the Polars optimizer. - """ + @staticmethod + def _windows(lower, upper, right, every): + import portion + + from polars_io_tools.io_sources.partitions import _date_windows + + return _date_windows(portion.Interval.from_atomic(portion.CLOSED, lower, upper, right), every) + + def test_exclusive_upper_drops_boundary_window(self): + import portion + + w = self._windows(datetime(2024, 1, 2), datetime(2024, 1, 10), portion.OPEN, "1d") + assert len(w) == 8 # Jan 2..Jan 9; the Jan 10 window is excluded + assert w[-1] == (datetime(2024, 1, 9), datetime(2024, 1, 10)) + assert datetime(2024, 1, 10) not in {lo for lo, _ in w} + + def test_inclusive_upper_keeps_boundary_window(self): + import portion - def test_closed_left_removes_upper_bound_partition(self): - """ - closed="left" produces [Jan 2, Jan 10). _partition_specs includes - a partition for Jan 10, but [Jan 10, Jan 11) ∩ [Jan 2, Jan 10) = ∅ - so it must be removed entirely. - """ - from polars_io_tools.io_sources.lazy_ray import _partition_specs, _trim_partition_specs - - start = date(2024, 1, 2) - end = date(2024, 1, 10) - - predicate = pl.col("ts").is_between(start, end, closed="left") - date_interval = convert_expr_to_datetime_range(predicate, "ts", get_enclosure=False) - assert not date_interval.empty - - specs = _partition_specs(date_interval.lower, date_interval.upper, "daily") - assert len(specs) == 9 # Jan 2..10 inclusive before trimming - - trimmed = _trim_partition_specs(specs, date_interval, pl.Date) - - assert len(trimmed) == 8, f"Expected 8 partitions (Jan 2-9), got {len(trimmed)}. The Jan 10 partition should have been removed." - trimmed_starts = {s for s, _ in trimmed} - assert datetime(2024, 1, 10) not in trimmed_starts - - def test_closed_both_keeps_all_partitions(self): - """closed='both' (default) must NOT remove any partitions.""" - from polars_io_tools.io_sources.lazy_ray import _partition_specs, _trim_partition_specs - - start = date(2024, 1, 2) - end = date(2024, 1, 10) - - predicate = pl.col("ts").is_between(start, end) - date_interval = convert_expr_to_datetime_range(predicate, "ts", get_enclosure=False) - specs = _partition_specs(date_interval.lower, date_interval.upper, "daily") - - trimmed = _trim_partition_specs(specs, date_interval, pl.Date) - - assert len(trimmed) == len(specs), f"closed='both' should not remove any partitions, but {len(specs) - len(trimmed)} were removed" - - def test_monthly_both_bounds_trimmed(self): - """ - Both start and end are trimmed to the intersection with the user's - interval. Closed upper bounds are extended via _extend_interval so - that _execute_partition's `col < end` still includes the boundary value. - """ - from polars_io_tools.io_sources.lazy_ray import _partition_specs, _trim_partition_specs - - # User wants Jan 15 through Mar 10 (Datetime column, us precision) - start = datetime(2024, 1, 15) - end = datetime(2024, 3, 10) - col_type = pl.Datetime("us") - - predicate = pl.col("ts").is_between(start, end) - date_interval = convert_expr_to_datetime_range(predicate, "ts", get_enclosure=False) - specs = _partition_specs(date_interval.lower, date_interval.upper, "monthly") - - # Raw specs: [Jan 1, Feb 1), [Feb 1, Mar 1), [Mar 1, Apr 1) - assert len(specs) == 3 - - trimmed = _trim_partition_specs(specs, date_interval, col_type) - - assert len(trimmed) == 3 # all three months have overlap - # First partition start trimmed: Jan 15, not Jan 1 - assert trimmed[0][0] == datetime(2024, 1, 15) - # First partition end: Feb 1 (intersection is [Jan 15, Feb 1), open upper → kept) - assert trimmed[0][1] == datetime(2024, 2, 1) - # Middle partition unchanged (intersection is [Feb 1, Mar 1), open upper) - assert trimmed[1] == (datetime(2024, 2, 1), datetime(2024, 3, 1)) - # Last partition: intersection is [Mar 1, Mar 10] (closed upper from user's filter) - # → end extended by _extend_interval: +1µs for Datetime("us") - assert trimmed[2][0] == datetime(2024, 3, 1) - assert trimmed[2][1] == datetime(2024, 3, 10, 0, 0, 0, 1) + w = self._windows(datetime(2024, 1, 2), datetime(2024, 1, 10), portion.CLOSED, "1d") + assert len(w) == 9 # Jan 2..Jan 10 inclusive + assert w[-1] == (datetime(2024, 1, 10), datetime(2024, 1, 11)) + + def test_monthly_floors_to_calendar_boundary(self): + import portion + + w = self._windows(datetime(2024, 1, 15), datetime(2024, 3, 10), portion.CLOSED, "1mo") + assert len(w) == 3 + assert w[0] == (datetime(2024, 1, 1), datetime(2024, 2, 1)) # floored to Jan 1, not Jan 15 + assert w[1] == (datetime(2024, 2, 1), datetime(2024, 3, 1)) + assert w[-1] == (datetime(2024, 3, 1), datetime(2024, 4, 1)) + + def test_disjoint_interval_skips_the_gaps(self): + import portion + + from polars_io_tools.io_sources.partitions import _date_windows + + # January and December only: two monthly windows, not twelve across the gap between them. + iv = portion.closedopen(datetime(2024, 1, 1), datetime(2024, 1, 15)) | portion.closedopen(datetime(2024, 12, 1), datetime(2024, 12, 15)) + w = _date_windows(iv, "1mo") + assert w == [ + (datetime(2024, 1, 1), datetime(2024, 2, 1)), + (datetime(2024, 12, 1), datetime(2025, 1, 1)), + ]