Skip to content

Commit f50710e

Browse files
committed
feat(sql): add opt-in partitioned reads to scan_db
Add a `partitions=` argument to `scan_db` that splits a read into independent slices pulled over parallel connections and concatenated in order, speeding up large scan-like extracts whose single-cursor transfer is the bottleneck. `scan_db` consumes the shared read-partition vocabulary from `io_sources.partitions` (the same `ReadPartition` / `by_time` / `by_value` / `by_range` used by the distributed executor). Each slice's predicate is translated to a SQL `WHERE` and injected as an innermost subquery so it survives outer projection; a predicate that does not translate to SQL raises rather than silently reading the whole slice. Concurrency is capped by the Polars thread pool (`POLARS_MAX_THREADS`), with an optional `max_concurrency` to throttle below it. The path is fully opt-in and falls back to a single query when a partitioner cannot derive a bounded split. By default slices are concatenated in partition order, matching an unpartitioned read. Pass `preserve_partition_order=False` to yield each slice as it finishes instead, lowering time-to-first-row and avoiding head-of-line blocking behind a slow slice, at the cost of a nondeterministic row order. Add unit tests for the partitioned reader and a how-to section in the Reading and Writing Data wiki. Signed-off-by: Pascal Tomecek <40371786+ptomecek@users.noreply.github.com>
1 parent 7e83ded commit f50710e

3 files changed

Lines changed: 510 additions & 45 deletions

File tree

docs/wiki/Reading-and-Writing-Data.md

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,52 @@ example SQL Server can still seek on `CAST(datetime AS date)`) while others fall
6767
full scan. When a filtered column is indexed and on a hot path, prefer filtering the
6868
physical column directly over its cast form.
6969

70+
## Speed up a large SQL read by partitioning it
71+
72+
When a scan-like extract is bounded by an indexed column, a single ODBC cursor is often the
73+
bottleneck. Pass `partitions=` and `scan_db` splits the read into independent slices and pulls
74+
them over parallel connections, concatenating the results in order. Build the slices from the
75+
filter you push down with `by_time`, `by_value`, or `by_range`:
76+
77+
```python
78+
from polars_io_tools import scan_db, by_time
79+
80+
lf = scan_db(
81+
"SELECT * FROM daily_prices",
82+
connection="Driver={PostgreSQL};Server=db.example.com;Database=mkt;Uid=reader;******",
83+
partitions=by_time("price_date", every="1mo"),
84+
)
85+
86+
# One slice per month, taken from the pushed-down date range:
87+
result = lf.filter(
88+
(pl.col("price_date") >= pl.date(2025, 1, 1)) & (pl.col("price_date") < pl.date(2025, 7, 1))
89+
).collect()
90+
```
91+
92+
- `by_time(column, every=)` — calendar windows; `every` is an interval string (`"1mo"`, `"2w"`,
93+
`"5d"`, `"1q"`, `"1y"`) or an integer number of days.
94+
- `by_value(column, values=None)` — one slice per discrete value; with `values=None` the values
95+
are read from the `IN` filter you push down.
96+
- `by_range(column, every=)` — fixed-width numeric buckets over the pushed-down range.
97+
98+
How many slices run at once is capped by a process-wide SQL connection budget
99+
(`POLARS_IO_TOOLS_MAX_SQL_CONNECTIONS`; default `min(pl.thread_pool_size(), 8)` — a modest 8 on a
100+
normal machine, self-throttling to 1 in fan-out clusters that pin `POLARS_MAX_THREADS=1`, since
101+
these reads are IO-bound and the cap is a connection budget, not the CPU thread pool); pass
102+
`max_concurrency=` to throttle below it on a busy server. Partitioning is fully opt-in — with no
103+
`partitions=`, or when a partitioner cannot derive a bounded split, `scan_db` runs the query over
104+
a single connection. Prefer a handful of medium slices over many tiny ones: each slice is a
105+
separate query with its own planning and round-trip cost.
106+
107+
For hand-built slices, pass an iterable of `ReadPartition(predicate, key)`. Predicates that
108+
translate to SQL are pushed to the database; any part that cannot (for example an arbitrary
109+
Python UDF) is still enforced client-side, so each slice stays exact.
110+
111+
By default the slices are concatenated in partition order, so a partitioned read has the same row
112+
order as the unpartitioned one. Pass `preserve_partition_order=False` to yield each slice as soon
113+
as it finishes instead — this lowers time-to-first-row and avoids a slow slice blocking the rest,
114+
at the cost of a nondeterministic row order (and, with a row limit, a nondeterministic subset).
115+
70116
## Read from ClickHouse
71117

72118
`scan_clickhouse` streams query results over ClickHouse's HTTP interface as Arrow IPC.

0 commit comments

Comments
 (0)