[vuanhphung]

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–7% in two runs here (median 840 → 891 MB on 1.6 GB) and +13% in an earlier prototype run.
• 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: close() now sets a per-manager cancel event. Downloads check it before the GET and again before decompressing, and shutdown(wait=False, cancel_futures=True) drops queued work. Bodies already being read still finish. cancel() only cancels the statement on the server: until the next link fetch fails, up to max_download_threads // 2 already-queued links can still start downloading. The downloader.py changes stay away from the decompression lines that 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.
• 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). Cancel-on-close stays on either way. It isn't documented publicly.

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

Workload	FetchResults	fetch p50	fetch p95	peak RSS
lineitem, 4.9 GB, fetchall_arrow, n=5	26 → 41	13.96 → 9.63 s	18.06 → 10.63 s	5577 → 5616 MB
⅓ of lineitem, 1.6 GB, fetchall_arrow, n=5	8 → 13	4.59 → 3.41 s	5.55 → 3.78 s	2374 → 2379 MB
⅓ of lineitem, 1.6 GB, fetchmany_arrow(1M), n=3	8 → 13	4.25 → 3.65 s	5.38 → 3.68 s	1244 → 1282 MB

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

	main	#1030	this PR
downloads in flight at close	2	8	8
MB read in the 5 s after close	4.5	47	58
RSS 5 s after close	310 MB	523 MB	387 MB
close()	0.11 s	0.11 s	0.11 s
exit lag when the process exits right after close	0.09 s	0.35 s	0.28 s

Most of the close cost comes from #1030's larger batches; this PR adds little to it. What remains after close is the bodies already being read (the interpreter waits for them at exit). The cancel check skips their decompression, which is why RSS after close drops. cancel() 1.5 s into fetchall_arrow raised RequestError 0.20 / 0.36 / 0.27 s later on main / #1030 / this PR. This PR also started 3 downloads from queued links after the cancel, where the others started none.

With #1032 applied on top (locally, n=3, 1.6 GB): peak RSS 1283 → 685 MB for fetchmany_arrow(1M) and 2392 → 1862 MB for fetchall_arrow, with fetch time 3.85 → 3.67 s and 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, the kill switch, and cancel-on-close (shutdown sets the event and cancels queued futures; a cancelled download skips the GET or the decompression). #1030's per-batch tests now run with the kill switch. The full unit suite passes, and #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.

[peco-review-bot]

Verdict: 1 Medium · 2 Low

Looks good overall, with one medium concern: the prefetch path discards any part of a FetchResults reply that isn't a CloudFetch link. Ordering, the memory bound and cancel-on-close look right. The kill switch is wired through every result-set constructor and has tests.

[peco-review-bot]

Verdict: 2 Low

Looks good overall, with two low-severity notes: a re-poll edge case in _schedule_downloads, and no end-to-end unit test of the result set handing link fetching to the queue. Ordering, the memory bound, cancel-on-close and the kill switch all behave as the PR describes. Nit: the link source is attached only after ThriftCloudFetchQueue.__init__ waits for the first file, so the first prefetch can't start until that download finishes.

[vuanhphung]

On the review-body nit (the first prefetch waits until ThriftCloudFetchQueue.__init__ has the first file): true. The cost is one file's download time (~250 ms here), once per result, and the initial batch keeps downloading in the meantime. The live numbers in the description include it. Removing it would mean passing the link source through ThriftResultSetQueueFactory.build_queue into the queue constructor, which is more change than it is worth, so I left it.

[peco-review-bot]

Verdict: 1 Low

Looks good overall, with one low concern: the manager does not check whether it was cancelled after the link-fetch call returns. Prefetch is wired correctly through _start_prefetch and the has_more_rows hand-off. The kill switch keeps the old per-batch path, and the tests cover ordering, empty replies, errors, budget clamping and cancel-on-close.

[peco-review-bot]

Verdict: 1 Medium · 2 Low

The design is solid and well tested. I found 1 medium issue (a link-fetch error mid-read silently drops the current file's rows, despite a test saying fetching can resume) and 2 low ones (a close/submit race, and the kill switch using bool() on a value that may be a string). Nit: in downloader.py, stdlib threading sits under the TYPE_CHECKING block next to a comment that's about keeping thrift out, which reads as if it also needs to be kept out.

[vuanhphung]

On the review-body nit (import threading under TYPE_CHECKING in downloader.py): it is annotation-only, so the block is a fine place for it. I put it there on purpose so this PR doesn't touch the top import block that #1032 rewrites; the two PRs then merge without conflicts. The comment above the block is about the thrift import and still holds.

[peco-review-bot]

Verdict: 1 Medium · 1 Low

The prefetch design looks sound, with one medium concern: if a link fetch fails and the caller retries, rows can be skipped with no error, where the old per-batch path raised a DataError. Nit: import threading in downloader.py sits under TYPE_CHECKING, next to a comment about keeping thrift out of the import path. It works because annotations are deferred, but threading is stdlib and could be a normal import.

[vuanhphung]

The import threading nit in this review body is answered in #1038 (comment). It stays under TYPE_CHECKING so this PR doesn't touch the top import block that #1032 rewrites.