GeezerVison

Stacked on #1030. The base is perf/thrift-cf-batch, so the diff shows only this change. I'll retarget to main once #1030 merges.

What type of PR is this?

• Other (performance)

Description

With #1030, every FetchResults batch holds about one CloudFetch link per download thread, but each batch still has to finish before the next one starts. The next FetchResults goes out only after the last file of the current batch is consumed, so the slowest file holds everything up, and a new download manager and thread pool are built for each batch. Download timelines on a 4.9 GB result show an average of 5.3 of 10 downloads in flight.

The download manager can now pull links itself. Once CloudFetch starts, ThriftResultSet hands the rest of the result to the queue's download manager. Whenever a download slot is free and no links are pending, the manager asks for more links: one per free slot plus half a pool queued ahead ((free + max_download_threads // 2) × largest link, clamped to 1–4× buffer_size_bytes as in #1030). The request runs on the consumer thread, as FetchResults does today, so there's no new concurrency on the Thrift client. One thread pool now serves the whole result, and files are still returned in row order (same offset check as before). A variant that refilled only when half the slots were free kept the average at 5.3 in flight and gave no consistent gain, so slots are refilled one at a time.

Bounds and trade-offs:

• Memory: downloaded but unconsumed files are still capped at max_download_threads, the existing bound. The pool now stays at that bound more of the time, so a slow consumer's peak RSS is higher: +5% to +16% across three runs on 1.6 GB (latest pair 831 → 918 MB).
• Link expiry: up to max_download_threads // 2 links (5 by default) wait for a free slot. They start once that many more files are consumed. A link expires 900 s after it is issued, so a consumer has to spend more than about 180 s per ~20 MB file before a queued link expires. With a slow consumer (0.3 s per 100k rows) the oldest link still had 898 s left when its download started (perf(thrift): request about one CloudFetch file per download thread #1030: 899.5 s). Thrift still has no link refresh.
• Close and cancel: closing a result set sets a per-manager cancel event and calls shutdown(wait=False, cancel_futures=True). Downloads now read the body in 1 MiB chunks (preload_content=False) and check the event before the GET and between chunks, so a download in flight stops within about 1 MiB and its connection is closed rather than returned to the pool. A fully read body releases its connection for reuse. urllib3 retried a failed body read only while it read the body itself, so a GET interrupted mid-body is now retried here (3 attempts). Cursor.cancel() also cancels the active result set's downloads. Later reads then raise OperationalError instead of returning a result that ends early; main raised RequestError, a subclass, at the next FetchResults. The cancellation is raised outside the HTTP client so it isn't logged as an HTTP error. downloader.py stays clear of the lines perf(cloudfetch): cut peak memory of large CloudFetch results #1032 edits, and perf(cloudfetch): cut peak memory of large CloudFetch results #1032 applies cleanly on top. The one adjustment needed after both merge is that perf(cloudfetch): cut peak memory of large CloudFetch results #1032's test_run_returns_arrow_buffer_for_lz4_file mock needs a stream(), a one-line change.
• Kill switch: the internal connection kwarg _cloudfetch_prefetch=False restores perf(thrift): request about one CloudFetch file per download thread #1030's per-batch fetching exactly (verified live: same batches and maxBytes). String values from a URL query also work ("false"/"0"/"no"/"off"), and an unrecognized string falls back to off with a warning. Cancel-on-close stays on either way. It isn't documented publicly.
• Errors: if a mid-stream link fetch fails, files already downloading are still returned. The read that needs rows past them raises the original error, and later reads raise, so a retry can't skip rows.

Live A/B against a PRO warehouse (Thrift; fresh processes, #1030 and this PR interleaved, warm-up discarded; fetch = time after execute() returns):

The 4.9 GB p95 comes from one 16.4 s run; the other four were 9.4–10.8 s. An earlier run of the same A/B, before the chunked reads, measured 13.96 → 9.63 s p50 and 18.06 → 10.63 s p95. Chunked reads keep connections in reuse: draining 1.6 GB (85 files) opened 6 new HTTPS connections on both #1030 and this PR.

Early close after 500k rows of the 1.6 GB result (n=3 medians):

cancel() 1.5 s into fetchall_arrow made the fetch fail 0.29 / 0.44 / 0.19 s later on main / #1030 / this PR. This PR raises OperationalError, or RequestError when the next FetchResults fails first.

With #1032 applied on top (locally, n=3, 1.6 GB): peak RSS 1312 → 642 MB for fetchmany_arrow(1M) and 2376 → 1860 MB for fetchall_arrow. Fetch time stayed within noise: 4.69 → 3.34 s and 3.56 → 5.71 s, the latter with one 10 s outlier; an earlier run measured 3.96 → 3.13 s.

How is this tested?

• Unit tests
• E2E Tests
• Manually

New unit tests in test_cloudfetch_prefetch.py cover the refill rule: links are requested only when a slot is free and nothing is pending, there are never more than max_download_threads tasks, and row order is kept across refills. They also check that an empty reply doesn't cause a busy loop, the link budget and its clamp, the offset check, the handoff, and the kill switch. On the cancel side they cover:

• chunked reads that release the connection;
• a cancel mid-body that stops reading and skips decompression;
• an interrupted body that is retried, and one that gives up after 3 attempts;
• cancel() during a link fetch, and a scheduled download cancelled before it starts;
• reads after a cancel that raise instead of ending the result;
• Cursor.cancel() cancelling the active result set's downloads. perf(thrift): request about one CloudFetch file per download thread #1030's per-batch tests now run with the kill switch. The full unit suite passes, and perf(cloudfetch): cut peak memory of large CloudFetch results #1032's tests pass with both PRs applied. Against a live warehouse, the test_driver.py e2e tests for large wide/narrow results, async large results, fetchone/fetchall/fetchmany/iteration, Arrow fetches, row limits and cancel pass (31 passed, Thrift and SEA).

Manual live checks on a 30M-row, 32-link result verified exact row counts and order across batches for fetchall_arrow, fetchmany_arrow, mixed fetchmany/fetchone, and fetchall. They also covered close() and cancel() mid-stream, and use_cloud_fetch=False, with prefetch on and with the kill switch.

Related Tickets & Documents

Stacked on #1030. Related: #1032 (CloudFetch decompression into private mmaps).

This pull request and its description were written by Isaac.