Add opt-in partitioned reads to scan_db

[ptomecek]

Description

Adds a partitions= argument to scan_db that splits a read into independent slices pulled over parallel connections and streamed batch-by-batch, speeding up large scan-like extracts whose single-cursor transfer is the bottleneck. Peak memory stays proportional to the connection budget × the Arrow batch size — a slice is never fully buffered — regardless of total row count.

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:

• by_time(column, every) — calendar windows from the pushed-down date range ("1mo"/"2w"/"5d"/"1q"/"1y" or an int of days)
• by_value(column, values=None) — one slice per discrete value (derived from an IN filter when omitted)
• by_range(column, every) — fixed-width numeric buckets
• an explicit iterable of ReadPartition(predicate, key)

How it works

Each slice's bound is ANDed into the pushed-down predicate and routed through the existing apply_polars_io_source_exprs (so MSSQL ORDER BY/OPTION hoisting and identifier quoting are reused). Columns the predicates need are retained through projection pushdown and the combined predicate is re-applied client-side, so slices stay exact even if only part of the predicate translates to SQL. Workers stream into a shared bounded queue and are relaunched as slices complete, keeping at most k connections live; n_rows/head pushdown stops the fan-out early and joins any in-flight workers before returning. A partitioner that cannot derive a bounded split falls back to a single query; a known-empty split yields an empty frame; an over-limit split raises.

Concurrency is governed by a shared connection budget (POLARS_IO_TOOLS_MAX_SQL_CONNECTIONS, default min(pl.thread_pool_size(), 8) — self-throttling to 1 in fan-out clusters that pin POLARS_MAX_THREADS=1), with an optional max_concurrency= to throttle below it.

Ordering / sorted output

