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

* ARROW-15 Add support for BSON Decimal128 type

* add builder test

* ARROW-15 Add support for BSON Decimal128 type

* Add PandasDecimal128 and update tests

* fix typo

* wip remove test workarounds

* fix auto schema

* remove unused import

* clean up struct

* address review

* test on big endian linux

* debug

* remove other builds

* run checkout first

* fix build-libbson

* install cmake

* make apt non-interactive

* fix syntax

* use python3

* install pip

* install pip

* fix syntax

* remove support for big-endian decimal128

* add note about big-endian systems

* add comment

Author: blink1073

bindings/python/README.rst

@@ -61,6 +61,9 @@ installed::
 
      $ python -m pip install pandas
 
+Note:  ``pymongoarrow`` is not supported or tested on big-endian systems
+(e.g. Linux s390x).
+
 Development Install
 ===================
 

bindings/python/docs/source/supported_types.rst

@@ -6,15 +6,13 @@ Supported Types
 PyMongoArrow currently supports a small subset of all BSON types.
 Support for additional types will be added in subsequent releases.
 
-.. note:: PyMongoArrow does not currently fully support extension types with Pandas/NumPy or Arrow.
-   However, they can be used in schemas.
-   This means that ObjectId and Decimal128 are not fully supported in Pandas DataFrames or Arrow Tables.
-   Instead, the schema type will be converted to a string or object representation of the type.
-   For more information see :doc:`extension_types`.
 
 .. note:: For more information about BSON types, see the
    `BSON specification <http://bsonspec.org/spec.html>`_.
 
+.. note:: ``Decimal128`` types are only supported on little-endian systems.
+   On big-endian systems, ``null`` will be used.
+
 .. list-table::
    :widths: auto
    :header-rows: 1
@@ -28,7 +26,9 @@ Support for additional types will be added in subsequent releases.
    * - Embedded array
      - :class:`py.list`, an instance of :class:`pyarrow.list_`,
    * - ObjectId
-     - :class:`py.bytes`, :class:`bson.ObjectId`, an instance of :class:`pymongoarrow.types.ObjectIdType`, an instance of :class:`pyarrow.FixedSizeBinaryScalar`
+     - :class:`py.bytes`, :class:`bson.ObjectId`, an instance of :class:`pymongoarrow.types.ObjectIdType`, an instance of :class:`pymongoarrow.pandas_types.PandasObjectId`
+   * - Decimal128
+     - :class:`bson.Decimal128`, an instance of :class:`pymongoarrow.types.Decimal128Type`, an instance of :class:`pymongoarrow.pandas_types.PandasDecimal128`.
    * - Boolean
      - an instance of :class:`~pyarrow.bool_`, :class:`~py.bool`
    * - 64-bit binary floating point

bindings/python/pymongoarrow/api.py

@@ -338,7 +338,7 @@ def write(collection, tabular):
     tab_size = len(tabular)
     if isinstance(tabular, Table):
         _validate_schema(tabular.schema.types)
