ARROW-16 BSON array support (#119)

Author: juliusgeo

bindings/python/docs/source/changelog.rst

@@ -4,6 +4,8 @@ Changelog
 Changes in Version 0.7.0
 ------------------------
 - Added support for BSON Embedded Document type.
+- Added support for BSON Array type.
+
 
 Changes in Version 0.6.3
 ------------------------

bindings/python/docs/source/quickstart.rst

@@ -76,6 +76,10 @@ Nested data (embedded documents) are also supported::
   from pymongoarrow.api import Schema
   schema = Schema({'_id': int, 'amount': float, 'account': { 'name': str, 'account_number': int}})
 
+Arrays (and nested arrays) are also supported::
+
+  from pymongoarrow.api import Schema
+  schema = Schema({'_id': int, 'amount': float, 'account': list_(int32())})
 
 .. note::
 
@@ -112,6 +116,12 @@ Nested data (embedded documents) are also supported::
   schema = Schema({'_id': int, 'amount': float, 'account': { 'name': str, 'account_number': int}})
   arrow_table = client.db.data.find_arrow_all({'amount': {'$gt': 0}}, schema=schema)
 
+Arrays (and nested arrays) are also supported::
+
+  from pymongoarrow.api import Schema
+  schema = Schema({'_id': int, 'amount': float, 'account': list_(int32())})
+  arrow_table = client.db.data.find_arrow_all({'amount': {'$gt': 0}}, schema=schema)
+
 Aggregate operations
 --------------------
 Running ``aggregate`` operations is similar to ``find``. Here is an example of

bindings/python/docs/source/supported_types.rst

@@ -19,6 +19,8 @@ Support for additional types will be added in subsequent releases.
      - :class:`py.str`, an instance of :class:`pyarrow.string`
    * - Embedded document
      - :class:`py.dict`, and instance of :class:`pyarrow.struct`
+   * - 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`
    * - Boolean

bindings/python/pymongoarrow/context.py

@@ -20,6 +20,7 @@
     DoubleBuilder,
     Int32Builder,
     Int64Builder,
+    ListBuilder,
     ObjectIdBuilder,
     StringBuilder,
 )
@@ -35,6 +36,7 @@
     _BsonArrowTypes.string: StringBuilder,
     _BsonArrowTypes.bool: BoolBuilder,
     _BsonArrowTypes.document: DocumentBuilder,
+    _BsonArrowTypes.array: ListBuilder,
 }
 
 
@@ -71,7 +73,6 @@ def from_schema(cls, schema, codec_options=DEFAULT_CODEC_OPTIONS):
 
         builder_map = {}
         tzinfo = codec_options.tzinfo
-
         str_type_map = _get_internal_typemap(schema.typemap)
         for fname, ftype in str_type_map.items():
             builder_cls = _TYPE_TO_BUILDER_CLS[ftype]
@@ -86,6 +87,10 @@ def from_schema(cls, schema, codec_options=DEFAULT_CODEC_OPTIONS):
             elif builder_cls == DocumentBuilder:
                 arrow_type = schema.typemap[fname]
                 builder_map[encoded_fname] = DocumentBuilder(arrow_type, tzinfo)
+            elif builder_cls == ListBuilder:
+                arrow_type = schema.typemap[fname]
+                builder_map[encoded_fname] = ListBuilder(arrow_type, tzinfo)
+
             else:
                 builder_map[encoded_fname] = builder_cls()
         return cls(schema, builder_map)

bindings/python/pymongoarrow/lib.pyx

@@ -26,8 +26,9 @@ import bson
 import numpy as np
 from pyarrow import timestamp, struct, field
 from pyarrow.lib import (
-    tobytes, StructType, int32, int64, float64, string, bool_
+    tobytes, StructType, int32, int64, float64, string, bool_, list_
 )
+
 from pymongoarrow.errors import InvalidBSON, PyMongoArrowError
 from pymongoarrow.context import PyMongoArrowContext
 from pymongoarrow.types import _BsonArrowTypes, _atypes, ObjectIdType, Decimal128StringType
@@ -65,7 +66,8 @@ _builder_type_map = {
     BSON_TYPE_UTF8: StringBuilder,
     BSON_TYPE_BOOL: BoolBuilder,
     BSON_TYPE_DOCUMENT: DocumentBuilder,
-    BSON_TYPE_DECIMAL128: StringBuilder
+    BSON_TYPE_DECIMAL128: StringBuilder,
+    BSON_TYPE_ARRAY: ListBuilder,
 }
 
 _field_type_map = {
@@ -75,9 +77,26 @@ _field_type_map = {
     BSON_TYPE_OID: ObjectIdType(),
     BSON_TYPE_UTF8: string(),
     BSON_TYPE_BOOL: bool_(),
-    BSON_TYPE_DECIMAL128: Decimal128StringType()
+    BSON_TYPE_DECIMAL128: Decimal128StringType(),
 }
 
+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"""
+    if value_t in _field_type_map:
+        field_type = _field_type_map[value_t]
+    elif value_t == BSON_TYPE_ARRAY:
+        bson_iter_recurse(doc_iter, child_iter)
+        list_dtype = extract_array_dtype(child_iter, context)
+        field_type = list_(list_dtype)
+    elif value_t == BSON_TYPE_DOCUMENT:
+        bson_iter_recurse(doc_iter, child_iter)
+        field_type = extract_document_dtype(child_iter, context)
+    elif value_t == BSON_TYPE_DATE_TIME:
+        field_type = timestamp('ms', tz=context.tzinfo)
+    else:
+        raise PyMongoArrowError('unknown value type {}'.format(value_t))
+    return field_type
+
 
 cdef extract_document_dtype(bson_iter_t * doc_iter, context):
     """Get the appropropriate data type for a sub document"""
@@ -88,19 +107,21 @@ cdef extract_document_dtype(bson_iter_t * doc_iter, context):
     while bson_iter_next(doc_iter):
         key = bson_iter_key(doc_iter)
         value_t = bson_iter_type(doc_iter)
-        if value_t in _field_type_map:
-            field_type = _field_type_map[value_t]
-        elif value_t == BSON_TYPE_DOCUMENT:
-            bson_iter_recurse(doc_iter, &child_iter)
-            field_type = extract_document_dtype(&child_iter, context)
-        elif value_t == BSON_TYPE_DATE_TIME:
-            field_type = timestamp('ms', tz=context.tzinfo)
-
+        field_type = extract_field_dtype(doc_iter, &child_iter, value_t, context)
         fields.append(field(key.decode('utf-8'), field_type))
     return struct(fields)
 
+cdef extract_array_dtype(bson_iter_t * doc_iter, context):
+    """Get the appropropriate data type for a sub array"""
+    cdef const char* key
+    cdef bson_type_t value_t
+    cdef bson_iter_t child_iter
+    fields = []
+    first_item = bson_iter_next(doc_iter)
+    value_t = bson_iter_type(doc_iter)
+    return extract_field_dtype(doc_iter, &child_iter, value_t, context)
 
-def process_bson_stream(bson_stream, context):
+def process_bson_stream(bson_stream, context, arr_value_builder=None):
     """Process a bson byte stream using a PyMongoArrowContext"""
     cdef const uint8_t* docstream = <const uint8_t *>bson_stream
     cdef size_t length = <size_t>PyBytes_Size(bson_stream)
@@ -110,6 +131,8 @@ def process_bson_stream(bson_stream, context):
     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 bson_decimal128_t dec128
     cdef bson_type_t value_t
     cdef const char * bson_str
@@ -131,12 +154,13 @@ def process_bson_stream(bson_stream, context):
     t_string = _BsonArrowTypes.string
     t_bool = _BsonArrowTypes.bool
     t_document = _BsonArrowTypes.document
+    t_array = _BsonArrowTypes.array
+
 
     # initialize count to current length of builders
     for _, builder in builder_map.items():
         count = len(builder)
         break
-
     try:
         while True:
             doc = bson_reader_read_safe(stream_reader)
@@ -146,7 +170,10 @@ def process_bson_stream(bson_stream, context):
                 raise InvalidBSON("Could not read BSON document")
             while bson_iter_next(&doc_iter):
                 key = bson_iter_key(&doc_iter)
-                builder = builder_map.get(key)
+                if arr_value_builder is not None:
+                    builder = arr_value_builder
+                else:
+                    builder = builder_map.get(key)
                 if builder is None:
                     builder = builder_map.get(key)
                 if builder is None and context.schema is None:
@@ -165,10 +192,15 @@ def process_bson_stream(bson_stream, context):
                         bson_iter_recurse(&doc_iter, &child_iter)
                         struct_dtype = extract_document_dtype(&child_iter, context)
                         builder = DocumentBuilder(struct_dtype, context.tzinfo)
+                    elif builder_type == ListBuilder:
+                        bson_iter_recurse(&doc_iter, &child_iter)
+                        list_dtype = extract_array_dtype(&child_iter, context)
+                        list_dtype = list_(list_dtype)
+                        builder = ListBuilder(list_dtype, context.tzinfo, value_builder=arr_value_builder)
                     else:
                         builder = builder_type()
-
-                    builder_map[key] = builder
+                    if arr_value_builder is None:
+                        builder_map[key] = builder
                     for _ in range(count):
                         builder.append_null()
 
@@ -231,6 +263,14 @@ def process_bson_stream(bson_stream, context):
                         builder.append(<bytes>doc_buf[:doc_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:
+                            raise ValueError("Subarray is invalid")
+                        builder.append(<bytes>doc_buf[:doc_buf_len])
+                    else:
+                        builder.append_null()
                 else:
                     raise PyMongoArrowError('unknown ftype {}'.format(ftype))
             count += 1
@@ -467,7 +507,11 @@ cdef class BoolBuilder(_ArrayBuilderBase):
 cdef object get_field_builder(field, tzinfo):
     """"Find the appropriate field builder given a pyarrow field"""
     cdef object field_builder
-    field_type = field.type
+    cdef DataType field_type
+    if isinstance(field, DataType):
+        field_type = field
+    else:
+        field_type = field.type
     if _atypes.is_int32(field_type):
         field_builder = Int32Builder()
     elif _atypes.is_int64(field_type):
@@ -484,6 +528,8 @@ cdef object get_field_builder(field, tzinfo):
         field_builder = BoolBuilder()
     elif _atypes.is_struct(field_type):
         field_builder = DocumentBuilder(field_type, tzinfo)
+    elif _atypes.is_list(field_type):
+        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:
@@ -549,3 +595,55 @@ cdef class DocumentBuilder(_ArrayBuilderBase):
 
     cdef shared_ptr[CStructBuilder] unwrap(self):
         return self.builder
+
+cdef class ListBuilder(_ArrayBuilderBase):
+    type_marker = _BsonArrowTypes.array
+
+    cdef:
+        shared_ptr[CListBuilder] builder
+        _ArrayBuilderBase child_builder
+        object dtype
+        object context
+
+    def __cinit__(self, DataType dtype, tzinfo=None, MemoryPool memory_pool=None, value_builder=None):
+        cdef StringBuilder field_builder
+        cdef CMemoryPool* pool = maybe_unbox_memory_pool(memory_pool)
+        cdef shared_ptr[CArrayBuilder] grandchild_builder
+        self.dtype = dtype
+        if not _atypes.is_list(dtype):
+            raise ValueError("dtype must be a list_()")
+        self.context = context = PyMongoArrowContext(None, {})
+        self.context.tzinfo = tzinfo
+        field_builder = <StringBuilder>get_field_builder(self.dtype.value_type, tzinfo)
+        grandchild_builder = <shared_ptr[CArrayBuilder]>field_builder.builder
+        self.child_builder = field_builder
+        self.builder.reset(new CListBuilder(pool, grandchild_builder, pyarrow_unwrap_data_type(dtype)))
+
+
+    @property
+    def dtype(self):
+        return self.dtype
+
+    cpdef append_null(self):
+        self.builder.get().AppendNull()
+
+    def __len__(self):
+        return self.builder.get().length()
+
+    cpdef append(self, value):
+        if not isinstance(value, bytes):
+            value = bson.encode(value)
+        # Append an element to the array.
+        # arr_value_builder will be appended to by process_bson_stream.
+        self.builder.get().Append(True)
+        process_bson_stream(value, self.context, arr_value_builder=self.child_builder)
+
+
+    cpdef finish(self):
+        cdef shared_ptr[CArray] out
+        with nogil:
+            self.builder.get().Finish(&out)
+        return pyarrow_wrap_array(out)
+
+    cdef shared_ptr[CListBuilder] unwrap(self):
+        return self.builder

bindings/python/pymongoarrow/libarrow.pxd

@@ -37,6 +37,14 @@ cdef extern from "arrow/builder.h" namespace "arrow" nogil:
         int32_t num_fields()
         shared_ptr[CDataType] type()
 
+    cdef cppclass CListBuilder" arrow::ListBuilder"(CArrayBuilder):
+        CListBuilder(CMemoryPool* pool,
+                       shared_ptr[CArrayBuilder] value_builder, shared_ptr[CDataType] dtype)
+        CStatus Append(uint8_t is_valid)
+        CArrayBuilder* value_builder(int i)
+        int32_t num_values()
+        shared_ptr[CDataType] type()
+
 
 cdef extern from "arrow/type_fwd.h" namespace "arrow" nogil:
     shared_ptr[CDataType] fixed_size_binary(int32_t byte_width)

bindings/python/pymongoarrow/libbson.pxd

@@ -133,6 +133,11 @@ cdef extern from "<bson/bson.h>":
                     uint32_t *document_len,  #  OUT
                     const uint8_t **document) # OUT
 
+    void bson_iter_array (const bson_iter_t *iter, # IN
+                 uint32_t *array_len,     # OUT
+                 const uint8_t **array)
+
+
 # bson_reader_t API
 cdef extern from "<bson/bson.h>":
     bson_reader_t * bson_reader_new_from_data(const uint8_t *data, size_t length)

bindings/python/pymongoarrow/types.py

@@ -26,6 +26,7 @@
     bool_,
     float64,
     int64,
+    list_,
     string,
     struct,
     timestamp,
@@ -42,6 +43,7 @@ class _BsonArrowTypes(enum.Enum):
     bool = 7
     decimal128_str = 8
     document = 9
+    array = 10
 
 
 # Custom Extension Types.
@@ -141,6 +143,7 @@ def get_numpy_type(type):
     _atypes.is_string: _BsonArrowTypes.string,
     _atypes.is_boolean: _BsonArrowTypes.bool,
     _atypes.is_struct: _BsonArrowTypes.document,
+    _atypes.is_list: _BsonArrowTypes.array,
 }
 
 
@@ -156,6 +159,8 @@ def _normalize_typeid(typeid, field_name):
         for sub_field_name, sub_typeid in typeid.items():
             fields.append((sub_field_name, _normalize_typeid(sub_typeid, sub_field_name)))
         return struct(fields)
+    elif isinstance(typeid, list):
+        return list_(_normalize_typeid(type(typeid[0]), "0"))
     elif _is_typeid_supported(typeid):
         normalizer = _TYPE_NORMALIZER_FACTORY[typeid]
         return normalizer(typeid)