refactor(ray): unify distributed partitioning on a shared vocabulary

Collapse the three distributed entry points into a single `execute_on_ray(partitions=...)`
and introduce a backend-neutral partition vocabulary in `io_sources/partitions.py`.

A partition is a `ReadPartition(predicate, key)`; `RayPartition` subclasses it to add
per-task `remote_options`. Partitions come from builders:

- `by_time(column, every)`  — calendar windows from the pushed-down date range
- `by_value(column, values)` — discrete values (derived from an IN filter when omitted)
- `by_range(column, every)`  — fixed-width numeric buckets
- `by_key(partitions, by)`   — equality on caller-enumerated keys (was execute_on_ray_by)
- `discrete_partitions` / `cartesian_partitions` — member lists and date x bucket products

`execute_on_ray` accepts a partitioner or an explicit ReadPartition list; the former
requires a bounded pushed-down predicate. The legacy `date_column`/`time_unit` calendar
call is kept as a shortcut. `execute_on_ray_by` and `execute_on_ray_partitions` are removed
(the first becomes `by_key`, the second is just an explicit list).

The vocabulary is intentionally backend-neutral so a SQL reader can consume the same
partitions. Update the ray tests to the unified entry point and add parity tests for the
builders; refresh the README and API-Reference.

Signed-off-by: Pascal Tomecek <40371786+ptomecek@users.noreply.github.com>

Author: ptomecek

README.md

@@ -90,9 +90,10 @@ 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
-  across an existing Ray cluster; `execute_on_ray_by` / `execute_on_ray_partitions`
-  partition by an arbitrary key or explicit specs. (For multi-stage distributed pipelines,
+- **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

docs/wiki/API-Reference.md

@@ -110,40 +110,32 @@ 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`.
-
-### `execute_on_ray_by`
-
-```python
-lf.piot.execute_on_ray_by(partitions, by, *, remote_options=None,
-                          partition_remote_options=None, return_as="arrow",
-                          max_concurrency=100, preserve_partition_order=False)
-```
-
-Distribute the LazyFrame across Ray by equality on caller-enumerated partition keys.
-`partitions` is a small frame (or `Series`/sequence) with one row per partition; `by` is a
-column name, list of names, selector, or a `pl.Expr` (e.g. `pl.col("id").hash() % N`). Each
-row runs one Ray task filtered to `col == key`. `partition_remote_options` sets per-partition
-Ray options from a struct column in `partitions` or a `{key: dict}` mapping.
-
-### `execute_on_ray_partitions`
-
-```python
-lf.piot.execute_on_ray_partitions(specs, *, remote_options=None, return_as="arrow",
-                                  max_concurrency=100, preserve_partition_order=False)
-```
-
-Distribute the LazyFrame using an explicit list of `RayPartition(predicate, key, remote_options)` specs — the general primitive for ranges, discrete member lists
-(`col.is_in(...)`), and compound predicates. Build common spec sets with `discrete_partitions`
-and `cartesian_partitions`.
-
-Chaining multiple `execute_on_ray*` calls relies on predicate pushdown surviving intervening
+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/).
 

polars_io_tools/io_sources/__init__.py

@@ -18,6 +18,7 @@
 from .lazy_iter_rows import *
 from .lazy_narwhals_reader import *
 from .lazy_sql_reader import *
+from .partitions import *
 from .pushdown_combine import *
 from .pushdown_pivot import *
 from .pushdown_unpivot import *
@@ -37,6 +38,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,17 +52,11 @@
     # 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,
-        execute_on_ray_by as _execute_on_ray_by_proto,
-        execute_on_ray_partitions as _execute_on_ray_partitions_proto,
     )
 else:
 
     def _execute_on_ray_proto(*_a, **_kw): ...
 
-    def _execute_on_ray_by_proto(*_a, **_kw): ...
-
-    def _execute_on_ray_partitions_proto(*_a, **_kw): ...
-
 
 @pl.api.register_lazyframe_namespace("piot")
 class PIOTOperations:
@@ -99,20 +95,6 @@ def execute_on_ray(self, *args, **kwargs) -> pl.LazyFrame:
 
         return _execute_on_ray(self._lf, *args, **kwargs)
 
-    @functools.wraps(_execute_on_ray_by_proto)
-    def execute_on_ray_by(self, *args, **kwargs) -> pl.LazyFrame:
-        # heavy import happens only when the user calls the method
-        from .lazy_ray import execute_on_ray_by as _execute_on_ray_by
-
-        return _execute_on_ray_by(self._lf, *args, **kwargs)
-
-    @functools.wraps(_execute_on_ray_partitions_proto)
-    def execute_on_ray_partitions(self, *args, **kwargs) -> pl.LazyFrame:
-        # heavy import happens only when the user calls the method
-        from .lazy_ray import execute_on_ray_partitions as _execute_on_ray_partitions
-
-        return _execute_on_ray_partitions(self._lf, *args, **kwargs)
-
     @functools.wraps(filtered_join_asof)
     def filtered_join_asof(self, *args, **kwargs) -> pl.LazyFrame:
         if self._DISABLE_OPTIMIZATIONS: