feat: fetch() add output_fields param (#358)

python/tests/detail/test_collection_dql.py

@@ -273,6 +273,118 @@ class TestCollectionFetch:
             f"Expected 0 results for empty ID list, but got {len(result)}"
         )
 
+    @pytest.mark.parametrize("doc_num", [3])
+    def test_fetch_with_output_fields(self, full_collection: Collection, doc_num):
+        """Test that fetch respects output_fields parameter."""
+        multiple_docs = [
+            generate_doc(i, full_collection.schema) for i in range(doc_num)
+        ]
+        result = full_collection.insert(multiple_docs)
+        for item in result:
+            assert item.ok(), f"Insert failed: {item.code()}"
+
+        doc_id = multiple_docs[0].id
+
+        # Case 1: output_fields=None -> all scalar fields returned
+        fetched_all = full_collection.fetch(ids=[doc_id], output_fields=None)
+        assert doc_id in fetched_all
+        doc_all = fetched_all[doc_id]
+        assert doc_all is not None
+        assert doc_all.has_field("int32_field"), (
+            "int32_field should be present when output_fields=None"
+        )
+        assert doc_all.has_field("string_field"), (
+            "string_field should be present when output_fields=None"
+        )
+
+        # Case 2: output_fields=["int32_field"] -> only int32_field returned
+        fetched_partial = full_collection.fetch(
+            ids=[doc_id], output_fields=["int32_field"]
+        )
+        assert doc_id in fetched_partial
+        doc_partial = fetched_partial[doc_id]
+        assert doc_partial is not None
+        assert doc_partial.has_field("int32_field"), "int32_field should be present"
+        assert not doc_partial.has_field("string_field"), (
+            'string_field should not be present when output_fields=["int32_field"]'
+        )
+        assert not doc_partial.has_field("float_field"), (
+            'float_field should not be present when output_fields=["int32_field"]'
+        )
+
+        # Case 3: output_fields=[] (empty) -> no scalar fields returned
+        fetched_empty = full_collection.fetch(ids=[doc_id], output_fields=[])
+        assert doc_id in fetched_empty
+        doc_empty = fetched_empty[doc_id]
+        assert doc_empty is not None
+        assert doc_empty.id == doc_id, "pk should still be set"
+        assert not doc_empty.has_field("int32_field"), (
+            "int32_field should not be present when output_fields=[]"
+        )
+        assert not doc_empty.has_field("string_field"), (
+            "string_field should not be present when output_fields=[]"
+        )
+
+        # Case 4: multiple output_fields
+        fetched_multi = full_collection.fetch(
+            ids=[doc_id], output_fields=["int32_field", "float_field"]
+        )
+        assert doc_id in fetched_multi
+        doc_multi = fetched_multi[doc_id]
+        assert doc_multi is not None
+        assert doc_multi.has_field("int32_field")
+        assert doc_multi.has_field("float_field")
+        assert not doc_multi.has_field("string_field")
+
+    @pytest.mark.parametrize("doc_num", [3])
+    def test_fetch_with_include_vector(self, full_collection: Collection, doc_num):
+        """Test that fetch respects include_vector parameter."""
+        multiple_docs = [
+            generate_doc(i, full_collection.schema) for i in range(doc_num)
+        ]
+        result = full_collection.insert(multiple_docs)
+        for item in result:
+            assert item.ok(), f"Insert failed: {item.code()}"
+
+        doc_id = multiple_docs[0].id
+
+        # Case 1: include_vector=True (default) -> vector data returned
+        fetched_with_vec = full_collection.fetch(ids=[doc_id])
+        assert doc_id in fetched_with_vec
+        doc_with_vec = fetched_with_vec[doc_id]
+        assert doc_with_vec is not None
+        assert doc_with_vec.has_field("int32_field"), (
+            "scalar fields should still be present"
+        )
+        assert doc_with_vec.vector("vector_fp32_field"), (
+            "vector should be present when include_vector=True (default)"
+        )
+
+        # Case 2: include_vector=False -> no vector data returned
+        fetched_no_vec = full_collection.fetch(ids=[doc_id], include_vector=False)
+        assert doc_id in fetched_no_vec
+        doc_no_vec = fetched_no_vec[doc_id]
+        assert doc_no_vec is not None
+        assert doc_no_vec.has_field("int32_field"), (
+            "scalar fields should still be present"
+        )
+        assert not doc_no_vec.vector("vector_fp32_field"), (
+            "vector should not be present when include_vector=False"
+        )
+
+        # Case 3: include_vector=False with output_fields
+        fetched_combo = full_collection.fetch(
+            ids=[doc_id], output_fields=["int32_field"], include_vector=False
+        )
+        assert doc_id in fetched_combo
+        doc_combo = fetched_combo[doc_id]
+        assert doc_combo is not None
+        assert doc_combo.has_field("int32_field")
+        assert not doc_combo.has_field("string_field")
+        assert not doc_combo.vector("vector_fp32_field"), (
+            "vector should not be present when include_vector=False"
+        )
+
 
 class TestCollectionQuery:
     @pytest.mark.parametrize("doc_num", [5])

python/zvec/__init__.pyi

@@ -115,7 +115,12 @@ class _Collection:
     def Destroy(self) -> None: ...
     def DropColumn(self, arg0: str) -> None: ...
     def DropIndex(self, arg0: str) -> None: ...
-    def Fetch(self, arg0: collections.abc.Sequence[str]) -> dict[str, _Doc]: ...
+    def Fetch(
+        self,
+        pks: collections.abc.Sequence[str],
+        output_fields: list[str] | None = None,
+        include_vector: bool = True,
+    ) -> dict[str, _Doc]: ...
     def Flush(self) -> None: ...
     def GroupByQuery(self, arg0: ...) -> list[...]: ...
     def Insert(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ...

python/zvec/model/collection.py

@@ -337,17 +337,27 @@ class Collection:
         self._obj.DeleteByFilter(filter)
 
     # ========== Collection DQL-fetch Methods ==========
-    def fetch(self, ids: Union[str, list[str]]) -> dict[str, Doc]:
+    def fetch(
+        self,
+        ids: Union[str, list[str]],
+        *,
+        output_fields: Optional[list[str]] = None,
+        include_vector: bool = True,
+    ) -> dict[str, Doc]:
         """Retrieve documents by ID.
 
         Args:
             ids (Union[str, list[str]]): Document IDs to fetch.
+            output_fields (Optional[list[str]], optional): Scalar fields to
+                include. If None, all fields are returned. Defaults to None.
+            include_vector (bool, optional): Whether to include vector data in
+                results. Defaults to True.
 
         Returns:
             dict[str, Doc]: Mapping from ID to document. Missing IDs are omitted.
         """
         ids = [ids] if isinstance(ids, str) else ids
-        docs = self._obj.Fetch(ids)
+        docs = self._obj.Fetch(ids, output_fields, include_vector)
         return {
             doc_id: py_doc
             for doc_id, core_doc in docs.items()