ARROW-15 Add support for BSON Decimal128 type (#72)

Author: blink1073

bindings/python/docs/source/changelog.rst

@@ -1,14 +1,19 @@
 Changelog
 =========
 
+Changes in Version 0.4.0
+------------------------
+- Support for :class:`~bson.decimal128.Decimal128` type.
+- Support for macOS arm64 architecture on Python 3.9+.
+
 Changes in Version 0.3.0
 ------------------------
 - Support for `PyArrow` 7.0.
-- Support for `ObjectId` `bson` type.
+- Support for :class:`~bson.objectid.ObjectId` type.
 - Improve error message when schema contains an unsupported type.
 - Add support for BSON string type.
 - Add support for BSON boolean type.
-- Upgraded to bundle ``libbson`` 1.21.1. If installing from source, the minimum supported ``libbson`` version is now 1.21.0.
+- Upgraded to bundle `libbson <http://mongoc.org/libbson/current/index.html>`_ 1.21.1. If installing from source, the minimum supported ``libbson`` version is now 1.21.0.
 - Dropped Python 3.6 support (it was dropped in `PyArrow` 7.0).
 
 Changes in Version 0.2.0
@@ -17,8 +22,8 @@ Changes in Version 0.2.0
 - Support for PyMongo 4.0.
 - Support for Python 3.10.
 - Support for Windows.
-- ``find_arrow_all`` now accepts a user-provided ``projection``.
-- ``find_raw_batches`` now accepts a ``session`` object.
+- :meth:`~pymongoarrow.api.find_arrow_all` now accepts a user-provided ``projection``.
+- :meth:`~pymongoarrow.api.find_arrow_all` now accepts a ``session`` object.
 - Note: The supported version of ``pyarrow`` is now ``>=6,<6.1``.
 
 Changes in Version 0.1.1

bindings/python/docs/source/conf.py

@@ -74,4 +74,5 @@
     "pandas": ("https://pandas.pydata.org/docs/", None),
     "numpy": ("https://numpy.org/doc/1.20/", None),
     "pymongo": ("https://pymongo.readthedocs.io/en/stable/", None),
+    "bson": ("https://pymongo.readthedocs.io/en/stable/", None),
 }

bindings/python/pymongoarrow/api.py

@@ -36,6 +36,7 @@
     "find_pandas_all",
     "aggregate_numpy_all",
     "find_numpy_all",
+    "write",
     "Schema",
 ]
 

bindings/python/pymongoarrow/context.py

@@ -30,6 +30,7 @@
     _BsonArrowTypes.double: DoubleBuilder,
     _BsonArrowTypes.datetime: DatetimeBuilder,
     _BsonArrowTypes.objectid: ObjectIdBuilder,
+    _BsonArrowTypes.decimal128_str: StringBuilder,
     _BsonArrowTypes.string: StringBuilder,
     _BsonArrowTypes.bool: BoolBuilder,
 }

bindings/python/pymongoarrow/lib.pyx

@@ -33,6 +33,7 @@ from cython.operator cimport dereference
 from libcpp cimport bool as cbool
 from libcpp.map cimport map
 from libcpp.string cimport string
+from libc.stdlib cimport malloc, free
 from pyarrow.lib cimport *
 from pyarrow.lib import tobytes
 from pymongoarrow.libarrow cimport *
@@ -59,11 +60,14 @@ def process_bson_stream(bson_stream, context):
     cdef bson_reader_t* stream_reader = bson_reader_new_from_data(docstream, length)
     cdef const bson_t * doc = NULL
     cdef bson_iter_t doc_iter
+    cdef bson_decimal128_t dec128
     cdef const char* key
     cdef bson_type_t value_t
     cdef Py_ssize_t count = 0
     cdef const char * bson_str
     cdef uint32_t str_len
+    cdef char *decimal128_str = <char *> malloc(
+        BSON_DECIMAL128_STRING * sizeof(char))
 
     # Localize types for better performance.
     t_int32 = _BsonArrowTypes.int32
@@ -115,6 +119,10 @@ def process_bson_stream(bson_stream, context):
                         if value_t == BSON_TYPE_UTF8:
                             bson_str = bson_iter_utf8(&doc_iter, &str_len)
                             builder.append(<bytes>(bson_str)[:str_len])
+                        elif value_t == BSON_TYPE_DECIMAL128:
+                            bson_iter_decimal128(&doc_iter, &dec128)
+                            bson_decimal128_to_string(&dec128, decimal128_str)
+                            builder.append(<bytes>(decimal128_str))
                         else:
                             builder.append_null()
                     elif ftype == t_double:
