Merge pull request #1671 from jstol/feature/extract-nested-object-array-props-from-data-models

Add support for extracting nested object array properties from data models

Author: dirkkul

integration/test_collection_nested.py

@@ -516,34 +516,77 @@ def test_nested_return_specific_properties(
     assert out.properties["nested"] == expected
 
 
-def test_nested_return_generic_properties(collection_factory: CollectionFactory) -> None:
-    class Child(TypedDict):
-        name: str
-        age: int
+class Child(TypedDict):
+    name: str
+    age: int
 
-    class Parent(TypedDict):
-        child: Nested[Child]
 
-    collection = collection_factory(
-        properties=[
-            Property(
-                name="child",
-                data_type=DataType.OBJECT,
-                nested_properties=[
-                    Property(
-                        name="name",
-                        data_type=DataType.TEXT,
-                    ),
-                    Property(
-                        name="age",
-                        data_type=DataType.INT,
-                    ),
+@pytest.mark.parametrize(
+    "data_model_cls,data_object,properties",
+    [
+        (
+            Child,
+            {"name": "Timmy", "age": 10},
+            [
+                Property(name="name", data_type=DataType.TEXT),
+                Property(name="age", data_type=DataType.INT),
+            ],
+        ),
+        (
+            TypedDict("ParentWithChild", {"name": str, "age": int, "child": Nested[Child]}),
+            {
+                "name": "Bob",
+                "age": 39,
+                "child": {"name": "Timmy", "age": 10},
+            },
+            [
+                Property(name="name", data_type=DataType.TEXT),
+                Property(name="age", data_type=DataType.INT),
+                Property(
+                    name="child",
+                    data_type=DataType.OBJECT,
+                    nested_properties=[
+                        Property(name="name", data_type=DataType.TEXT),
+                        Property(name="age", data_type=DataType.INT),
+                    ],
+                ),
+            ],
+        ),
+        (
+            TypedDict(
+                "ParentWithChildList", {"name": str, "age": int, "children": Nested[list[Child]]}
+            ),
+            {
+                "name": "Bob",
+                "age": 39,
+                "children": [
+                    {"name": "Timmy", "age": 10},
+                    {"name": "Daisy", "age": 8},
                 ],
-            )
-        ],
-        data_model_properties=Parent,
+            },
+            [
+                Property(name="name", data_type=DataType.TEXT),
+                Property(name="age", data_type=DataType.INT),
+                Property(
+                    name="children",
+                    data_type=DataType.OBJECT_ARRAY,
+                    nested_properties=[
+                        Property(name="name", data_type=DataType.TEXT),
+                        Property(name="age", data_type=DataType.INT),
+                    ],
+                ),
+            ],
+        ),
+    ],
+)
+def test_nested_return_generic_properties(
+    collection_factory, data_model_cls, data_object, properties
+):
+    collection = collection_factory(
+        properties=properties,
+        data_model_properties=data_model_cls,
     )
+    collection.data.insert(data_object)
+    results = collection.query.fetch_objects(return_properties=data_model_cls)
 
-    collection.data.insert(Parent(child=Child(name="Timmy", age=10)))
-    results = collection.query.fetch_objects(return_properties=Parent)
-    assert results.objects[0].properties["child"] == {"name": "Timmy", "age": 10}
+    assert results.objects[0].properties == data_object

weaviate/collections/aggregations/hybrid/executor.py

@@ -12,8 +12,8 @@
 from weaviate.connect import executor
 from weaviate.connect.v4 import ConnectionType
 from weaviate.exceptions import WeaviateUnsupportedFeatureError
-from weaviate.types import NUMBER
 from weaviate.proto.v1 import aggregate_pb2
+from weaviate.types import NUMBER
 
 
 class _HybridExecutor(Generic[ConnectionType], _BaseExecutor[ConnectionType]):

weaviate/collections/aggregations/near_image/executor.py

@@ -11,9 +11,9 @@
 from weaviate.collections.filters import _FilterToGRPC
 from weaviate.connect import executor
 from weaviate.connect.v4 import ConnectionType
+from weaviate.proto.v1 import aggregate_pb2
 from weaviate.types import BLOB_INPUT, NUMBER
 from weaviate.util import parse_blob
-from weaviate.proto.v1 import aggregate_pb2
 
 
 class _NearImageExecutor(Generic[ConnectionType], _BaseExecutor[ConnectionType]):

weaviate/collections/aggregations/near_object/executor.py

@@ -11,8 +11,8 @@
 from weaviate.collections.filters import _FilterToGRPC
 from weaviate.connect import executor
 from weaviate.connect.v4 import ConnectionType
-from weaviate.types import NUMBER, UUID
 from weaviate.proto.v1 import aggregate_pb2
+from weaviate.types import NUMBER, UUID
 
 
 class _NearObjectExecutor(Generic[ConnectionType], _BaseExecutor[ConnectionType]):

weaviate/collections/aggregations/near_text/executor.py

@@ -12,8 +12,8 @@
 from weaviate.collections.filters import _FilterToGRPC
 from weaviate.connect import executor
 from weaviate.connect.v4 import ConnectionType
-from weaviate.types import NUMBER
 from weaviate.proto.v1 import aggregate_pb2
+from weaviate.types import NUMBER
 
 
 class _NearTextExecutor(Generic[ConnectionType], _BaseExecutor[ConnectionType]):

weaviate/collections/aggregations/near_vector/executor.py

@@ -16,8 +16,8 @@
 from weaviate.connect import executor
 from weaviate.connect.v4 import ConnectionType
 from weaviate.exceptions import WeaviateInvalidInputError
-from weaviate.types import NUMBER
 from weaviate.proto.v1 import aggregate_pb2
+from weaviate.types import NUMBER
 
 
 class _NearVectorExecutor(Generic[ConnectionType], _BaseExecutor[ConnectionType]):

weaviate/collections/classes/internal.py

@@ -438,6 +438,10 @@ def __is_nested(value: Any) -> bool:
 
 def __create_nested_property_from_nested(name: str, value: Any) -> QueryNested:
     inner_type = get_args(value)[0]
+    # If this nested property contains an object array, use the element type
+    if get_origin(inner_type) is list:
+        inner_type = get_args(inner_type)[0]
+
     return QueryNested(
         name=name,
         properties=[