ARROW-12 Support ObjectId datatype (#44)

Author: blink1073

bindings/python/docs/source/changelog.rst

@@ -1,6 +1,10 @@
 Changelog
 =========
 
+Changes in Version 0.3.0
+------------------------
+- Support for `ObjectId` `bson` type.
+
 Changes in Version 0.2.0
 ------------------------
 

bindings/python/docs/source/supported_types.rst

@@ -15,6 +15,8 @@ Support for additional types will be added in subsequent releases.
 
    * - BSON Type
      - Type Identifiers
+   * - ObjectId
+     - :class:`py.bytes`, :class:`bson.ObjectId`, an instance of :class:`pyarrow.FixedSizeBinaryScalar`
    * - 64-bit binary floating point
      - :class:`py.float`, an instance of :meth:`pyarrow.float64`
    * - 32-bit integer
@@ -27,6 +29,10 @@ Support for additional types will be added in subsequent releases.
 Type identifiers can be used to specify that a field is of a certain type
 during :class:`pymongoarrow.api.Schema` declaration. For example, if your data
 has fields 'f1' and 'f2' bearing types 32-bit integer and UTC datetime
-respectively, your schema can be defined as::
+respectively, and '_id' that is an `ObjectId`, your schema can be defined as::
 
-  schema = Schema({'f1': pyarrow.int32(), 'f2': pyarrow.timestamp('ms')})
+  schema = Schema({
+    '_id': ObjectId,
+    'f1': pyarrow.int32(),
+    'f2': pyarrow.timestamp('ms')
+})

bindings/python/pymongoarrow/bson.pyi

@@ -36,6 +36,7 @@ def process_bson_stream(bson_stream, context):
     t_int64 = _BsonArrowTypes.int64
     t_double = _BsonArrowTypes.double
     t_datetime = _BsonArrowTypes.datetime
+    t_oid = _BsonArrowTypes.objectid
     builder_map = context.builder_map
 
     # initialize count to current length of builders
@@ -69,6 +70,11 @@ def process_bson_stream(bson_stream, context):
                             builder.append(bson_iter_as_int64(&doc_iter))
                         else:
                             builder.append_null()
+                    elif ftype == t_oid:
+                        if value_t == BSON_TYPE_OID:
+                            builder.append(<bytes>(<uint8_t*>bson_iter_oid(&doc_iter))[:12])
+                        else:
+                            builder.append_null()
                     elif ftype == t_double:
                         if (value_t == BSON_TYPE_DOUBLE or
                                 value_t == BSON_TYPE_BOOL or

bindings/python/pymongoarrow/builders.pyi

@@ -19,6 +19,40 @@ cdef class _ArrayBuilderBase:
             self.append(value)
 
 
+cdef class ObjectIdBuilder(_ArrayBuilderBase):
+    type_marker = _BsonArrowTypes.objectid
+    cdef:
+        shared_ptr[CFixedSizeBinaryBuilder] builder
+
+    def __cinit__(self, MemoryPool memory_pool=None):
+        cdef shared_ptr[CDataType] dtype = fixed_size_binary(12)
+        cdef CMemoryPool* pool = maybe_unbox_memory_pool(memory_pool)
+        self.builder.reset(new CFixedSizeBinaryBuilder(dtype, pool))
+
+    cpdef append_null(self):
+        self.builder.get().AppendNull()
+
+    def __len__(self):
+        return self.builder.get().length()
+
+    cpdef append(self, value):
+        if value is None or value is np.nan:
+            self.builder.get().AppendNull()
+        elif isinstance(value, bytes):
+            self.builder.get().Append(value)
+        else:
+            raise TypeError('ObjectIdBuilder only accepts bytes objects')
+
+    cpdef finish(self):
+        cdef shared_ptr[CArray] out
+        with nogil:
+            self.builder.get().Finish(&out)
+        return pyarrow_wrap_array(out)
+
+    cdef shared_ptr[CFixedSizeBinaryBuilder] unwrap(self):
+        return self.builder
+
+
 cdef class Int32Builder(_ArrayBuilderBase):
     type_marker = _BsonArrowTypes.int32
     cdef:

bindings/python/pymongoarrow/context.py

@@ -14,15 +14,16 @@
 from bson.codec_options import DEFAULT_CODEC_OPTIONS
 
 from pyarrow import timestamp, Table
-from pymongoarrow.lib import Int32Builder, Int64Builder, DoubleBuilder, DatetimeBuilder
+from pymongoarrow.lib import Int32Builder, Int64Builder, DoubleBuilder, DatetimeBuilder, ObjectIdBuilder
 from pymongoarrow.types import _get_internal_typemap, _BsonArrowTypes
 
 
 _TYPE_TO_BUILDER_CLS = {
     _BsonArrowTypes.int32: Int32Builder,
     _BsonArrowTypes.int64: Int64Builder,
     _BsonArrowTypes.double: DoubleBuilder,
-    _BsonArrowTypes.datetime: DatetimeBuilder
+    _BsonArrowTypes.datetime: DatetimeBuilder,
+    _BsonArrowTypes.objectid: ObjectIdBuilder
 }
 
 

bindings/python/pymongoarrow/lib.pyx

@@ -34,6 +34,7 @@ from libcpp cimport bool as cbool
 from libcpp.map cimport map
 from libcpp.string cimport string
 from pyarrow.lib cimport *
+from pymongoarrow.libarrow cimport *
 from pymongoarrow.libbson cimport *
 from pymongoarrow.types import _BsonArrowTypes
 

bindings/python/pymongoarrow/libarrow.pxd

@@ -0,0 +1,32 @@
+# Copyright 2021-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.
+
+# Cython compiler directives
+# cython: language_level=3
+# distutils: language=c
+from libc.stdint cimport int32_t, uint8_t
+from pyarrow.lib cimport *
+from pyarrow.includes.libarrow cimport (CStatus, CMemoryPool)
+
+
+# libarrow type wrappings
+
+cdef extern from "arrow/builder.h" namespace "arrow" nogil:
+    cdef cppclass CFixedSizeBinaryBuilder" arrow::FixedSizeBinaryBuilder"(CArrayBuilder):
+        CFixedSizeBinaryBuilder(shared_ptr[CDataType], CMemoryPool* pool)
+        CStatus Append(const uint8_t* value)
+
+
+cdef extern from "arrow/type_fwd.h" namespace "arrow" nogil:
+    shared_ptr[CDataType] fixed_size_binary(int32_t byte_width);

bindings/python/pymongoarrow/libbson.pxd

@@ -108,6 +108,8 @@ cdef extern from "<bson/bson.h>":
 
     int64_t bson_iter_date_time(const bson_iter_t *iter)
 
+    const bson_oid_t * bson_iter_oid (const bson_iter_t *iter)
+
     # TODO: add decimal128
 
     double bson_iter_double(const bson_iter_t *iter)

bindings/python/pymongoarrow/types.py

@@ -14,7 +14,7 @@
 from datetime import datetime
 import enum
 
-from bson import Int64
+from bson import Int64, ObjectId
 
 from pyarrow import timestamp, float64, int64, int32
 from pyarrow import DataType as _ArrowDataType
@@ -26,21 +26,28 @@ class _BsonArrowTypes(enum.Enum):
     double = 2
     int32 = 3
     int64 = 4
+    objectid = 5
 
 
 _TYPE_NORMALIZER_FACTORY = {
     Int64: lambda _: int64(),
     float: lambda _: float64(),
     int: lambda _: int64(),
-    datetime: lambda _: timestamp('ms')     # TODO: add tzinfo support
+    datetime: lambda _: timestamp('ms'),     # TODO: add tzinfo support
+    ObjectId: lambda _: ObjectId
 }
 
 
+def _is_objectid(obj):
+    return obj == ObjectId
+
+
 _TYPE_CHECKER_TO_INTERNAL_TYPE = {
     _atypes.is_int32: _BsonArrowTypes.int32,
     _atypes.is_int64: _BsonArrowTypes.int64,
     _atypes.is_float64: _BsonArrowTypes.double,
-    _atypes.is_timestamp: _BsonArrowTypes.datetime
+    _atypes.is_timestamp: _BsonArrowTypes.datetime,
+    _is_objectid: _BsonArrowTypes.objectid
 }
 
 
@@ -64,8 +71,14 @@ def _get_internal_typemap(typemap):
     internal_typemap = {}
     for fname, ftype in typemap.items():
         for checker, internal_id in _TYPE_CHECKER_TO_INTERNAL_TYPE.items():
-            if checker(ftype):
-                internal_typemap[fname] = internal_id
-                continue
+            # 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
     assert len(internal_typemap) == len(typemap)
     return internal_typemap

bindings/python/test/test_bson.py

@@ -18,12 +18,12 @@
 from pymongoarrow.context import PyMongoArrowContext
 from pymongoarrow.lib import process_bson_stream
 from pymongoarrow.schema import Schema
-from pymongoarrow.types import int32, int64
+from pymongoarrow.types import int32, int64, ObjectId
 
 
 class TestBsonToArrowConversionBase(TestCase):
     def setUp(self):
-        self.schema = Schema({'_id': int32(),
+        self.schema = Schema({'_id': ObjectId,
                               'data': int64()})
         self.context = PyMongoArrowContext.from_schema(
             self.schema)
@@ -47,26 +47,29 @@ def _run_test(self, doclist, as_dict):
 
 
 class TestValidBsonToArrowConversion(TestBsonToArrowConversionBase):
+
     def test_simple(self):
-        docs = [{'_id': 1, 'data': 10},
-                {'_id': 2, 'data': 20},
-                {'_id': 3, 'data': 30},
-                {'_id': 4, 'data': 40}]
+        ids = [ObjectId() for i in range(4)]
+        docs = [{'_id': ids[0], 'data': 10},
+                {'_id': ids[1], 'data': 20},
+                {'_id': ids[2], 'data': 30},
+                {'_id': ids[3], 'data': 40}]
         as_dict = {
-            '_id': [1, 2, 3, 4],
+            '_id': [oid.binary for oid in ids] ,
             'data': [10, 20, 30, 40]}
 
         self._run_test(docs, as_dict)
 
     def test_with_nulls(self):
-        docs = [{'_id': 1, 'data': 10},
-                {'_id': 2, 'data': 20},
-                {'_id': 3},
-                {'_id': 4, 'data': 40},
+        ids = [ObjectId() for i in range(4)]
+        docs = [{'_id': ids[0], 'data': 10},
+                {'_id': ids[1], 'data': 20},
+                {'_id': ids[2]},
+                {'_id': ids[3], 'data': 40},
                 {'foo': 1},
                 {}]
         as_dict = {
-            '_id': [1, 2, 3, 4, None, None],
+            '_id': [oid.binary for oid in ids] + [None, None],
             'data': [10, 20, None, 40, None, None]}
 
         self._run_test(docs, as_dict)
@@ -79,12 +82,13 @@ def _generate_payload(doclist):
             doclist)[:-2]
 
     def test_simple(self):
-        docs = [{'_id': 1, 'data': 10},
-                {'_id': 2, 'data': 20},
-                {'_id': 3, 'data': 30},
-                {'_id': 4, 'data': 40}]
+        ids = [ObjectId() for i in range(4)]
+        docs = [{'_id': ids[0], 'data': 10},
+                {'_id': ids[1], 'data': 20},
+                {'_id': ids[2], 'data': 30},
+                {'_id': ids[3], 'data': 40}]
         as_dict = {
-            '_id': [1, 2, 3, 4],
+            '_id': [oid.binary for oid in ids],
             'data': [10, 20, 30, 40]}
 
         with self.assertRaisesRegex(