/
record_batch_reader.go
263 lines (229 loc) · 7.44 KB
/
record_batch_reader.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package flight
import (
"bytes"
"fmt"
"sync/atomic"
"github.com/apache/arrow/go/v15/arrow"
"github.com/apache/arrow/go/v15/arrow/array"
"github.com/apache/arrow/go/v15/arrow/arrio"
"github.com/apache/arrow/go/v15/arrow/internal/debug"
"github.com/apache/arrow/go/v15/arrow/ipc"
"github.com/apache/arrow/go/v15/arrow/memory"
)
// DataStreamReader is an interface for receiving flight data messages on a stream
// such as via grpc with Arrow Flight.
type DataStreamReader interface {
Recv() (*FlightData, error)
}
type dataMessageReader struct {
rdr DataStreamReader
peeked *FlightData
refCount int64
msg *ipc.Message
lastAppMetadata []byte
descr *FlightDescriptor
}
func (d *dataMessageReader) Message() (*ipc.Message, error) {
var (
fd *FlightData
err error
)
if d.peeked != nil {
fd = d.peeked
d.peeked = nil
} else {
fd, err = d.rdr.Recv()
}
if err != nil {
if d.msg != nil {
// clear the previous message in the error case
d.msg.Release()
d.msg = nil
}
d.lastAppMetadata = nil
d.descr = nil
return nil, err
}
d.lastAppMetadata = fd.AppMetadata
d.descr = fd.FlightDescriptor
d.msg = ipc.NewMessage(memory.NewBufferBytes(fd.DataHeader), memory.NewBufferBytes(fd.DataBody))
return d.msg, nil
}
func (d *dataMessageReader) Retain() {
atomic.AddInt64(&d.refCount, 1)
}
func (d *dataMessageReader) Release() {
debug.Assert(atomic.LoadInt64(&d.refCount) > 0, "too many releases")
if atomic.AddInt64(&d.refCount, -1) == 0 {
if d.msg != nil {
d.msg.Release()
d.msg = nil
}
d.lastAppMetadata = nil
}
}
// Reader is an ipc.Reader which also keeps track of the metadata from
// the FlightData messages as they come in, calling LatestAppMetadata
// will return the metadata bytes from the most recently read message.
type Reader struct {
*ipc.Reader
dmr *dataMessageReader
}
// Retain increases the reference count for the underlying message reader
// and ipc.Reader which are utilized by this Reader.
func (r *Reader) Retain() {
r.Reader.Retain()
r.dmr.Retain()
}
// Release reduces the reference count for the underlying message reader
// and ipc.Reader, when the reference counts become zero, the allocated
// memory is released for the stored record and metadata.
func (r *Reader) Release() {
r.Reader.Release()
r.dmr.Release()
}
// LatestAppMetadata returns the bytes from the AppMetadata field of the
// most recently read FlightData message that was processed by calling
// the Next function. The metadata returned would correspond to the record
// retrieved by calling Record().
func (r *Reader) LatestAppMetadata() []byte {
return r.dmr.lastAppMetadata
}
// LatestFlightDescriptor returns a pointer to the last FlightDescriptor object
// that was received in the most recently read FlightData message that was
// processed by calling the Next function. The descriptor returned would correspond
// to the record retrieved by calling Record().
func (r *Reader) LatestFlightDescriptor() *FlightDescriptor {
return r.dmr.descr
}
// Chunk is a convenience function to return a chunk of the flight stream
// returning the RecordBatch along with the FlightDescriptor and any AppMetadata.
// Each of these can be retrieved separately with their respective functions,
// this is just a convenience to retrieve all three with one function call.
func (r *Reader) Chunk() StreamChunk {
return StreamChunk{
Data: r.Record(),
Desc: r.dmr.descr,
AppMetadata: r.dmr.lastAppMetadata,
}
}
// NewRecordReader constructs an ipc reader using the flight data stream reader
// as the source of the ipc messages, opts passed will be passed to the underlying
// ipc.Reader such as ipc.WithSchema and ipc.WithAllocator
func NewRecordReader(r DataStreamReader, opts ...ipc.Option) (*Reader, error) {
// peek the first message for a descriptor
data, err := r.Recv()
if err != nil {
return nil, err
}
rdr := &Reader{dmr: &dataMessageReader{rdr: r, refCount: 1}}
rdr.dmr.descr = data.FlightDescriptor
if len(data.DataHeader) > 0 {
rdr.dmr.peeked = data
}
rdr.dmr.Retain()
if rdr.Reader, err = ipc.NewReaderFromMessageReader(rdr.dmr, opts...); err != nil {
return nil, fmt.Errorf("arrow/flight: could not create flight reader: %w", err)
}
return rdr, nil
}
// DeserializeSchema takes the schema bytes from FlightInfo or SchemaResult
// and returns the deserialized arrow schema.
func DeserializeSchema(info []byte, mem memory.Allocator) (*arrow.Schema, error) {
// even though the Flight proto file says that the bytes should be the
// flatbuffer message as per Schema.fbs, the current implementations send
// a serialized recordbatch with no body rows rather than just the
// schema message. So let's make sure to follow that.
rdr, err := ipc.NewReader(bytes.NewReader(info), ipc.WithAllocator(mem))
if err != nil {
return nil, err
}
defer rdr.Release()
return rdr.Schema(), nil
}
// StreamChunk represents a single chunk of a FlightData stream
type StreamChunk struct {
Data arrow.Record
Desc *FlightDescriptor
AppMetadata []byte
Err error
}
// MessageReader is an interface representing a RecordReader
// that also provides StreamChunks and/or the ability to retrieve
// FlightDescriptors and AppMetadata from the flight stream
type MessageReader interface {
array.RecordReader
arrio.Reader
Err() error
Chunk() StreamChunk
LatestFlightDescriptor() *FlightDescriptor
LatestAppMetadata() []byte
}
type haserr interface {
Err() error
}
// StreamChunksFromReader is a convenience function to populate a channel
// from a record reader. It is intended to be run using a separate goroutine
// by calling `go flight.StreamChunksFromReader(rdr, ch)`.
//
// If the record reader panics, an error chunk will get sent on the channel.
//
// This will close the channel and release the reader when it completes.
func StreamChunksFromReader(rdr array.RecordReader, ch chan<- StreamChunk) {
defer close(ch)
defer func() {
if err := recover(); err != nil {
ch <- StreamChunk{Err: fmt.Errorf("panic while reading: %s", err)}
}
}()
defer rdr.Release()
for rdr.Next() {
rec := rdr.Record()
rec.Retain()
ch <- StreamChunk{Data: rec}
}
if e, ok := rdr.(haserr); ok {
if e.Err() != nil {
ch <- StreamChunk{Err: e.Err()}
}
}
}
func ConcatenateReaders(rdrs []array.RecordReader, ch chan<- StreamChunk) {
defer close(ch)
defer func() {
for _, r := range rdrs {
r.Release()
}
if err := recover(); err != nil {
ch <- StreamChunk{Err: fmt.Errorf("panic while reading: %s", err)}
}
}()
for _, r := range rdrs {
for r.Next() {
rec := r.Record()
rec.Retain()
ch <- StreamChunk{Data: rec}
}
if e, ok := r.(haserr); ok {
if e.Err() != nil {
ch <- StreamChunk{Err: e.Err()}
return
}
}
}
}