Skip to content

Commit 3373c79

Browse files
fix: allow null auto_id primary column in DataFrame insert
Co-authored-by: chatgpt-codex-connector[bot] <199175422+chatgpt-codex-connector[bot]@users.noreply.github.com> Signed-off-by: MiaoXing <73048950+VectorPeak@users.noreply.github.com>
1 parent 0ee0227 commit 3373c79

2 files changed

Lines changed: 13 additions & 1 deletion

File tree

pymilvus/orm/schema.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1328,7 +1328,7 @@ def check_insert_schema(schema: CollectionSchema, data: Union[List[List], pd.Dat
13281328
msg = f"Expect no data for auto_id primary field: {schema.primary_field.name}"
13291329
raise DataNotMatchException(message=msg)
13301330
columns = list(data.columns)
1331-
columns.remove(schema.primary_field)
1331+
columns.remove(schema.primary_field.name)
13321332
data = data[columns]
13331333

13341334
tmp_fields = list(

tests/unit/orm/test_schema.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1756,6 +1756,18 @@ def test_check_insert_schema_auto_id_with_data(self):
17561756
with pytest.raises(DataNotMatchException, match="auto_id"):
17571757
check_insert_schema(schema, insert_data)
17581758

1759+
def test_check_insert_schema_auto_id_with_null_pk_column(self):
1760+
"""Test auto_id DataFrame inserts tolerate an all-null pk column."""
1761+
schema = CollectionSchema(
1762+
[
1763+
FieldSchema("id", DataType.INT64, is_primary=True, auto_id=True),
1764+
FieldSchema("vec", DataType.FLOAT_VECTOR, dim=2),
1765+
]
1766+
)
1767+
insert_data = pd.DataFrame({"id": [None, None], "vec": [[1.0, 2.0], [3.0, 4.0]]})
1768+
1769+
check_insert_schema(schema, insert_data)
1770+
17591771

17601772
class TestCheckUpsertSchema:
17611773
"""Tests for check_upsert_schema function."""

0 commit comments

Comments
 (0)