@@ -144,6 +152,7 @@ def process_bson_stream(bson_stream, context):
                     builder.append_null()
     finally:
         bson_reader_destroy(stream_reader)
+        free(decimal128_str)
 
 
 # Builders

bindings/python/pymongoarrow/libbson.pxd

@@ -124,6 +124,7 @@ cdef extern from "<bson/bson.h>":
 
     const char * bson_iter_utf8(const bson_iter_t *iter, uint32_t *length)
 
+    bint bson_iter_decimal128 (const bson_iter_t *iter, bson_decimal128_t *dec)
 
 # bson_reader_t API
 cdef extern from "<bson/bson.h>":
@@ -133,6 +134,10 @@ cdef extern from "<bson/bson.h>":
 
     void bson_reader_destroy(bson_reader_t *reader)
 
+# bson_decimal128 API
+cdef extern from "<bson/bson.h>":
+    cdef int32_t BSON_DECIMAL128_STRING
+    void bson_decimal128_to_string (const bson_decimal128_t *dec, char *str)
 
 # runtime version checking API
 cdef extern from "<bson/bson.h>":

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 Int64, ObjectId
+from bson import Decimal128, Int64, ObjectId
 from pyarrow import DataType as _ArrowDataType
 from pyarrow import PyExtensionType, binary, bool_, float64, int64, string, timestamp
 
@@ -30,6 +30,7 @@ class _BsonArrowTypes(enum.Enum):
     objectid = 5
     string = 6
     bool = 7
+    decimal128_str = 8
 
 
 # Custom Extension Types.
@@ -47,6 +48,16 @@ def __reduce__(self):
         return ObjectIdType, ()
 
 
+class Decimal128StringType(PyExtensionType):
+    _type_marker = _BsonArrowTypes.decimal128_str
+
+    def __init__(self):
+        super().__init__(string())
+
+    def __reduce__(self):
+        return Decimal128StringType, ()
+
+
 # Internal Type Handling.
 
 
@@ -55,14 +66,20 @@ def _is_objectid(obj):
     return type_marker == ObjectIdType._type_marker
 
 
+def _is_decimal128_str(obj):
+    type_marker = getattr(obj, "_type_marker", "")
+    return type_marker == Decimal128StringType._type_marker
+
+
 _TYPE_NORMALIZER_FACTORY = {
     Int64: lambda _: int64(),
     float: lambda _: float64(),
     int: lambda _: int64(),
     datetime: lambda _: timestamp("ms"),  # TODO: add tzinfo support
     ObjectId: lambda _: ObjectIdType(),
-    str: lambda: string(),
-    bool: lambda: bool_(),
+    Decimal128: lambda _: Decimal128StringType(),
+    str: lambda _: string(),
+    bool: lambda _: bool_(),
 }
 
 
@@ -90,6 +107,7 @@ def get_numpy_type(type):
     _atypes.is_float64: _BsonArrowTypes.double,
     _atypes.is_timestamp: _BsonArrowTypes.datetime,
     _is_objectid: _BsonArrowTypes.objectid,
+    _is_decimal128_str: _BsonArrowTypes.decimal128_str,
     _atypes.is_string: _BsonArrowTypes.string,
     _atypes.is_boolean: _BsonArrowTypes.bool,
 }

bindings/python/test/test_arrow.py

@@ -18,7 +18,8 @@
 from test.utils import AllowListEventListener
 
 import pymongo
-from pyarrow import Table, bool_, decimal256, float64, int32, int64
+from bson import Decimal128, ObjectId
+from pyarrow import Table, binary, bool_, decimal256, float64, int32, int64
 from pyarrow import schema as ArrowSchema
 from pyarrow import string, timestamp
 from pyarrow.parquet import read_table, write_table
@@ -27,6 +28,7 @@
 from pymongoarrow.api import Schema, aggregate_arrow_all, find_arrow_all, write
 from pymongoarrow.errors import ArrowWriteError
 from pymongoarrow.monkey import patch_all
+from pymongoarrow.types import Decimal128StringType, ObjectIdType
 
 
 class TestArrowApiMixin:
@@ -292,3 +294,42 @@ def run_find(self, *args, **kwargs):
 
     def run_aggregate(self, *args, **kwargs):
         return self.coll.aggregate_arrow_all(*args, **kwargs)
