ARROW-12050: [C++][Python][FlightRPC] Make Flight operations interruptible in Python

[lidavidm]

This uses a stop token to let interactive users interrupt a long-running Flight operation. It's not perfect: the operation won't be cancelled until the server delivers a message, so this doesn't protect against very slow servers. (In that case, we'd need some way for the stop source to call TryCancel() on the gRPC RPC object, which would be tricky.) But so long as the server is being responsive, this means Ctrl-C should do what people expect in Python.

[github-actions]

https://issues.apache.org/jira/browse/ARROW-12050

[cyb70289]

Not quite sure of the issue addressed.
Is it possible to reproduce the issue on my local host and verify the fix?

[lidavidm]

It solves a problem for Python users: if you Ctrl-C in something like read_all(), nothing will happen because control is in C++/Flight code, and then you're stuck waiting for the call to complete.

If you apply this diff to python/examples/flight/server.py:

diff --git a/python/examples/flight/server.py b/python/examples/flight/server.py
index 7a6b6697e..bc1df21d1 100644
--- a/python/examples/flight/server.py
+++ b/python/examples/flight/server.py
@@ -73,6 +73,7 @@ class FlightServer(pyarrow.flight.FlightServerBase):
             yield self._make_flight_info(key, descriptor, table)
 
     def get_flight_info(self, context, descriptor):
+        return self._make_flight_info("", descriptor, pyarrow.schema([]).empty_table())
         key = FlightServer.descriptor_to_key(descriptor)
         if key in self.flights:
             table = self.flights[key]
@@ -86,10 +87,10 @@ class FlightServer(pyarrow.flight.FlightServerBase):
         print(self.flights[key])
 
     def do_get(self, context, ticket):
-        key = ast.literal_eval(ticket.ticket.decode())
-        if key not in self.flights:
-            return None
-        return pyarrow.flight.RecordBatchStream(self.flights[key])
+        import itertools
+        schema = pyarrow.schema([])
+        rb = pyarrow.RecordBatch.from_arrays([], schema=schema)
+        return pyarrow.flight.GeneratorStream(schema, itertools.repeat(rb))
 
     def list_actions(self, context):
         return [

Then you can start the server:

arrow/python$ env PYTHONPATH=$(pwd) python examples/flight/server.py --port 2000 &
arrow/python$ env PYTHONPATH=(pwd) python examples/flight/client.py get localhost:2000 -c foo

Without this patch, if you Ctrl-C, your client will still be stuck forever. With this patch you'll interrupt the call as expected.

[pitrou]

This is really nice!

[pitrou]

Is there a reason for the Table version to be virtual? It seems it can simply be implemented in the base class...?

Of course, it doesn't matter much as long as we have only one implementation of this.

[lidavidm]

I made this non-virtual.

[pitrou]

Am I reading this wrong, or is this twice the same condition?

[lidavidm]

I meant for the second to be KeyboardInterrupt.

[pitrou]

Would this raise an exception if the context is cancelled, or just succeed immediately?

[lidavidm]

It should raise an exception, but either way, the client is gone at this point so there's not much the server can do.

[pitrou]

Ok, I was just wondering if the server could end up spending more time than necessary here :-)

[pitrou]

As a sidenote, I think we'll need to allow adding a callback on a StopToken, so that arbitrary callables such as gRPC's TryCancel can be called. I opened a JIRA for it: https://issues.apache.org/jira/browse/ARROW-12938