diff --git a/docarray/array/storage/memory/getsetdel.py b/docarray/array/storage/memory/getsetdel.py index f48a4f3246e..17fc2d98607 100644 --- a/docarray/array/storage/memory/getsetdel.py +++ b/docarray/array/storage/memory/getsetdel.py @@ -15,6 +15,8 @@ class GetSetDelMixin(BaseGetSetDelMixin): @needs_id2offset_rebuild def _del_docs_by_mask(self, mask: Sequence[bool]): + if len(mask) < len(self._data): + mask = mask + [False for _ in range(len(self._data) - len(mask))] self._data = list(itertools.compress(self._data, (not _i for _i in mask))) @needs_id2offset_rebuild diff --git a/tests/unit/array/mixins/test_del.py b/tests/unit/array/mixins/test_del.py index dffd573e85d..52f132624fa 100644 --- a/tests/unit/array/mixins/test_del.py +++ b/tests/unit/array/mixins/test_del.py @@ -42,7 +42,7 @@ def test_del_all(docs, to_delete): ], ) def test_del_boolean_mask(docs, to_delete, missing_id): - + all_ids = docs[:, 'id'] # assert each missing_id is present before deleting for m_id in missing_id: assert m_id in docs[:, 'id'] @@ -52,6 +52,8 @@ def test_del_boolean_mask(docs, to_delete, missing_id): # assert each missing_id is NOT present AFTER deleting for m_id in missing_id: assert m_id not in docs[:, 'id'] + for m_id in filter(lambda id: id not in missing_id, all_ids): + assert m_id in docs[:, 'id'] @pytest.mark.parametrize(