from docarray import Document, DocumentArray
import numpy as np

n_docs = 500
n_features = 128

# works in sqlite
storage='sqlite'
da = DocumentArray(storage=storage).empty(n_docs)
da[:, 'embedding'] = [x for x in np.random.random((n_docs,n_features))]
assert da[:,'embedding'].shape == (n_docs, n_features)
da.embeddings = np.random.random((n_docs,n_features))
assert da[:,'embedding'].shape == (n_docs, n_features)

storage='pqlite'
da_pqlite = DocumentArray(storage=storage, config={'n_dim': n_features}).empty(n_docs)
da_pqlite[:, 'embedding'] = [x for x in np.random.random((len(da_pqlite),n_features))]
assert da_pqlite[:,'embedding'].shape == (500, 128)
da_pqlite.embeddings = np.random.random((len(da_pqlite),n_features))
assert da_pqlite[:,'embedding'].shape == (500, 128)

storage='pqlite'
da_pqlite = DocumentArray(storage=storage).empty(n_docs, config={'n_dim': n_features})
da_pqlite[:, 'embedding'] = [x for x in np.random.random((len(da_pqlite),n_features))]
assert da_pqlite[:,'embedding'].shape == (500, 128)
da_pqlite.embeddings = np.random.random((len(da_pqlite),n_features))
assert da_pqlite[:,'embedding'].shape == (500, 128)