Rows are yielded in completion order (whichever slice's batches arrive first), which is not deterministic — apply a Polars .sort() after the read if you need a specific order (this matches connector-x, which likewise has no ordering knob). A top-level ORDER BY is honoured only for an unpartitioned read; under partitioning it can't be re-established across independent slices, so it's dropped with a warning (the row set is still correct). Clauses that would change the result set under partitioning — a top-level LIMIT/OFFSET/TOP/FETCH, QUALIFY, or a window function — raise instead; apply them in Polars after the read.

Tests / docs

Unit tests for the builders and the partitioned reader (DuckDB-backed), including bounded-memory / head-pushdown and early-stop worker-join coverage, plus a how-to section in the Reading and Writing Data wiki. make lint-py / lint-docs clean.

[github-actions]

Test Results

1 590 tests  +66   1 542 ✅ +66   57s ⏱️ +17s
    2 suites ± 0      48 💤 ± 0 
    2 files   ± 0       0 ❌ ± 0 

Results for commit c172725. ± Comparison against base commit 2a17d11.

♻️ This comment has been updated with latest results.

[feussy]

The prefetch here isn't bounded by data volume, only by slice count, and that undoes streaming for the sink paths.

I measured peak RSS in isolated subprocesses, scaling the source data 8x:

                      120MB   240MB   480MB   960MB
unpartitioned sink :   318     329     350     330
partitioned (k=12) :   439     556     725    1105

The unpartitioned sink_parquet peak is flat across the whole range -- it really does stream. The partitioned path tracks base + k * slice_size, which I confirmed by holding the data fixed and varying the budget (max_concurrency 1/2/4/8 -> 349/413/481/564 MB).

collect() is only about +19% and I wouldn't worry about it, since Polars materializes anyway. It's specifically sink_* where this hurts, which is the one place a user has explicitly asked for bounded memory -- and there we go from a flat ceiling to something that grows with both k and slice size, with nothing capping it.

A bounded queue of depth d between the workers and the generator would put a ceiling on it (d * slice instead of k * slice) and let k stay tuned for connection concurrency. Also worth noting even at k=1 we buffer a whole slice before emitting anything, so time-to-first-row becomes time-to-first-slice.

Can we get backpressure in before this merges? Happy to talk through what d should default to.

[ptomecek]

Good call — this is reworked. Workers now stream batch-by-batch into a shared bounded queue (maxsize=k) rather than buffering a whole slice, so peak memory is O(k × arrow_batch_size), independent of slice size and total row count. I confirmed the ceiling stays roughly flat at scale (held ~1.1–1.3 GB across 2.4M→5.5M rows). Time-to-first-row is now first-batch, not first-slice, and the sink_* growth-with-slice-size you measured should be gone.

[feussy]

The unpartitioned branch passes n_rows into _build_sql but this one drops it:

.head(5) unpartitioned : 1 query with LIMIT 5   ->   5 rows
.head(5) partitioned   : 8 unbounded queries    -> 244 rows

Narrower than it looks in practice -- n_rows only arrives non-None for .head(n) with no filter, and the derived partitioners need a filter to derive bounds, so you need enumerated partitions (by_value(values=[...]) or an explicit ReadPartition list) to get here. Still reachable though.

Pushing LIMIT n into each slice should be safe: you can only ever come back with more than n overall, never fewer, and the client-side counter already trims. Was there a reason to leave it out that I'm missing?

[ptomecek]

The wrong row count is fixed, just not via LIMIT pushdown: the fan-out streams and the generator enforces the global n_rows cap and stops early, so .head(5) returns 5, not 244 (covered by test_unordered_row_limit). I left the per-slice LIMIT n push out since streaming early-stop already bounds how much we pull — but you're right that pushing LIMIT n would also shrink the server-side scan, so happy to add it as an optimization if you think it's worth it.

[feussy]

This docstring is accurate and it's good that it's called out, including the head() consequence. My concern is that it's the only thing standing between a user and a wrong answer, and it lives under a different parameter.

Since we wrap the query as SELECT * FROM (<query>) AS __cpl_subq WHERE ..., a subquery's ORDER BY isn't binding on the outer query, so with k slices nothing re-establishes the global order:

unpartitioned: [47,46,45, ... ,1,0]
partitioned  : [12,11,...,0 | 24,...,13 | 36,...,25 | 47,...,37]

head(5) unpartitioned: [47, 46, 45, 44, 43]
head(5) partitioned  : [12, 11, 10,  9,  8]

That second pair is what bothers me -- it's not a different order, it's different rows, silently. "Give me the 5 most recent" quietly returns 5 arbitrary rows. It also fails in a nasty way: when the sort direction happens to line up with ascending partition-key order the result is byte-identical and correctly sorted, so it looks fine right up until someone flips to DESC.

Worth saying a polars .sort() is completely fine -- sort isn't in the register_io_source pushdown protocol, so it lands after the source and gets the right answer even with preserve_partition_order=False. This is only about literal ORDER BY text in the query.

Related, and same root cause: LIMIT/TOP/OFFSET and window functions cost k x the scan. Rows actually read from the base table:

plain SELECT *        1 query  48 | 4 slices  48   (1.0x)
DISTINCT              1 query  48 | 4 slices  48   (1.0x)
LIMIT 20              1 query  48 | 4 slices 192   (4.0x)
window ROW_NUMBER()   1 query  48 | 4 slices 192   (4.0x)

Plain reads and DISTINCT are fine -- the optimizer pushes the predicate into the scan and each slice touches only its own range. But you can't legally push a filter through a limit or a window frame, so every slice re-reads the whole input. Answers stay correct, you just get no speedup and k x the work. That's SQL semantics, so it'll hold on MSSQL too, not just the duckdb I tested on.

Nobody should partition a query like that, but I think it's a realistic footgun for anyone who wraps scan_db in a house helper with a default partitions= and applies it everywhere.

Suggestion: we already do the right thing when a partitioner can't derive bounds (test_unbounded_predicate_falls_back_to_single_query), so route these into the same path -- parsed_query.find(exp.Order, exp.Limit, exp.Offset, exp.Window) and fall back to a single query. One caveat: falling back is only safe for derived partitioners. For enumerated ones it'd change the row set, so those should raise instead. I'd log it at warning rather than debug -- the user asked for partitioning and isn't getting it.

[ptomecek]

Addressed, and the riskiest part is gone. Under partitioning we now hard-raise on the result-changing clauses — top-level LIMIT/OFFSET/TOP/FETCH, QUALIFY, and window functions — so you can't silently get wrong rows or k× the scan; you get a clear error pointing you to apply them in Polars. A literal ORDER BY warns (row set correct, just not globally ordered) and points at .sort(). Unpartitioned reads pass ORDER BY straight through to the engine. I went with raise rather than a silent single-query fallback so the user is actually told partitioning isn't happening. And the whole preserve_partition_order "looks-sorted-until-DESC" mode you flagged has been removed — output is completion-order, sort in Polars if you need order.

[feussy]

read_arrow_batches_from_odbc(connection_string=...) opens and tears down a connection per call, so we do N connects for N slices rather than one per worker. With max_partitions defaulting to 512 and auth/TLS on each connect, that's potentially 512 connects where we wanted 8.

arrow_odbc has Connection / connect() and an enable_odbc_connection_pooling() we're not using. Would a connection per pool worker be workable here, or is there a thread-affinity problem with the ODBC handles I'm not seeing?

[ptomecek]

This one still stands — we open/tear down a connection per slice; I haven't switched to pooling in this PR. Our at-scale benchmarks actually back up the concern: connect/auth overhead is real enough that coarse partitioning (few big slices ≈ k) clearly beats many small ones, so the practical guidance is to keep slice count low. I'd rather do proper connection reuse (one per pool worker, or enable_odbc_connection_pooling) as a focused follow-up than fold it in here — I'll file an issue. I don't see a thread-affinity blocker, but it needs a quick check that an arrow_odbc connection is safe to reuse across successive reads on a pool thread.

[ptomecek]

Following up: after digging into arrow_odbc, connection reuse is really two separate things, and neither is a good hidden default inside scan_db.

• Process-global pooling — arrow_odbc.enable_odbc_connection_pooling() lets the ODBC driver manager recycle physical connections so each connect skips the TLS/Kerberos handshake. But it's a process-wide arrow-odbc / driver-manager setting: it affects all ODBC use in the process and reuses physical connections (so session-scoped state like temp tables can persist across them), and it has to be enabled before the first ODBC use. So rather than wrap it in our own flag/env — which would just be re-exporting someone else's global switch with timing baggage — I've documented it as the escape hatch in the Reading & Writing wiki; callers opt in themselves before their first read.
• A scoped per-scan connection pool (create ≤k connections, check out/return per slice, drop refs in the generator's finally) is the safe-by-default, no-global-side-effect version and would be the "right" fix. But it interacts with the shared process-global executor's lifetime (Connection has no close() — it's GC/RAII, so thread-locals on a permanent pool would leak k connections) and needs the cross-thread-reuse check you flagged verified against the real driver. That's a focused change with its own testing, so I'm leaving it as a possible future improvement rather than expanding this PR.

Walking back my earlier "I'll file an issue" — holding off for now; correctness is unaffected today (this is purely connect overhead), and the practical guidance (prefer a few larger slices) is in the docs. Easy to revisit if connect cost bites in practice.

[feussy]

When this works out to a single slice we still take the partitioned path, so we pay for full-slice buffering and an extra connect for no benefit. Falling back to the unpartitioned path at k == 1 would fix both.

One trap: that's only safe for the derived partitioners. by_value("g", [[0],[1],[2]]) returns 300 rows where the unpartitioned query returns 1200, so short-circuiting on count alone would silently change results for enumerated partitions. Same derived-vs-enumerated split the query-shape comment needs -- might be worth an exhaustive flag on the Partitioner protocol so both can key off one thing, though I realize that's #25's code and may be out of scope here.

[ptomecek]

Fixed — when the work resolves to a single slice we now take a dedicated single-connection streaming path (if len(work) == 1:), so no thread-pool fan-out and no whole-slice buffering. It still applies the slice predicate, so enumerated single partitions like by_value("g", [[0]]) correctly return just their subset — no silent result change. The exhaustive flag on the Partitioner protocol is a nice cleanup but agreed, that's #25's code and out of scope here.

[feussy]

The pre-PR version put {final_sql} in here and this one drops it. Since we synthesize the WHERE clause, the user never sees the query that actually ran, which makes a failing slice pretty hard to debug.

Separately, there's a latent bug in this handler: work is bound partway through the generator but the except also covers the code above that point, so an early failure hits UnboundLocalError and buries the real exception.

I've got a small patch for both if it's useful -- hoist work to the top, and have the worker stash the SQL on the exception before re-raising (concurrent.futures re-raises the same exception object, so the attribute survives .result()), falling back to work[0][0] as a representative query. About 13 lines, suite stays green.

[ptomecek]

Both handled. The worker now stashes the slice SQL on the exception before re-raising (RuntimeError(f"...slice {idx}:\n{sql}")), and concurrent.futures preserves that object through .result(), so a failing slice surfaces the exact query it ran (asserted in test_partition_worker_error_surfaces). And the outer handler only references predicate/with_columns (both function params, always bound), never work — so an early failure can't hit UnboundLocalError and bury the real error.

[feussy]

This runs at import and isn't guarded, so a bad value takes down import polars_io_tools entirely:

' '     -> ValueError: invalid literal for int() with base 10: ' '
'auto'  -> ValueError: invalid literal for int() with base 10: 'auto'
'8.5'   -> ValueError: invalid literal for int() with base 10: '8.5'

It's a SQL tuning knob, but a typo in it breaks the package for people who only use the Parquet or Arrow paths and never touch scan_db. The traceback lands in module import, so it reads like the library is broken rather than like a config problem.

The ' ' case is the one I'd actually expect to see -- empty string is handled since it's falsy, but a single space gets past that and straight into int(), which is what you get from FOO: " " in a Helm values file or a trailing space in a .env.

Either catch it, warn, and fall back to the default, or move the parse into _get_sql_executor() (already lazy and lock-guarded) so it fails on first use with a message naming the variable and the value. I'd lean toward the second so we're not silently ignoring something the user explicitly asked for. Adding .strip() seems worth doing regardless.

[ptomecek]

Fixed the way you suggested (option 2). The parse moved into _get_sql_executor() — lazy and lock-guarded — so a bad value fails on first parallel use with a message naming the variable, not at import polars_io_tools. Added .strip() so ' ' is treated as unset, and blank/whitespace/auto/8.5/0/-1 are all covered by test_connection_budget_env_parse. Import stays clean for Parquet/Arrow-only users.