Skip to content

fix(arrow/flight): limit descriptors to next payload - #1042

Open
fallintoplace wants to merge 2 commits into
apache:mainfrom
fallintoplace:fix/flight-descriptor-next-payload
Open

fix(arrow/flight): limit descriptors to next payload#1042
fallintoplace wants to merge 2 commits into
apache:mainfrom
fallintoplace:fix/flight-descriptor-next-payload

Conversation

@fallintoplace

@fallintoplace fallintoplace commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

What changed

Clear a configured Flight descriptor immediately after sending one IPC payload rather than after the entire record-batch write. On the read side, retain the latest non-nil descriptor so it remains associated with the first record batch.

Why

The descriptor remained set for the duration of the initial Write. In the common schema-plus-record path, this attached it to both the schema message and the first record batch even though SetFlightDescriptor documents that it applies only to the next payload.

Clearing after the send attempt keeps the descriptor on that attempted payload but prevents it from being reused if the caller encounters the send error. Since the stream has failed, retrying through a later payload would not preserve the original ordering contract.

Sending the descriptor only with the schema exposed an existing reader assumption: a following record message with no descriptor overwrote the stored value. Latching non-nil descriptors preserves the documented LatestFlightDescriptor and Chunk().Desc behavior for the corresponding first record batch.

The FlightSQL prepared-statement mock now expects the descriptor only on the schema payload and nil on the data payload.

Validation

go test ./arrow/flight/...

@fallintoplace
fallintoplace requested a review from zeroshade as a code owner July 28, 2026 10:54

@zeroshade zeroshade left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The underlying change is spec-correct and I want it to land — but it breaks an existing test in a package the PR didn't run, and the description misstates the bug, so please fix both first.

The direction is right, and I verified it against the other implementations. The authoritative proto comment is in-tree at arrow/flight/gen/flight/Flight.pb.go:1407-1408: "The descriptor of the data. This is only relevant when a client is starting a new DoPut stream." And both reference implementations attach it once at stream start:

  • Arrow C++ cpp/src/arrow/flight/client.cc:339 uses a first_payload_ flag and attaches the descriptor to the first payload only — it even requires that payload be MessageType::SCHEMA.
  • Arrow Java FlightClient.java:409 (// Send the descriptor to start.) sends a dedicated descriptor-only first message and never repeats it.

So arrow-go is currently the outlier and this reduces interop divergence. No concerns there.

Blocker: it breaks arrow/flight/flightsql.

WITH this PR:  FAIL arrow/flight/flightsql
                 --- FAIL: TestFlightSqlClient/TestPreparedStatementExecuteParamBinding
BASELINE (only record_batch_writer.go reverted):  ok arrow/flight/flightsql

I isolated it to the writer change by reverting that one file. It is the only new failure across a full-repo run. The cause is arrow/flight/flightsql/client_test.go:475-477, which encodes the old behavior explicitly, comment and all:

})).Return(nil).Twice() // first sends schema message, second sends data

This escaped because the PR's stated validation was go test ./arrow/flight — which excludes the flightsql subpackage. go test ./arrow/flight/... would have caught it. I confirmed the fix is bounded: replacing that single .Twice() with one .Once() for the schema message plus one .Once() asserting fd.FlightDescriptor == nil for the data message makes the package pass.

Please also correct the description. It says "every payload carried the descriptor". That isn't what happens — measured on baseline with one SetFlightDescriptor and three writes:

BASELINE:  msg0=true  msg1=true   msg2=false  msg3=false
PATCHED:   msg0=true  msg1=false  msg2=false  msg3=false

The old code cleared the descriptor after the whole Write(), but one Write emits several payloads (schema, dictionaries, record batch), so it leaked onto exactly one extra message — the first record batch — not all of them. The narrower claim is still a real bug worth fixing; it just should be stated accurately since the body becomes the squash commit message.

One more thing to address, a public-API behavior change the PR doesn't mention. dataMessageReader.Message() (record_batch_reader.go:76) overwrites descr on every message, so Reader.LatestFlightDescriptor() / Chunk().Desc now return nil for the first record batch:

BASELINE: after Next() #0 -> Chunk().Desc != nil = true
PATCHED:  after Next() #0 -> Chunk().Desc != nil = false

That contradicts the documented contract at record_batch_reader.go:130-133 ("The descriptor returned would correspond to the record batch retrieved by calling RecordBatch()"). Both in-tree consumers happen to be safe because they read the descriptor before Next() (flightsql/server.go:966, flight_integration/scenario.go:327), but an out-of-tree DoExchange server routing on rdr.Chunk().Desc would silently start seeing nil. Either latch the first descriptor on the reader side or update that doc comment.

The new writer test itself is good — it fails on unpatched main and passes with the fix. Once client_test.go is updated and CI is green on ./arrow/flight/..., I'm happy to merge this.

Comment on lines +50 to +52
err := f.w.Send(&f.fd)
f.fd.FlightDescriptor = nil
return err

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Clearing per-payload rather than per-Write is the correct granularity, and it matches C++ (client.cc:339, first_payload_ flag) and Java (FlightClient.java:409).

Two consequences to handle:

  1. arrow/flight/flightsql/client_test.go:475-477 asserts the old behavior via .Twice() and now fails. Needs updating in this PR.
  2. On the read side, record_batch_reader.go:76 overwrites descr for every message, so LatestFlightDescriptor() / Chunk().Desc now return nil for the first record batch — which contradicts the contract documented at record_batch_reader.go:130-133. Consider latching the first non-nil descriptor there, or adjusting that doc.

Also worth noting err is now returned after mutating f.fd, so the descriptor is cleared even when Send fails. That's almost certainly what you want (the stream is broken anyway), but it's a deliberate choice worth a comment.

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.

2 participants