Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions python/pyarrow/_flight.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -324,6 +324,9 @@ cdef class Action(_Weakrefable):
CAction.Deserialize(tobytes(serialized)))
return action

def __reduce__(self):
return Action.deserialize, (self.serialize(),)

def __eq__(self, Action other):
return self.action == other.action

Expand Down Expand Up @@ -391,6 +394,9 @@ cdef class Result(_Weakrefable):
CFlightResult.Deserialize(tobytes(serialized)))))
return result

def __reduce__(self):
return Result.deserialize, (self.serialize(),)

def __eq__(self, Result other):
return deref(self.result.get()) == deref(other.result.get())

Expand Down Expand Up @@ -437,6 +443,9 @@ cdef class BasicAuth(_Weakrefable):
def serialize(self):
return GetResultValue(self.basic_auth.get().SerializeToString())

def __reduce__(self):
return BasicAuth.deserialize, (self.serialize(),)

def __eq__(self, BasicAuth other):
return deref(self.basic_auth.get()) == deref(other.basic_auth.get())

Expand Down Expand Up @@ -596,6 +605,9 @@ cdef class FlightDescriptor(_Weakrefable):
CFlightDescriptor.Deserialize(tobytes(serialized)))
return descriptor

def __reduce__(self):
return FlightDescriptor.deserialize, (self.serialize(),)

def __eq__(self, FlightDescriptor other):
return self.descriptor == other.descriptor

Expand Down Expand Up @@ -635,6 +647,9 @@ cdef class Ticket(_Weakrefable):
CTicket.Deserialize(tobytes(serialized)))
return ticket

def __reduce__(self):
return Ticket.deserialize, (self.serialize(),)

def __eq__(self, Ticket other):
return self.c_ticket == other.c_ticket

Expand Down Expand Up @@ -665,6 +680,10 @@ cdef class Location(_Weakrefable):
return NotImplemented
return self.location.Equals((<Location> other).location)

def __reduce__(self):
# Location has no protobuf serialize helper; URI is sufficient.
return Location, (self.uri,)

@staticmethod
def for_grpc_tcp(host, port):
"""Create a Location for a TCP-based gRPC service."""
Expand Down Expand Up @@ -829,6 +848,9 @@ cdef class FlightEndpoint(_Weakrefable):
CFlightEndpoint.Deserialize(tobytes(serialized)))
return endpoint

def __reduce__(self):
return FlightEndpoint.deserialize, (self.serialize(),)

def __repr__(self):
return (f"<pyarrow.flight.FlightEndpoint ticket={self.ticket!r} "
f"locations={self.locations!r} "
Expand Down Expand Up @@ -888,6 +910,9 @@ cdef class SchemaResult(_Weakrefable):
CSchemaResult.Deserialize(tobytes(serialized)))))
return result

def __reduce__(self):
return SchemaResult.deserialize, (self.serialize(),)

def __eq__(self, SchemaResult other):
return deref(self.result.get()) == deref(other.result.get())

Expand Down Expand Up @@ -1032,6 +1057,9 @@ cdef class FlightInfo(_Weakrefable):
CFlightInfo.Deserialize(tobytes(serialized))))
return info

def __reduce__(self):
return FlightInfo.deserialize, (self.serialize(),)

def __eq__(self, FlightInfo other):
return deref(self.info.get()) == deref(other.info.get())

Expand Down
62 changes: 62 additions & 0 deletions python/pyarrow/tests/test_flight.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import itertools
import json
import os
import pickle
import pathlib
import signal
import struct
Expand Down Expand Up @@ -2059,6 +2060,67 @@ def test_roundtrip_types():
assert endpoint == flight.FlightEndpoint.deserialize(endpoint.serialize())


def test_pickle_flight_types():
"""Flight value types should round-trip through pickle (GH-35150)."""
action = flight.Action("action1", b"action1-body")
assert action == pickle.loads(pickle.dumps(action))

ticket = flight.Ticket("foo")
assert ticket == pickle.loads(pickle.dumps(ticket))

result = flight.Result(b"result1")
assert result == pickle.loads(pickle.dumps(result))

basic_auth = flight.BasicAuth("username1", "password1")
assert basic_auth == pickle.loads(pickle.dumps(basic_auth))

schema_result = flight.SchemaResult(pa.schema([('a', pa.int32())]))
assert schema_result == pickle.loads(pickle.dumps(schema_result))

desc = flight.FlightDescriptor.for_command("test")
assert desc == pickle.loads(pickle.dumps(desc))

desc = flight.FlightDescriptor.for_path("a", "b", "test.arrow")
assert desc == pickle.loads(pickle.dumps(desc))

location = flight.Location.for_grpc_tcp('localhost', 5005)
assert location == pickle.loads(pickle.dumps(location))

endpoint = flight.FlightEndpoint(
ticket,
['grpc://test', location],
pa.scalar("2023-04-05T12:34:56").cast(pa.timestamp("s")),
b'endpoint app metadata'
)
assert endpoint == pickle.loads(pickle.dumps(endpoint))

info = flight.FlightInfo(
pa.schema([('a', pa.int32())]),
desc,
[
flight.FlightEndpoint(b'', ['grpc://test']),
flight.FlightEndpoint(
b'',
[location],
pa.scalar("2023-04-05T12:34:56.789012345").cast(pa.timestamp("ns")),
b'endpoint app metadata'
),
],
1,
42,
True,
b'test app metadata'
)
info2 = pickle.loads(pickle.dumps(info))
assert info.schema == info2.schema
assert info.descriptor == info2.descriptor
assert info.total_bytes == info2.total_bytes
assert info.total_records == info2.total_records
assert info.ordered == info2.ordered
assert info.app_metadata == info2.app_metadata
assert info.endpoints == info2.endpoints


def test_roundtrip_errors():
"""Ensure that Flight errors propagate from server to client."""
with ErrorFlightServer() as server, \
Expand Down
Loading