Fix Flight Timestamp type, revert workaround from apache#43537

EnricoMi · Nov 8, 2024 · 13630cd · 13630cd
1 parent b73e6f9
commit 13630cd
Show file tree

Hide file tree

Showing 6 changed files with 8 additions and 66 deletions.
diff --git a/python/pyarrow/_flight.pyx b/python/pyarrow/_flight.pyx
@@ -750,11 +750,8 @@ cdef class FlightEndpoint(_Weakrefable):
 
         if expiration_time is not None:
             if isinstance(expiration_time, lib.TimestampScalar):
-                # Convert into OS-dependent std::chrono::system_clock::time_point from
-                # std::chrono::time_point<std::chrono::system_clock, std::chrono::nanoseconds>
-                # See Timestamp in cpp/src/arrow/flight/types.h
-                self.endpoint.expiration_time = TimePoint_to_system_time(TimePoint_from_ns(
-                    expiration_time.cast(timestamp("ns")).value))
+                self.endpoint.expiration_time = TimePoint_from_ns(
+                    expiration_time.cast(timestamp("ns")).value)
             else:
                 raise TypeError("Argument expiration_time must be a TimestampScalar, "
                                 "not '{}'".format(type(expiration_time)))
@@ -786,11 +783,7 @@ cdef class FlightEndpoint(_Weakrefable):
         cdef:
             int64_t time_since_epoch
         if self.endpoint.expiration_time.has_value():
-            time_since_epoch = TimePoint_to_ns(
-                # Convert from OS-dependent std::chrono::system_clock::time_point into
-                # std::chrono::time_point<std::chrono::system_clock, std::chrono::nanoseconds>
-                # See Timestamp in cpp/src/arrow/flight/types.h
-                TimePoint_from_system_time(self.endpoint.expiration_time.value()))
+            time_since_epoch = TimePoint_to_ns(self.endpoint.expiration_time.value())
             return lib.scalar(time_since_epoch, timestamp("ns", "UTC"))
         return None
 

diff --git a/python/pyarrow/includes/chrono.pxd b/python/pyarrow/includes/chrono.pxd
diff --git a/python/pyarrow/includes/libarrow_flight.pxd b/python/pyarrow/includes/libarrow_flight.pxd
@@ -19,7 +19,7 @@
 
 from pyarrow.includes.common cimport *
 from pyarrow.includes.libarrow cimport *
-from pyarrow.includes.chrono cimport time_point
+from pyarrow.includes.libarrow_python cimport CTimePoint
 
 
 cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
@@ -135,7 +135,7 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
 
         CTicket ticket
         vector[CLocation] locations
-        optional[time_point] expiration_time
+        optional[CTimePoint] expiration_time
         c_string app_metadata
 
         bint operator==(CFlightEndpoint)

diff --git a/python/pyarrow/includes/libarrow_python.pxd b/python/pyarrow/includes/libarrow_python.pxd
@@ -17,7 +17,6 @@
 
 # distutils: language = c++
 
-from pyarrow.includes.chrono cimport time_point
 from pyarrow.includes.common cimport *
 from pyarrow.includes.libarrow cimport *
 
@@ -245,9 +244,6 @@ cdef extern from "arrow/python/api.h" namespace "arrow::py::internal" nogil:
     CTimePoint TimePoint_from_s(double val)
     CTimePoint TimePoint_from_ns(int64_t val)
 
-    CTimePoint TimePoint_from_system_time(time_point val)
-    time_point TimePoint_to_system_time(CTimePoint val)
-
     CResult[c_string] TzinfoToString(PyObject* pytzinfo)
     CResult[PyObject*] StringToTzinfo(c_string)
 

diff --git a/python/pyarrow/src/arrow/python/datetime.h b/python/pyarrow/src/arrow/python/datetime.h
@@ -144,20 +144,6 @@ inline TimePoint TimePoint_from_ns(int64_t val) {
   return TimePoint(TimePoint::duration(val));
 }
 
-ARROW_PYTHON_EXPORT
-// Note: Needed by FlightEndpoint.expiration_time, which is an OS-dependent
-// std::chrono::system_clock::time_point
-inline std::chrono::system_clock::time_point TimePoint_to_system_time(TimePoint val) {
-  return std::chrono::time_point_cast<std::chrono::system_clock::duration>(val);
-}
-
-ARROW_PYTHON_EXPORT
-// Note: Needed by FlightEndpoint.expiration_time, which is an OS-dependent
-// std::chrono::system_clock::time_point
-inline TimePoint TimePoint_from_system_time(std::chrono::system_clock::time_point val) {
-  return std::chrono::time_point_cast<TimePoint::duration>(val);
-}
-
 ARROW_PYTHON_EXPORT
 inline int64_t PyDelta_to_s(PyDateTime_Delta* pytimedelta) {
   return (PyDateTime_DELTA_GET_DAYS(pytimedelta) * 86400LL +

diff --git a/python/pyarrow/tests/test_flight.py b/python/pyarrow/tests/test_flight.py
@@ -1179,19 +1179,9 @@ def test_flight_get_info():
         assert info.endpoints[0].expiration_time is None
         assert info.endpoints[0].app_metadata == b""
         assert info.endpoints[0].locations[0] == flight.Location('grpc://test')
-        # on macOS, system_clock::duration is milliseconds
-        # on Windows, system_clock::duration is 100 nanoseconds
-        # on Linux, system_clock::duration is nanoseconds
-        ts = None
-        if pa._platform.system() == 'Darwin':
-            ts = "2023-04-05T12:34:56.789012000+00:00"
-        elif pa._platform.system() == 'Windows':
-            ts = "2023-04-05T12:34:56.789012300+00:00"
-        elif pa._platform.system() == 'Linux':
-            ts = "2023-04-05T12:34:56.789012345+00:00"
-        if ts is not None:
-            assert info.endpoints[1].expiration_time == \
-                pa.scalar(ts).cast(pa.timestamp("ns", "UTC"))
+        assert info.endpoints[1].expiration_time == \
+            pa.scalar("2023-04-05T12:34:56.789012345+00:00") \
+              .cast(pa.timestamp("ns", "UTC"))
         assert info.endpoints[1].app_metadata == b"endpoint app metadata"
         assert info.endpoints[1].locations[0] == \
             flight.Location.for_grpc_tcp('localhost', 5005)