Skip to content

[PureGo] Accept Arrow IPC bytes as ingest input - #768

Draft
zlata-stefanovic-db wants to merge 2 commits into
mainfrom
purego-arrow-ipc-preflight
Draft

[PureGo] Accept Arrow IPC bytes as ingest input#768
zlata-stefanovic-db wants to merge 2 commits into
mainfrom
purego-arrow-ipc-preflight

Conversation

@zlata-stefanovic-db

Copy link
Copy Markdown
Contributor

What changes are proposed in this pull request?

Accepts caller-owned Arrow IPC bytes as an ingestion input, alongside the live
RecordBatch that #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. EncodeIPC
takes those bytes and produces the same canonical payload the typed path
produces, so everything downstream is unchanged.

The two pieces worth review attention:

  • The input is reserialized, not forwarded. The batch is read back and
    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.
  • Admission cannot use the wire length. A compressed batch weighs a fraction
    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
    preflightIPCExpansion walks the message headers directly — without building
    a 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 shape
into 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/flatbuffers becomes 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:

  1. Flight frame encoding: the 2 MiB chunk plan measured by binary searching
    actual encoded protobuf size, and the frame emitter.
  2. Flight acknowledgment and batch metadata parsing in internal/transport.
  3. stream.EncoderHooks and the stream-core wiring.

#732 also carries DecodeIPCRecordBatch and exported EncodeSchemaIPC /
DecodeSchemaIPC. Nothing outside their own tests consumes them even there, so
as with #758 they ride with the change that gives them a caller. #732's version
of EncodeIPC additionally computes a Flight chunk plan; that arrives with the
frame encoder.

How is this tested?

18 unit tests in internal/arrowproto (33 cases counting subtests), none of
which need a server or a network. go test -race ./... passes across the module
(586 cases, 10 packages), as does CGO_ENABLED=0 go test ./..., and gofmt,
go vet, and go mod tidy -diff are clean.

Coverage worth calling out:

  • The compressed-admission case asserts the estimate reaches at least the
    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.NewCheckedAllocator and asserts zero allocations, which is the other
    half of the claim: the expansion is measured without ever being performed. Both
    LZ4 and Zstd are covered.
  • A dictionary-encoded compressed stream is charged for its dictionary values,
    the case that distinguishes walking every message from walking only the record
    batch.
  • Declared sizes that overflow int64 are rejected, driven through a
    hand-built flatbuffer rather than a real Arrow stream, since arrow-go will not
    produce one.
  • Malformed inputs — empty, garbage, truncated mid-message, trailing bytes,
    schema with no batch — all return an error and report zero expanded bytes, so
    a rejected stream cannot leave a usable-looking size behind.
  • EncodeIPC canonicalizes and copies: the caller's buffer is zeroed
    immediately after the call and the payload still decodes to the original rows.
  • Admission covers the materialized payload for both an uncompressed and a
    dictionary-encoded input.

Not covered here because it needs the changes that follow: the core rejecting an
oversized compressed stream with ErrPayloadTooLarge before the decoder runs
(#732 tests this through EncoderHooks, which arrives with the wiring), Flight
frame chunking, and anything end-to-end against a Flight server.

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant