@@ -86,6 +86,62 @@ def __getitem__(self, item):
8686 return super ().__getitem__ (item )
8787
8888 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+
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+ """
89145
90146 if not (self .is_stacked ()):
91147
0 commit comments