From de8fb760ef61af12420fc1f0e7053ea23d5e3f71 Mon Sep 17 00:00:00 2001 From: lichen2015 Date: Fri, 29 May 2026 17:36:30 +0800 Subject: [PATCH] fix: sync querier schema after column DDL to fix empty query fields (#429) * fix: sync querier schema after column DDL to fix empty query fields (#426) * fix: sync querier schema after create_index/drop_index --- python/tests/detail/test_collection_ddl.py | 163 +++++++++++++++++++++ python/zvec/model/collection.py | 5 + 2 files changed, 168 insertions(+) diff --git a/python/tests/detail/test_collection_ddl.py b/python/tests/detail/test_collection_ddl.py index 2f61d26..a2c682f 100644 --- a/python/tests/detail/test_collection_ddl.py +++ b/python/tests/detail/test_collection_ddl.py @@ -1555,3 +1555,166 @@ class TestColumnDDL: result = basic_collection.insert(doc) assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value) + + def test_add_column_then_query_returns_new_field( + self, basic_collection: Collection + ): + """Regression test for issue #426: query() should return fields added via add_column().""" + basic_collection.add_column( + field_schema=FieldSchema("score", DataType.INT64, nullable=True), + ) + + docs = [ + Doc( + id="1", + fields={"id": 1, "name": "alice", "weight": 60.0, "score": 100}, + vectors={ + "dense": generate_constant_vector(1, 128), + "sparse": generate_sparse_vector(1), + }, + ), + Doc( + id="2", + fields={"id": 2, "name": "bob", "weight": 70.0, "score": 200}, + vectors={ + "dense": generate_constant_vector(2, 128), + "sparse": generate_sparse_vector(2), + }, + ), + ] + result = basic_collection.insert(docs) + assert all(r.ok() for r in result) + + # Query with explicit output_fields + query_result = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + output_fields=["score"], + ) + assert len(query_result) == 2 + for doc in query_result: + assert "score" in doc.fields, ( + f"Doc {doc.id} missing 'score' field after add_column (output_fields explicit)" + ) + assert doc.fields["score"] in (100, 200) + + # Query with select-all (no output_fields) + query_result_all = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + ) + assert len(query_result_all) == 2 + for doc in query_result_all: + assert "score" in doc.fields, ( + f"Doc {doc.id} missing 'score' field after add_column (select all)" + ) + + def test_alter_column_rename_then_query_returns_new_name( + self, basic_collection: Collection + ): + """Regression test: query() should use the new field name after alter_column rename.""" + docs = [ + Doc( + id="1", + fields={"id": 1, "name": "alice", "weight": 60.0}, + vectors={ + "dense": generate_constant_vector(1, 128), + "sparse": generate_sparse_vector(1), + }, + ), + Doc( + id="2", + fields={"id": 2, "name": "bob", "weight": 70.0}, + vectors={ + "dense": generate_constant_vector(2, 128), + "sparse": generate_sparse_vector(2), + }, + ), + ] + result = basic_collection.insert(docs) + assert all(r.ok() for r in result) + + # Rename 'weight' -> 'mass' + basic_collection.alter_column("weight", new_name="mass") + + # Query with explicit output_fields using new name + query_result = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + output_fields=["mass"], + ) + assert len(query_result) == 2 + for doc in query_result: + assert "mass" in doc.fields, ( + f"Doc {doc.id} missing 'mass' field after alter_column rename" + ) + assert "weight" not in doc.fields, ( + f"Doc {doc.id} still has old name 'weight' after rename" + ) + + # Query with select-all + query_result_all = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + ) + assert len(query_result_all) == 2 + for doc in query_result_all: + assert "mass" in doc.fields, ( + f"Doc {doc.id} missing 'mass' in select-all after alter_column rename" + ) + assert "weight" not in doc.fields, ( + f"Doc {doc.id} still has old name 'weight' in select-all after rename" + ) + + def test_drop_column_then_query_excludes_dropped_field( + self, basic_collection: Collection + ): + """Regression test: query() should not return fields removed via drop_column().""" + basic_collection.add_column( + field_schema=FieldSchema("score", DataType.INT64, nullable=True), + ) + + docs = [ + Doc( + id="1", + fields={"id": 1, "name": "alice", "weight": 60.0, "score": 100}, + vectors={ + "dense": generate_constant_vector(1, 128), + "sparse": generate_sparse_vector(1), + }, + ), + Doc( + id="2", + fields={"id": 2, "name": "bob", "weight": 70.0, "score": 200}, + vectors={ + "dense": generate_constant_vector(2, 128), + "sparse": generate_sparse_vector(2), + }, + ), + ] + result = basic_collection.insert(docs) + assert all(r.ok() for r in result) + + # Verify field exists before drop + query_before = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + ) + assert all("score" in doc.fields for doc in query_before) + + # Drop the column + basic_collection.drop_column("score") + + # Query after drop - 'score' should not appear + query_after = basic_collection.query( + Query(field_name="dense", vector=generate_constant_vector(1, 128)), + topk=2, + ) + assert len(query_after) == 2 + for doc in query_after: + assert "score" not in doc.fields, ( + f"Doc {doc.id} still has 'score' after drop_column" + ) + assert "name" in doc.fields, ( + f"Doc {doc.id} missing 'name' - other fields should still be present" + ) diff --git a/python/zvec/model/collection.py b/python/zvec/model/collection.py index d513a43..6262581 100644 --- a/python/zvec/model/collection.py +++ b/python/zvec/model/collection.py @@ -130,6 +130,7 @@ class Collection: """ self._obj.CreateIndex(field_name, index_param, option) self._schema = CollectionSchema._from_core(self._obj.Schema()) + self._querier._schema = self._schema def drop_index(self, field_name: str) -> None: """Remove the index from a field. @@ -139,6 +140,7 @@ class Collection: """ self._obj.DropIndex(field_name) self._schema = CollectionSchema._from_core(self._obj.Schema()) + self._querier._schema = self._schema def optimize(self, option: OptimizeOption = OptimizeOption()) -> None: """Optimize the collection (e.g., merge segments, rebuild index). @@ -168,6 +170,7 @@ class Collection: """ self._obj.AddColumn(field_schema._get_object(), expression, option) self._schema = CollectionSchema._from_core(self._obj.Schema()) + self._querier._schema = self._schema def drop_column(self, field_name: str) -> None: """Remove a column from the collection. @@ -177,6 +180,7 @@ class Collection: """ self._obj.DropColumn(field_name) self._schema = CollectionSchema._from_core(self._obj.Schema()) + self._querier._schema = self._schema def alter_column( self, @@ -224,6 +228,7 @@ class Collection: option, ) self._schema = CollectionSchema._from_core(self._obj.Schema()) + self._querier._schema = self._schema # ========== Collection DDL Methods ========== @overload