handle status and switch to shared mem pool

Author: blink1073

bindings/python/pymongoarrow/lib.pyx

@@ -22,7 +22,7 @@ import sys
 # Python imports
 import bson
 import numpy as np
-from pyarrow import timestamp
+from pyarrow import timestamp, default_memory_pool
 
 from pymongoarrow.errors import InvalidBSON
 from pymongoarrow.types import ObjectIdType, Decimal128Type as Decimal128Type_, BinaryType, CodeType
@@ -62,22 +62,24 @@ cdef class BuilderManager:
         uint64_t count
         bint has_schema
         object tzinfo
+        object pool
 
     def __cinit__(self, dict schema_map, bint has_schema, object tzinfo):
         self.has_schema = has_schema
         self.tzinfo = tzinfo
         self.count = 0
         self.builder_map = {}
+        self.pool = default_memory_pool()
         # Unpack the schema map.
         for fname, (ftype, arrow_type) in schema_map.items():
             name = fname.encode('utf-8')
             # special-case initializing builders for parameterized types
             if ftype == BSON_TYPE_DATE_TIME:
                 if tzinfo is not None and arrow_type.tz is None:
                     arrow_type = timestamp(arrow_type.unit, tz=tzinfo)  # noqa: PLW2901
-                self.builder_map[name] = DatetimeBuilder(dtype=arrow_type)
+                self.builder_map[name] = DatetimeBuilder(dtype=arrow_type, memory_pool=self.pool)
             elif ftype == BSON_TYPE_BINARY:
-                self.builder_map[name] = BinaryBuilder(arrow_type.subtype)
+                self.builder_map[name] = BinaryBuilder(arrow_type.subtype, memory_pool=self.pool)
             else:
                 # We only use the doc_iter for binary arrays, which are handled already.
                 self.get_builder(name, ftype, <bson_iter_t *>nullptr)
@@ -100,39 +102,39 @@ cdef class BuilderManager:
         if value_t == BSON_TYPE_DATE_TIME:
             if self.tzinfo is not None:
                 arrow_type = timestamp('ms', tz=self.tzinfo)
-                builder = DatetimeBuilder(dtype=arrow_type)
+                builder = DatetimeBuilder(dtype=arrow_type, memory_pool=self.pool)
             else:
-                builder = DatetimeBuilder()
+                builder = DatetimeBuilder(memory_pool=self.pool)
         elif value_t == BSON_TYPE_DOCUMENT:
             builder = DocumentBuilder()
         elif value_t == BSON_TYPE_ARRAY:
-            builder = ListBuilder()
+            builder = ListBuilder(memory_pool=self.pool)
         elif value_t == BSON_TYPE_BINARY:
             if doc_iter == NULL:
                 raise ValueError('Did not pass a doc_iter!')
             bson_iter_binary (doc_iter, &subtype,
                 &val_buf_len, &val_buf)
-            builder = BinaryBuilder(subtype)
+            builder = BinaryBuilder(subtype, memory_pool=self.pool)
         elif value_t == ARROW_TYPE_DATE32:
-            builder = Date32Builder()
+            builder = Date32Builder(memory_pool=self.pool)
         elif value_t == ARROW_TYPE_DATE64:
-            builder = Date64Builder()
+            builder = Date64Builder(memory_pool=self.pool)
         elif value_t == BSON_TYPE_INT32:
-            builder = Int32Builder()
+            builder = Int32Builder(memory_pool=self.pool)
         elif value_t == BSON_TYPE_INT64:
-            builder = Int64Builder()
+            builder = Int64Builder(memory_pool=self.pool)
         elif value_t == BSON_TYPE_DOUBLE:
-            builder = DoubleBuilder()
+            builder = DoubleBuilder(memory_pool=self.pool)
         elif value_t == BSON_TYPE_OID:
-            builder = ObjectIdBuilder()
+            builder = ObjectIdBuilder(memory_pool=self.pool)
         elif value_t == BSON_TYPE_UTF8:
-            builder = StringBuilder()
+            builder = StringBuilder(memory_pool=self.pool)
         elif value_t == BSON_TYPE_BOOL:
-            builder = BoolBuilder()
+            builder = BoolBuilder(memory_pool=self.pool)
         elif value_t == BSON_TYPE_DECIMAL128:
-            builder = Decimal128Builder()
+            builder = Decimal128Builder(memory_pool=self.pool)
         elif value_t == BSON_TYPE_CODE:
-            builder = CodeBuilder()
+            builder = CodeBuilder(memory_pool=self.pool)
 
         self.builder_map[key] = builder
         return builder
@@ -175,12 +177,14 @@ cdef class BuilderManager:
             # For lists, the nulls are stored in the parent.
             if parent_type != BSON_TYPE_ARRAY:
                 if count > builder.length():
-                    builder.append_nulls(count - builder.length())
+                    status = builder.append_nulls_raw(count - builder.length())
+                    if not status.ok():
+                        raise ValueError("Failed to append nulls to", full_key.decode('utf8'))
 
             # Append the next value.
             status = builder.append_raw(doc_iter, value_t)
             if not status.ok():
-                raise ValueError("Could not append raw value")
+                raise ValueError("Could not append raw value to", full_key.decode('utf8'))
 
             # Recurse into documents.
             if value_t == BSON_TYPE_DOCUMENT:
