From 235efa7f316065fc8ae6a0abe5c67fef2b8a6ac5 Mon Sep 17 00:00:00 2001
From: Vincent <1729065730@qq.com>
Date: Thu, 29 Oct 2026 19:14:52 +0800
Subject: [PATCH] GH-35150: [Python] Make Flight value types pickleable

Add __reduce__ for Action, Result, BasicAuth, FlightDescriptor, Ticket,
Location, FlightEndpoint, SchemaResult, and FlightInfo so they round-trip
through pickle (via existing serialize/deserialize, or URI for Location).
---
 python/pyarrow/_flight.pyx          | 28 +++++++++++++
 python/pyarrow/tests/test_flight.py | 62 +++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+)

diff --git a/python/pyarrow/_flight.pyx b/python/pyarrow/_flight.pyx
index 2ab0ab618aec..b9a54483edd8 100644
--- a/python/pyarrow/_flight.pyx
+++ b/python/pyarrow/_flight.pyx
@@ -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
 
@@ -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())
 
@@ -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())
 
@@ -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
 
@@ -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
 
@@ -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."""
@@ -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} "
@@ -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())
 
@@ -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())
 
diff --git a/python/pyarrow/tests/test_flight.py b/python/pyarrow/tests/test_flight.py
index e1c404d81e72..8a3a9524f0c5 100644
--- a/python/pyarrow/tests/test_flight.py
+++ b/python/pyarrow/tests/test_flight.py
@@ -22,6 +22,7 @@
 import itertools
 import json
 import os
+import pickle
 import pathlib
 import signal
 import struct
@@ -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, \