Skip to content

Commit dcbaaa9

Browse files
David Lipitrou
David Li
authored and
pitrou
committed
ARROW-3162: [Python] [Flight] Enable implementing Flight servers in Python
Author: David Li <David.M.Li@twosigma.com> Closes apache#3712 from lihalite/flight-python-server and squashes the following commits: 10dccb6 <David Li> Use 'except *' to handle exceptions in Cython 445eb87 <David Li> Copy Flight out parameter values instead of moving in Python 9839e0f <David Li> Lint Python Flight code f33b51f <David Li> Guard against Flight server methods not filling out parameter 422b4bc <David Li> Bind Flight action endpoints to Python c09b138 <David Li> Bind Flight ListFlights server call to Python 4a5cbaf <David Li> Guard against Flight server returning null listing 033734d <David Li> Bind Flight DoGet server call to Python 35c62f6 <David Li> Treat null FlightDataStream as error
1 parent 1d0b369 commit dcbaaa9

7 files changed

Lines changed: 301 additions & 49 deletions

File tree

‎cpp/src/arrow/flight/server.cc‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,9 @@ class FlightServiceImpl : public FlightService::Service {
108108

109109
template <typename UserType, typename Iterator, typename ProtoType>
110110
grpc::Status WriteStream(Iterator* iterator, ServerWriter<ProtoType>* writer) {
111+
if (!iterator) {
112+
return grpc::Status(grpc::StatusCode::INTERNAL, "No items to iterate");
113+
}
111114
// Write flight info to stream until listing is exhausted
112115
ProtoType pb_value;
113116
std::unique_ptr<UserType> value;
@@ -153,6 +156,10 @@ class FlightServiceImpl : public FlightService::Service {
153156
GRPC_RETURN_NOT_OK(internal::FromProto(*request, &criteria));
154157
}
155158
GRPC_RETURN_NOT_OK(server_->ListFlights(&criteria, &listing));
159+
if (!listing) {
160+
// Treat null listing as no flights available
161+
return grpc::Status::OK;
162+
}
156163
return WriteStream<FlightInfo>(listing.get(), writer);
157164
}
158165

@@ -166,6 +173,11 @@ class FlightServiceImpl : public FlightService::Service {
166173
std::unique_ptr<FlightInfo> info;
167174
GRPC_RETURN_NOT_OK(server_->GetFlightInfo(descr, &info));
168175

176+
if (!info) {
177+
// Treat null listing as no flights available
178+
return grpc::Status(grpc::StatusCode::NOT_FOUND, "Flight not found");
179+
}
180+
169181
GRPC_RETURN_NOT_OK(internal::ToProto(*info, response));
170182
return grpc::Status::OK;
171183
}
@@ -180,6 +192,10 @@ class FlightServiceImpl : public FlightService::Service {
180192
std::unique_ptr<FlightDataStream> data_stream;
181193
GRPC_RETURN_NOT_OK(server_->DoGet(ticket, &data_stream));
182194

195+
if (!data_stream) {
196+
return grpc::Status(grpc::StatusCode::NOT_FOUND, "No data in this flight");
197+
}
198+
183199
// Write the schema as the first message in the stream
184200
FlightPayload schema_payload;
185201
MemoryPool* pool = default_memory_pool();
@@ -218,7 +234,7 @@ class FlightServiceImpl : public FlightService::Service {
218234
if (!message || message->type() != ipc::Message::Type::SCHEMA) {
219235
return internal::ToGrpcStatus(
220236
Status(StatusCode::Invalid, "DoPut must start with schema/descriptor"));
221-
} else if (data.descriptor == nullptr) {
237+
} else if (!data.descriptor) {
222238
return internal::ToGrpcStatus(
223239
Status(StatusCode::Invalid, "DoPut must start with non-null descriptor"));
224240
} else {
@@ -253,14 +269,18 @@ class FlightServiceImpl : public FlightService::Service {
253269
std::unique_ptr<ResultStream> results;
254270
GRPC_RETURN_NOT_OK(server_->DoAction(action, &results));
255271

256-
std::unique_ptr<Result> result;
257-
pb::Result pb_result;
272+
if (!results) {
273+
return grpc::Status::CANCELLED;
274+
}
275+
258276
while (true) {
277+
std::unique_ptr<Result> result;
259278
GRPC_RETURN_NOT_OK(results->Next(&result));
260279
if (!result) {
261280
// No more results
262281
break;
263282
}
283+
pb::Result pb_result;
264284
GRPC_RETURN_NOT_OK(internal::ToProto(*result, &pb_result));
265285
if (!writer->Write(pb_result)) {
266286
// Stream may be closed

‎cpp/src/arrow/python/flight.cc‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,9 @@ PyFlightServer::PyFlightServer(PyObject* server, PyFlightServerVtable vtable)
3333
Status PyFlightServer::ListFlights(
3434
const arrow::flight::Criteria* criteria,
3535
std::unique_ptr<arrow::flight::FlightListing>* listings) {
36-
return Status::NotImplemented("NYI");
36+
PyAcquireGIL lock;
37+
vtable_.list_flights(server_.obj(), criteria, listings);
38+
return CheckPyError();
3739
}
3840

3941
Status PyFlightServer::GetFlightInfo(const arrow::flight::FlightDescriptor& request,
@@ -58,11 +60,28 @@ Status PyFlightServer::DoPut(std::unique_ptr<arrow::flight::FlightMessageReader>
5860

5961
Status PyFlightServer::DoAction(const arrow::flight::Action& action,
6062
std::unique_ptr<arrow::flight::ResultStream>* result) {
61-
return Status::NotImplemented("NYI");
63+
PyAcquireGIL lock;
64+
vtable_.do_action(server_.obj(), action, result);
65+
return CheckPyError();
6266
}
6367

6468
Status PyFlightServer::ListActions(std::vector<arrow::flight::ActionType>* actions) {
65-
return Status::NotImplemented("NYI");
69+
PyAcquireGIL lock;
70+
vtable_.list_actions(server_.obj(), actions);
71+
return CheckPyError();
72+
}
73+
74+
PyFlightResultStream::PyFlightResultStream(PyObject* generator,
75+
PyFlightResultStreamCallback callback)
76+
: callback_(callback) {
77+
Py_INCREF(generator);
78+
generator_.reset(generator);
79+
}
80+
81+
Status PyFlightResultStream::Next(std::unique_ptr<arrow::flight::Result>* result) {
82+
PyAcquireGIL lock;
83+
callback_(generator_.obj(), result);
84+
return CheckPyError();
6685
}
6786

6887
Status CreateFlightInfo(const std::shared_ptr<arrow::Schema>& schema,

‎cpp/src/arrow/python/flight.h‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,14 +35,21 @@ namespace flight {
3535
/// Python.
3636
class ARROW_PYTHON_EXPORT PyFlightServerVtable {
3737
public:
38+
std::function<void(PyObject*, const arrow::flight::Criteria*,
39+
std::unique_ptr<arrow::flight::FlightListing>*)>
40+
list_flights;
3841
std::function<void(PyObject*, const arrow::flight::FlightDescriptor&,
3942
std::unique_ptr<arrow::flight::FlightInfo>*)>
4043
get_flight_info;
41-
std::function<void(PyObject*, std::unique_ptr<arrow::flight::FlightMessageReader>)>
42-
do_put;
4344
std::function<void(PyObject*, const arrow::flight::Ticket&,
4445
std::unique_ptr<arrow::flight::FlightDataStream>*)>
4546
do_get;
47+
std::function<void(PyObject*, std::unique_ptr<arrow::flight::FlightMessageReader>)>
48+
do_put;
49+
std::function<void(PyObject*, const arrow::flight::Action&,
50+
std::unique_ptr<arrow::flight::ResultStream>*)>
51+
do_action;
52+
std::function<void(PyObject*, std::vector<arrow::flight::ActionType>*)> list_actions;
4653
};
4754

4855
class ARROW_PYTHON_EXPORT PyFlightServer : public arrow::flight::FlightServerBase {
@@ -65,6 +72,22 @@ class ARROW_PYTHON_EXPORT PyFlightServer : public arrow::flight::FlightServerBas
6572
PyFlightServerVtable vtable_;
6673
};
6774

75+
/// \brief A callback that obtains the next result from a Flight action.
76+
typedef std::function<void(PyObject*, std::unique_ptr<arrow::flight::Result>*)>
77+
PyFlightResultStreamCallback;
78+
79+
/// \brief A ResultStream built around a Python callback.
80+
class ARROW_PYTHON_EXPORT PyFlightResultStream : public arrow::flight::ResultStream {
81+
public:
82+
explicit PyFlightResultStream(PyObject* generator,
83+
PyFlightResultStreamCallback callback);
84+
Status Next(std::unique_ptr<arrow::flight::Result>* result) override;
85+
86+
private:
87+
OwnedRefNoGIL generator_;
88+
PyFlightResultStreamCallback callback_;
89+
};
90+
6891
ARROW_PYTHON_EXPORT
6992
Status CreateFlightInfo(const std::shared_ptr<arrow::Schema>& schema,
7093
const arrow::flight::FlightDescriptor& descriptor,

‎python/examples/flight/server.py‎

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@
1717

1818
"""An example Flight Python server."""
1919

20+
import ast
21+
import threading
22+
import time
23+
2024
import pyarrow
2125
import pyarrow.flight
2226

@@ -28,15 +32,35 @@ def __init__(self):
2832

2933
@classmethod
3034
def descriptor_to_key(self, descriptor):
31-
return (descriptor.descriptor_type, descriptor.command,
35+
return (descriptor.descriptor_type.value, descriptor.command,
3236
tuple(descriptor.path or tuple()))
3337

38+
def list_flights(self, criteria):
39+
for key, table in self.flights.items():
40+
if key[1] is not None:
41+
descriptor = \
42+
pyarrow.flight.FlightDescriptor.for_command(key[1])
43+
else:
44+
descriptor = pyarrow.flight.FlightDescriptor.for_path(*key[2])
45+
46+
endpoints = [
47+
pyarrow.flight.FlightEndpoint(repr(key),
48+
[('localhost', 5005)]),
49+
]
50+
yield pyarrow.flight.FlightInfo(table.schema,
51+
descriptor, endpoints,
52+
table.num_rows, 0)
53+
3454
def get_flight_info(self, descriptor):
3555
key = FlightServer.descriptor_to_key(descriptor)
3656
if key in self.flights:
3757
table = self.flights[key]
58+
endpoints = [
59+
pyarrow.flight.FlightEndpoint(repr(key),
60+
[('localhost', 5005)]),
61+
]
3862
return pyarrow.flight.FlightInfo(table.schema,
39-
descriptor, [],
63+
descriptor, endpoints,
4064
table.num_rows, 0)
4165
raise KeyError('Flight not found.')
4266

@@ -46,6 +70,33 @@ def do_put(self, descriptor, reader):
4670
self.flights[key] = reader.read_all()
4771
print(self.flights[key])
4872

73+
def do_get(self, ticket):
74+
key = ast.literal_eval(ticket.ticket.decode())
75+
if key not in self.flights:
76+
return None
77+
return pyarrow.flight.RecordBatchStream(self.flights[key])
78+
79+
def list_actions(self):
80+
return [
81+
("clear", "Clear the stored flights."),
82+
("shutdown", "Shut down this server."),
83+
]
84+
85+
def do_action(self, action):
86+
if action.type == "clear":
87+
raise NotImplementedError(
88+
"{} is not implemented.".format(action.type))
89+
else:
90+
yield pyarrow.flight.Result(pyarrow.py_buffer(b'Shutdown!'))
91+
# Shut down on background thread to avoid blocking current
92+
# request
93+
threading.Thread(target=self._shutdown).start()
94+
95+
def _shutdown(self):
96+
"""Shut down after a delay."""
97+
time.sleep(2)
98+
self.shutdown()
99+
49100

50101
def main():
51102
server = FlightServer()

0 commit comments

Comments
 (0)