ARROW-144 Support BSON Code Type (#145)

Author: blink1073

bindings/python/pymongoarrow/context.py

@@ -16,6 +16,7 @@
 from pymongoarrow.lib import (
     BinaryBuilder,
     BoolBuilder,
+    CodeBuilder,
     DatetimeBuilder,
     Decimal128Builder,
     DocumentBuilder,
@@ -40,6 +41,7 @@
     _BsonArrowTypes.document: DocumentBuilder,
     _BsonArrowTypes.array: ListBuilder,
     _BsonArrowTypes.binary: BinaryBuilder,
+    _BsonArrowTypes.code: CodeBuilder,
 }
 
 

bindings/python/pymongoarrow/lib.pyx

@@ -33,13 +33,14 @@ from pyarrow.lib import (
 
 from pymongoarrow.errors import InvalidBSON, PyMongoArrowError
 from pymongoarrow.context import PyMongoArrowContext
-from pymongoarrow.types import _BsonArrowTypes, _atypes, ObjectIdType, Decimal128Type as Decimal128Type_, BinaryType
+from pymongoarrow.types import _BsonArrowTypes, _atypes, ObjectIdType, Decimal128Type as Decimal128Type_, BinaryType, CodeType
 
 # Cython imports
 from cpython cimport PyBytes_Size, object
 from cython.operator cimport dereference
 from libcpp cimport bool as cbool
 from libcpp.map cimport map
+from libc.string cimport strlen
 from libcpp.vector cimport vector
 from pyarrow.lib cimport *
 from pymongoarrow.libarrow cimport *
@@ -69,7 +70,8 @@ _builder_type_map = {
     BSON_TYPE_DOCUMENT: DocumentBuilder,
     BSON_TYPE_DECIMAL128: Decimal128Builder,
     BSON_TYPE_ARRAY: ListBuilder,
-    BSON_TYPE_BINARY: BinaryBuilder
+    BSON_TYPE_BINARY: BinaryBuilder,
+    BSON_TYPE_CODE: CodeBuilder,
 }
 
 _field_type_map = {
@@ -80,8 +82,10 @@ _field_type_map = {
     BSON_TYPE_UTF8: string(),
     BSON_TYPE_BOOL: bool_(),
     BSON_TYPE_DECIMAL128: Decimal128Type_(),
+    BSON_TYPE_CODE: CodeType(),
 }
 
+
 cdef extract_field_dtype(bson_iter_t * doc_iter, bson_iter_t * child_iter, bson_type_t value_t, context):
     """Get the appropriate data type for a specific field"""
     cdef const uint8_t *val_buf = NULL
@@ -148,7 +152,7 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
     cdef uint32_t val_buf_len = 0
     cdef bson_decimal128_t dec128
     cdef bson_type_t value_t
-    cdef const char * bson_str
+    cdef const char * bson_str = NULL
     cdef StructType struct_dtype
     cdef const bson_t * doc = NULL
     cdef bson_iter_t doc_iter
@@ -171,6 +175,7 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
     t_array = _BsonArrowTypes.array
     t_binary = _BsonArrowTypes.binary
     t_decimal128 = _BsonArrowTypes.decimal128
+    t_code = _BsonArrowTypes.code
 
     # initialize count to current length of builders
     for _, builder in builder_map.items():
@@ -256,6 +261,12 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
                         builder.append(<bytes>(bson_str)[:str_len])
                     else:
                         builder.append_null()
+                elif ftype == t_code:
+                    if value_t == BSON_TYPE_CODE:
+                        bson_str = bson_iter_code(&doc_iter, &str_len)
+                        builder.append(<bytes>(bson_str)[:str_len])
+                    else:
+                        builder.append_null()
                 elif ftype == t_decimal128:
                     if value_t == BSON_TYPE_DECIMAL128:
                         bson_iter_decimal128(&doc_iter, &dec128)
@@ -359,6 +370,16 @@ cdef class StringBuilder(_ArrayBuilderBase):
         return self.builder
 
 
+cdef class CodeBuilder(StringBuilder):
+    type_marker = _BsonArrowTypes.code
+
+    cpdef finish(self):
+        cdef shared_ptr[CArray] out
+        with nogil:
+            self.builder.get().Finish(&out)
+        return pyarrow_wrap_array(out).cast(CodeType())
+
+
 cdef class ObjectIdBuilder(_ArrayBuilderBase):
     type_marker = _BsonArrowTypes.objectid
     cdef:

bindings/python/pymongoarrow/libbson.pxd

@@ -130,6 +130,9 @@ cdef extern from "<bson/bson.h>":
 
     bint bson_iter_decimal128 (const bson_iter_t *iter, bson_decimal128_t *dec)
 
+    const char * bson_iter_code (const bson_iter_t *iter, # IN
+                                 uint32_t *length)        # OUT
+
     bint  bson_iter_recurse (const bson_iter_t *iter, # IN
                              bson_iter_t *child)  # OUT
 

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, Decimal128, ObjectId
+from bson import Binary, Code, Decimal128, ObjectId
 from pandas.api.extensions import (
     ExtensionArray,
     ExtensionDtype,
@@ -262,7 +262,7 @@ def __arrow_array__(self, type=None):
 
 @register_extension_dtype
 class PandasDecimal128(PandasBSONDtype):
-    """A pandas extension type for BSON ObjectId data type."""
+    """A pandas extension type for BSON Decimal128 data type."""
 
     type = Decimal128
 
@@ -282,3 +282,33 @@ def __arrow_array__(self, type=None):
         from pymongoarrow.types import Decimal128Type
 
         return pa.array(self.data, type=Decimal128Type())
+
+
+@register_extension_dtype
+class PandasCode(PandasBSONDtype):
+    """A pandas extension type for BSON Code data type."""
+
+    type = Code
+
+    @classmethod
+    def construct_array_type(cls) -> Type["PandasCodeArray"]:
+        return PandasCodeArray
+
+
+class PandasCodeArray(PandasBSONExtensionArray):
+    """A pandas extension type for BSON Code data arrays."""
+
+    @property
+    def _default_dtype(self):
+        return PandasCode()
+
+    def __eq__(self, other):
+        # Code types do not support element-wise comparison.
+        if isinstance(other, Code):
+            other = np.array(other, dtype=object)
+        return super().__eq__(other)
+
+    def __arrow_array__(self, type=None):
+        from pymongoarrow.types import CodeType
+
+        return pa.array(self.data, type=CodeType())

bindings/python/pymongoarrow/types.py

@@ -17,7 +17,7 @@
 import numpy as np
 import pyarrow as pa
 import pyarrow.types as _atypes
-from bson import Binary, Decimal128, Int64, ObjectId
+from bson import Binary, Code, Decimal128, Int64, ObjectId
 from pyarrow import DataType as _ArrowDataType
 from pyarrow import (
     ExtensionScalar,
@@ -31,7 +31,12 @@
     struct,
     timestamp,
 )
-from pymongoarrow.pandas_types import PandasBinary, PandasDecimal128, PandasObjectId
+from pymongoarrow.pandas_types import (
+    PandasBinary,
+    PandasCode,
+    PandasDecimal128,
+    PandasObjectId,
+)
 
 
 class _BsonArrowTypes(enum.Enum):
@@ -46,18 +51,23 @@ class _BsonArrowTypes(enum.Enum):
     document = 9
     array = 10
     binary = 11
+    code = 12
 
 
 # Custom Extension Types.
 # See https://arrow.apache.org/docs/python/extending_types.html#defining-extension-types-user-defined-types
 # for details.
 
 
-class ObjectIdScalar(ExtensionScalar):
+class BSONExtensionScalar(ExtensionScalar):
     def as_py(self):
         if self.value is None:
             return None
-        return ObjectId(self.value.as_py())
+        return self._bson_class(self.value.as_py())
+
+
+class ObjectIdScalar(BSONExtensionScalar):
+    _bson_class = ObjectId
 
 
 class ObjectIdType(PyExtensionType):
@@ -128,6 +138,26 @@ def to_pandas_dtype(self):
         return PandasBinary(self.subtype)
 
 
+class CodeScalar(BSONExtensionScalar):
+    _bson_class = Code
+
+
+class CodeType(PyExtensionType):
+    _type_marker = _BsonArrowTypes.code
+
+    def __init__(self):
+        super().__init__(string())
+
+    def __reduce__(self):
+        return CodeType, ()
+
+    def __arrow_ext_scalar_class__(self):
+        return CodeScalar
+
+    def to_pandas_dtype(self):
+        return PandasCode()
+
+
 # Internal Type Handling.
 
 
@@ -146,6 +176,11 @@ def _is_binary(obj):
     return type_marker == BinaryType._type_marker
 
 
+def _is_code(obj):
+    type_marker = getattr(obj, "_type_marker", "")
+    return type_marker == CodeType._type_marker
+
+
 _TYPE_NORMALIZER_FACTORY = {
     Int64: lambda _: int64(),
     float: lambda _: float64(),
@@ -159,6 +194,7 @@ def _is_binary(obj):
     str: lambda _: string(),
     bool: lambda _: bool_(),
     Binary: lambda subtype: BinaryType(subtype),
+    Code: lambda _: CodeType(),
 }
 
 
@@ -188,6 +224,7 @@ def get_numpy_type(type):
     _is_objectid: _BsonArrowTypes.objectid,
     _is_decimal128: _BsonArrowTypes.decimal128,
     _is_binary: _BsonArrowTypes.binary,
+    _is_code: _BsonArrowTypes.code,
     _atypes.is_string: _BsonArrowTypes.string,
     _atypes.is_boolean: _BsonArrowTypes.bool,
     _atypes.is_struct: _BsonArrowTypes.document,

bindings/python/test/pandas_types/test_code.py

@@ -0,0 +1,105 @@
+# 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 Code
+from pandas.tests.extension import base
+from pymongoarrow.pandas_types import PandasCode, PandasCodeArray
+
+try:
+    base.BaseIndexTests
+except AttributeError:
+    pytest.skip("Not available", allow_module_level=True)
+
+
+def make_datum():
+    return Code(str(np.random.rand()))
+
+
+@pytest.fixture
+def dtype():
+    return PandasCode()
+
+
+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 PandasCodeArray(np.array(make_data(), dtype=object), dtype=dtype)
+
+
+@pytest.fixture
+def data_missing(dtype):
+    return PandasCodeArray(np.array([np.nan, make_datum()], dtype=object), dtype=dtype)
+
+
+@pytest.fixture
+def data_for_sorting(dtype):
+    return PandasCodeArray(
+        np.array([make_datum(), make_datum(), make_datum()], dtype=object), dtype=dtype
+    )
+
+
+@pytest.fixture
+def data_missing_for_sorting(dtype):
+    return PandasCodeArray(
+        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):
+    def test_contains(self):
+        # We cannot compare a Code object to an array.
+        pass
+
+
+class TestConstructors(base.BaseConstructorsTests):
+    pass
+
+
+class TestGetitem(base.BaseGetitemTests):
+    pass
+
+
+class TestSetitem(base.BaseSetitemTests):
+    def test_setitem_frame_2d_values(self):
+        # Results in passing an integer as a value, which
+        # cannot be converted to Code type.
+        pass
+
+
+class TestIndex(base.BaseIndexTests):
+    pass
+
+
+class TestMissing(base.BaseMissingTests):
+    pass