From 0db7ef95dba12cc41ef7d4c31f5e1fb877f2f972 Mon Sep 17 00:00:00 2001 From: Sami Jaghouar Date: Wed, 11 Jan 2023 12:04:01 +0100 Subject: [PATCH 1/2] refactor: better management of stack mod Signed-off-by: Sami Jaghouar --- docarray/array/array_stacked.py | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/docarray/array/array_stacked.py b/docarray/array/array_stacked.py index acc9906f400..147bfc9ae28 100644 --- a/docarray/array/array_stacked.py +++ b/docarray/array/array_stacked.py @@ -152,7 +152,6 @@ def _create_columns( if is_tensor_union(type_): val = tensor_type.get_comp_backend().none_value() columns_to_stack[field_to_stack].append(val) - delattr(doc, field_to_stack) for field_to_stack, to_stack in columns_to_stack.items(): @@ -168,6 +167,10 @@ def _create_columns( elif issubclass(type_, AbstractTensor): columns[field_to_stack] = type_.__docarray_stack__(to_stack) # type: ignore # noqa: E501 + for field, column in columns.items(): + for doc, val in zip(docs, column): + setattr(doc, field, val) + return columns def _get_array_attribute( @@ -224,18 +227,6 @@ def __iter__(self): for i in range(len(self)): yield self[i] - def __getitem_without_columns__(self, item): # note this should handle slices - """Return the document at the given index with the columns item put to None""" - doc = self._docs[item] - for field in self._columns.keys(): - if hasattr(doc, field): - delattr(doc, field) - return doc - - def __iter_without_columns__(self): - for i in range(len(self)): - yield self.__getitem_without_columns__(i) - def __len__(self): return len(self._docs) @@ -262,7 +253,7 @@ def to_protobuf(self) -> 'DocumentArrayStackedProto': ) da_proto = DocumentArrayProto() - for doc in self.__iter_without_columns__(): + for doc in self: da_proto.docs.append(doc.to_protobuf()) columns_proto: Dict[str, UnionArrayProto] = dict() From f9b4af6c6b721e702b007705b7e503046817095b Mon Sep 17 00:00:00 2001 From: Sami Jaghouar Date: Wed, 11 Jan 2023 14:23:00 +0100 Subject: [PATCH 2/2] fix: mypy Signed-off-by: Sami Jaghouar --- docarray/array/array_stacked.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docarray/array/array_stacked.py b/docarray/array/array_stacked.py index 147bfc9ae28..71d5b3547a8 100644 --- a/docarray/array/array_stacked.py +++ b/docarray/array/array_stacked.py @@ -167,9 +167,9 @@ def _create_columns( elif issubclass(type_, AbstractTensor): columns[field_to_stack] = type_.__docarray_stack__(to_stack) # type: ignore # noqa: E501 - for field, column in columns.items(): + for field_name, column in columns.items(): for doc, val in zip(docs, column): - setattr(doc, field, val) + setattr(doc, field_name, val) return columns