ARROW-54 Make ObjectId an Extension Type (#50)

Author: blink1073

bindings/python/docs/source/supported_types.rst

@@ -18,7 +18,7 @@ Support for additional types will be added in subsequent releases.
    * - String
      - :class:`py.str`, an instance of :class:`pyarrow.string`
    * - ObjectId
-     - :class:`py.bytes`, :class:`bson.ObjectId`, an instance of :class:`pyarrow.FixedSizeBinaryScalar`
+     - :class:`py.bytes`, :class:`bson.ObjectId`, an instance of :class:`pymongoarrow.types.ObjectIdType`, an instance of :class:`pyarrow.FixedSizeBinaryScalar`
    * - 64-bit binary floating point
      - :class:`py.float`, an instance of :meth:`pyarrow.float64`
    * - 32-bit integer

bindings/python/pymongoarrow/types.py

@@ -16,7 +16,8 @@
 
 from bson import Int64, ObjectId
 
-from pyarrow import timestamp, float64, int64, int32, string
+from pyarrow import timestamp, binary, float64, int64, int32, string
+from pyarrow import PyExtensionType
 from pyarrow import DataType as _ArrowDataType
 import pyarrow.types as _atypes
 
@@ -30,20 +31,37 @@ class _BsonArrowTypes(enum.Enum):
     string = 6
 
 
+# Custom Extension Types.
+# See https://arrow.apache.org/docs/python/extending_types.html#defining-extension-types-user-defined-types
+# for details.
+
+class ObjectIdType(PyExtensionType):
+    _type_marker = _BsonArrowTypes.objectid
+
+    def __init__(self):
+        super().__init__(binary(12))
+
+    def __reduce__(self):
+        return ObjectIdType, ()
+
+
+# Internal Type Handling.
+
+def _is_objectid(obj):
+    type_marker = getattr(obj, '_type_marker', '')
+    return type_marker == ObjectIdType._type_marker
+
+
 _TYPE_NORMALIZER_FACTORY = {
     Int64: lambda _: int64(),
     float: lambda _: float64(),
     int: lambda _: int64(),
     datetime: lambda _: timestamp('ms'),     # TODO: add tzinfo support
-    ObjectId: lambda _: ObjectId,
+    ObjectId: lambda _: ObjectIdType(),
     str: lambda: string(),
 }
 
 
-def _is_objectid(obj):
-    return obj == ObjectId
-
-
 _TYPE_CHECKER_TO_INTERNAL_TYPE = {
     _atypes.is_int32: _BsonArrowTypes.int32,
     _atypes.is_int64: _BsonArrowTypes.int64,
@@ -74,15 +92,9 @@ def _get_internal_typemap(typemap):
     internal_typemap = {}
     for fname, ftype in typemap.items():
         for checker, internal_id in _TYPE_CHECKER_TO_INTERNAL_TYPE.items():
-            # Catch error where the pyarrow checkers are looking for an `id`
-            # attribute that might not exist on non-pyarrow types
-            # (like ObjectId).  For example, `is_int32()` checks for
-            # `t.id == lib.Type_INT32`.
-            try:
-                if checker(ftype):
-                    internal_typemap[fname] = internal_id
-            except AttributeError:
-                pass
+            if checker(ftype):
+                internal_typemap[fname] = internal_id
+
         if fname not in internal_typemap:
             raise ValueError('Unsupported data type in schema for ' +
                 f'field "{fname}" of type "{ftype}"')

bindings/python/test/test_bson.py

@@ -15,11 +15,11 @@
 
 from bson import encode, InvalidBSON
 
-import pyarrow
+import pyarrow as pa
 from pymongoarrow.context import PyMongoArrowContext
 from pymongoarrow.lib import process_bson_stream
 from pymongoarrow.schema import Schema
-from pymongoarrow.types import int32, int64, string, ObjectId
+from pymongoarrow.types import int32, int64, string, ObjectId, ObjectIdType
 
 
 class TestBsonToArrowConversionBase(TestCase):
@@ -105,7 +105,7 @@ class TestUnsupportedDataType(TestBsonToArrowConversionBase):
     def test_simple(self):
         schema = Schema({'_id': ObjectId,
                          'data': int64(),
-                         'fake':  pyarrow.float16() })
+                         'fake':  pa.float16() })
         msg =  ("Unsupported data type in schema for field " +
                 '"fake" of type "halffloat"')
         with self.assertRaisesRegex(ValueError, msg):
@@ -115,7 +115,7 @@ def test_simple(self):
 class TestNonAsciiFieldName(TestBsonToArrowConversionBase):
 
     def setUp(self):
-        self.schema = Schema({'_id': ObjectId,
+        self.schema = Schema({'_id': ObjectIdType(),
                               'dätá': int64()})
         self.context = PyMongoArrowContext.from_schema(
             self.schema)
@@ -132,3 +132,25 @@ def test_simple(self):
         }
 
         self._run_test(docs, as_dict)
+
+
+class TestSerializeExtensions(TestCase):
+    # Follows example in
+    # https://arrow.apache.org/docs/python/extending_types.html#defining-extension-types-user-defined-types
+
+    def serialize_array(self, arr):
+        batch = pa.RecordBatch.from_arrays([arr], ["ext"])
+        sink = pa.BufferOutputStream()
+        with pa.RecordBatchStreamWriter(sink, batch.schema) as writer:
+            writer.write_batch(batch)
+        buf = sink.getvalue()
+        with pa.ipc.open_stream(buf) as reader:
+            result = reader.read_all()
+        return result.column('ext')
+
+    def test_object_id_type(self):
+        oids = [ObjectId().binary for _ in range(4)]
+        storage_array = pa.array(oids, pa.binary(12))
+        arr = pa.ExtensionArray.from_storage(ObjectIdType(), storage_array)
+        result = self.serialize_array(arr)
+        assert result.type._type_marker == ObjectIdType._type_marker