further progress on buildermanager and context

Author: blink1073

bindings/python/pymongoarrow/context.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from pyarrow import ListArray, StructArray, Table
-from pyarrow.types import is_struct
 
 from pymongoarrow.types import _BsonArrowTypes, _get_internal_typemap
 
@@ -47,45 +46,45 @@ def process_bson_stream(self, stream):
         self.manager.process_bson_stream(stream, len(stream))
 
     def finish(self):
-        array_map = _parse_array_map(self.manager.finish())
+        array_map = _parse_builder_map(self.manager.finish())
         arrays = list(array_map.values())
         if self.schema is not None:
             return Table.from_arrays(arrays=arrays, schema=self.schema.to_arrow())
         return Table.from_arrays(arrays=arrays, names=list(array_map.keys()))
 
 
-def _parse_array_map(array_map):
+def _parse_builder_map(builder_map):
     # Handle nested builders.
     to_remove = []
     # Traverse the builder map right to left.
-    for key, value in reversed(array_map.items()):
-        field = key.decode("utf-8")
-        if value.type_marker == _BsonArrowTypes.document:
-            full_names = [f"{field}.{name.decode('utf-8')}" for name in value]
-            arrs = [array_map[c.encode("utf-8")] for c in full_names]
-            array_map[field] = StructArray.from_arrays(arrs, names=value)
+    for key, value in reversed(builder_map.items()):
+        if value.type_marker == _BsonArrowTypes.document.value:
+            names = value.finish()
+            full_names = [f"{key}.{name}" for name in names]
+            arrs = [builder_map[c] for c in full_names]
+            builder_map[key] = StructArray.from_arrays(arrs, names=names)
             to_remove.extend(full_names)
-        elif value.type_marker == _BsonArrowTypes.array:
-            child_name = field + "[]"
+        elif value.type_marker == _BsonArrowTypes.array.value:
+            child_name = key + "[]"
             to_remove.append(child_name)
-            child = array_map[child_name.encode("utf-8")]
-            array_map[key] = ListArray.from_arrays(value, child)
+            child = builder_map[child_name]
+            builder_map[key] = ListArray.from_arrays(value.finish(), child)
+        else:
+            builder_map[key] = value.finish()
 
-    for field in to_remove:
-        key = field.encode("utf-8")
-        if key in array_map:
-            del array_map[key]
+    for key in to_remove:
+        if key in builder_map:
+            del builder_map[key]
 
-    return array_map
+    return builder_map
 
 
 def _parse_types(str_type_map, schema_map, tzinfo):
     for fname, (ftype, arrow_type) in str_type_map.items():
-        encoded_fname = fname.encode("utf-8")
-        schema_map[encoded_fname] = ftype, arrow_type
+        schema_map[fname] = ftype, arrow_type
 
         # special-case nested builders
-        if ftype == _BsonArrowTypes.document:
+        if ftype == _BsonArrowTypes.document.value:
             # construct a sub type map here
             sub_type_map = {}
             for i in range(arrow_type.num_fields):
@@ -94,13 +93,10 @@ def _parse_types(str_type_map, schema_map, tzinfo):
                 sub_type_map[sub_name] = field.type
             sub_type_map = _get_internal_typemap(sub_type_map)
             _parse_types(sub_type_map, schema_map, tzinfo)
-        elif ftype == _BsonArrowTypes.array:
-            if is_struct(arrow_type.value_type):
-                # construct a sub type map here
-                sub_type_map = {}
-                for i in range(arrow_type.value_type.num_fields):
-                    field = arrow_type.value_type[i]
-                    sub_name = f"{fname}[].{field.name}"
-                    sub_type_map[sub_name] = field.type
-                sub_type_map = _get_internal_typemap(sub_type_map)
-                _parse_types(sub_type_map, schema_map, tzinfo)
+        elif ftype == _BsonArrowTypes.array.value:
+            sub_type_map = {}
+            sub_name = f"{fname}[]"
+            sub_value_type = arrow_type.value_type
+            sub_type_map[sub_name] = sub_value_type
+            sub_type_map = _get_internal_typemap(sub_type_map)
+            _parse_types(sub_type_map, schema_map, tzinfo)

