@@ -86,6 +86,62 @@ def __getitem__(self, item): |
86 | 86 | return super().__getitem__(item) |
87 | 87 |
|
88 | 88 | def stack(self): |
| 89 | + """ |
| 90 | + Calling this method will make the DocumentArray enter in stacked mode |
| 91 | +
|
| 92 | + This mode will stack all the Tensor and Nested fields of the documents in the |
| 93 | + array into columns. This is useful when you want to perform operations on the |
| 94 | + whole array at once. Especially when it involves working on Tensor directly |
| 95 | + where you want to avoid for loops and work on batch directly. |
| 96 | +
|
| 97 | + By default, a DocumentArray is not in stack mode, it means that each Document |
| 98 | + holds its own data and that the DocumentArray is just a python list of |
| 99 | + Documents. When you access fields of the DocumentArray you actually loop over |
| 100 | + the list of Documents and access the field of each Document. |
| 101 | +
|
| 102 | + In stack mode, accessing or setting fields of the DocumentArray will access or |
| 103 | + set the column of the array. |
| 104 | +
|
| 105 | + When entering stack mode the DocumentArray will create a column for each field |
| 106 | + of the Document that are Tensor or that are Nested Document that contains at |
| 107 | + least one Tensor field. |
| 108 | +
|
| 109 | + IMPORTANT: in stacked_mode you cannot add or remove Document from the array. |
| 110 | + You can only modify the fields of the DocumentArray. This is intentially done to |
| 111 | + avoid inconsistencies between the columns and the list of Documents and because |
| 112 | + adding to a column is slow. If you need to do so please use the unstack method |
| 113 | + to exit stacked mode. |
| 114 | +
|
| 115 | + EXAMPLE USAGE |
| 116 | + .. code-block:: python |
| 117 | + from docarray import Document, DocumentArray |
| 118 | + from docarray.typing import NdArray |
| 119 | +
|
| 120 | +
|
| 121 | + class Image(Document): |
| 122 | + tensor: NdArray[100] |
| 123 | +
|
| 124 | +
|
| 125 | + batch = DocumentArray[Image]([Image(tensor=np.zeros((100))) \ |
| 126 | + for _ in range(10)]) |
| 127 | +
|
| 128 | + batch.stack() |
| 129 | +
|
| 130 | + print(batch[0].tensor[0]) |
| 131 | + # >>> 0 |
| 132 | +
|
| 133 | + print(batch.tensor.shape) |
| 134 | + # >>> (10, 3, 224, 224) |
| 135 | +
|
| 136 | + batch.tensor = np.ones((10, 100)) |
| 137 | +
|
| 138 | + print(batch[0].tensor[0]) |
| 139 | + # >>> 1 |
| 140 | +
|
| 141 | + batch.append(Image(tensor=np.zeros((100)))) |
| 142 | + # >>> raise RuntimeError('Cannot call append when the document array is in |
| 143 | + # >>> stack mode' |
| 144 | + """ |
89 | 145 |
|
90 | 146 | if not (self.is_stacked()): |
91 | 147 |
|
|