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.

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>

Author: ptomecek

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

@@ -41,6 +41,46 @@ The dialect is detected from the ODBC connection, so the generated SQL matches y
 database. Pass `fetch_size=` to control the batch size used when Polars does not
 request one.
 
+## Speed up a large SQL read by partitioning it
+
+When a scan-like extract is bounded by an indexed column, a single ODBC cursor is often the
+bottleneck. Pass `partitions=` and `scan_db` splits the read into independent slices and pulls
+them over parallel connections, concatenating the results in order. Build the slices from the
+filter you push down with `by_time`, `by_value`, or `by_range`:
+
+```python
+from polars_io_tools import scan_db, by_time
+
+lf = scan_db(
+    "SELECT * FROM daily_prices",
+    connection="Driver={PostgreSQL};Server=db.example.com;Database=mkt;Uid=reader;******",
+    partitions=by_time("price_date", every="1mo"),
+)
+
+# One slice per month, taken from the pushed-down date range:
+result = lf.filter(
+    (pl.col("price_date") >= pl.date(2025, 1, 1)) & (pl.col("price_date") < pl.date(2025, 7, 1))
+).collect()
+```
+
+- `by_time(column, every=)` — calendar windows; `every` is an interval string (`"1mo"`, `"2w"`,
+  `"5d"`, `"1q"`, `"1y"`) or an integer number of days.
+- `by_value(column, values=None)` — one slice per discrete value; with `values=None` the values
+  are read from the `IN` filter you push down.
+- `by_range(column, every=)` — fixed-width numeric buckets over the pushed-down range.
+
+How many slices run at once is capped by a process-wide SQL connection budget
+(`POLARS_IO_TOOLS_MAX_SQL_CONNECTIONS`, default 8 — these reads are IO-bound, so the cap is a
+connection budget, not the CPU thread pool); pass `max_concurrency=` to throttle below it on a
+busy server. Partitioning is fully opt-in — with no `partitions=`, or when a partitioner cannot
+derive a bounded split, `scan_db` runs the query over a single connection. Prefer a handful of
+medium slices over many tiny ones: each slice is a separate query with its own planning and
+round-trip cost.
+
+For hand-built slices, pass an iterable of `ReadPartition(predicate, key)`. Predicates that
+translate to SQL are pushed to the database; any part that cannot (for example an arbitrary
+Python UDF) is still enforced client-side, so each slice stays exact.
+
 ## Read from ClickHouse
 
 `scan_clickhouse` streams query results over ClickHouse's HTTP interface as Arrow IPC.

polars_io_tools/io_sources/lazy_sql_reader.py

@@ -1,11 +1,16 @@
+import concurrent.futures
 import logging
+import os
+import threading
+from collections.abc import Iterable
 from functools import lru_cache
 from typing import Any
 
 import polars as pl
 from sqlglot import exp, parse_one
 from sqlglot.dialects.dialect import Dialect
 
