ARROW-53 Add Support for Boolean Type (#51)

Author: juliusgeo

bindings/python/docs/source/changelog.rst

@@ -6,6 +6,7 @@ Changes in Version 0.3.0
 - Support for `ObjectId` `bson` type.
 - Improve error message when schema contains an unsupported type.
 - Add support for BSON string type.
+- Add support for BSON boolean type.
 
 Changes in Version 0.2.0
 ------------------------

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`
    * - ObjectId
      - :class:`py.bytes`, :class:`bson.ObjectId`, an instance of :class:`pymongoarrow.types.ObjectIdType`, an instance of :class:`pyarrow.FixedSizeBinaryScalar`
+   * - Boolean
+     - an instance of :class:`~pyarrow.bool_`, :class:`~py.bool`
    * - 64-bit binary floating point
      - :class:`py.float`, an instance of :meth:`pyarrow.float64`
    * - 32-bit integer

bindings/python/pymongoarrow/context.py

@@ -14,7 +14,9 @@
 from bson.codec_options import DEFAULT_CODEC_OPTIONS
 
 from pyarrow import timestamp, Table
-from pymongoarrow.lib import Int32Builder, Int64Builder, DoubleBuilder, DatetimeBuilder, ObjectIdBuilder, StringBuilder
+from pymongoarrow.lib import (Int32Builder, Int64Builder, DoubleBuilder,
+                              DatetimeBuilder, ObjectIdBuilder,
+                              StringBuilder, BoolBuilder)
 from pymongoarrow.types import _get_internal_typemap, _BsonArrowTypes
 
 
@@ -25,6 +27,7 @@
     _BsonArrowTypes.datetime: DatetimeBuilder,
     _BsonArrowTypes.objectid: ObjectIdBuilder,
     _BsonArrowTypes.string: StringBuilder,
+    _BsonArrowTypes.bool: BoolBuilder,
 }
 
 

bindings/python/pymongoarrow/lib.pyx

@@ -72,6 +72,7 @@ def process_bson_stream(bson_stream, context):
     t_datetime = _BsonArrowTypes.datetime
     t_oid = _BsonArrowTypes.objectid
     t_string = _BsonArrowTypes.string
+    t_bool = _BsonArrowTypes.bool
     builder_map = context.builder_map
 
     # initialize count to current length of builders
@@ -129,6 +130,11 @@ def process_bson_stream(bson_stream, context):
                             builder.append(bson_iter_date_time(&doc_iter))
                         else:
                             builder.append_null()
+                    elif ftype == t_bool:
+                        if value_t == BSON_TYPE_BOOL:
+                            builder.append(bson_iter_bool(&doc_iter))
+                        else:
+                            builder.append_null()
                     else:
                         raise PyMongoArrowError('unknown ftype {}'.format(ftype))
             count += 1
@@ -330,3 +336,31 @@ cdef class DatetimeBuilder(_ArrayBuilderBase):
 
     cdef shared_ptr[CTimestampBuilder] unwrap(self):
         return self.builder
+
+
+cdef class BoolBuilder(_ArrayBuilderBase):
+    type_marker = _BsonArrowTypes.bool
+    cdef:
+        shared_ptr[CBooleanBuilder] builder
+
+    def __cinit__(self, MemoryPool memory_pool=None):
+        cdef CMemoryPool* pool = maybe_unbox_memory_pool(memory_pool)
+        self.builder.reset(new CBooleanBuilder(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(<c_bool>value)
+
+    cpdef finish(self):
+        cdef shared_ptr[CArray] out
+        with nogil:
+            self.builder.get().Finish(&out)
+        return pyarrow_wrap_array(out)
+
+    cdef shared_ptr[CBooleanBuilder] unwrap(self):
+        return self.builder

bindings/python/pymongoarrow/types.py

@@ -16,7 +16,7 @@
 
 from bson import Int64, ObjectId
 
-from pyarrow import timestamp, binary, float64, int64, int32, string
+from pyarrow import timestamp, binary, float64, int64, int32, string, bool_
 from pyarrow import PyExtensionType
 from pyarrow import DataType as _ArrowDataType
 import pyarrow.types as _atypes
@@ -29,6 +29,7 @@ class _BsonArrowTypes(enum.Enum):
     int64 = 4
     objectid = 5
     string = 6
+    bool = 7
 
 
 # Custom Extension Types.
@@ -59,6 +60,7 @@ def _is_objectid(obj):
     datetime: lambda _: timestamp('ms'),     # TODO: add tzinfo support
     ObjectId: lambda _: ObjectIdType(),
     str: lambda: string(),
+    bool: lambda: bool_(),
 }
 
 
@@ -69,6 +71,7 @@ def _is_objectid(obj):
     _atypes.is_timestamp: _BsonArrowTypes.datetime,
     _is_objectid: _BsonArrowTypes.objectid,
     _atypes.is_string: _BsonArrowTypes.string,
+    _atypes.is_boolean: _BsonArrowTypes.bool,
 }
 
 

bindings/python/test/test_bson.py

@@ -19,7 +19,7 @@
 from pymongoarrow.context import PyMongoArrowContext
 from pymongoarrow.lib import process_bson_stream
 from pymongoarrow.schema import Schema
-from pymongoarrow.types import int32, int64, string, ObjectId, ObjectIdType
+from pymongoarrow.types import int32, int64, string, ObjectId, ObjectIdType, bool_
 
 
 class TestBsonToArrowConversionBase(TestCase):
@@ -154,3 +154,22 @@ def test_object_id_type(self):
         arr = pa.ExtensionArray.from_storage(ObjectIdType(), storage_array)
         result = self.serialize_array(arr)
         assert result.type._type_marker == ObjectIdType._type_marker
+
+
+class TestBooleanType(TestBsonToArrowConversionBase):
+    def setUp(self):
+        self.schema = Schema({'data': bool_()})
+        self.context = PyMongoArrowContext.from_schema(
+            self.schema)
+
+    def test_simple(self):
+        docs = [{'data': True},
+                {'data': False},
+                {'data': 19},
+                {'data': "string"},
+                {'data': False},
+                {'data': True}]
+        as_dict = {
+            'data': [True, False, None, None, False, True]
+        }
+        self._run_test(docs, as_dict)

bindings/python/test/test_builders.py

@@ -16,11 +16,11 @@
 from unittest import TestCase
 from bson.objectid import ObjectId
 
-from pyarrow import Array, timestamp, int32, int64
+from pyarrow import Array, timestamp, int32, int64, bool_
 
 from pymongoarrow.lib import (
     DatetimeBuilder, DoubleBuilder, Int32Builder, Int64Builder,
-    ObjectIdBuilder, StringBuilder)
+    ObjectIdBuilder, StringBuilder, BoolBuilder)
 
 
 class TestIntBuildersMixin:
@@ -146,3 +146,26 @@ def test_simple(self):
         self.assertEqual(len(arr), 5)
         self.assertEqual(
             arr.to_pylist(), values + [None])
+
+
+class TestBoolBuilderMixin:
+    def test_simple(self):
+        builder = BoolBuilder()
+        builder.append(False)
+        builder.append_values([True, False, True, False, True, False])
+        builder.append_null()
+        arr = builder.finish()
+
+        self.assertIsInstance(arr, Array)
+        self.assertEqual(arr.null_count, 1)
+        self.assertEqual(len(arr), 8)
+        self.assertEqual(
+            arr.to_pylist(), [False, True, False, True, False, True, False,
+                              None])
+        self.assertEqual(arr.type, self.data_type)
+
+
+class TestBoolBuilder(TestCase, TestBoolBuilderMixin):
+    def setUp(self):
+        self.builder_cls = BoolBuilder
+        self.data_type = bool_()