GeezerVison

What type of PR is this?

• Other (performance)

Description

This PR changes how convert_arrow_based_set_to_arrow_table handles inline Arrow results on Thrift, with two effects: more statements per second for many-batch results under many threads, and large inline results that release their memory when freed. It supersedes #1036. Decoding through pyarrow alone made large-result memory worse than main, so both changes have to land together.

Throughput. The warehouse returns small inline results as dozens of LZ4-compressed Arrow batches: range(100) and range(1000) both arrive as 72 batches. Each batch was decompressed with lz4.frame.decompress, which releases and reacquires the GIL on every call. When many threads run statements, each reacquire waits behind the other runnable threads. Against a mock that replays the captured reply with 50 ms RTT, a 1,000-row statement at 64 threads made ~600 voluntary context switches and used 5.3 ms of CPU, half of it sys time (1.4 ms with one thread), and throughput stalled at ~390 statements/s. Replaying the decompressed bytes from a dict, with no C call, removed the effect, which confirmed the cause. For three or more compressed batches, the frames are now concatenated and decoded with one pyarrow.CompressedInputStream read. Arrow decodes concatenated LZ4 frames natively; python-lz4's one-shot decompress silently returns only the first frame. One- and two-batch results such as SELECT 1, uncompressed results, and pyarrow builds without the lz4 codec keep the per-batch path.

Memory. Each FetchResults row set (~3.7 MB for a large result) used to be concatenated into a bytearray on the glibc heap. The Arrow table reads it zero-copy, so these buffers live as long as the result. glibc couldn't give the space back afterwards: a 1.7 GB inline fetchall_arrow() left ~1.6 GB resident on main. Decoding through pyarrow alone made this worse, because its large read temporaries pushed the result buffers onto the brk heap (peak 2909 MB, 2859 MB left after free). Each row set's stream is now built in its own private anonymous mmap. The map is reserved at 4x the compressed size (untouched pages cost no memory), filled by readinto() straight from the decode stream (or batch by batch on the per-batch path), grown via mremap if needed, and shrunk to size before the zero-copy export. It returns to the OS when the table is released. A stream moves to a map only once it holds at least 1 MiB, because each held map is one VMA (vm.max_map_count); reaching the default limit would take over 64 GB of inline results held at once. Smaller results, mmap failures and non-Linux platforms keep the bytearray, and whenever no map was created (other platforms, or a failed mmap), results whose compressed size reaches 1 MiB keep the per-batch decode, as on main.

CloudFetch, SEA and the kernel backend are unchanged.

Trade-off: Arrow takes the stream's decode buffers from pyarrow's memory pool, which keeps a heap per thread. In a 64-thread mock run, peak RSS was 198 MB on main, 228 MB with the decode change alone and 231 MB with both, about 0.5 MB per busy thread. The memory map itself adds nothing measurable. An earlier mock measured +50 MB at 200 threads.

How is this tested?

• Unit tests
• E2E Tests
• Manually

Unit tests decode real LZ4-compressed Arrow batches: 1, 2, 3 and 72 batches; compressed and not; a highly compressible result that needs several reads; and the no-codec fallback. They check that 72 batches no longer call lz4.frame.decompress. Corrupt and truncated frames raise RuntimeError on both paths, and MemoryError passes through unchanged. The memory-map tests cover:

• 72 batches decoding as one stream into one map, and two large batches taking the per-batch path into the map;
• growth when the estimate is too small, including the retry at the exact size;
• the final shrink, and slicing when the shrink fails;
• heap fallbacks when mmap or growth fails;
• a compressible stream moving to a map once it crosses 1 MiB, while smaller streams stay on the heap;
• non-Linux platforms;
• the table keeping its map alive and releasing it with the table.

They pass on pyarrow 14.0.2, 23.0.1 and 26.0.1, and the full unit suite passes. Live e2e: the TestPySQLCoreSuite fetch/Arrow tests, TestPySQLLargeNarrowResultSet and test_complex_types.py pass (52). Earlier on the decode commits, the test_driver.py core, large wide/narrow and async suites passed (132); the 6 test_staging_ingestion_* failures there were a missing ingestion user in that environment and failed the same way on main. Inline results (lz4 on and off) matched CloudFetch results for 2M rows of mixed types, and 40 captured live row sets decode identically.

Memory, live Thrift warehouse with use_cloud_fetch=False: inline fetchall_arrow, 3 fresh processes per cell in rotating order, median RSS increase over the pre-query baseline. "Decode only" is this PR without the memory-map commit.

Throughput, 64 threads, a connection and a cursor per statement per thread, 1,000-row / 72-batch inline result, mock with 50 ms RTT, 3 interleaved runs: 528 → 1109 statements/s and p50 120 → 56 ms on main → this PR (decode only: 1113 statements/s; the memory map doesn't affect small results).

Results from the decode commits, measured with the original harness (median of 5 interleaved runs, 1,000-row result). The small-result path they exercise is unchanged by the memory-map commit.

On the live warehouse, 64-thread fetchall_arrow and all 1- and 16-thread runs were capped by the shared warehouse, and the higher live p99 at 200 threads coincides with ~30% more load on it. SELECT 1 (one batch, unchanged path) and the kernel backend measured the same on both sides.

Related Tickets & Documents

Supersedes #1036.

This pull request and its description were written by Isaac.