[PureGo] Accept Arrow IPC bytes as ingest input - #768
Draft
zlata-stefanovic-db wants to merge 2 commits into
Draft
[PureGo] Accept Arrow IPC bytes as ingest input#768zlata-stefanovic-db wants to merge 2 commits into
zlata-stefanovic-db wants to merge 2 commits into
Conversation
A caller that already holds Arrow IPC bytes had no way in: only a live RecordBatch was accepted. EncodeIPC reads the stream back and reserializes it, so the payload carries its own schema and dictionary state rather than depending on frames the caller sent earlier and will not send again after a reconnect. Admission cannot use the wire length for this input. A compressed stream weighs a fraction of what Arrow allocates when it materializes, so the IPC metadata is walked directly and the uncompressed size every buffer declares is charged instead. On a 4,096-row batch of highly compressible strings the declared sizes are 41x the doubled wire length, which is the difference between refusing the stream and expanding it past the buffer limit. The dictionary message is charged the same way, since its values are the large buffers while the record batch holds only indices. The walk parses flatbuffers by hand, so it validates as it goes and converts a malformed shape into an error rather than a panic. Nothing here is reachable from a public API yet. Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>
The dictionaryBatch helper built values from rune('a'+i%distinct), which
runs past ASCII at distinct=64, so a third of the values encoded as
two-byte UTF-8 and the dictionary was 540 KB larger than the helper's own
distinct*valueBytes arithmetic implied. Draw from a 64-character ASCII
alphabet instead so that product is exact and the admission assertion
compares against the real footprint.
Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>
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
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.
What changes are proposed in this pull request?
Accepts caller-owned Arrow IPC bytes as an ingestion input, alongside the live
RecordBatchthat #758 added. This is the second of five changes split out of#732. Nothing in the package is reachable from a public API yet.
A caller that already holds IPC bytes — read from a file, forwarded from another
Arrow producer, or received over Flight — previously had no way in.
EncodeIPCtakes those bytes and produces the same canonical payload the typed path
produces, so everything downstream is unchanged.
The two pieces worth review attention:
written out again so the payload carries its own schema and dictionary state.
A caller's stream may reference dictionaries declared in frames it sent
earlier and will not send again, and a payload has to survive a reconnect on
its own. Reserializing also puts the input through the same exactly-one-batch
validation as the typed path, which matters more here because this is the
first input a caller supplies in wire form: a stream with trailing bytes, two
batches, or none is rejected rather than silently truncated.
of what Arrow allocates when it materializes, and the core admits before
encoding, so charging the compressed size would let a stream through and then
expand past the buffer limit while Arrow allocates for it. Every compressed
buffer declares its uncompressed size in the IPC metadata, so
preflightIPCExpansionwalks the message headers directly — without buildinga single Arrow array — and charges what those declarations add up to. On the
4,096-row string batch in the tests that number is 41x the doubled wire
length.
The walk is hand-rolled flatbuffer parsing rather than a call into arrow-go,
because arrow-go's own readers reach these fields only while materializing the
batch, which is the allocation this is trying to get ahead of. It validates as
it goes — message order, header type, body bounds, per-buffer ranges, negative
and overflowing declared sizes — and a
recover()converts a malformed shapeinto an error rather than a panic escaping into the caller's goroutine.
The dictionary message is charged on the same path. A dictionary-encoded stream
keeps its values in a separate IPC message while the record batch holds only
indices, so measuring the record batch alone would miss the bulk of it.
Dependencies.
github.com/google/flatbuffersbecomes a direct dependency.It was already in the module graph as an indirect dependency through arrow-go,
so this adds no new code to the build — only an explicit
require.Deliberately not here, each following as its own change:
actual encoded protobuf size, and the frame emitter.
internal/transport.stream.EncoderHooksand the stream-core wiring.#732 also carries
DecodeIPCRecordBatchand exportedEncodeSchemaIPC/DecodeSchemaIPC. Nothing outside their own tests consumes them even there, soas with #758 they ride with the change that gives them a caller. #732's version
of
EncodeIPCadditionally computes a Flight chunk plan; that arrives with theframe encoder.
How is this tested?
18 unit tests in
internal/arrowproto(33 cases counting subtests), none ofwhich need a server or a network.
go test -race ./...passes across the module(586 cases, 10 packages), as does
CGO_ENABLED=0 go test ./..., andgofmt,go vet, andgo mod tidy -diffare clean.Coverage worth calling out:
uncompressed footprint the buffers declare, which doubling the compressed
length cannot reach — so it fails if the preflight stops reading declarations
rather than merely looking conservative. It runs on a
memory.NewCheckedAllocatorand asserts zero allocations, which is the otherhalf of the claim: the expansion is measured without ever being performed. Both
LZ4 and Zstd are covered.
the case that distinguishes walking every message from walking only the record
batch.
int64are rejected, driven through ahand-built flatbuffer rather than a real Arrow stream, since arrow-go will not
produce one.
schema with no batch — all return an error and report zero expanded bytes, so
a rejected stream cannot leave a usable-looking size behind.
EncodeIPCcanonicalizes and copies: the caller's buffer is zeroedimmediately after the call and the payload still decodes to the original rows.
dictionary-encoded input.
Not covered here because it needs the changes that follow: the core rejecting an
oversized compressed stream with
ErrPayloadTooLargebefore the decoder runs(#732 tests this through
EncoderHooks, which arrives with the wiring), Flightframe chunking, and anything end-to-end against a Flight server.