feat: support nested array

examples/nested_array.py

@@ -0,0 +1,110 @@
+"""Store integer and string arrays inside an array of structs.
+
+Requires a Milvus server that supports array sub-fields in structs.
+Run from the repository root: PYTHONPATH=. python examples/nested_array.py
+"""
+
+from pprint import pprint
+
+from pymilvus import DataType, MilvusClient
+
+URI = "http://localhost:19530"
+COLLECTION_NAME = "nested_array_example"
+
+
+def main() -> None:
+    client = MilvusClient(URI)
+    try:
+        if client.has_collection(COLLECTION_NAME):
+            client.drop_collection(COLLECTION_NAME)
+
+        schema = client.create_schema(auto_id=False)
+        schema.add_field("pk", DataType.INT64, is_primary=True)
+        schema.add_field("vector", DataType.FLOAT_VECTOR, dim=4)
+
+        # Each Struct element contains Array<Int32> and Array<VarChar> fields.
+        group_schema = client.create_struct_field_schema()
+        group_schema.add_field(
+            "values", DataType.ARRAY, element_type=DataType.INT32, max_capacity=4
+        )
+        group_schema.add_field(
+            "labels",
+            DataType.ARRAY,
+            element_type=DataType.VARCHAR,
+            max_capacity=4,
+            max_length=64,
+        )
+        schema.add_field(
+            "groups",
+            DataType.ARRAY,
+            element_type=DataType.STRUCT,
+            struct_schema=group_schema,
+            max_capacity=8,  # Up to eight Struct elements per row.
+            nullable=True,  # Nullability applies to the whole Struct array.
+        )
+
+        index_params = client.prepare_index_params()
+        index_params.add_index(field_name="vector", index_type="AUTOINDEX", metric_type="COSINE")
+        client.create_collection(
+            collection_name=COLLECTION_NAME,
+            schema=schema,
+            index_params=index_params,
+            consistency_level="Strong",
+        )
+
+        # Describe returns the same logical field types used above.
+        print("Describe collection:")
+        pprint(client.describe_collection(COLLECTION_NAME), sort_dicts=False)
+
+        rows = [
+            {
+                "pk": 0,
+                "vector": [1.0, 0.0, 0.0, 0.0],
+                "groups": [
+                    {"values": [1, 2], "labels": ["alpha", "shared"]},
+                    {"values": [-3, 4], "labels": ["beta", "shared"]},
+                ],
+            },
+            {
+                "pk": 1,
+                "vector": [0.0, 1.0, 0.0, 0.0],
+                "groups": None,  # A null Struct array.
+            },
+            {
+                "pk": 2,
+                "vector": [0.0, 0.0, 1.0, 0.0],
+                "groups": [{"values": [], "labels": []}],  # Empty array sub-fields.
+            },
+            {
+                "pk": 3,
+                "vector": [0.0, 0.0, 0.0, 1.0],
+                "groups": [],  # An empty Struct array.
+            },
+        ]
+        result = client.insert(collection_name=COLLECTION_NAME, data=rows)
+        print(f"Inserted {result['insert_count']} rows")
+
+        print("Query inserted rows:")
+        result = client.query(
+            collection_name=COLLECTION_NAME,
+            filter="pk in [0, 1, 2, 3]",
+            output_fields=["pk", "groups"],
+        )
+        pprint(sorted(result, key=lambda row: row["pk"]), sort_dicts=False)
+
+        # Match rows with at least one Struct element whose labels contain "alpha".
+        print("Rows containing the label 'alpha' (pk=0):")
+        result = client.query(
+            collection_name=COLLECTION_NAME,
+            filter='MATCH_ANY(groups, array_contains($[labels], "alpha"))',
+            output_fields=["pk", "groups"],
+        )
+        pprint(result, sort_dicts=False)
+
+        client.drop_collection(COLLECTION_NAME)
+    finally:
+        client.close()
+
+
+if __name__ == "__main__":
+    main()

pymilvus/client/abstract.py

@@ -18,6 +18,35 @@
 logger = logging.getLogger(__name__)
 
 
+def _type_schema_params_to_dict(raw_params: Any) -> Dict[str, Any]:
+    params = {}
+    for type_param in raw_params:
+        key = "mmap_enabled" if type_param.key == "mmap.enabled" else type_param.key
+        value: Any = type_param.value
+        if key == "mmap_enabled":
+            value = value.lower() != "false"
+        elif key in ("dim", "max_capacity", Config.MaxVarCharLengthKey):
+            value = int(value)
+        params[key] = value
+    return params
+
+
+def _type_schema_to_dict(raw: Any) -> Dict[str, Any]:
+    kind = raw.WhichOneof("kind")
+    result = (
+        {"array_element": _type_schema_to_dict(raw.array_element)}
+        if kind == "array_element"
+        else {"leaf_type": DataType(raw.leaf_type)}
+    )
+    if raw.nullable:
+        result["nullable"] = True
+
+    params = _type_schema_params_to_dict(raw.type_params)
+    if params:
+        result["type_params"] = params
+    return result
+
+
 class FieldSchema:
     def __init__(self, raw: Any):
         self._raw = raw
