ARROW-181 Revert nested document behavior and document limitation (#173)

* Revert "ARROW-176 Nested extension objects are not handled in auto schema (#166)"

This reverts commit 1cde144.

* ARROW-181 Revert nested document behavior and document limitation

Author: blink1073

bindings/python/docs/source/data_types.rst

@@ -162,3 +162,13 @@ with the code below (taken from `here <https://arrow.apache.org/docs/python/pand
    ... del arrow_table
 
 Defining a conversion for `pa.string()` in addition converts Arrow strings to NumPy strings, and not objects.
+
+Nested Extension Types
+----------------------
+
+Pending ``ARROW-179``, extension types such as ``ObjectId`` that appear in nested documents will not
+be converted to the corresponding PyMongoArrow extension type, but will
+instead have the raw Arrow type (``FixedSizeBinaryType(fixed_size_binary[12])``).
+
+These values can either be consumed as-is or converted individually to the
+desired extension type, e.g. ``_id = out['nested'][0]['_id'].cast(ObjectIdType())``.

bindings/python/pymongoarrow/lib.pyx

@@ -490,6 +490,7 @@ cdef class ObjectIdBuilder(_ArrayBuilderBase):
     cdef shared_ptr[CFixedSizeBinaryBuilder] unwrap(self):
         return self.builder
 
+
 cdef class Int32Builder(_ArrayBuilderBase):
     cdef:
         shared_ptr[CInt32Builder] builder
@@ -721,8 +722,6 @@ cdef object get_field_builder(object field, object tzinfo):
         field_builder = Decimal128Builder()
     elif getattr(field_type, '_type_marker') == _BsonArrowTypes.binary:
         field_builder = BinaryBuilder(field_type.subtype)
-    elif getattr(field_type, '_type_marker') == _BsonArrowTypes.code:
-        field_builder = CodeBuilder()
     else:
         field_builder = StringBuilder()
     return field_builder
@@ -733,7 +732,6 @@ cdef class DocumentBuilder(_ArrayBuilderBase):
         shared_ptr[CStructBuilder] builder
         object dtype
         object context
-        object builder_map
 
     def __cinit__(self, StructType dtype, tzinfo=None, MemoryPool memory_pool=None):
         cdef StringBuilder field_builder
@@ -746,11 +744,11 @@ cdef class DocumentBuilder(_ArrayBuilderBase):
 
         self.context = context = PyMongoArrowContext(None, {})
         context.tzinfo = tzinfo
-        self.builder_map = context.builder_map
+        builder_map = context.builder_map
 
         for field in dtype:
             field_builder = <StringBuilder>get_field_builder(field, tzinfo)
-            self.builder_map[field.name.encode('utf-8')] = field_builder
+            builder_map[field.name.encode('utf-8')] = field_builder
             c_field_builders.push_back(<shared_ptr[CArrayBuilder]>field_builder.builder)
 
         self.builder.reset(new CStructBuilder(pyarrow_unwrap_data_type(dtype), pool, c_field_builders))
@@ -783,30 +781,7 @@ cdef class DocumentBuilder(_ArrayBuilderBase):
         cdef shared_ptr[CArray] out
         with nogil:
             self.builder.get().Finish(&out)
-
-        struct_array = pyarrow_wrap_array(out)
-        for struct_def in struct_array:
-            new_types = []
-            new_names = list(struct_def.keys())
-            for fname, ftype in struct_def.items():
-                builder_instance = self.builder_map[fname.encode('utf-8')]
-                if isinstance(builder_instance, ObjectIdBuilder): # ObjectIdType
-                    new_ftype = ObjectIdType()
-                    new_types.append(new_ftype)
-                elif isinstance(builder_instance, Decimal128Builder): # Decimal128Type
-                   new_ftype = Decimal128Type_()
-                   new_types.append(new_ftype)
-                elif isinstance(builder_instance, BinaryBuilder): # BinaryType
-                    new_ftype = BinaryType(self.dtype.field(fname).type.subtype)
-                    new_types.append(new_ftype)
-                elif isinstance(builder_instance, CodeBuilder): # CodeType
-                    new_ftype = CodeType()
-                    new_types.append(new_ftype)
-                else:
-                    new_types.append(ftype.type)
-
-            new_dtype = struct(zip(new_names, new_types))
-        return struct_array.cast(new_dtype)
+        return pyarrow_wrap_array(out)
 
     cdef shared_ptr[CStructBuilder] unwrap(self):
         return self.builder

bindings/python/test/test_arrow.py

@@ -22,7 +22,18 @@
 import pyarrow
 import pymongo
 from bson import Binary, Code, CodecOptions, Decimal128, ObjectId
-from pyarrow import Table, bool_, csv, decimal256, field, int32, int64, list_
+from pyarrow import (
+    DataType,
+    FixedSizeBinaryType,
+    Table,
+    bool_,
+    csv,
+    decimal256,
+    field,
+    int32,
+    int64,
+    list_,
+)
 from pyarrow import schema as ArrowSchema
 from pyarrow import string, struct, timestamp
 from pyarrow.parquet import read_table, write_table
@@ -665,10 +676,18 @@ 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, 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)
+        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)
 
 
 class TestArrowExplicitApi(ArrowApiTestMixin, unittest.TestCase):