ARROW-175 fix Incorrect projection on nested fields (#165)

Author: NoahStapp

bindings/python/pymongoarrow/schema.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import collections.abc as abc
 
+from pyarrow import ListType, StructType
 from pymongoarrow.types import _normalize_typeid
 
 
@@ -62,10 +63,21 @@ def _normalize_mapping(mapping):
 
     def _get_projection(self):
         projection = {"_id": False}
-        for fname, _ in self.typemap.items():
-            projection[fname] = True
+        for fname, ftype in self.typemap.items():
+            projection[fname] = self._get_field_projection_value(ftype)
         return projection
 
+    def _get_field_projection_value(self, ftype):
+        value = True
+        if isinstance(ftype, ListType):
+            return self._get_field_projection_value(ftype.value_field.type)
+        elif isinstance(ftype, StructType):
+            projection = {}
+            for nested_ftype in ftype:
+                projection[nested_ftype.name] = True
+            value = projection
+        return value
+
     def __eq__(self, other):
         if isinstance(other, type(self)):
             return self.typemap == other.typemap

bindings/python/test/test_arrow.py

@@ -640,6 +640,16 @@ def test_mixed_types_int64(self):
         with self.assertRaises(OverflowError):
             find_arrow_all(self.coll, {}, projection={"_id": 0}, schema=Schema({"a": int32()}))
 
+    def test_nested_contradicting_unused_schema(self):
+        data = [{"obj": {"a": 1, "b": 1000000000}}, {"obj": {"a": 2, "b": 1.0e50}}]
+        schema = Schema({"obj": {"a": int32()}})
+
+        self.coll.drop()
+        self.coll.insert_many(data)
+        for func in [find_arrow_all, aggregate_arrow_all]:
+            out = func(self.coll, {} if func == find_arrow_all else [], schema=schema)
+            self.assertEqual(out["obj"].to_pylist(), [{"a": 1}, {"a": 2}])
+
 
 class TestArrowExplicitApi(ArrowApiTestMixin, unittest.TestCase):
     def run_find(self, *args, **kwargs):

bindings/python/test/test_schema.py

@@ -15,9 +15,9 @@
 from unittest import TestCase
 
 from bson import Binary, Code, Decimal128, Int64, ObjectId
-from pyarrow import Table, float64, int64
+from pyarrow import Table, field, float64, int64, list_
 from pyarrow import schema as ArrowSchema
-from pyarrow import timestamp
+from pyarrow import struct, timestamp
 from pymongoarrow.schema import Schema
 from pymongoarrow.types import _TYPE_NORMALIZER_FACTORY
 
@@ -70,3 +70,22 @@ def test_from_bson_units(self):
     def test_from_arrow_units(self):
         schema = Schema({"field1": int64(), "field2": timestamp("s")})
         self.assertEqual(schema.typemap, {"field1": int64(), "field2": timestamp("s")})
+
+    def test_nested_projection(self):
+        schema = Schema({"_id": int64(), "obj": {"a": int64(), "b": int64()}})
+        self.assertEqual(schema._get_projection(), {"_id": True, "obj": {"a": True, "b": True}})
+
+    def test_list_projection(self):
+        schema = Schema(
+            {"_id": int64(), "list": list_(struct([field("a", int64()), field("b", int64())]))}
+        )
+        self.assertEqual(schema._get_projection(), {"_id": True, "list": {"a": True, "b": True}})
+
+    def test_list_of_list_projection(self):
+        schema = Schema(
+            {
+                "_id": int64(),
+                "list": list_(list_(struct([field("a", int64()), field("b", int64())]))),
+            }
+        )
+        self.assertEqual(schema._get_projection(), {"_id": True, "list": {"a": True, "b": True}})