bindings/python/pymongoarrow/lib.pyx

@@ -67,22 +67,22 @@ cdef class BuilderManager:
         self.has_schema = has_schema
         self.tzinfo = tzinfo
         self.count = 0
-        self.builder_map = builder_map = {}
+        self.builder_map = {}
         # Unpack the schema map.
         for fname, (ftype, arrow_type) in schema_map.items():
-            encoded_fname = fname.encode("utf-8")
+            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
-                builder_map[encoded_fname] = DatetimeBuilder(dtype=arrow_type)
+                self.builder_map[name] = DatetimeBuilder(dtype=arrow_type)
             elif ftype == BSON_TYPE_BINARY:
-                subtype = arrow_type.subtype
-                builder_map[encoded_fname] = BinaryBuilder(subtype)
+                self.builder_map[name] = BinaryBuilder(arrow_type.subtype)
             else:
-                self.get_builder(encoded_fname, ftype, <bson_iter_t *>nullptr)
+                # We only use the doc_iter for binary arrays, which are handled already.
+                self.get_builder(name, ftype, <bson_iter_t *>nullptr)
 
-    cdef _ArrayBuilderBase get_builder(self, cstring key, bson_type_t value_t, bson_iter_t * doc_iter):
+    cdef _ArrayBuilderBase get_builder(self, cstring key, bson_type_t value_t, bson_iter_t * doc_iter) except *:
         cdef _ArrayBuilderBase builder = None
         cdef bson_subtype_t subtype
         cdef const uint8_t *val_buf = NULL
@@ -108,6 +108,8 @@ cdef class BuilderManager:
         elif value_t == BSON_TYPE_ARRAY:
             builder = ListBuilder()
         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)
@@ -147,7 +149,6 @@ cdef class BuilderManager:
             # Get the key and and value.
             key = bson_iter_key(doc_iter)
             value_t = bson_iter_type(doc_iter)
-            print('handling', key, value_t)
 
             # Get the appropriate full key.
             if parent_type == BSON_TYPE_ARRAY:
@@ -163,6 +164,8 @@ cdef class BuilderManager:
             else:
                 full_key = key
 
+            print('handling', full_key, value_t)
+
             # Get the builder.
             builder = <_ArrayBuilderBase>self.builder_map.get(full_key, None)
             if builder is None and not self.has_schema:
@@ -193,9 +196,9 @@ cdef class BuilderManager:
             if parent_type == BSON_TYPE_ARRAY:
                 (<ListBuilder>self.builder_map[base_key]).append_count()
 
-            # Update our count.
-            if builder.length() > self.count:
-                self.count = builder.length()
+        # Update our count for top level documents.
+        if parent_type == 0:
+            self.count += 1
 
     cpdef void process_bson_stream(self, const uint8_t* bson_stream, size_t length):
         """Process a bson byte stream."""
@@ -214,35 +217,35 @@ cdef class BuilderManager:
                 bson_reader_destroy(stream_reader)
 
     cpdef finish(self):
-        """Finish building the arrays."""
-        cdef dict builder_map = self.builder_map
-        cdef dict array_map = {}
+        """Finish appending to the builders."""
+        cdef dict return_map = {}
         cdef bytes key
         cdef str field
         cdef _ArrayBuilderBase value
 
+        # Move the builders to a new dict with string keys.
+        for key, value in self.builder_map.items():
+            return_map[key.decode('utf-8')] = value
+
         # Insert null fields.
-        for key in list(builder_map):
-            if builder_map[key] is None:
-                builder_map[key] = NullBuilder(self.count)
+        for field in list(return_map):
+            if return_map[field] is None:
+                return_map[field] = NullBuilder(self.count)
 
         # Pad fields as needed.
-        for key, value in builder_map.items():
-            field = key.decode("utf-8")
-
+        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())
 
-            array_map[field] = value.finish()
-        return array_map
+        return return_map
 
 
 cdef class _ArrayBuilderBase:
     cdef:
