ARROW-52 Add support for BSON binary type (#128)

Author: blink1073

bindings/python/pymongoarrow/api.py

@@ -18,20 +18,10 @@
 from bson import encode
 from bson.codec_options import TypeEncoder, TypeRegistry
 from bson.raw_bson import RawBSONDocument
+from numpy import ndarray
+from pandas import NA, DataFrame
 from pyarrow import Schema as ArrowSchema
 from pyarrow import Table
-
-try:
-    from numpy import ndarray
-except ImportError:
-    ndarray = None
-
-try:
-    from pandas import NA, DataFrame
-except ImportError:
-    DataFrame = None
-    NA = None
-
 from pymongo.bulk import BulkWriteError
 from pymongo.common import MAX_WRITE_BATCH_SIZE
 from pymongoarrow.context import PyMongoArrowContext

bindings/python/pymongoarrow/context.py

@@ -14,6 +14,7 @@
 from bson.codec_options import DEFAULT_CODEC_OPTIONS
 from pyarrow import Table, timestamp
 from pymongoarrow.lib import (
+    BinaryBuilder,
     BoolBuilder,
     DatetimeBuilder,
     DocumentBuilder,
@@ -37,6 +38,7 @@
     _BsonArrowTypes.bool: BoolBuilder,
     _BsonArrowTypes.document: DocumentBuilder,
     _BsonArrowTypes.array: ListBuilder,
+    _BsonArrowTypes.binary: BinaryBuilder,
 }
 
 
@@ -90,7 +92,9 @@ def from_schema(cls, schema, codec_options=DEFAULT_CODEC_OPTIONS):
             elif builder_cls == ListBuilder:
                 arrow_type = schema.typemap[fname]
                 builder_map[encoded_fname] = ListBuilder(arrow_type, tzinfo)
-
+            elif builder_cls == BinaryBuilder:
+                subtype = schema.typemap[fname].subtype
+                builder_map[encoded_fname] = BinaryBuilder(subtype)
             else:
                 builder_map[encoded_fname] = builder_cls()
         return cls(schema, builder_map)

bindings/python/pymongoarrow/lib.pyx

@@ -31,7 +31,7 @@ from pyarrow.lib import (
 
 from pymongoarrow.errors import InvalidBSON, PyMongoArrowError
 from pymongoarrow.context import PyMongoArrowContext
-from pymongoarrow.types import _BsonArrowTypes, _atypes, ObjectIdType, Decimal128StringType
+from pymongoarrow.types import _BsonArrowTypes, _atypes, ObjectIdType, Decimal128StringType, BinaryType
 
 # Cython imports
 from cpython cimport PyBytes_Size, object
@@ -68,6 +68,7 @@ _builder_type_map = {
     BSON_TYPE_DOCUMENT: DocumentBuilder,
     BSON_TYPE_DECIMAL128: StringBuilder,
     BSON_TYPE_ARRAY: ListBuilder,
+    BSON_TYPE_BINARY: BinaryBuilder
 }
 
 _field_type_map = {
@@ -82,6 +83,10 @@ _field_type_map = {
 
 cdef extract_field_dtype(bson_iter_t * doc_iter, bson_iter_t * child_iter, bson_type_t value_t, context):
     """Get the appropropriate data type for a specific field"""
+    cdef const uint8_t *val_buf = NULL
+    cdef uint32_t val_buf_len = 0
+    cdef bson_subtype_t subtype
+
     if value_t in _field_type_map:
         field_type = _field_type_map[value_t]
     elif value_t == BSON_TYPE_ARRAY:
@@ -93,6 +98,9 @@ cdef extract_field_dtype(bson_iter_t * doc_iter, bson_iter_t * child_iter, bson_
         field_type = extract_document_dtype(child_iter, context)
     elif value_t == BSON_TYPE_DATE_TIME:
         field_type = timestamp('ms', tz=context.tzinfo)
+    elif value_t == BSON_TYPE_BINARY:
+        bson_iter_binary (doc_iter, &subtype, &val_buf_len, &val_buf)
+        field_type = BinaryType(subtype)
     else:
         raise PyMongoArrowError('unknown value type {}'.format(value_t))
     return field_type
@@ -129,10 +137,8 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
     cdef char *decimal128_str = <char *> malloc(
         BSON_DECIMAL128_STRING * sizeof(char))
     cdef uint32_t str_len
-    cdef const uint8_t *doc_buf = NULL
-    cdef uint32_t doc_buf_len = 0;
-    cdef const uint8_t *arr_buf = NULL
-    cdef uint32_t arr_buf_len = 0;
+    cdef const uint8_t *val_buf = NULL
+    cdef uint32_t val_buf_len = 0
     cdef bson_decimal128_t dec128
     cdef bson_type_t value_t
     cdef const char * bson_str
@@ -142,6 +148,7 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
     cdef bson_iter_t child_iter
     cdef const char* key
     cdef Py_ssize_t count = 0
+    cdef bson_subtype_t subtype
 
     builder_map = context.builder_map
 
@@ -155,6 +162,7 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
     t_bool = _BsonArrowTypes.bool
     t_document = _BsonArrowTypes.document
     t_array = _BsonArrowTypes.array
+    t_binary = _BsonArrowTypes.binary
 
 
     # initialize count to current length of builders
@@ -197,6 +205,10 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
                         list_dtype = extract_array_dtype(&child_iter, context)
                         list_dtype = list_(list_dtype)
                         builder = ListBuilder(list_dtype, context.tzinfo, value_builder=arr_value_builder)
+                    elif builder_type == BinaryBuilder:
+                        bson_iter_binary (&doc_iter, &subtype,
+                            &val_buf_len, &val_buf)
+                        builder = BinaryBuilder(subtype)
                     else:
                         builder = builder_type()
                     if arr_value_builder is None:
@@ -257,20 +269,28 @@ def process_bson_stream(bson_stream, context, arr_value_builder=None):
                         builder.append_null()
                 elif ftype == t_document:
                     if value_t == BSON_TYPE_DOCUMENT:
-                        bson_iter_document(&doc_iter, &doc_buf_len, &doc_buf)
-                        if doc_buf_len <= 0:
+                        bson_iter_document(&doc_iter, &val_buf_len, &val_buf)
+                        if val_buf_len <= 0:
                             raise ValueError("Subdocument is invalid")
-                        builder.append(<bytes>doc_buf[:doc_buf_len])
+                        builder.append(<bytes>val_buf[:val_buf_len])
                     else:
                         builder.append_null()
                 elif ftype == t_array:
                     if value_t == BSON_TYPE_ARRAY:
-                        bson_iter_array(&doc_iter, &doc_buf_len, &doc_buf)
-                        if doc_buf_len <= 0:
+                        bson_iter_array(&doc_iter, &val_buf_len, &val_buf)
+                        if val_buf_len <= 0:
                             raise ValueError("Subarray is invalid")
-                        builder.append(<bytes>doc_buf[:doc_buf_len])
+                        builder.append(<bytes>val_buf[:val_buf_len])
                     else:
                         builder.append_null()
+                elif ftype == t_binary:
+                    if value_t == BSON_TYPE_BINARY:
+                        bson_iter_binary (&doc_iter, &subtype,
+                            &val_buf_len, &val_buf)
+                        if subtype != builder.subtype:
+                            builder.append_null()
+                        else:
+                            builder.append(<bytes>val_buf[:val_buf_len])
                 else:
                     raise PyMongoArrowError('unknown ftype {}'.format(ftype))
             count += 1
@@ -534,6 +554,8 @@ cdef object get_field_builder(field, tzinfo):
         field_builder = ObjectIdBuilder()
     elif getattr(field_type, '_type_marker') == _BsonArrowTypes.decimal128_str:
         field_builder = StringBuilder()
+    elif getattr(field_type, '_type_marker') == _BsonArrowTypes.binary:
+        field_builder = BinaryBuilder(field_type.subtype)
     else:
         field_builder = StringBuilder()
     return field_builder
@@ -596,6 +618,7 @@ cdef class DocumentBuilder(_ArrayBuilderBase):
     cdef shared_ptr[CStructBuilder] unwrap(self):
         return self.builder
 
+
 cdef class ListBuilder(_ArrayBuilderBase):
     type_marker = _BsonArrowTypes.array
 
@@ -647,3 +670,36 @@ cdef class ListBuilder(_ArrayBuilderBase):
 
     cdef shared_ptr[CListBuilder] unwrap(self):
         return self.builder
+
+
+cdef class BinaryBuilder(_ArrayBuilderBase):
+    type_marker = _BsonArrowTypes.binary
+    cdef:
+        shared_ptr[CBinaryBuilder] builder
+        uint8_t _subtype
+
+    def __cinit__(self, uint8_t subtype):
+        self._subtype = subtype
+        self.builder.reset(new CBinaryBuilder())
+
+    cpdef append_null(self):
+        self.builder.get().AppendNull()
+
+    @property
+    def subtype(self):
+        return self._subtype
+
+    def __len__(self):
+        return self.builder.get().length()
+
+    cpdef append(self, value):
+        self.builder.get().Append(<bytes>value, len(value))
+
+    cpdef finish(self):
+        cdef shared_ptr[CArray] out
+        with nogil:
+            self.builder.get().Finish(&out)
+        return pyarrow_wrap_array(out).cast(BinaryType(self._subtype))
+
+    cdef shared_ptr[CBinaryBuilder] unwrap(self):
+        return self.builder

bindings/python/pymongoarrow/libbson.pxd

@@ -48,6 +48,9 @@ cdef extern from "<bson/bson.h>":
     ctypedef struct bson_iter_t:
         pass
 
+    ctypedef enum bson_subtype_t:
+        pass
+
     ctypedef struct bson_oid_t:
         uint8_t bytes[12]
 
@@ -137,6 +140,10 @@ cdef extern from "<bson/bson.h>":
                  uint32_t *array_len,     # OUT
                  const uint8_t **array)
 
+    void bson_iter_binary (const bson_iter_t *iter, # IN
+                  bson_subtype_t *subtype, # OUT
+                  uint32_t *binary_len,    # OUT
+                  const uint8_t **binary)  # OUT
 
 # bson_reader_t API
 cdef extern from "<bson/bson.h>":