@@ -218,6 +222,7 @@ cdef class BuilderManager:
         cdef dict return_map = {}
         cdef bytes key
         cdef str field
+        cdef CStatus status
         cdef _ArrayBuilderBase value
 
         # Move the builders to a new dict with string keys.
@@ -227,15 +232,17 @@ cdef class BuilderManager:
         # Insert null fields.
         for field in list(return_map):
             if return_map[field] is None:
-                return_map[field] = NullBuilder(self.count)
+                return_map[field] = NullBuilder(memory_pool=self.pool)
 
         # Pad fields as needed.
         for field, value in return_map.items():
             # If it isn't a list item, append nulls as needed.
             # For lists, the nulls are stored in the parent.
             if not field.endswith('[]'):
                 if value.length() < self.count:
-                    value.append_nulls(self.count - value.length())
+                    status = value.append_nulls_raw(self.count - value.length())
+                    if not status.ok():
+                        raise ValueError("Failed to append nulls to", field)
 
         return return_map
 
@@ -281,21 +288,36 @@ cdef class _ArrayBuilderBase:
     def __len__(self):
         return self.length()
 
-    cpdef void append_null(self):
-        self.get_builder().get().AppendNull()
+    cpdef append_null(self):
+        cdef CStatus status = self.append_null_raw()
+        if not status.ok():
+            raise ValueError("Could not append null value")
 
     cpdef void append_nulls(self, uint64_t count):
         for _ in range(count):
             self.append_null()
 
+    cdef CStatus append_null_raw(self):
+        return self.get_builder().get().AppendNull()
+
+    cdef CStatus append_nulls_raw(self, uint64_t count):
+        cdef CStatus status
+        for _ in range(count):
+            status = self.append_null_raw()
+            if not status.ok():
+                return status
+
     cpdef uint64_t length(self):
         return self.get_builder().get().length()
 
     def finish(self):
         cdef shared_ptr[CArray] out
+        cdef CStatus status
         cdef shared_ptr[CArrayBuilder] builder = self.get_builder()
         with nogil:
-            builder.get().Finish(&out)
+            status = builder.get().Finish(&out)
+        if not status.ok():
+            raise ValueError("Failed to convert value to array")
         return pyarrow_wrap_array(out)
 
 
@@ -526,13 +548,10 @@ cdef class Date32Builder(_ArrayBuilderBase):
 cdef class NullBuilder(_ArrayBuilderBase):
     cdef shared_ptr[CArrayBuilder] builder
 
-    def __cinit__(self, uint64_t count, MemoryPool memory_pool=None):
+    def __cinit__(self, MemoryPool memory_pool=None):
         cdef CMemoryPool* pool = maybe_unbox_memory_pool(memory_pool)
-        cdef uint64_t i
         self.builder.reset(new CNullBuilder(pool))
         self.type_marker = ARROW_TYPE_NULL
-        for i in range(count):
-            self.append_null()
 
     cdef CStatus append_raw(self, bson_iter_t * doc_iter, bson_type_t value_t):
         return self.builder.get().AppendNull()
@@ -596,11 +615,12 @@ cdef class Decimal128Builder(_ArrayBuilderBase):
 cdef class BinaryBuilder(_ArrayBuilderBase):
     cdef:
         uint8_t _subtype
-        shared_ptr[CBinaryBuilder] builder
+        shared_ptr[CStringBuilder] builder
 
-    def __cinit__(self, uint8_t subtype):
+    def __cinit__(self, uint8_t subtype, MemoryPool memory_pool=None):
+        cdef CMemoryPool* pool = maybe_unbox_memory_pool(memory_pool)
         self._subtype = subtype
-        self.builder.reset(new CBinaryBuilder())
+        self.builder.reset(new CStringBuilder(pool))
         self.type_marker = BSON_TYPE_BINARY
 
     @property
@@ -630,7 +650,7 @@ cdef class DocumentBuilder(_ArrayBuilderBase):
     """The document builder stores a map of field names that can be retrieved as a set."""
     cdef:
         dict field_map
-        int32_t count
+        int64_t count
 
     def __cinit__(self):
         self.type_marker = BSON_TYPE_DOCUMENT
@@ -643,8 +663,9 @@ cdef class DocumentBuilder(_ArrayBuilderBase):
     cpdef uint64_t length(self):
         return self.count
 
-    cpdef void append_null(self):
+    cdef CStatus append_null_raw(self):
         self.count += 1
+        return CStatus_OK()
 
     cpdef void add_field(self, cstring field_name):
         self.field_map[field_name] = 1
@@ -657,7 +678,7 @@ cdef class DocumentBuilder(_ArrayBuilderBase):
 cdef class ListBuilder(_ArrayBuilderBase):
     """The list builder stores an int32 list of offsets and a counter with the current value."""
     cdef:
-        int32_t count
+        int64_t count
         shared_ptr[CInt32Builder] builder
 
     def __cinit__(self, MemoryPool memory_pool=None):
@@ -672,8 +693,8 @@ cdef class ListBuilder(_ArrayBuilderBase):
     cpdef void append_count(self):
         self.count += 1
 
-    cpdef void append_null(self):
-        self.builder.get().Append(self.count)
+    cdef CStatus append_null_raw(self):
+        return self.builder.get().Append(self.count)
 
     cdef shared_ptr[CArrayBuilder] get_builder(self):
         return <shared_ptr[CArrayBuilder]>self.builder