Stream CloudFetch chunk decompression to bound memory and prevent OOM on large results

[msrathore-db]

Summary

Fixes #1508.

Downloading a large query result via CloudFetch on the 3.x driver could exhaust the JVM heap and throw java.lang.OutOfMemoryError: Java heap space (surfaced as Failed to ready chunk / Download failed for chunk index 0), where the 2.x driver succeeded at the same heap. This was a 2.x → 3.x regression.

Root cause

The OutOfMemoryError is on-heap. Each chunk was decompressed by materializing the entire decompressed Arrow payload into an on-heap byte[] (several times the compressed size) before parsing. With up to cloudFetchThreadPoolSize (default 16) chunks decompressing in parallel, transient on-heap usage scaled with download concurrency rather than with available heap.

(The parsed Arrow vectors themselves are off-heap, so they were not the cause — the transient on-heap decompression buffers were.)

Fix

Streamed decompression. Decompression is now streamed directly into the Arrow reader via DecompressionUtil.decompressToInputStream (a lazy LZ4FrameInputStream), so the decompressed payload is never materialized on-heap alongside the compressed bytes. Peak per-chunk on-heap cost drops from compressed + full-decompressed to compressed + one LZ4 block. Chunks are still downloaded, decompressed, and parsed in parallel ahead of consumption — throughput is unaffected.

This is the same LZ4FrameInputStream the 2.x Simba driver feeds its decompression through; this change simply pipes it into Arrow instead of buffering the output.

For constrained heaps

If a workload still hits OOM on a very small heap (many concurrent chunks × per-chunk cost), lower the existing cloudFetchThreadPoolSize connection property (default 16) to reduce the number of chunks buffered concurrently — this is the memory-vs-throughput lever. No new connection property is introduced.

Scope note

An earlier revision of this PR also added an in-memory byte-budget scheduler (cloudFetchMaxBytesInMemory). It was removed after testing showed it was redundant with the existing cloudFetchThreadPoolSize limit and its heuristic default did not reliably bind (the manifest reports compressed size, while the OOM is driven by decompressed + parse transients). Streaming decompression is the robust fix; concurrency is already tunable via cloudFetchThreadPoolSize.

Testing

Verified against a SQL warehouse on the issue's 26-column, 169,769-row repro table:

• Memory: the unpatched 3.x driver OOMs at -Xmx128m; the patched driver streams the full result at -Xmx128m (peak ~99 MB), and down to -Xmx48m with cloudFetchThreadPoolSize=1 (peak ~42 MB) — below where the 2.7.5 / 2.8.1 Simba drivers OOM.
• Correctness: a cell checksum over the full result is byte-identical between the patched driver, the unpatched driver, and the 2.7.5 / 2.8.1 Simba drivers, across SEA/Thrift paths, the LZ4_FRAME and NONE codecs, complex types (struct/array/map, decimal, binary, NULLs), and heap sizes — no truncation or corruption. Verified with a position-sensitive (ordered, chained) row digest.
• Unit tests: added streaming-decompression coverage (LZ4 frame + NONE) in DecompressionUtilTest; the jdbc-core unit suite passes.

NO_CHANGELOG=true

This pull request and its description were written by Isaac.

[gopalldb]

will it get attempted again?

[msrathore-db]

Yes. When canScheduleChunkDownload(...) returns false we break without advancing nextChunkToDownload, so the chunk is deferred, not dropped. downloadNextChunks() is re-invoked from releaseChunk() every time a consumed chunk frees budget, and it retries the same nextChunkToDownload. The always-allow-one rule (totalChunksInMemory == 0) guarantees forward progress even if a single chunk is larger than the whole budget. I've added a comment at the break documenting this.

[gopalldb]

Had added similar handling: https://github.com/databricks/databricks-jdbc/pull/1470/changes#diff-b1b86f2824a50564183d80c1022f233ecb2c960037b6a51f59dbac4f72fe7f91

[msrathore-db]

Good catch — #1470's decompressLazy and this decompressToInputStream are the same idea: wrap the compressed bytes in LZ4FrameInputStream and let the Arrow reader pull decompressed bytes lazily instead of materializing the full payload. Since #1470 is still open (and stacked on #1468), I didn't want to introduce a merge dependency here. Proposal: whichever lands first, the other consolidates onto the single shared helper — they're behaviourally identical so the merge is trivial. Happy to be the one to converge.

(Side note from testing: LZ4FrameInputStream already reads concatenated LZ4 frames on demand, so this also covers the >2 GB-decompressed case without a custom multi-frame reader.)

[gopalldb]

⚠️ [MAJOR] Byte budget silently inert when chunk sizes are unknown (0), disabling OOM protection
Line 331 | bug | logical_claude_agent | [1/3 Claude flagged]

In-memory byte budget is silently disabled whenever chunk sizes are unknown (0). chunkSizeInBytes defaults to 0 when the SEA manifest reports a null getByteCount() (ArrowResultChunk.java:222-223) or when the Thrift getBytesNum() is unset (line 256). With 0-sized chunks, totalBytesInMemory never grows, so canScheduleChunkDownload (AbstractRemoteChunkProvider.java:331-334) only ever enforces the count limit and the byte budget this PR adds is inactive. On exactly the large-result workloads where byte sizes may be missing, the heap bound silently does not apply, reintroducing the OOM the PR targets.

