This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new 9e6dadb3 fix(arrow/flight): limit descriptors to next payload (#1042)
9e6dadb3 is described below
commit 9e6dadb33f6c2fafcbba11cceb2da1e313efa947
Author: Minh Vu <[email protected]>
AuthorDate: Wed Aug 5 19:07:22 2026 +0200
fix(arrow/flight): limit descriptors to next payload (#1042)
## 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/...`
---
arrow/flight/flight_test.go | 59 +++++++++++++++++++++++++++++++++++
arrow/flight/flightsql/client_test.go | 5 ++-
arrow/flight/record_batch_reader.go | 4 ++-
arrow/flight/record_batch_writer.go | 9 ++----
4 files changed, 69 insertions(+), 8 deletions(-)
diff --git a/arrow/flight/flight_test.go b/arrow/flight/flight_test.go
index e6b8a25b..a95db2bc 100644
--- a/arrow/flight/flight_test.go
+++ b/arrow/flight/flight_test.go
@@ -481,6 +481,65 @@ func TestWriterInconsistentSchema(t *testing.T) {
require.NoError(t, w.Close())
}
+func TestWriterFlightDescriptorOnlyAppliesToNextPayload(t *testing.T) {
+ recs, ok := arrdata.Records["primitives"]
+ require.True(t, ok)
+
+ fs := collectingFlightStreamWriter{}
+ w := flight.NewRecordWriter(&fs, ipc.WithSchema(recs[0].Schema()))
+ w.SetFlightDescriptor(&flight.FlightDescriptor{Path:
[]string{"dataset"}})
+ require.NoError(t, w.Write(recs[0]))
+ require.Greater(t, len(fs.hasDescriptor), 1)
+ require.True(t, fs.hasDescriptor[0])
+ for _, hasDescriptor := range fs.hasDescriptor[1:] {
+ require.False(t, hasDescriptor)
+ }
+ require.NoError(t, w.Close())
+}
+
+func TestReaderRetainsDescriptorForFirstRecordBatch(t *testing.T) {
+ recs, ok := arrdata.Records["primitives"]
+ require.True(t, ok)
+
+ descriptor := &flight.FlightDescriptor{Path: []string{"dataset"}}
+ fs := collectingFlightStreamWriter{}
+ w := flight.NewRecordWriter(&fs, ipc.WithSchema(recs[0].Schema()))
+ w.SetFlightDescriptor(descriptor)
+ require.NoError(t, w.Write(recs[0]))
+ require.NoError(t, w.Close())
+
+ r, err := flight.NewRecordReader(&fs)
+ require.NoError(t, err)
+ defer r.Release()
+ require.True(t, r.Next())
+ require.Equal(t, descriptor, r.Chunk().Desc)
+}
+
+type collectingFlightStreamWriter struct {
+ hasDescriptor []bool
+ payloads []*flight.FlightData
+ next int
+}
+
+func (f *collectingFlightStreamWriter) Send(data *flight.FlightData) error {
+ f.hasDescriptor = append(f.hasDescriptor, data.FlightDescriptor != nil)
+ f.payloads = append(f.payloads, &flight.FlightData{
+ FlightDescriptor: data.FlightDescriptor,
+ DataHeader: append([]byte(nil), data.DataHeader...),
+ DataBody: append([]byte(nil), data.DataBody...),
+ })
+ return nil
+}
+
+func (f *collectingFlightStreamWriter) Recv() (*flight.FlightData, error) {
+ if f.next == len(f.payloads) {
+ return nil, io.EOF
+ }
+ payload := f.payloads[f.next]
+ f.next++
+ return payload, nil
+}
+
type flightStreamWriter struct{}
// Send implements flight.DataStreamWriter.
diff --git a/arrow/flight/flightsql/client_test.go
b/arrow/flight/flightsql/client_test.go
index 858526ea..af690fb1 100644
--- a/arrow/flight/flightsql/client_test.go
+++ b/arrow/flight/flightsql/client_test.go
@@ -474,7 +474,10 @@ func (s *FlightSqlClientSuite)
TestPreparedStatementExecuteParamBinding() {
s.mockClient.On("DoPut", s.callOpts).Return(mockedPut, nil)
mockedPut.On("Send", mock.MatchedBy(func(fd *flight.FlightData) bool {
return proto.Equal(expectedDesc, fd.FlightDescriptor)
- })).Return(nil).Twice() // first sends schema message, second sends data
+ })).Return(nil).Once()
+ mockedPut.On("Send", mock.MatchedBy(func(fd *flight.FlightData) bool {
+ return fd.FlightDescriptor == nil
+ })).Return(nil).Once()
mockedPut.On("CloseSend").Return(nil)
mockedPut.On("Recv").Return(putResult, nil)
diff --git a/arrow/flight/record_batch_reader.go
b/arrow/flight/record_batch_reader.go
index 8c9d3876..dac7a352 100644
--- a/arrow/flight/record_batch_reader.go
+++ b/arrow/flight/record_batch_reader.go
@@ -75,7 +75,9 @@ func (d *dataMessageReader) Message() (*ipc.Message, error) {
}
d.lastAppMetadata = fd.AppMetadata
- d.descr = fd.FlightDescriptor
+ if fd.FlightDescriptor != nil {
+ d.descr = fd.FlightDescriptor
+ }
d.msg = ipc.NewMessage(memory.NewBufferBytes(fd.DataHeader),
memory.NewBufferBytes(fd.DataBody))
return d.msg, nil
}
diff --git a/arrow/flight/record_batch_writer.go
b/arrow/flight/record_batch_writer.go
index 9cfb6bd8..124ceb36 100644
--- a/arrow/flight/record_batch_writer.go
+++ b/arrow/flight/record_batch_writer.go
@@ -47,7 +47,9 @@ func (f *flightPayloadWriter) WritePayload(payload
ipc.Payload) error {
payload.SerializeBody(&f.buf)
f.fd.DataBody = f.buf.Bytes()
- return f.w.Send(&f.fd)
+ err := f.w.Send(&f.fd)
+ f.fd.FlightDescriptor = nil
+ return err
}
func (f *flightPayloadWriter) Close() error { return nil }
@@ -75,11 +77,6 @@ func (w *Writer) SetFlightDescriptor(descr
*FlightDescriptor) {
// Write writes a recordbatch payload and returns any error, implementing the
arrio.Writer interface
func (w *Writer) Write(rec arrow.RecordBatch) error {
- if w.pw.fd.FlightDescriptor != nil {
- defer func() {
- w.pw.fd.FlightDescriptor = nil
- }()
- }
return w.Writer.Write(rec)
}