+from .partitions import Partitioner, ReadPartition, as_partition_list, retained_columns
 from .sql_dialects import MSSQL
 from .sql_utils import (
     apply_polars_io_source_exprs,
@@ -19,6 +24,24 @@
 # Configure logging
 log = logging.getLogger(__name__)
 
+# Process-wide budget for concurrent SQL connections. Partition reads are IO-bound (each worker
+# is mostly blocked on the database), so the right cap is an IO concurrency budget, not the CPU
+# compute pool: `POLARS_MAX_THREADS` / `pl.thread_pool_size()` governs Rayon compute and is the
+# wrong dimension here. The compute portion (Arrow decode + client-side filter) stays bounded by
+# Rayon regardless. Tune with the `POLARS_IO_TOOLS_MAX_SQL_CONNECTIONS` env var (default 8); the
+# shared executor is the global governor so concurrent scans cannot collectively oversubscribe.
+_DEFAULT_SQL_CONNECTIONS = max(1, int(os.environ.get("POLARS_IO_TOOLS_MAX_SQL_CONNECTIONS", "8")))
+_SQL_EXECUTOR: concurrent.futures.ThreadPoolExecutor | None = None
+_SQL_EXECUTOR_LOCK = threading.Lock()
+
+
+def _get_sql_executor() -> concurrent.futures.ThreadPoolExecutor:
+    global _SQL_EXECUTOR
+    with _SQL_EXECUTOR_LOCK:
+        if _SQL_EXECUTOR is None:
+            _SQL_EXECUTOR = concurrent.futures.ThreadPoolExecutor(max_workers=_DEFAULT_SQL_CONNECTIONS)
+    return _SQL_EXECUTOR
+
 
 @lru_cache(None)
 def get_sqlglot_dialect_odbc(conn_string: str) -> str | type[Dialect] | None:
@@ -90,7 +113,16 @@ def get_schema_from_query_odbc(
         raise ValueError(f"Could not determine schema for query: {query}, with error: {e}") from e
 
 
-def scan_db(query: str, connection: str, fetch_size: int = 10000, **kwargs) -> pl.LazyFrame:
+def scan_db(
+    query: str,
+    connection: str,
+    fetch_size: int = 10000,
+    *,
+    partitions: "Partitioner | Iterable[ReadPartition] | None" = None,
+    max_partitions: int = 512,
+    max_concurrency: int | None = None,
+    **kwargs,
+) -> pl.LazyFrame:
     """
     Create a LazyFrame from a SQL query with predicate pushdown support.
 
@@ -100,19 +132,38 @@ def scan_db(query: str, connection: str, fetch_size: int = 10000, **kwargs) -> p
     connection with optimized predicate pushdown. Filters applied to the LazyFrame will
     be translated back to SQL and pushed to the database.
 
+    When ``partitions`` is set, the reader splits the query into independent slices and pulls
+    them over parallel connections, concatenating the results in order. This can dramatically
+    speed up large, scan-like extracts whose single-cursor transfer is the bottleneck. It is
+    fully opt-in: with ``partitions=None`` the behaviour is identical to a plain single scan.
+
     Args:
         query (str): The SQL query to execute
         connection (str): A connection string (*not* a database connection object)
         fetch_size (int, default 10000): Number of rows to fetch at a time. This is a default needed by the \
             source generator function that scan_db wraps (because it is required \
             by the Polars IO plugins API). This value will only be used if Polars \
             does not pass a value for batch size; if it does, that will be used instead.
+        partitions (Partitioner | Iterable[ReadPartition] | None, default None): How to split the read. \
+            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` for hand-built slices. Each slice becomes one query on its own \
+            connection; any part of a slice's predicate that cannot be pushed to SQL is enforced \
+            client-side. When a partitioner cannot derive a bounded split, the query runs unpartitioned.
+        max_partitions (int, default 512): Guardrail -- if partitioning would produce more than this \
+            many slices, raise (raise this limit or coarsen the partitions).
+        max_concurrency (int | None, default None): Optional cap on the number of partitions pulled \
+            simultaneously. Hard-capped at the process-wide SQL connection budget \
+            (``POLARS_IO_TOOLS_MAX_SQL_CONNECTIONS``, default 8); use it to throttle *below* that to be \
+            gentle on a shared server. None means use the full budget.
         **kwargs: Additional arguments for the database connector
 
     Returns:
         pl.LazyFrame: A Polars LazyFrame with predicate pushdown support
     """
 
+    conn_string = connection if isinstance(connection, str) else str(connection)
+
     def _fetch_info_needing_connection() -> tuple[
         dict[str, pl.DataType],
         exp.Expression,
@@ -132,6 +183,52 @@ def _fetch_info_needing_connection() -> tuple[
 
     schema, parsed_query, dialect = _fetch_info_needing_connection()
 
+    if isinstance(partitions, Partitioner) and partitions.on not in schema:
+        raise ValueError(f"partition column {partitions.on!r} is not in the query schema {list(schema)}")
+
+    def _select_cols(df: pl.DataFrame, with_columns: list[str] | None) -> pl.DataFrame:
+        if with_columns is not None:
+            wanted = set(with_columns)
+            return df.select(col for col in schema if col in wanted)
+        return df
+
+    def _build_sql(
+        predicate: pl.Expr | None,
+        with_columns: list[str] | None,
+        n_rows: int | None,
+        batch_size: int | None,
+    ) -> str:
+        # Reuse the shared subquery machinery (MSSQL ORDER BY / OPTION hoisting, identifier
+        # quoting) for both the pushed predicate and any partition bound folded into it.
+        final_query_expr = apply_polars_io_source_exprs(parsed_query.copy(), dialect, with_columns, predicate, n_rows, batch_size)
+        return final_query_expr.transform(fix_three_part_identifiers).sql(dialect=dialect)
+
+    def _read(sql: str, batch_size: int | None):
+        from arrow_odbc import read_arrow_batches_from_odbc
+
+        return read_arrow_batches_from_odbc(
+            query=sql,
+            batch_size=fetch_size if batch_size is None else batch_size,
+            connection_string=conn_string,
+            **kwargs,
+        )
+
+    def _fetch_all(sql: str, client_predicate: pl.Expr | None, keep_columns: list[str] | None, batch_size: int | None) -> list[pl.DataFrame]:
+        """Run one partition on its own connection and return all its DataFrames (worker body).
+
+        ``client_predicate`` (the pushed predicate ANDed with the partition bound) is reapplied
+        here so the slice is exact even if only part of it translated to SQL -- this is what
+        guarantees partitions stay disjoint. ``keep_columns`` are the caller's requested columns;
+        any extra columns retained only to evaluate the predicate are dropped.
+        """
+        out: list[pl.DataFrame] = []
+        for record_batch in _read(sql, batch_size):
+            df = pl.DataFrame(record_batch)
+            if client_predicate is not None:
+                df = df.filter(client_predicate)
+            out.append(_select_cols(df, keep_columns))
+        return out
+
     # Create the generator function for our custom IO source
     def source_generator(
         with_columns: list[str] | None,
@@ -142,57 +239,119 @@ def source_generator(
         # Short-circuit: if the caller already knows zero rows are needed
         # (e.g. from head(0) on a contradictory filter), skip the query entirely.
         if n_rows == 0:
-            empty = pl.DataFrame({}, schema=schema)
-            if with_columns is not None:
-                empty = empty.select(col for col in schema if col in set(with_columns))
-            yield empty
+            yield _select_cols(pl.DataFrame({}, schema=schema), with_columns)
             return
 
-        # Generate a new SQL query by combining the original query with the predicate
-        query_copy = parsed_query.copy()
-        final_query_expr = apply_polars_io_source_exprs(query_copy, dialect, with_columns, predicate, n_rows, batch_size)
-        # Convert back to SQL string
-        final_sql = final_query_expr.sql(dialect=dialect)
-        log.debug(f"Executing SQL with pushdown: {final_sql}")
-
-        # Create a connection string if needed
-        conn_string = connection if isinstance(connection, str) else str(connection)
-        try:
-            from arrow_odbc import read_arrow_batches_from_odbc
-
-            # Use arrow_odbc directly to fetch results
-            batch_reader = read_arrow_batches_from_odbc(
-                query=final_sql,
-                batch_size=fetch_size if batch_size is None else batch_size,
-                connection_string=conn_string,
-                # Pass through additional connection options
-                # that the user specified in the parent function
-                **kwargs,
+        # Resolve the partition slices. ``as_partition_list`` distinguishes three cases:
+        #   None -> a partitioner could not derive a bounded split (fall back to one query),
+        #   []   -> a known-empty partition set (yield nothing),
+        #   list -> concrete slices.
+        partition_list = as_partition_list(partitions, predicate) if partitions is not None else None
+
+        if partition_list is not None and len(partition_list) > max_partitions:
+            # A concrete partition set may intentionally select a subset (e.g. by_value with an
+            # explicit value list), so it cannot be silently replaced by a single unpartitioned
+            # query -- raise rather than risk returning extra rows.
+            raise ValueError(
+                f"Partition count {len(partition_list)} exceeds max_partitions={max_partitions}; raise max_partitions or coarsen the partitions."
             )
 
-            # Track if we've yielded any batches yet
-            # This is necessary in case the query yields
-            # no records
-            count = 0
-
-            def select_cols(df) -> pl.DataFrame:
-                if with_columns is not None:
-                    with_cols_set = set(with_columns)
-                    return df.select(col for col in schema if col in with_cols_set)
-                return df
+        if partition_list is not None and len(partition_list) == 0:
+            yield _select_cols(pl.DataFrame({}, schema=schema), with_columns)
+            return
 
-            for record_batch in batch_reader:
-                df = pl.DataFrame(record_batch)
-                if predicate is not None:
-                    df = df.filter(predicate)
-                yield select_cols(df)
-                count += 1
+        # Build (sql, client_predicate) per slice. For a partitioned read we AND the partition
+        # bound into the pushed predicate (reused by both the server-side SQL and the client-side
+        # safety filter) and retain the columns those predicates need through projection so the
+        # filter can be evaluated; n_rows is applied client-side across the ordered stream.
+        if partition_list:
+            preds = [part.predicate for part in partition_list]
+            if predicate is not None:
+                preds.append(predicate)
+            effective_wc = retained_columns(preds, with_columns)
+            work: list[tuple[str, pl.Expr | None]] = []
+            for part in partition_list:
+                combined = part.predicate if predicate is None else (predicate & part.predicate)
+                work.append((_build_sql(combined, effective_wc, None, batch_size), combined))
+        else:
+            work = [(_build_sql(predicate, with_columns, n_rows, batch_size), predicate)]
+
+        for sql, _ in work:
+            log.debug("Executing SQL with pushdown: %s", sql)
+        log.debug("scan_db running %d partition(s)", len(work))
+
+        yielded_rows = 0
+
+        def _emit(df: pl.DataFrame):
+            """Yield a frame honouring the global n_rows cap across the ordered stream."""
+            nonlocal yielded_rows
+            if n_rows is not None:
+                remaining = n_rows - yielded_rows
+                if remaining <= 0:
+                    return
+                if df.height > remaining:
+                    df = df.head(remaining)
+            yielded_rows += df.height
+            yield df
 
-            if count == 0:
-                yield select_cols(pl.DataFrame({}, schema=schema))
+        try:
+            if len(work) == 1:
+                # Single-connection streaming path (preserves original behaviour incl. empty result).
+                sql, client_predicate = work[0]
+                count = 0
+                for record_batch in _read(sql, batch_size):
+                    df = pl.DataFrame(record_batch)
+                    if client_predicate is not None:
+                        df = df.filter(client_predicate)
+                    yield from _emit(_select_cols(df, with_columns))
+                    count += 1
+                    if n_rows is not None and yielded_rows >= n_rows:
+                        break
+                if count == 0:
+                    yield _select_cols(pl.DataFrame({}, schema=schema), with_columns)
+                return
+
+            # Partitioned path: bounded-concurrency fan-out with strict in-order yield.
+            executor = _get_sql_executor()
+            k = _DEFAULT_SQL_CONNECTIONS if max_concurrency is None else max(1, min(max_concurrency, _DEFAULT_SQL_CONNECTIONS))
+            futures: dict[concurrent.futures.Future, int] = {}
+            completed: dict[int, list[pl.DataFrame]] = {}
+            next_submit = 0
+            next_yield = 0
+
+            def _submit_more():
+                nonlocal next_submit
+                while len(futures) < k and next_submit < len(work):
+                    sql, client_predicate = work[next_submit]
+                    fut = executor.submit(_fetch_all, sql, client_predicate, with_columns, batch_size)
+                    futures[fut] = next_submit
+                    next_submit += 1
+
+            _submit_more()
+            while next_yield < len(work):
+                while next_yield in completed:
+                    for df in completed.pop(next_yield):
+                        yield from _emit(df)
+                    next_yield += 1
+                    _submit_more()
+                    if n_rows is not None and yielded_rows >= n_rows:
+                        for fut in futures:
+                            fut.cancel()
+                        futures.clear()
+                        completed.clear()
+                        return
+                if next_yield >= len(work) or not futures:
+                    break
+                done, _ = concurrent.futures.wait(list(futures.keys()), return_when=concurrent.futures.FIRST_COMPLETED)
+                for fut in done:
+                    completed[futures.pop(fut)] = fut.result()
+                _submit_more()
+
+            if yielded_rows == 0:
+                yield _select_cols(pl.DataFrame({}, schema=schema), with_columns)
 
         except Exception as e:
-            err_msg = f"Failed to execute SQL query: {final_sql}\nPredicate:\n{predicate}\n The `with_columns` used: {with_columns}\n"
+            err_msg = f"Failed to execute SQL query.\nPredicate:\n{predicate}\n The `with_columns` used: {with_columns}\n"
             err_msg += f"\n\nWhile running the above, received error: {e.__class__.__name__}:{e}"
             raise RuntimeError(err_msg) from e