Repository navigation
SNOW-2499292 Implement serialisable arrow streams - #1834
Draft
sfc-gh-pfus wants to merge 1 commit into
Draft
sfc-gh-pfus wants to merge 1 commit into
sfc-gh-pfus wants to merge 1 commit into
Conversation
zeroshade
suggested changes
Aug 7, 2026
Comment on lines
+148
to
+150
| if client == nil { | ||
| return nil, fmt.Errorf("gosnowflake: no http client provided for chunk download") | ||
| } |
Contributor
There was a problem hiding this comment.
is there any reason not to allow using a default http.Client{} if one wasn't set? that would be preferable to erroring here.
| } | ||
| if resp.StatusCode != http.StatusOK { | ||
| defer func() { _ = resp.Body.Close() }() | ||
| b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) |
Contributor
There was a problem hiding this comment.
why ignore the error here?
Comment on lines
+172
to
+183
| bufStream := bufio.NewReader(resp.Body) | ||
| // Peek is best-effort: a short or empty payload simply skips gzip detection. | ||
| gzipMagic, _ := bufStream.Peek(2) | ||
| if len(gzipMagic) == 2 && gzipMagic[0] == 0x1f && gzipMagic[1] == 0x8b { | ||
| gz, err := gzip.NewReader(bufStream) | ||
| if err != nil { | ||
| _ = resp.Body.Close() | ||
| return nil, err | ||
| } | ||
| return &streamWrapReader{Reader: gz, wrapped: resp.Body}, nil | ||
| } | ||
| return &streamWrapReader{Reader: bufStream, wrapped: resp.Body}, nil |
Contributor
There was a problem hiding this comment.
instead of replicating this here, can we make this a function and call it in both places? (here and in the original chunk handler)
Comment on lines
+176
to
+181
| gz, err := gzip.NewReader(bufStream) | ||
| if err != nil { | ||
| _ = resp.Body.Close() | ||
| return nil, err | ||
| } | ||
| return &streamWrapReader{Reader: gz, wrapped: resp.Body}, nil |
Contributor
There was a problem hiding this comment.
You can simplify this by doing bufStream = gz instead so that you only have a single return at the bottom
Comment on lines
+21
to
+23
| // The descriptor carries only the chunk address and payload format; it deliberately does not | ||
| // carry column metadata or timezone. A consumer that needs to decode the bytes obtains | ||
| // RowTypes() and Location() from the ArrowStreamLoader and carries them alongside. |
Contributor
There was a problem hiding this comment.
why do this instead of including them in the serialized data already?
In addition, is there a way to know whether the result set is sorted (order matters) versus when order doesn't matter?
Comment on lines
+48
to
+69
| // ToSerializable converts a stream batch into a portable descriptor for distributed fetch. | ||
| // Call it before GetStream. The batch must have been produced by ArrowStreamLoader.GetBatches | ||
| // (i.e. be backed by a downloader). | ||
| func (asb *ArrowStreamBatch) ToSerializable() (SerializableArrowStreamBatch, error) { | ||
| if asb.scd == nil { | ||
| return SerializableArrowStreamBatch{}, fmt.Errorf( | ||
| "arrow stream batch is not backed by a downloader and cannot be serialized") | ||
| } | ||
| s := SerializableArrowStreamBatch{ | ||
| Format: asb.scd.queryResultFormat, | ||
| NumRows: asb.numrows, | ||
| } | ||
| if asb.inlineData != nil { | ||
| s.InlineData = asb.inlineData | ||
| return s, nil | ||
| } | ||
| meta := asb.scd.ChunkMetas[asb.idx] | ||
| s.URL = meta.URL | ||
| s.UncompressedSize = meta.UncompressedSize | ||
| s.Headers = streamChunkHeaders(asb.scd.ChunkHeader, asb.scd.Qrmk) | ||
| return s, nil | ||
| } |
Contributor
There was a problem hiding this comment.
why a separate ToSerializable method instead of just implementing the UnmarshalJSON directly on ArrowStreamBatch?
sfc-gh-pfus
force-pushed
the
pfus-08-06-snow-2499292_implement_serialisable_arrow_streams
branch
from
August 25, 2026 09:56
fd626ef to
661618d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
TL;DR
Adds
SerializableArrowStreamBatchto enable distributed Arrow result-set fetching without a Snowflake connection on worker nodes.What changed?
Introduced
SerializableArrowStreamBatch, a JSON-serializable descriptor that captures everything needed to download and read a single Arrow result chunk independently of the originating Snowflake connection. It carries the presigned chunk URL, required HTTP headers (including SSE-C decryption keys for S3, or provider-specific headers for GCS/Azure), inline data for the first batch, row count, and format.ArrowStreamBatch.ToSerializable()converts a batch into this portable descriptor, handling both inline (first-batch) and remote (presigned URL) cases.SerializableArrowStreamBatch.Reader(opts...)reconstructs anArrowStreamPartitionReaderfrom the descriptor. Remote batches require an*http.Clientsupplied viaWithStreamHTTPClient.ArrowStreamPartitionReader.GetStream(ctx)returns anio.ReadCloserover the chunk bytes, downloading and gzip-inflating as needed, with no Snowflake session involved.openStandaloneChunkStreamperforms the HTTP GET against the presigned URL, mirrors the gzip-detection logic from the existing downloader, and surfaces HTTP errors asSnowflakeError.How to test?
Unit tests cover:
ToSerializablefor remote SSE-C, remote ChunkHeader (GCS/Azure), inline, and missing-downloader error cases.ArrowStreamPartitionReaderfor plain and gzip-compressed remote payloads, inline payloads, missing-client errors, and expired (403) URLs.An integration test (
TestQueryArrowStreamSerializableRoundTrip) runs a 50,000-row query, serializes each batch descriptor to JSON, deserializes it, fetches the chunk via a bare*http.Client, and asserts that the reconstructed bytes are byte-for-byte identical to the original and decode to the correct total row count viaipc.NewReader.Why make this change?
Distributed query engines (e.g. ADBC
ExecutePartitions/ReadPartition) need to split a Snowflake result set on a coordinator and download chunks in parallel across worker nodes that have no Snowflake connection. Previously,ArrowStreamBatchwas tightly coupled to the driver's internal downloader and could not be shipped across process or machine boundaries.SerializableArrowStreamBatchdecouples the chunk descriptor from the connection, enabling fan-out fetch patterns while keeping the security-sensitive presigned URL and SSE-C key contained within the serialized payload.