-    elif DataFrame is not None and isinstance(tabular, DataFrame):
+    elif isinstance(tabular, DataFrame):
         _validate_schema(ArrowSchema.from_pandas(tabular).types)
     elif (
         isinstance(tabular, dict)
@@ -359,9 +359,8 @@ def write(collection, tabular):
 
     # Handle Pandas NA objects.
     codec_options = collection.codec_options
-    if DataFrame is not None:
-        type_registry = TypeRegistry([_PandasNACodec()])
-        codec_options = codec_options.with_options(type_registry=type_registry)
+    type_registry = TypeRegistry([_PandasNACodec()])
+    codec_options = codec_options.with_options(type_registry=type_registry)
 
     while cur_offset < tab_size:
         cur_size = 0

bindings/python/pymongoarrow/context.py

@@ -17,6 +17,7 @@
     BinaryBuilder,
     BoolBuilder,
     DatetimeBuilder,
+    Decimal128Builder,
     DocumentBuilder,
     DoubleBuilder,
     Int32Builder,
@@ -33,7 +34,7 @@
     _BsonArrowTypes.double: DoubleBuilder,
     _BsonArrowTypes.datetime: DatetimeBuilder,
     _BsonArrowTypes.objectid: ObjectIdBuilder,
-    _BsonArrowTypes.decimal128_str: StringBuilder,
+    _BsonArrowTypes.decimal128: Decimal128Builder,
     _BsonArrowTypes.string: StringBuilder,
     _BsonArrowTypes.bool: BoolBuilder,
     _BsonArrowTypes.document: DocumentBuilder,

bindings/python/pymongoarrow/lib.pyx

@@ -20,6 +20,8 @@
 import copy
 import datetime
 import enum
+import struct as pystruct
+import sys
 
 # Python imports
 import bson
@@ -31,15 +33,14 @@ from pyarrow.lib import (
 
 from pymongoarrow.errors import InvalidBSON, PyMongoArrowError
 from pymongoarrow.context import PyMongoArrowContext
-from pymongoarrow.types import _BsonArrowTypes, _atypes, ObjectIdType, Decimal128StringType, BinaryType
+from pymongoarrow.types import _BsonArrowTypes, _atypes, ObjectIdType, Decimal128Type as Decimal128Type_, BinaryType
 
 # 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 libcpp.vector cimport vector
-from libc.stdlib cimport malloc, free
 from pyarrow.lib cimport *
 from pymongoarrow.libarrow cimport *
 from pymongoarrow.libbson cimport *
@@ -66,7 +67,7 @@ _builder_type_map = {
     BSON_TYPE_UTF8: StringBuilder,
     BSON_TYPE_BOOL: BoolBuilder,
     BSON_TYPE_DOCUMENT: DocumentBuilder,
-    BSON_TYPE_DECIMAL128: StringBuilder,
+    BSON_TYPE_DECIMAL128: Decimal128Builder,
     BSON_TYPE_ARRAY: ListBuilder,
     BSON_TYPE_BINARY: BinaryBuilder
 }
@@ -78,7 +79,7 @@ _field_type_map = {
     BSON_TYPE_OID: ObjectIdType(),
     BSON_TYPE_UTF8: string(),
     BSON_TYPE_BOOL: bool_(),
-    BSON_TYPE_DECIMAL128: Decimal128StringType(),
+    BSON_TYPE_DECIMAL128: Decimal128Type_(),
 }
 
 cdef extract_field_dtype(bson_iter_t * doc_iter, bson_iter_t * child_iter, bson_type_t value_t, context):
@@ -134,8 +135,6 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
     cdef const uint8_t* docstream = <const uint8_t *>bson_stream
     cdef size_t length = <size_t>PyBytes_Size(bson_stream)
     cdef bson_reader_t* stream_reader = bson_reader_new_from_data(docstream, length)
-    cdef char *decimal128_str = <char *> malloc(
-        BSON_DECIMAL128_STRING * sizeof(char))
     cdef uint32_t str_len
     cdef const uint8_t *val_buf = NULL
     cdef uint32_t val_buf_len = 0
@@ -163,7 +162,7 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
     t_document = _BsonArrowTypes.document
     t_array = _BsonArrowTypes.array
     t_binary = _BsonArrowTypes.binary
-
+    t_decimal128 = _BsonArrowTypes.decimal128
 
     # initialize count to current length of builders
     for _, builder in builder_map.items():
@@ -243,10 +242,17 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
                     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:
+                    else:
+                        builder.append_null()
+                elif ftype == t_decimal128:
+                    if value_t == BSON_TYPE_DECIMAL128:
                         bson_iter_decimal128(&doc_iter, &dec128)
-                        bson_decimal128_to_string(&dec128, decimal128_str)
-                        builder.append(<bytes>(decimal128_str))
+                        if sys.byteorder == 'little':
+                            val = pystruct.pack('<QQ', dec128.low, dec128.high)
+                            builder.append(val)
+                        else:
+                            # We do not support big-endian systems.
+                            builder.append_null()
                     else:
                         builder.append_null()
                 elif ftype == t_double:
@@ -300,7 +306,6 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
                     builder.append_null()
     finally:
         bson_reader_destroy(stream_reader)
-        free(decimal128_str)
 
 
 # Builders
@@ -523,6 +528,35 @@ cdef class BoolBuilder(_ArrayBuilderBase):
         return self.builder
 
 
+cdef class Decimal128Builder(_ArrayBuilderBase):
+    type_marker = _BsonArrowTypes.decimal128
+
+    cdef:
+        shared_ptr[CFixedSizeBinaryBuilder] builder
+
+    def __cinit__(self, MemoryPool memory_pool=None):
+        cdef shared_ptr[CDataType] dtype = fixed_size_binary(16)
+        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):
+        self.builder.get().Append(value)
+
+    cpdef finish(self):
+        cdef shared_ptr[CArray] out
+        with nogil:
+            self.builder.get().Finish(&out)
+        return pyarrow_wrap_array(out).cast(Decimal128Type_())
+
+    cdef shared_ptr[CFixedSizeBinaryBuilder] unwrap(self):
+        return self.builder
+
 
 cdef object get_field_builder(field, tzinfo):
     """"Find the appropriate field builder given a pyarrow field"""
@@ -552,8 +586,8 @@ cdef object get_field_builder(field, tzinfo):
         field_builder = ListBuilder(field_type, tzinfo)
     elif getattr(field_type, '_type_marker') == _BsonArrowTypes.objectid:
         field_builder = ObjectIdBuilder()
-    elif getattr(field_type, '_type_marker') == _BsonArrowTypes.decimal128_str:
-        field_builder = StringBuilder()
+    elif getattr(field_type, '_type_marker') == _BsonArrowTypes.decimal128:
+        field_builder = Decimal128Builder()
     elif getattr(field_type, '_type_marker') == _BsonArrowTypes.binary:
         field_builder = BinaryBuilder(field_type.subtype)
     else:

bindings/python/pymongoarrow/libbson.pxd

@@ -38,7 +38,8 @@ cdef extern from "<bson/bson.h>":
         pass
 
     ctypedef struct bson_decimal128_t:
-        pass
+        uint64_t high
+        uint64_t low
 
     ctypedef struct bson_error_t:
         uint32_t domain

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, ObjectId
+from bson import Binary, Decimal128, ObjectId
 from pandas.api.extensions import (
     ExtensionArray,
     ExtensionDtype,
@@ -56,7 +56,10 @@ def __from_arrow__(self, array: Union[pa.Array, pa.ChunkedArray]) -> ExtensionAr
             typ = self.type
             for val in np.array(arr):
                 if not pd.isna(val) and not isinstance(val, typ):
-                    val = typ(val)
+                    if typ == Decimal128:
+                        val = Decimal128.from_bid(val)
+                    else:
+                        val = typ(val)
                 vals.append(val)
             arr = np.array(vals, dtype=object)
             # using _from_sequence to ensure None is converted to NA
@@ -249,3 +252,27 @@ def __arrow_array__(self, type=None):
         from pymongoarrow.types import ObjectIdType
 
         return pa.array(self.data, type=ObjectIdType())
+
+
+@register_extension_dtype
+class PandasDecimal128(PandasBSONDtype):
+    """A pandas extension type for BSON ObjectId data type."""
+
+    type = Decimal128
+
+    @classmethod
+    def construct_array_type(cls) -> Type["PandasDecimal128Array"]:
+        return PandasDecimal128Array
+
+
+class PandasDecimal128Array(PandasBSONExtensionArray):
+    """A pandas extension type for BSON Binary data arrays."""
+
+    @property
+    def _default_dtype(self):
+        return PandasDecimal128()
+
+    def __arrow_array__(self, type=None):
+        from pymongoarrow.types import Decimal128Type
+
+        return pa.array(self.data, type=Decimal128Type())

bindings/python/pymongoarrow/types.py

@@ -31,7 +31,7 @@
     struct,
     timestamp,
 )
-from pymongoarrow.pandas_types import PandasBinary, PandasObjectId
+from pymongoarrow.pandas_types import PandasBinary, PandasDecimal128, PandasObjectId
 
 
 class _BsonArrowTypes(enum.Enum):
@@ -42,7 +42,7 @@ class _BsonArrowTypes(enum.Enum):
     objectid = 5
     string = 6
     bool = 7
-    decimal128_str = 8
+    decimal128 = 8
     document = 9
     array = 10
     binary = 11
@@ -78,21 +78,26 @@ def to_pandas_dtype(self):
 
 class Decimal128Scalar(ExtensionScalar):
     def as_py(self):
-        return Decimal128(self.value.as_py())
+        if self.value is None:
+            return None
+        return Decimal128.from_bid(self.value.as_py())
 
 
-class Decimal128StringType(PyExtensionType):
-    _type_marker = _BsonArrowTypes.decimal128_str
+class Decimal128Type(PyExtensionType):
+    _type_marker = _BsonArrowTypes.decimal128
 
     def __init__(self):
-        super().__init__(string())
+        super().__init__(binary(16))
 
     def __reduce__(self):
-        return Decimal128StringType, ()
+        return Decimal128Type, ()
 
     def __arrow_ext_scalar_class__(self):
         return Decimal128Scalar
 
+    def to_pandas_dtype(self):
+        return PandasDecimal128()
+
 
 class BinaryScalar(ExtensionScalar):
     def as_py(self):
@@ -131,9 +136,9 @@ def _is_objectid(obj):
     return type_marker == ObjectIdType._type_marker
 
 
-def _is_decimal128_str(obj):
+def _is_decimal128(obj):
     type_marker = getattr(obj, "_type_marker", "")
-    return type_marker == Decimal128StringType._type_marker
+    return type_marker == Decimal128Type._type_marker
 
 
 def _is_binary(obj):
@@ -150,7 +155,7 @@ def _is_binary(obj):
     # must be used directly.
     datetime: lambda _: timestamp("ms"),
     ObjectId: lambda _: ObjectIdType(),
-    Decimal128: lambda _: Decimal128StringType(),
+    Decimal128: lambda _: Decimal128Type(),
     str: lambda _: string(),
     bool: lambda _: bool_(),
     Binary: lambda subtype: BinaryType(subtype),
@@ -181,7 +186,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,
+    _is_decimal128: _BsonArrowTypes.decimal128,
     _is_binary: _BsonArrowTypes.binary,
     _atypes.is_string: _BsonArrowTypes.string,
     _atypes.is_boolean: _BsonArrowTypes.bool,