+
+
+class TestBSONTypes(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        if not client_context.connected:
+            raise unittest.SkipTest("cannot connect to MongoDB")
+        cls.cmd_listener = AllowListEventListener("find", "aggregate")
+        cls.getmore_listener = AllowListEventListener("getMore")
+        cls.client = client_context.get_client(
+            event_listeners=[cls.getmore_listener, cls.cmd_listener]
+        )
+
+    def test_find_decimal128(self):
+        oids = list(ObjectId() for i in range(4))
+        decs = [Decimal128(i) for i in ["0.1", "1.0", "1e-5"]]
+        schema = Schema({"_id": ObjectIdType(), "data": Decimal128StringType()})
+        expected = Table.from_pydict(
+            {
+                "_id": [i.binary for i in oids],
+                "data": [str(decs[0]), str(decs[1]), str(decs[2]), None],
+            },
+            ArrowSchema([("_id", binary(12)), ("data", string())]),
+        )
+        coll = self.client.pymongoarrow_test.get_collection(
+            "test", write_concern=WriteConcern(w="majority")
+        )
+
+        coll.drop()
+        coll.insert_many(
+            [
+                {"_id": oids[0], "data": decs[0]},
+                {"_id": oids[1], "data": decs[1]},
+                {"_id": oids[2], "data": decs[2]},
+                {"_id": oids[3]},
+            ]
+        )
+        table = find_arrow_all(coll, {}, schema=schema)
+        self.assertEqual(table, expected)

bindings/python/test/test_bson.py

@@ -14,11 +14,11 @@
 from unittest import TestCase
 
 import pyarrow as pa
-from bson import InvalidBSON, encode
+from bson import Decimal128, Int64, InvalidBSON, encode
 from pymongoarrow.context import PyMongoArrowContext
 from pymongoarrow.lib import process_bson_stream
 from pymongoarrow.schema import Schema
-from pymongoarrow.types import ObjectId, ObjectIdType, bool_, int64, string
+from pymongoarrow.types import ObjectId, ObjectIdType, int64, string
 
 
 class TestBsonToArrowConversionBase(TestCase):
@@ -147,9 +147,24 @@ def test_object_id_type(self):
         assert result.type._type_marker == ObjectIdType._type_marker
 
 
+class TestInt64Type(TestBsonToArrowConversionBase):
+    def setUp(self):
+        self.schema = Schema({"data": Int64})
+        self.context = PyMongoArrowContext.from_schema(self.schema)
+
+    def test_simple(self):
+        docs = [
+            {"data": Int64(1e15)},
+            {"data": Int64(-1e5)},
+            {"data": Int64(10)},
+        ]
+        as_dict = {"data": [1e15, -1e5, 10]}
+        self._run_test(docs, as_dict)
+
+
 class TestBooleanType(TestBsonToArrowConversionBase):
     def setUp(self):
-        self.schema = Schema({"data": bool_()})
+        self.schema = Schema({"data": bool})
         self.context = PyMongoArrowContext.from_schema(self.schema)
 
     def test_simple(self):
@@ -163,3 +178,32 @@ def test_simple(self):
         ]
         as_dict = {"data": [True, False, None, None, False, True]}
         self._run_test(docs, as_dict)
+
+
+class TestStringType(TestBsonToArrowConversionBase):
+    def setUp(self):
+        self.schema = Schema({"data": str})
+        self.context = PyMongoArrowContext.from_schema(self.schema)
+
+    def test_simple(self):
+        docs = [
+            {"data": "a"},
+            {"data": "dätá"},
+        ]
+        as_dict = {"data": ["a", "dätá"]}
+        self._run_test(docs, as_dict)
+
+
+class TestDecimal128StringType(TestBsonToArrowConversionBase):
+    def setUp(self):
+        self.schema = Schema({"data": Decimal128})
+        self.context = PyMongoArrowContext.from_schema(self.schema)
+
+    def test_simple(self):
+        docs = [
+            {"data": Decimal128("0.01")},
+            {"data": Decimal128("1e-5")},
+            {"data": Decimal128("1.0e+5")},
+        ]
+        as_dict = {"data": ["0.01", "0.00001", "1.0E+5"]}
+        self._run_test(docs, as_dict)

bindings/python/test/test_monkey.py

@@ -18,7 +18,7 @@
 
 class TestMonkey(TestCase):
     def test_all_methods_listed_for_patching(self):
-        unpatched_apis = {"Schema"}
+        unpatched_apis = {"Schema", "write"}
         all_methods = set(api.__all__) - unpatched_apis
         methods_to_patch = set(api._PATCH_METHODS)
         self.assertEqual(all_methods, methods_to_patch)