Generalize execute_on_ray partitioning (arbitrary key, explicit specs)

[ptomecek]

Description

execute_on_ray currently partitions only on a single datetime column by calendar unit
(daily / monthly / yearly). This generalizes the distributed source to partition on
arbitrary keys and explicit specs, and refactors the calendar path onto a shared execution
core (its public behavior is unchanged). Closes #24.

Two new LazyFrame entry points on the piot namespace:

• execute_on_ray_by(partitions, by, ...) — partition by equality on caller-enumerated
  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.
• execute_on_ray_partitions(specs, ...) — the general primitive taking an explicit list
  of RayPartition(predicate, key, remote_options) specs. Helpers discrete_partitions
  (member lists via is_in) and cartesian_partitions (date_windows × buckets) build
  common spec sets.

Changes

• New shared core that distributes a LazyFrame across Ray from a list of partition
  predicates; the calendar execute_on_ray and both new entry points build specs over it.
• Columns a partition predicate needs are retained through projection pushdown and dropped
  from results; empty results use the requested output schema.
• Per-partition remote_options — a uniform dict, a struct column, a {key: dict} mapping,
  or per-RayPartition.
• Results stream in completion order by default (preserve_partition_order=True restores
  spec order); execute_on_ray keeps chronological order.
• Duplicate partition keys are rejected; return_as / max_concurrency are validated up
  front; outstanding tasks are cancelled when n_rows is satisfied or a partition fails.
• Each partition executes inner.filter(<pushed-down predicate>).filter(<partition predicate>), so a partition predicate reaches inner sources (e.g. a lookback-expanding
  filter) unchanged.
• Optional pruning: execute_on_ray_by skips partitions whose key a pushed-down predicate on
  the key column(s) cannot match, evaluated exactly with Polars and restricted to non-float
  keys and non-UDF predicates so a matching partition can never be dropped. The explicit
  execute_on_ray_partitions spec list is never pruned.

Docs

API-Reference entries for the new functions and a note pointing multi-stage distributed
pipelines at Polars Cloud.

Tests

Partitioning by column and by expression, discrete member lists, compound
date_window × bucket, projection / column retention, empty-result schema, duplicate-key
rejection, completion vs preserved ordering, per-partition options, pruning (including
no-false-prune cases for NaN/is_in, float coercion, -0.0, casts, UDFs, and multi-column
nulls), and byte-identical single-frame vs distributed results through a lookback source.
test_lazy_ray (38) and the full polars_io_tools suite pass; ruff and mdformat clean.

Type of Change

• New feature

Checklist

• Linting passes (make lint-py, make lint-docs)
• Tests pass (make test-py)
• New tests added for new functionality
• Documentation updated (if applicable)
• Changelog / version bump (if applicable) — version bumps are handled separately

[github-actions]

Test Results

1 524 tests  +54   1 476 ✅ +54   50s ⏱️ +17s
    2 suites ± 0      48 💤 ± 0 
    2 files   ± 0       0 ❌ ± 0 

Results for commit 836f22c. ± Comparison against base commit 36d2260.

This pull request removes 3 and adds 57 tests. Note that renamed tests count towards both.

polars_io_tools.tests.io_sources.test_range_visitor.TestPartitionPruningLogic ‑ test_closed_both_keeps_all_partitions
polars_io_tools.tests.io_sources.test_range_visitor.TestPartitionPruningLogic ‑ test_closed_left_removes_upper_bound_partition
polars_io_tools.tests.io_sources.test_range_visitor.TestPartitionPruningLogic ‑ test_monthly_both_bounds_trimmed

polars_io_tools.tests.io_sources.test_lazy_ray ‑ test_by_column_scalar_keys
polars_io_tools.tests.io_sources.test_lazy_ray ‑ test_by_expr_cast_no_false_prune
polars_io_tools.tests.io_sources.test_lazy_ray ‑ test_by_expr_hash_bucket
polars_io_tools.tests.io_sources.test_lazy_ray ‑ test_by_expr_null_key
polars_io_tools.tests.io_sources.test_lazy_ray ‑ test_by_selector_multi_column
polars_io_tools.tests.io_sources.test_lazy_ray ‑ test_by_series_name_mismatch
polars_io_tools.tests.io_sources.test_lazy_ray ‑ test_cartesian_date_x_bucket
polars_io_tools.tests.io_sources.test_lazy_ray ‑ test_cartesian_no_pruning_correct_under_filter
polars_io_tools.tests.io_sources.test_lazy_ray ‑ test_cartesian_none_bucket_matches_nulls
polars_io_tools.tests.io_sources.test_lazy_ray ‑ test_column_retention_under_projection
…

♻️ This comment has been updated with latest results.

[gauglertodd]

so i THINK that preserve_partition_order=True can buffer ~O(total output) on the driver instead of O(max_concurrency). The backpressure parameters should handle the number of buffers that polars/rust is (pre)fetching AND the pending tasks...i think.

[ptomecek]

so i THINK that preserve_partition_order=True can buffer ~O(total output) on the driver instead of O(max_concurrency). The backpressure parameters should handle the number of buffers that polars/rust is (pre)fetching AND the pending tasks...i think.

You are right, that was a good catch. IT's now deferring materializing things on the driver until actually needed and controlling the number of pending tasks better.

[gauglertodd]

lgtm!