@@ -38,6 +67,7 @@ def __init__(self, raw: Any):
         self.external_field = ""
         # For array field
         self.element_type = None
+        self.type_schema = None
         self.is_clustering_key = False
         self.__pack(self._raw)
 
@@ -50,6 +80,9 @@ def __pack(self, raw: Any):
         self.type = DataType(raw.data_type)
         self.is_partition_key = raw.is_partition_key
         self.element_type = DataType(raw.element_type)
+        if self.type == DataType.ARRAY and self.element_type == DataType.ARRAY:
+            self.type_schema = _type_schema_to_dict(raw.type_schema)
+            self.params.update(self.type_schema.get("type_params", {}))
         self.is_clustering_key = raw.is_clustering_key
         self.default_value = raw.default_value
         if raw.default_value is not None and raw.default_value.WhichOneof("data") is None:
@@ -122,6 +155,8 @@ def dict(self):
 
         if self.element_type:
             _dict["element_type"] = self.element_type
+        if self.type_schema is not None:
+            _dict["type_schema"] = self.type_schema
 
         if self.is_partition_key:
             _dict["is_partition_key"] = True

pymilvus/client/entity_helper.py

@@ -337,6 +337,25 @@ def convert_to_array(obj: List[Any], field_info: Any):
         obj = obj.tolist()
 
     field_data = schema_types.ScalarField()
+    type_schema = field_info.get("type_schema")
+    if type_schema is not None:
+        element = type_schema["array_element"]
+        if "leaf_type" in element:
+            return convert_to_array(
+                obj, {"name": field_info.get("name"), "element_type": element["leaf_type"]}
+            )
+        array_data = field_data.array_data
+        array_data.element_type = (
+            DataType.ARRAY
+            if "array_element" in element["array_element"]
+            else element["array_element"]["leaf_type"]
+        )
+        array_data.data.extend(
+            convert_to_array(value, {"name": field_info.get("name"), "type_schema": element})
+            for value in obj
+        )
+        return field_data
+
     element_type = field_info.get("element_type", None)
     attr_name = type_info.get_array_element_attr(element_type)
     if attr_name is not None:
@@ -497,6 +516,8 @@ def _pack_scalar_row(
 ):
     try:
         payload = _get_protobuf_payload(field_data, dtype)
+        if dtype == DataType.ARRAY and field_info.get("type_schema") is not None:
+            payload.element_type = DataType.ARRAY
         if value is None:
             payload.data.extend([])
             return
@@ -667,6 +688,8 @@ def entity_to_field_data(entity: Dict, field_info: Any, num_rows: int) -> schema
             elif entity_type == DataType.JSON:
                 entity_values = entity_to_json_arr(entity_values, field_info)
             elif entity_type == DataType.ARRAY:
+                if field_info.get("type_schema") is not None:
+                    field_data.scalars.array_data.element_type = DataType.ARRAY
                 entity_values = entity_to_array_arr(entity_values, field_info)
             getattr(field_data.scalars, attr_name).data.extend(entity_values)
         elif entity_type == DataType.FLOAT_VECTOR:

pymilvus/client/field_data_extractors.py

@@ -69,9 +69,22 @@ def decode_array_value(array_cell: Any, index: int) -> Any:
     data = getattr(array_cell, attr).data
     if len(data) <= index:
         return None
+    if attr == "array_data":
+        return decode_array(data[index], array_cell.array_data.element_type)
     return data[index]
 
 
+def decode_array(array_cell: Any, element_type: DataType) -> Any:
+    """Recursively decode one ScalarField containing an ARRAY value."""
+    if element_type == DataType.ARRAY:
+        array_data = array_cell.array_data
+        return [decode_array(value, array_data.element_type) for value in array_data.data]
+    attr = type_info.get_array_element_attr(element_type)
+    if attr is None:
+        raise MilvusException(message=f"Unsupported data type: {element_type}")
+    return list(getattr(array_cell, attr).data)
+
+
 def vector_array_length(vector_data: Any, element_type: DataType) -> int:
     """Return the number of vector values stored in one VectorArray row."""
     element_type = _resolve_vector_array_element_type(vector_data, element_type)
@@ -113,6 +126,8 @@ def decode_array_cell(field_data: Any, logical_index: int) -> Any:
     array_data = field_data.scalars.array_data
     if logical_index >= len(array_data.data):
         return None
+    if array_data.element_type == DataType.ARRAY:
+        return decode_array(array_data.data[logical_index], array_data.element_type)
     attr = type_info.get_array_element_attr(array_data.element_type)
     if attr is None:
         raise MilvusException(message=f"Unsupported data type: {array_data.element_type}")
@@ -241,8 +256,8 @@ def _resolve_vector_array_element_type(vector_data: Any, element_type: DataType)
 
 
 def _populated_array_attr(array_cell: Any) -> Optional[str]:
