fix: del remove attribute (#304)

davidbp · davidbp · commit 155c635d71b4 · 2022-04-28T12:08:53.000+02:00
* fix: del remove attribute

* test: add test that covers fix
diff --git a/docarray/array/mixins/delitem.py b/docarray/array/mixins/delitem.py
@@ -56,6 +56,8 @@ def __delitem__(self, index: 'DocumentArrayIndexType'):
                     for _d in self[index[0]]:
                         for _aa in _attrs:
                             self._set_doc_attr_by_id(_d.id, _aa, None)
+                            _d.pop(_aa)
+
             elif isinstance(index[0], bool):
                 self._del_docs_by_mask(index)
             elif isinstance(index[0], int):
diff --git a/tests/unit/array/mixins/test_del.py b/tests/unit/array/mixins/test_del.py
@@ -2,6 +2,7 @@
 
 from docarray import DocumentArray, Document
 from docarray.array.weaviate import DocumentArrayWeaviate
+import numpy as np
 
 
 @pytest.fixture()
@@ -81,3 +82,26 @@ def test_del_da_persist(da_cls, config, persist, docs, start_storage):
         assert len(da2) == len(docs)
     else:
         assert len(da2) == 0
+
+
+def test_del_da_attribute():
+
+    da = DocumentArray(
+        [
+            Document(embedding=np.array([1, 2, 3]), text='d1'),
+            Document(embedding=np.array([1, 2, 3]), text='d2'),
+        ]
+    )
+
+    q = DocumentArray(
+        [
+            Document(embedding=np.array([4, 5, 6]), text='q1'),
+            Document(embedding=np.array([2, 3, 4]), text='q1'),
+        ]
+    )
+
+    da.match(q)
+    del da[...][:, 'embedding']
+
+    for d in da:
+        assert d.embedding is None
