diff --git a/py/client/pydeephaven/_arrow_flight_service.py b/py/client/pydeephaven/_arrow_flight_service.py index 0f0834d081d..3ca1e2ca390 100644 --- a/py/client/pydeephaven/_arrow_flight_service.py +++ b/py/client/pydeephaven/_arrow_flight_service.py @@ -31,7 +31,9 @@ def import_table(self, data: pa.Table) -> Table: writer, reader = self._flight_client.do_put( pa.flight.FlightDescriptor.for_path("export", str(ticket)), dh_schema) writer.write_table(data) - writer.close() + # Note that pyarrow's write_table completes the gRPC. If we send another gRPC close + # it is possible that by the time the request arrives at the server that it no longer + # knows what it is for and sends a RST_STREAM causing a failure. _ = reader.read() flight_ticket = self.session.make_ticket(ticket) return Table(self.session, ticket=flight_ticket, size=data.num_rows, schema=dh_schema) @@ -54,7 +56,7 @@ def do_exchange(self): Returns: The corresponding Arrow FlightStreamWriter and FlightStreamReader. - """ + """ try: desc = pa.flight.FlightDescriptor.for_command(b"dphn") options = paflight.FlightCallOptions(headers=self.session.grpc_metadata) @@ -63,4 +65,4 @@ def do_exchange(self): except Exception as e: raise DHError("failed to perform a flight DoExchange on the table.") from e - +