INTPYTHON-538 Add support for PyArrow Decimal128 type

Author: blink1073

bindings/python/pymongoarrow/api.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import warnings
+from decimal import Decimal
 
 import numpy as np
 import pandas as pd
@@ -25,6 +26,7 @@
 import pymongo.errors
 from bson import encode
 from bson.codec_options import TypeEncoder, TypeRegistry
+from bson.decimal128 import Decimal128
 from bson.raw_bson import RawBSONDocument
 from numpy import ndarray
 from pyarrow import Schema as ArrowSchema
@@ -416,6 +418,18 @@ def transform_python(self, _):
         return
 
 
+class _DecimalCodec(TypeEncoder):
+    """A custom type codec for Decimal objects."""
+
+    @property
+    def python_type(self):
+        return Decimal
+
+    def transform_python(self, value):
+        """Transform an Decimal object into a BSON Decimal128 object"""
+        return Decimal128(value)
+
+
 def write(collection, tabular, *, exclude_none: bool = False):
     """Write data from `tabular` into the given MongoDB `collection`.
 
@@ -471,7 +485,7 @@ def write(collection, tabular, *, exclude_none: bool = False):
 
     # Handle Pandas NA objects.
     codec_options = collection.codec_options
-    type_registry = TypeRegistry([_PandasNACodec()])
+    type_registry = TypeRegistry([_PandasNACodec(), _DecimalCodec()])
     codec_options = codec_options.with_options(type_registry=type_registry)
 
     while cur_offset < tab_size:

bindings/python/pymongoarrow/types.py

@@ -277,6 +277,7 @@ def get_numpy_type(type):
     _atypes.is_date64: _BsonArrowTypes.date64.value,
     _atypes.is_large_string: _BsonArrowTypes.string.value,
     _atypes.is_large_list: _BsonArrowTypes.array.value,
+    _atypes.is_decimal128: _BsonArrowTypes.decimal128.value,
 }
 
 

bindings/python/test/test_arrow.py

@@ -966,6 +966,17 @@ def test_exclude_none(self):
         col_data = list(self.coll.find({}))
         assert "b" not in col_data[2]
 
+    def test_decimal128(self):
+        import decimal
+
+        a = decimal.Decimal("123.45")
+        arr = pa.array([a], pa.decimal128(5, 2))
+        data = Table.from_arrays([arr], names=["data"])
+        self.coll.drop()
+        write(self.coll, data)
+        coll_data = list(self.coll.find({}))
+        assert coll_data[0]["data"] == Decimal128(a)
+
 
 class TestArrowExplicitApi(ArrowApiTestMixin, unittest.TestCase):
     def run_find(self, *args, **kwargs):