Generalize execute_on_ray partitioning (arbitrary key, explicit specs)

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.

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`
 

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): ...