[gopalldb]

Mostly LG, some minor comments

[msrathore-db]

Fixed. You're right: with chunkSizeInBytes == 0 (SEA getByteCount() null / Thrift getBytesNum() unset) totalBytesInMemory never grew, so canScheduleChunkDownload enforced only the count limit and the byte budget was inert.

I added effectiveChunkSizeInBytes() in AbstractRemoteChunkProvider, which charges a size-less chunk a conservative fallback estimate (UNKNOWN_CHUNK_SIZE_ESTIMATE_BYTES), applied symmetrically on reserve (both RemoteChunkProvider and RemoteChunkProviderV2) and release, so the budget stays active. Added testUnknownChunkSizeStillConsumesBudget, which asserts a size-less chunk still consumes budget (previously inert).

One honest caveat worth flagging: the budget charges the manifest's (compressed) byteCount, which under-counts the real per-chunk on-heap footprint (compressed buffer + Arrow parse transients). On narrow results this fix + streaming decompression comfortably resolves #1508 (verified: streams the repro table at -Xmx48m, where unpatched 3.x and even 2.7.5/2.8.1 OOM). On very wide/many-chunk results the default budget can still be too loose and rely on cloudFetchMaxBytesInMemory being tuned down. I think tightening the default (or charging a decompressed-size estimate) is worth a follow-up rather than expanding this PR — open to your preference.

[gopalldb]

📄 src/main/java/com/databricks/jdbc/common/util/DecompressionUtil.java

💡 [MINOR] Possibly unused decompressToStream after switch to decompressToInputStream
Line 60 | style | deadcode_claude_agent

Potentially dead method — This PR replaces the ArrowResultChunk call site from DecompressionUtil.decompressToStream(...) (removed line 111 in ArrowResultChunk.java) to the new DecompressionUtil.decompressToInputStream(...). If decompressToStream had no other callers, it is now unused.

---

📄 src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingChunkProvider.java

⚠️ [MAJOR] Byte budget not applied to StreamingChunkProvider, leaving its OOM path open
Line 48 | bug | logical_claude_agent

PR's byte-budget bound is missing from StreamingChunkProvider. The PR description states the in-memory byte budget applies to both RemoteChunkProvider and StreamingChunkProvider, but StreamingChunkProvider (StreamingChunkProvider.java:48) implements ChunkProvider directly and does not extend AbstractRemoteChunkProvider, so it inherits none of the new gating (maxBytesInMemory, canScheduleChunkDownload, totalBytesInMemory). Its triggerDownloads() (lines 504-530) schedules downloads bounded only by the chunk count maxChunksInMemory (line 509). As a result, when this provider is selected (constructed in ArrowStreamResult.java:121/202), up to maxChunksInMemory large chunks can be downloaded and decompressed concurrently with no byte cap — the exact heap-exhaustion (OOM) scenario the PR set out to fix remains open for this path, and the PR description overstates coverage.

---

📄 src/main/java/com/databricks/jdbc/api/impl/arrow/incubator/RemoteChunkProviderV2.java

💡 [MINOR] In-memory budget charged before throwing link fetch in RemoteChunkProviderV2.downloadNextChunks leaks on failure
Line 123 | bug | logical_claude_agent

Byte/count budget charged before a throwing link fetch, leaked on error. In RemoteChunkProviderV2.downloadNextChunks(), totalChunksInMemory++ and the PR-added totalBytesInMemory += chunkSizeInBytes (lines 123-124) run before the linkDownloadService.getLinkForChunk(...).get() block (lines 125-142). That block re-throws DatabricksSQLException on ExecutionException/InterruptedException, exiting downloadNextChunks() before nextChunkToDownload++ (line 144). The reserved count and bytes are never reconciled on this path — a link-fetch failure leaves the budget over-charged for a chunk that was never scheduled. Unlike RemoteChunkProvider (lines 128-132), which increments only after a successful submit(), V2 reserves budget before the fallible operation. If scheduling continues after such a failure, the inflated totals can throttle or block further downloads.

[msrathore-db]

Heads-up on a scope change (force-pushed): I've removed the in-memory byte-budget (cloudFetchMaxBytesInMemory) and reduced this PR to streaming decompression only.

Why: real-warehouse testing showed the byte budget was redundant with the existing cloudFetchThreadPoolSize limit, and its default couldn't reliably bind — the manifest's byteCount is the compressed size, whereas the OOM is driven by the decompressed + Arrow-parse transients (~10–20×), so a compressed-bytes budget large enough to not throttle normal queries never engaged on the wide-result case anyway. Streaming decompression is the robust fix; for constrained heaps, cloudFetchThreadPoolSize (already a connection property) is the memory/throughput lever.

This makes the two earlier review comments obsolete:

• the decompressToInputStream vs Add SEA inline Arrow support (stacked on #1468) #1470 decompressLazy duplication — still one shared helper worth consolidating whenever the two land, but no longer entangled with budget code; and
• the size-0 byte-budget concern — moot, as there is no byte budget.

Verified byte-identical results (SEA/Thrift, LZ4/NONE, complex types, multiple heaps) vs unpatched and the 2.7.5/2.8.1 Simba drivers.