INTPYTHON-549 Rework the projection logic to allow reading list-of-struct data structures using a specified schema (#285)

Author: zpks

bindings/python/pymongoarrow/schema.py

@@ -68,19 +68,21 @@ def _normalize_mapping(mapping):
     def _get_projection(self):
         projection = {"_id": False}
         for fname, ftype in self.typemap.items():
-            projection[fname] = self._get_field_projection_value(ftype)
+            projection = self._get_field_projection_value(fname, ftype, projection)
         return projection
 
-    def _get_field_projection_value(self, ftype):
+    def _get_field_projection_value(self, fname, ftype, projection):
         value = True
         if isinstance(ftype, pa.ListType):
-            return self._get_field_projection_value(ftype.value_field.type)
+            return self._get_field_projection_value(fname, ftype.value_field.type, projection)
         if isinstance(ftype, pa.StructType):
-            projection = {}
             for nested_ftype in ftype:
-                projection[nested_ftype.name] = True
-            value = projection
-        return value
+                projection = self._get_field_projection_value(
+                    fname + "." + nested_ftype.name, nested_ftype.type, projection
+                )
+            return projection
+        projection[fname] = value
+        return projection
 
     def __eq__(self, other):
         if isinstance(other, type(self)):

bindings/python/test/test_arrow.py

@@ -32,6 +32,7 @@
     date64,
     decimal256,
     field,
+    float64,
     int32,
     int64,
     large_list,
@@ -415,6 +416,49 @@ def inner(i):
         raw_data["nested"] = [inner(i) for i in range(3)]
         return schema, Table.from_pydict(raw_data, ArrowSchema(schema))
 
+    def test_write_nested_schema_validation(self):
+        raw_data = {
+            "_id": [1, 2],
+            "top": [
+                {
+                    "middle": {
+                        "value": "string_1",
+                        "bottom": [
+                            {"event": datetime(2012, 1, 1), "value": 1.1},
+                            {"event": datetime(2014, 1, 1), "value": 1.2},
+                        ],
+                    }
+                },
+                {
+                    "middle": {
+                        "value": "string_2",
+                        "bottom": [
+                            {"event": datetime(2013, 1, 1), "value": 1.2},
+                            {"event": datetime(2019, 1, 1), "value": 1.5},
+                        ],
+                    }
+                },
+            ],
+        }
+
+        schema = {
+            "_id": int64(),
+            "top": struct(
+                {
+                    "middle": struct(
+                        {
+                            "value": string(),
+                            "bottom": list_(struct({"event": timestamp("ms"), "value": float64()})),
+                        }
+                    )
+                }
+            ),
+        }
+
+        data = Table.from_pydict(raw_data, ArrowSchema(schema))
+
+        self.round_trip(data, Schema(schema))
+
     def test_parquet(self):
         schema, data = self._create_nested_data()
         with tempfile.NamedTemporaryFile(suffix=".parquet") as f:

bindings/python/test/test_schema.py

@@ -76,7 +76,7 @@ def test_from_arrow_units(self):
 
     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}})
+        self.assertEqual(schema._get_projection(), {"_id": True, "obj.a": True, "obj.b": True})
 
     def test_list_projection(self):
         schema = Schema(
@@ -85,7 +85,7 @@ def test_list_projection(self):
                 "list": list_(struct([field("a", int64()), field("b", int64())])),
             }
         )
-        self.assertEqual(schema._get_projection(), {"_id": True, "list": {"a": True, "b": True}})
+        self.assertEqual(schema._get_projection(), {"_id": True, "list.a": True, "list.b": True})
 
     def test_list_of_list_projection(self):
         schema = Schema(
@@ -94,14 +94,14 @@ def test_list_of_list_projection(self):
                 "list": list_(list_(struct([field("a", int64()), field("b", int64())]))),
             }
         )
-        self.assertEqual(schema._get_projection(), {"_id": True, "list": {"a": True, "b": True}})
+        self.assertEqual(schema._get_projection(), {"_id": True, "list.a": True, "list.b": True})
 
     def test_py_list_projection(self):
         schema = Schema(
             {"_id": ObjectId, "list": [(struct([field("a", int64()), field("b", float64())]))]}
         )
 
-        self.assertEqual(schema._get_projection(), {"_id": True, "list": {"a": True, "b": True}})
+        self.assertEqual(schema._get_projection(), {"_id": True, "list.a": True, "list.b": True})
 
     def test_py_list_with_multiple_fields_raises(self):
         with pytest.raises(