[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% 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). 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.27 → 10.66 s	14.59 → 15.28 s	5648 → 5663 MB
⅓ of lineitem, 1.6 GB, fetchall_arrow, n=5	8 → 13	4.40 → 3.29 s	7.87 → 3.77 s	2399 → 2379 MB
⅓ of lineitem, 1.6 GB, fetchmany_arrow(1M), n=3	8 → 13	4.55 → 3.59 s	4.66 → 3.68 s	1244 → 1304 MB

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):

	main	#1030	this PR
downloads in flight at close	0	6	8
MB read in the 5 s after close	0	45	6
RSS 5 s after close	312 MB	514 MB	292 MB
close()	0.14 s	0.13 s	0.12 s
exit lag when the process exits right after close	0.20 s	0.47 s	0.09 s

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.

[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

Solid, well-tested change. One medium issue: once cancel has been requested, reads raise OperationalError even when the result had already been fully received and nothing was lost.

[peco-review-bot]

Verdict: 2 Low

Looks good overall, with two low-severity notes on the downloader's new chunked read and retry loop. The download manager's prefetch, cancel-race handling (keeping links whose submit failed, keeping the link source attached after a cancel) and row-order logic hold up when traced through, and the new unit tests cover the race paths.

[peco-review-bot]

Verdict: 1 Medium · 1 Low

The prefetch design looks solid, with one medium concern: an error from the mid-stream link fetch can silently drop rows when the caller retries. Cancel/close handling, the retry of interrupted body reads, and the kill switch all check out. Unit coverage is generally thorough, apart from the kill-switch forwarding noted in the second finding.

[peco-review-bot]

Verdict: 4 Low

Looks good overall. The prefetch, cancel and retry logic is careful and well tested, including the cancel and shutdown races. I found only low-severity issues: the meaning of has_more_rows has changed, the stored link error is re-raised as the same object, a secondary error in Cursor.cancel() can hide the original one, and body-read retries have no backoff.

[peco-review-bot]

Verdict: 1 Low

Looks good: one low-severity issue with how the kill switch reads _cloudfetch_prefetch=None. I checked the prefetch scheduling loop, what happens on cancel or shutdown (including the submit() racing cancel()), chunked body reads and retries in UnifiedHttpClient.request_context, and the has_more_rows handoff in ThriftResultSet. All of them hold up, and tests/unit covers them well.

[peco-review-bot]

Verdict: 1 Low · 2 Nit

Looks good overall, with 3 low findings: lenient kill-switch parsing, retry backoff counted in the download-speed metrics, and a new result-set/download-manager reference cycle. The prefetch, cancel and retry logic hangs together well and has thorough unit tests, including the cancel races and the case where a failed link fetch must not let a retry skip rows.

Other findings

• ⚪ Nit — start_time is taken before the retry loop. When a body read is retried, download_duration therefore also covers the failed attempt(s) and the 0.5 s / 1 s backoff waits. _log_download_metrics then reports a lower MB/s and can log the CloudFetch download slower than threshold warning for a file whose successful attempt was fast. To fix it, reset start_time at the start of each attempt, or log the retry separately.
• ⚪ Nit — set_link_source(self._fetch_more_links) creates a reference cycle: result set → results queue → download_manager._link_source (a bound method) → result set. Before this change, a ThriftResultSet dropped without close() was freed by refcounting, and its ThreadPoolExecutor (up to max_download_threads idle workers) was torn down right away. Now the whole graph, including the executor and its threads, stays alive until the cyclic GC runs. A weakref.WeakMethod (or clearing _link_source once has_more is False, which already happens, plus on close) would avoid holding the result set strongly while a stream is still in progress.

[vuanhphung]

On the two review-body nits:

• start_time across retries: fixed in 779182a. start_time is reset after the backoff and checks, so the speed metrics (and the slow-download warning) cover only the attempt that succeeded.
• Reference cycle through the bound _fetch_more_links: I'm leaving this as is. The cycle is broken explicitly in both normal endings. At end of stream the manager sets _link_source = None when the last reply has hasMoreRows=False, and on close _shutdown_manager() clears it. Cursor.close() and the next execute() both close the active result set. The cycle only lasts for a result set that is dropped mid-stream without being closed, and the cyclic GC then collects it, after which the idle executor threads exit as they do on main. A WeakMethod would change the LinkSource contract for that leak-only case.

[peco-review-bot]

Verdict: 1 Low

Looks good, with one low concern: after a failed link fetch, downloads already in flight keep running until close(). The hand-off, cancel/close, retry and kill-switch paths look correct to me. Tests cover them well, including the races with cancel and the sticky error after a failed link fetch.

[peco-review-bot]

Verdict: 1 Low

Looks good, with one low concern: what the internal kill switch does with an unrecognized value. I checked the cancel and close races, the order of link fetches against the shutdown and cancel events, the body-read retry and connection release in downloader.py, the handoff of has_more_rows, and how SEA and the per-batch path share the download manager. I found no correctness problems, and the tests cover the new states well.

[peco-review-bot]

Verdict: 1 Medium · 2 Low

Solid change overall: prefetching, ordering, the cancel/close paths and the kill switch are carefully done and well tested. One medium concern: a failed link fetch now aborts reads of files that are already downloaded. There are also two minor robustness gaps.

[peco-review-bot]

🟡 Medium — With prefetch on, every get_next_downloaded_file call can send a FetchResults request through _schedule_downloads() before it returns the file the caller asked for. If that request fails after the backend's retries, the exception comes out of this read, even though the file the caller needs (and possibly up to max_download_threads - 1 more) may already be downloaded. _link_error is also set, so every later read fails as well. Before this change, a FetchResults failure only showed up once the current batch had been fully consumed. A fetchmany/fetchone consumer could read all the rows that had already arrived. Now those rows are lost, and the failure lands inside CloudFetchQueue.next_n_rows, so the rows that call had already collected are dropped too.

The cancel path already handles this by failing only once the downloaded files run out. Consider doing the same for link errors: store the error, keep returning the tasks already in _download_tasks, and raise only when no downloaded file is left (the same branch as the _cancel_requested check). Add a test where the link source fails while downloaded files are still queued.

[vuanhphung]

I don't think this is a correctness issue, so I'm keeping the current behaviour. Once FetchResults fails after the backend's own retries, the result can't be finished on main or here, and both raise. The only difference is how many already-downloaded rows a streaming consumer reads before the error. A fetchall* caller gets no rows either way, since its collected chunks are dropped when the error propagates. A fetchmany/fetchone consumer on main reads the rest of the current batch, then fails at the next FetchResults and loses the rows that call had collected, exactly as here.

Deferring the error would mean keeping the manager running after the failure so the remaining files can be read, which undoes the round-13 fix (stop downloads nobody will read). It would also add a third deferred-failure state next to cancel and shutdown, for a path that already ends in an error. Raising at the failure, with the sticky error making every later read fail loudly, is the simpler contract: the result is unusable and the statement has to be re-executed. If we want to return more partial rows before failing, I'd do it as a separate change with its own tests.

[peco-review-bot]

Fair point that a full drain ends in an error either way, and a fetchall* caller loses its rows on both branches. The case that still differs is a caller that never needs the failing batch. Take a fetchmany(n) or fetchone() that reads a few rows and then closes, or breaks out of iteration, or a LIMIT-style consumer. On main that caller never sends the failing FetchResults, so it succeeds. With prefetch, the background request can fail during a read whose file is already downloaded, so a read that would have worked now raises. The fix doesn't need to keep downloads running. You could record _link_error, stop scheduling, keep returning the tasks already in _download_tasks, and raise in the same "nothing left" branch the cancel path uses. That keeps round-13's "stop downloads nobody will read" and only delays the error until the caller actually needs rows that can't be fetched. If you'd still rather do this in a follow-up, that's your call. I'm leaving the thread open so the regression for partial readers stays visible.

[peco-review-bot]

Verdict: 2 Low

Looks good overall: the prefetch hand-off, the cancel and close semantics, and the chunked body read with retries are carefully designed and well tested. Two low findings: the first link prefetch waits until the first file has downloaded and the caller reads again, and the hot-loop case (no downloads in flight, server keeps returning no links) has no test.

[peco-review-bot]

🔵 Low — _start_prefetch() only runs after the queue has been built, and ThriftCloudFetchQueue.__init__ already calls _create_next_table(). That call goes through get_next_downloaded_file → _schedule_downloads while _link_source is still None, and then blocks on task.result() for the first file. So the first prefetching FetchResults can't start until the first file has finished downloading and the caller makes its next read. If the initial batch has fewer links than max_download_threads (for example the direct-results batch), the free slots sit idle for the whole first download, which is exactly the gap this PR is meant to close.

This costs some performance but doesn't affect correctness. Possible fixes: pass the link source into the queue/manager before the first table is built, or have set_link_source call _schedule_downloads() right away. The second runs on the consumer thread, so it keeps the 'no new concurrency on the Thrift client' property. If you keep the current ordering on purpose, a comment would help.

[peco-review-bot]

🔵 Low — The break that avoids polling the server only fires when self._download_tasks is non-empty. If every download slot is empty and the server keeps replying ([], True), or with only zero-row links, the while loop calls _link_source back-to-back with no backoff. That case isn't tested. The new tests (test_empty_reply_is_retried_on_the_next_call_not_in_a_loop, test_reply_with_only_empty_links_is_retried_on_the_next_call) always have a download in flight. The old per-batch fetchall loop behaves the same way, so this probably isn't a regression. Still, a test that pins the empty-tasks case would document the intent, and it would also show whether the comment ('ask again on the next call instead of polling the server now') is meant to cover this case.