-        uint8_t type_marker
+        public uint8_t type_marker
 
     def append_values(self, values):
         for value in values:
@@ -656,7 +659,8 @@ cdef class DocumentBuilder(_ArrayBuilderBase):
         self.field_map[field_name] = 1
 
     def finish(self):
-        return set((f.decode('utf-8') for f in self.field_map))
+        # Fields must be in order if we were given a schema.
+        return list(f.decode('utf-8') for f in self.field_map)
 
 
 cdef class ListBuilder(_ArrayBuilderBase):

bindings/python/pymongoarrow/types.py

@@ -260,23 +260,23 @@ def get_numpy_type(type):
 
 
 _TYPE_CHECKER_TO_INTERNAL_TYPE = {
-    _atypes.is_int32: _BsonArrowTypes.int32,
-    _atypes.is_int64: _BsonArrowTypes.int64,
-    _atypes.is_float64: _BsonArrowTypes.double,
-    _atypes.is_timestamp: _BsonArrowTypes.datetime,
-    _atypes.is_null: _BsonArrowTypes.null,
-    _is_objectid: _BsonArrowTypes.objectid,
-    _is_decimal128: _BsonArrowTypes.decimal128,
-    _is_binary: _BsonArrowTypes.binary,
-    _is_code: _BsonArrowTypes.code,
-    _atypes.is_string: _BsonArrowTypes.string,
-    _atypes.is_boolean: _BsonArrowTypes.bool,
-    _atypes.is_struct: _BsonArrowTypes.document,
-    _atypes.is_list: _BsonArrowTypes.array,
-    _atypes.is_date32: _BsonArrowTypes.date32,
-    _atypes.is_date64: _BsonArrowTypes.date64,
-    _atypes.is_large_string: _BsonArrowTypes.string,
-    _atypes.is_large_list: _BsonArrowTypes.array,
+    _atypes.is_int32: _BsonArrowTypes.int32.value,
+    _atypes.is_int64: _BsonArrowTypes.int64.value,
+    _atypes.is_float64: _BsonArrowTypes.double.value,
+    _atypes.is_timestamp: _BsonArrowTypes.datetime.value,
+    _atypes.is_null: _BsonArrowTypes.null.value,
+    _is_objectid: _BsonArrowTypes.objectid.value,
+    _is_decimal128: _BsonArrowTypes.decimal128.value,
+    _is_binary: _BsonArrowTypes.binary.value,
+    _is_code: _BsonArrowTypes.code.value,
+    _atypes.is_string: _BsonArrowTypes.string.value,
+    _atypes.is_boolean: _BsonArrowTypes.bool.value,
+    _atypes.is_struct: _BsonArrowTypes.document.value,
+    _atypes.is_list: _BsonArrowTypes.array.value,
+    _atypes.is_date32: _BsonArrowTypes.date32.value,
+    _atypes.is_date64: _BsonArrowTypes.date64.value,
+    _atypes.is_large_string: _BsonArrowTypes.string.value,
+    _atypes.is_large_list: _BsonArrowTypes.array.value,
 }
 
 

bindings/python/test/test_arrow.py

@@ -23,8 +23,6 @@
 import pymongo
 from bson import Binary, Code, CodecOptions, Decimal128, ObjectId
 from pyarrow import (
-    DataType,
-    FixedSizeBinaryType,
     Table,
     bool_,
     csv,
@@ -405,7 +403,7 @@ def inner(i):
             if use_none:
                 inner_dict["null"] = None
             if nested_elem:
-                inner_dict["list"] = [nested_elem]
+                inner_dict["list_of_objs"] = [nested_elem]
             return inner_dict
 
         if nested_elem:
@@ -480,6 +478,18 @@ def test_auto_schema_nested(self):
             for name in out.column_names:
                 self.assertEqual(data[name], out[name].cast(data[name].type))
 
+    def test_auto_schema_nested_null(self):
+        # Create table with random data of various types.
+        _, data = self._create_nested_data(use_none=True)
+
+        self.coll.drop()
+        res = write(self.coll, data)
+        self.assertEqual(len(data), res.raw_result["insertedCount"])
+        for func in [find_arrow_all, aggregate_arrow_all]:
+            out = func(self.coll, {} if func == find_arrow_all else []).drop(["_id"])
+            for name in out.column_names:
+                self.assertEqual(data[name], out[name].cast(data[name].type))
+
     def test_schema_nested_null(self):
         schema, data = self._create_nested_data(use_none=True)
 
@@ -514,7 +524,13 @@ def test_auto_schema_first_list_null(self):
             {"a": ["str"]},
             {"a": []},
         ]
