ARROW-139 Fully support ObjectId extension type (#129)

Author: blink1073

bindings/python/pymongoarrow/lib.pyx

@@ -365,7 +365,7 @@ cdef class ObjectIdBuilder(_ArrayBuilderBase):
         cdef shared_ptr[CArray] out
         with nogil:
             self.builder.get().Finish(&out)
-        return pyarrow_wrap_array(out)
+        return pyarrow_wrap_array(out).cast(ObjectIdType())
 
     cdef shared_ptr[CFixedSizeBinaryBuilder] unwrap(self):
         return self.builder

bindings/python/pymongoarrow/pandas_types.py

@@ -21,7 +21,7 @@
 import numpy as np
 import pandas as pd
 import pyarrow as pa
-from bson import Binary
+from bson import Binary, ObjectId
 from pandas.api.extensions import (
     ExtensionArray,
     ExtensionDtype,
@@ -34,16 +34,9 @@ class PandasBSONDtype(ExtensionDtype):
 
     na_value = np.nan
 
-    def __init__(self, subtype):
-        self._subtype = subtype
-
-    @property
-    def subtype(self) -> int:
-        return self._subtype
-
     @property
     def name(self) -> str:
-        return f"bson_{self.type.__name__}[{self.subtype}]"
+        return f"bson_{self.__class__.__name__}"
 
     def __from_arrow__(self, array: Union[pa.Array, pa.ChunkedArray]) -> ExtensionArray:
 
@@ -65,7 +58,7 @@ def __from_arrow__(self, array: Union[pa.Array, pa.ChunkedArray]) -> ExtensionAr
                 if not pd.isna(val) and not isinstance(val, typ):
                     val = typ(val)
                 vals.append(val)
-            arr = np.array(vals)
+            arr = np.array(vals, dtype=object)
             # using _from_sequence to ensure None is converted to NA
             to_append = arr_type._from_sequence(arr, dtype=dtype)
             results.append(to_append)
@@ -75,13 +68,25 @@ def __from_arrow__(self, array: Union[pa.Array, pa.ChunkedArray]) -> ExtensionAr
         else:
             return arr_type(np.array([], dtype="object"))
 
+    @classmethod
+    def construct_from_string(cls, string):
+        if not isinstance(string, str):
+            raise TypeError(f"'construct_from_string' expects a string, got {type(string)}")
+        default = cls()
+        if string != default.name:
+            raise TypeError(f"Cannot construct a '{cls.__name__}' from '{string}'")
+        return default
+
 
 class PandasBSONExtensionArray(ExtensionArray):
     """The base class for Pandas BSON extension arrays."""
 
+    _default_dtype = None
+
     def __init__(self, values, dtype, copy=False) -> None:
         if not isinstance(values, np.ndarray):
             raise TypeError("Need to pass a numpy array as values")
+        dtype = dtype or self._default_dtype
         if dtype is None:
             raise ValueError("dtype must be a valid data type")
         for val in values:
@@ -132,7 +137,7 @@ def __len__(self) -> int:
     def isna(self):
         return np.array(
             [
-                x is not None and not isinstance(x, self.dtype.type) and np.isnan(x)
+                x is None or (x is not None and not isinstance(x, self.dtype.type) and np.isnan(x))
                 for x in self.data
             ],
             dtype=bool,
@@ -181,14 +186,25 @@ def _concat_same_type(cls, to_concat):
 
 
 @register_extension_dtype
-class PandasBSONBinary(PandasBSONDtype):
+class PandasBinary(PandasBSONDtype):
     """A pandas extension type for BSON Binary data type."""
 
     type = Binary
 
+    def __init__(self, subtype):
+        self._subtype = subtype
+
+    @property
+    def subtype(self) -> int:
+        return self._subtype
+
+    @property
+    def name(self) -> str:
+        return f"bson_{self.type.__name__}[{self.subtype}]"
+
     @classmethod
-    def construct_array_type(cls) -> Type["PandasBSONArray"]:
-        return PandasBSONArray
+    def construct_array_type(cls) -> Type["PandasBinaryArray"]:
+        return PandasBinaryArray
 
     @classmethod
     def construct_from_string(cls, string):
@@ -202,10 +218,34 @@ def construct_from_string(cls, string):
             raise TypeError(f"Cannot construct a '{cls.__name__}' from '{string}'")
 
 
-class PandasBSONArray(PandasBSONExtensionArray):
+class PandasBinaryArray(PandasBSONExtensionArray):
     """A pandas extension type for BSON Binary data arrays."""
 
     def __arrow_array__(self, type=None):
         from pymongoarrow.types import BinaryType
 
         return pa.array(self.data, type=BinaryType(self.dtype.subtype))
+
+
+@register_extension_dtype
+class PandasObjectId(PandasBSONDtype):
+    """A pandas extension type for BSON ObjectId data type."""
+
+    type = ObjectId
+
+    @classmethod
+    def construct_array_type(cls) -> Type["PandasObjectIdArray"]:
+        return PandasObjectIdArray
+
+
+class PandasObjectIdArray(PandasBSONExtensionArray):
+    """A pandas extension type for BSON Binary data arrays."""
+
+    @property
+    def _default_dtype(self):
+        return PandasObjectId()
+
+    def __arrow_array__(self, type=None):
+        from pymongoarrow.types import ObjectIdType
+
+        return pa.array(self.data, type=ObjectIdType())

bindings/python/pymongoarrow/types.py

@@ -31,7 +31,7 @@
     struct,
     timestamp,
 )
-from pymongoarrow.pandas_types import PandasBSONBinary
+from pymongoarrow.pandas_types import PandasBinary, PandasObjectId
 
 
 class _BsonArrowTypes(enum.Enum):
@@ -55,6 +55,8 @@ class _BsonArrowTypes(enum.Enum):
 
 class ObjectIdScalar(ExtensionScalar):
     def as_py(self):
+        if self.value is None:
+            return None
         return ObjectId(self.value.as_py())
 
 
@@ -70,6 +72,9 @@ def __reduce__(self):
     def __arrow_ext_scalar_class__(self):
         return ObjectIdScalar
 
+    def to_pandas_dtype(self):
+        return PandasObjectId()
+
 
 class Decimal128Scalar(ExtensionScalar):
     def as_py(self):
@@ -115,7 +120,7 @@ def __arrow_ext_scalar_class__(self):
         return BinaryScalar
 
     def to_pandas_dtype(self):
-        return PandasBSONBinary(self.subtype)
+        return PandasBinary(self.subtype)
 
 
 # Internal Type Handling.

bindings/python/test/pandas_types/test_binary.py

@@ -15,7 +15,7 @@
 import pytest
 from bson import Binary
 from pandas.tests.extension import base
-from pymongoarrow.pandas_types import PandasBSONArray, PandasBSONBinary
+from pymongoarrow.pandas_types import PandasBinary, PandasBinaryArray
 
 try:
     base.BaseIndexTests
@@ -30,7 +30,7 @@ def make_datum():
 
 @pytest.fixture
 def dtype():
-    return PandasBSONBinary(10)
+    return PandasBinary(10)
 
 
 def make_data():
@@ -45,24 +45,24 @@ def make_data():
 
 @pytest.fixture
 def data(dtype):
-    return PandasBSONArray(np.array(make_data(), dtype=object), dtype=dtype)
+    return PandasBinaryArray(np.array(make_data(), dtype=object), dtype=dtype)
 
 
 @pytest.fixture
 def data_missing(dtype):
-    return PandasBSONArray(np.array([np.nan, make_datum()], dtype=object), dtype=dtype)
+    return PandasBinaryArray(np.array([np.nan, make_datum()], dtype=object), dtype=dtype)
 
 
 @pytest.fixture
 def data_for_sorting(dtype):
-    return PandasBSONArray(
+    return PandasBinaryArray(
         np.array([make_datum(), make_datum(), make_datum()], dtype=object), dtype=dtype
     )
 
 
 @pytest.fixture
 def data_missing_for_sorting(dtype):
-    return PandasBSONArray(
+    return PandasBinaryArray(
         np.array([make_datum(), np.nan, make_datum()], dtype=object), dtype=dtype
     )
 

bindings/python/test/pandas_types/test_objectid.py

@@ -0,0 +1,100 @@
+# Copyright 2023-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import numpy as np
+import pytest
+from bson import ObjectId
+from pandas.tests.extension import base
+from pymongoarrow.pandas_types import PandasObjectId, PandasObjectIdArray
+
+try:
+    base.BaseIndexTests
+except AttributeError:
+    pytest.skip("Not available", allow_module_level=True)
+
+
+def make_datum():
+    return ObjectId()
+
+
+@pytest.fixture
+def dtype():
+    return PandasObjectId()
+
+
+def make_data():
+    return (
+        [make_datum() for _ in range(8)]
+        + [np.nan]
+        + [make_datum() for _ in range(88)]
+        + [np.nan]
+        + [make_datum(), make_datum()]
+    )
+
+
+@pytest.fixture
+def data(dtype):
+    return PandasObjectIdArray(np.array(make_data(), dtype=object), dtype=dtype)
+
+
+@pytest.fixture
+def data_missing(dtype):
+    return PandasObjectIdArray(np.array([np.nan, make_datum()], dtype=object), dtype=dtype)
+
+
+@pytest.fixture
+def data_for_sorting(dtype):
+    return PandasObjectIdArray(
+        np.array([make_datum(), make_datum(), make_datum()], dtype=object), dtype=dtype
+    )
+
+
+@pytest.fixture
+def data_missing_for_sorting(dtype):
+    return PandasObjectIdArray(
+        np.array([make_datum(), np.nan, make_datum()], dtype=object), dtype=dtype
+    )
+
+
+class TestDtype(base.BaseDtypeTests):
+    def test_is_not_string_type(self, data):
+        # Override to not return a value, which raises a warning.
+        super().test_is_not_string_type(data)
+
+    def test_is_not_object_type(self, data):
+        # Override to not return a value, which raises a warning.
+        super().test_is_not_object_type(data)
+
+
+class TestInterface(base.BaseInterfaceTests):
+    pass
+
+
+class TestConstructors(base.BaseConstructorsTests):
+    pass
+
+
+class TestGetitem(base.BaseGetitemTests):
+    pass
+
+
+class TestSetitem(base.BaseSetitemTests):
+    pass
+
+
+class TestIndex(base.BaseIndexTests):
+    pass
+
+
+class TestMissing(base.BaseMissingTests):
+    pass