-    for info in type_info.TYPE_INFO.values():
-        attr = info.array_element_attr
+    attrs = (info.array_element_attr for info in type_info.TYPE_INFO.values())
+    for attr in (*attrs, "array_data"):
         if attr is None:
             continue
         data = getattr(array_cell, attr, None)

pymilvus/client/prepare.py

@@ -1,4 +1,5 @@
 import base64
+import copy
 import datetime
 import json
 import re
@@ -86,6 +87,29 @@
 
 
 class Prepare:
+    @staticmethod
+    def _type_schema_from_dict(type_schema: Dict) -> schema_types.TypeSchema:
+        """Encode the complete recursive type tree without adding ARRAY levels."""
+        result = schema_types.TypeSchema(nullable=type_schema.get("nullable", False))
+        if "array_element" in type_schema:
+            result.array_element.CopyFrom(
+                Prepare._type_schema_from_dict(type_schema["array_element"])
+            )
+        else:
+            result.leaf_type = type_schema["leaf_type"]
+        for key, value in type_schema.get("type_params", {}).items():
+            result.type_params.append(
+                common_types.KeyValuePair(
+                    key=str(key) if key != "mmap_enabled" else "mmap.enabled",
+                    value=(
+                        orjson.dumps(value).decode(Config.EncodeProtocol)
+                        if not isinstance(value, str)
+                        else value
+                    ),
+                )
+            )
+        return result
+
     @classmethod
     def create_collection_request(
         cls,
@@ -204,6 +228,9 @@ def get_schema_from_collection_schema(
                 )
                 field_schema.type_params.append(kv_pair)
 
+            if f.type_schema is not None:
+                field_schema.type_schema.CopyFrom(cls._type_schema_from_dict(f.type_schema))
+
             schema.fields.append(field_schema)
 
         for struct in fields.struct_fields:
@@ -265,14 +292,48 @@ def get_struct_array_field_schema(
                 is_function_output=f.is_function_output,
             )
 
-            field_params = dict(f.params) if f.params else {}
+            logical_array_type = None
+            if f.dtype == DataType.ARRAY:
+                if f.type_schema is not None:
+                    logical_array_type = copy.deepcopy(f.type_schema)
+                else:
+                    array_params = {"max_capacity", "mmap_enabled", "warmup"}
+                    logical_array_type = {
+                        "array_element": {
+                            "leaf_type": f.element_type,
+                            "type_params": {
+                                key: value
+                                for key, value in f.params.items()
+                                if key not in array_params
+                            },
+                        },
+                        "type_params": {
+                            key: value for key, value in f.params.items() if key in array_params
+                        },
+                    }
+                logical_params = logical_array_type.get("type_params", {})
+                field_params = {
+                    key: logical_params.pop(key)
+                    for key in ("mmap_enabled", "warmup")
+                    if key in logical_params
+                }
+            else:
+                field_params = dict(f.params) if f.params else {}
             field_params["max_capacity"] = struct.max_capacity
 
             for k, v in field_params.items():
                 kv_pair = common_types.KeyValuePair(
                     key=str(k) if k != "mmap_enabled" else "mmap.enabled", value=json.dumps(v)
                 )
                 field_schema.type_params.append(kv_pair)
+            if logical_array_type is not None:
+                physical_array_type = {
+                    "array_element": logical_array_type,
+                    "type_params": field_params,
+                }
+                field_schema.type_schema.CopyFrom(
+                    Prepare._type_schema_from_dict(physical_array_type)
+                )
             struct_schema.fields.append(field_schema)
 
         return struct_schema
@@ -345,6 +406,15 @@ def get_field_schema(
         ]
         field_schema.type_params.extend(kvs)
 
+        if field.get("type_schema") is not None:
+            normalized = FieldSchema.construct_from_dict(field)
+            field_schema.element_type = normalized.element_type
+            field_schema.type_schema.CopyFrom(
+                Prepare._type_schema_from_dict(normalized.type_schema)
+            )
+            field_schema.ClearField("type_params")
+            field_schema.type_params.extend(field_schema.type_schema.type_params)
+
         return field_schema, primary_field, auto_id_field
 
     @classmethod
@@ -748,6 +818,8 @@ def _add_empty_struct_data(
                 continue
 
             if field_info["type"] == DataType.ARRAY:
+                if field_info.get("type_schema") is not None:
+                    field_data.scalars.array_data.element_type = DataType.ARRAY
                 field_data.scalars.array_data.data.append(convert_to_array([], field_info))
             elif field_info["type"] == DataType._ARRAY_OF_VECTOR:
                 field_data.vectors.vector_array.dim = Prepare._get_dim_value(field_info)
@@ -801,6 +873,8 @@ def _process_struct_values(
             field_info = struct_field_info[field_name]
 
             if field_info["type"] == DataType.ARRAY:
+                if field_info.get("type_schema") is not None:
+                    field_data.scalars.array_data.element_type = DataType.ARRAY
                 field_data.scalars.array_data.data.append(convert_to_array(values, field_info))
             elif field_info["type"] == DataType._ARRAY_OF_VECTOR:
                 field_data.vectors.vector_array.dim = Prepare._get_dim_value(field_info)