-        expected = pa.Table.from_pylist(docs)
+        expected = pa.Table.from_pylist(
+            [
+                {"a": []},
+                {"a": ["str"]},
+                {"a": []},
+            ]
+        )
         self._test_auto_schema_list(docs, expected)
 
     def test_auto_schema_first_list_empty(self):
@@ -525,7 +541,7 @@ def test_auto_schema_first_list_empty(self):
         ]
         expected = pa.Table.from_pylist(
             [
-                {"a": None},  # TODO: We incorrectly set the first empty list to null.
+                {"a": []},
                 {"a": ["str"]},
                 {"a": []},
             ]
@@ -538,20 +554,30 @@ def test_auto_schema_first_list_element_null(self):
             {"a": [None, None, "str"]},  # Inferred schema should use the first non-null element.
             {"a": []},
         ]
-        expected = pa.Table.from_pylist(docs)
+        expected = pa.Table.from_pylist(
+            [
+                {"a": []},
+                {"a": ["str"]},  # Inferred schema should use the first non-null element.
+                {"a": []},
+            ]
+        )
         self._test_auto_schema_list(docs, expected)
 
-    @unittest.expectedFailure  # TODO: Our inferred value for the first a.b field differs from pyarrow's.
     def test_auto_schema_first_embedded_list_null(self):
         docs = [
             {"a": {"b": None}},
             {"a": {"b": ["str"]}},
             {"a": {"b": []}},
         ]
-        expected = pa.Table.from_pylist(docs)
+        expected = pa.Table.from_pylist(
+            [
+                {"a": {"b": []}},
+                {"a": {"b": ["str"]}},
+                {"a": {"b": []}},
+            ]
+        )
         self._test_auto_schema_list(docs, expected)
 
-    @unittest.expectedFailure  # TODO: Our inferred value for the first a.b field differs from pyarrow's.
     def test_auto_schema_first_embedded_doc_null(self):
         docs = [
             {"a": {"b": None}},
@@ -747,18 +773,10 @@ def test_nested_bson_extension_types(self):
         out = find_arrow_all(self.coll, {})
         obj_schema_type = out.field("obj").type
 
-        self.assertIsInstance(obj_schema_type.field("obj_id").type, FixedSizeBinaryType)
-        self.assertIsInstance(obj_schema_type.field("dec_128").type, FixedSizeBinaryType)
-        self.assertIsInstance(obj_schema_type.field("binary").type, DataType)
-        self.assertIsInstance(obj_schema_type.field("code").type, DataType)
-
-        new_types = [ObjectIdType(), Decimal128Type(), BinaryType(0), CodeType()]
-        new_names = [f.name for f in out["obj"].type]
-        new_obj = out["obj"].cast(struct(zip(new_names, new_types)))
-        self.assertIsInstance(new_obj.type[0].type, ObjectIdType)
-        self.assertIsInstance(new_obj.type[1].type, Decimal128Type)
-        self.assertIsInstance(new_obj.type[2].type, BinaryType)
-        self.assertIsInstance(new_obj.type[3].type, CodeType)
+        self.assertIsInstance(obj_schema_type.field("obj_id").type, ObjectIdType)
+        self.assertIsInstance(obj_schema_type.field("dec_128").type, Decimal128Type)
+        self.assertIsInstance(obj_schema_type.field("binary").type, BinaryType)
+        self.assertIsInstance(obj_schema_type.field("code").type, CodeType)
 
     def test_large_string_type(self):
         """Tests pyarrow._large_string() DataType"""