ARROW-188 Support PyArrow 14 (#178)

* ARROW-188 Support PyArrow 14

* add setuptools

* fix py312 compat

* fix version spec

* address warning

* cleanup

Author: blink1073

.github/workflows/release-python.yml

@@ -32,7 +32,7 @@ jobs:
         - [ubuntu-20.04, manylinux_aarch64]
         - [macos-11, macosx_*]
         - [windows-2019, win_amd64]
-        python: ["cp38", "cp39", "cp310", "cp311"]
+        python: ["cp38", "cp39", "cp310", "cp311", "cp312"]
 
     steps:
       - name: Checkout pymongoarrow

.github/workflows/test-python.yml

@@ -33,7 +33,7 @@ jobs:
     strategy:
       matrix:
         os: ["ubuntu-latest", "macos-latest", "windows-latest"]
-        python-version: [3.8, 3.9, "3.10", "3.11"]
+        python-version: [3.8, 3.9, "3.10", "3.11", "3.12"]
       fail-fast: false
     name: CPython ${{ matrix.python-version }}-${{ matrix.os }}
     steps:
@@ -72,6 +72,7 @@ jobs:
       - name: Install libbson
         run: |
           pip install packaging  # needed for mongo-c-driver-1.24.4/build/calc_release_version.py
+          pip install setuptools  # needed for use of distutils
           ./build-libbson.sh
       - name: Install Python dependencies
         run: |

bindings/python/pymongoarrow/__init__.py

@@ -17,18 +17,10 @@
 
 # We must import pyarrow before attempting to load the Cython module.
 import pyarrow as pa  # noqa: F401
+from packaging.version import parse as _parse_version
 
 from pymongoarrow.version import _MIN_LIBBSON_VERSION, __version__  # noqa: F401
 
-try:
-    from pkg_resources import parse_version as _parse_version
-except ImportError:
-    from distutils.version import LooseVersion as _LooseVersion
-
-    def _parse_version(version):
-        return _LooseVersion(version)
-
-
 try:
     from pymongoarrow.lib import libbson_version
 except ImportError:

bindings/python/pymongoarrow/types.py

@@ -21,7 +21,7 @@
 from pyarrow import DataType as _ArrowDataType
 from pyarrow import (
     ExtensionScalar,
-    PyExtensionType,
+    ExtensionType,
     binary,
     bool_,
     float64,
@@ -71,11 +71,11 @@ class ObjectIdScalar(BSONExtensionScalar):
     _bson_class = ObjectId
 
 
-class ObjectIdType(PyExtensionType):
+class ObjectIdType(ExtensionType):
     _type_marker = _BsonArrowTypes.objectid
 
     def __init__(self):
-        super().__init__(binary(12))
+        super().__init__(binary(12), "pymongoarrow.objectid")
 
     def __reduce__(self):
         return ObjectIdType, ()
@@ -86,6 +86,13 @@ def __arrow_ext_scalar_class__(self):
     def to_pandas_dtype(self):
         return PandasObjectId()
 
+    def __arrow_ext_serialize__(self):
+        return b""
+
+    @classmethod
+    def __arrow_ext_deserialize__(self, storage_type, serialized):
+        return ObjectIdType()
+
 
 class Decimal128Scalar(ExtensionScalar):
     def as_py(self):
@@ -94,11 +101,11 @@ def as_py(self):
         return Decimal128.from_bid(self.value.as_py())
 
 
-class Decimal128Type(PyExtensionType):
+class Decimal128Type(ExtensionType):
     _type_marker = _BsonArrowTypes.decimal128
 
     def __init__(self):
-        super().__init__(binary(16))
+        super().__init__(binary(16), "pymongoarrow.decimal128")
 
     def __reduce__(self):
         return Decimal128Type, ()
@@ -109,6 +116,13 @@ def __arrow_ext_scalar_class__(self):
     def to_pandas_dtype(self):
         return PandasDecimal128()
 
+    def __arrow_ext_serialize__(self):
+        return b""
+
+    @classmethod
+    def __arrow_ext_deserialize__(self, storage_type, serialized):
+        return Decimal128Type()
+
 
 class BinaryScalar(ExtensionScalar):
     def as_py(self):
@@ -118,12 +132,12 @@ def as_py(self):
         return Binary(self.value.as_py(), self.type.subtype)
 
 
-class BinaryType(PyExtensionType):
+class BinaryType(ExtensionType):
     _type_marker = _BsonArrowTypes.binary
 
     def __init__(self, subtype):
         self._subtype = subtype
-        super().__init__(binary())
+        super().__init__(binary(), "pymongoarrow.binary")
 
     @property
     def subtype(self):
@@ -138,16 +152,26 @@ def __arrow_ext_scalar_class__(self):
     def to_pandas_dtype(self):
         return PandasBinary(self.subtype)
 
+    def __arrow_ext_serialize__(self):
+        return f"subtype={self.subtype}".encode()
+
+    @classmethod
+    def __arrow_ext_deserialize__(cls, storage_type, serialized):
+        serialized = serialized.decode()
+        assert serialized.startswith("subtype=")  # noqa: S101
+        subtype = int(serialized.split("=")[1])
+        return BinaryType(subtype)
+
 
 class CodeScalar(BSONExtensionScalar):
     _bson_class = Code
 
 
-class CodeType(PyExtensionType):
+class CodeType(ExtensionType):
     _type_marker = _BsonArrowTypes.code
 
     def __init__(self):
-        super().__init__(string())
+        super().__init__(string(), "pymongoarrow.code")
 
     def __reduce__(self):
         return CodeType, ()
@@ -158,6 +182,18 @@ def __arrow_ext_scalar_class__(self):
     def to_pandas_dtype(self):
         return PandasCode()
 
+    def __arrow_ext_serialize__(self):
+        return b""
+
+    @classmethod
+    def __arrow_ext_deserialize__(self, storage_type, serialized):
+        return CodeType()
+
+
+# Register all of the extension types.
+for dtype in [ObjectIdType, CodeType, Decimal128Type]:
+    pa.register_extension_type(dtype())
+pa.register_extension_type(BinaryType(0))
 
 # Internal Type Handling.
 

bindings/python/pyproject.toml

@@ -4,7 +4,7 @@ requires = [
     "wheel>=0.37",
     "cython>=0.29",
     # Must be kept in sync with "project.dependencies" below.
-    "pyarrow>=13.0,<13.1.0",
+    "pyarrow>=14.0,<14.1.0",
 ]
 
 [project]
@@ -33,9 +33,10 @@ classifiers = [
 requires-python = ">=3.8"
 dependencies = [
     # Must be kept in sync with "build_sytem.requires" above.
-    "pyarrow >=13.0,<13.1",
+    "pyarrow >=14.0,<14.1",
     "pymongo >=4.4,<5",
     "pandas >=1.3.5,<3",
+    "packaging >=23.2,<24"
 ]
 dynamic = ["version"]
 
@@ -100,7 +101,9 @@ norecursedirs = ["test/*"]
 faulthandler_timeout = 1500
 xfail_strict = true
 filterwarnings = [
-  "error"
+  "error",
+  # https://github.com/dateutil/dateutil/issues/1314
+  "module:datetime.datetime.utc:DeprecationWarning",
 ]
 
 [tool.ruff]

bindings/python/test/test_arrow.py

@@ -327,18 +327,14 @@ def test_write_batching(self, mock):
         self.assertEqual(mock.call_count, 2)
 
     def _create_nested_data(self, nested_elem=None):
-        schema = {k.__name__: v(True) for k, v in _TYPE_NORMALIZER_FACTORY.items()}
+        schema = {k.__name__: v(0) for k, v in _TYPE_NORMALIZER_FACTORY.items()}
         if nested_elem:
             schem_ent, nested_elem = nested_elem
             schema["list"] = list_(schem_ent)
 
         # PyArrow does not support from_pydict with nested extension types.
         schema["nested"] = struct(
-            [
-                field(a, b)
-                for (a, b) in list(schema.items())
-                if not isinstance(b, pa.PyExtensionType)
-            ]
+            [field(a, b) for (a, b) in list(schema.items()) if not isinstance(b, pa.ExtensionType)]
         )
         raw_data = {
             "str": [None] + [str(i) for i in range(2)],