fix(arrow/flight): limit descriptors to next payload - #1042
fix(arrow/flight): limit descriptors to next payload#1042fallintoplace wants to merge 2 commits into
Conversation
zeroshade
left a comment
There was a problem hiding this comment.
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:339uses afirst_payload_flag and attaches the descriptor to the first payload only — it even requires that payload beMessageType::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 dataThis 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.
| err := f.w.Send(&f.fd) | ||
| f.fd.FlightDescriptor = nil | ||
| return err |
There was a problem hiding this comment.
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:
arrow/flight/flightsql/client_test.go:475-477asserts the old behavior via.Twice()and now fails. Needs updating in this PR.- On the read side,
record_batch_reader.go:76overwritesdescrfor every message, soLatestFlightDescriptor()/Chunk().Descnow return nil for the first record batch — which contradicts the contract documented atrecord_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.
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 thoughSetFlightDescriptordocuments 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
LatestFlightDescriptorandChunk().Descbehavior 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/...