# create a Document Store
from docarray.storage.backends.HnswDocStore import HnswDocumentIndex
from pydantic import Field
from docarray.typing import NdArray
from docarray import DocumentArray, BaseDocument

class MyDocument(BaseDocument):
    tens: NdArray[10]

store = HnswDocumentIndex[MyDocument](work_dir='path/to/work/dir')

# index some data    
data = DocumentArray[MyDocument]([MyDocument(tens=np.random.randn(10)) for i in range(10)])
store.index(data)

# find by np array
result = store.find(np.random.random((10,)), embedding_field='tens', limit=5)
result = store.find_batched(np.random.random((2, 10)), embedding_field='tens', limit=5)

# find by Document
result = store.find(MyDocument(tens=np.random.randn(10)), search_field='tens', limit=5)
result = store.find_batched(DocumentArray[MyDocument]([MyDocument(tens=np.random.randn(10)) for i in range(10)]), search_field='tens', limit=5)

# delete data
del store[data.id]  # delete all

# advanced configs
class MyDocument(BaseDocument):
    tens_one: NdArray = Field(dim=10, space='cosine')
    tens_two: NdArray = Field(dim=10, space='l2')
store = HnswDocumentStore[MyDocument](work_dir='path/to/work/dir')

# select search field
data = DocumentArray[MyDocument]([MyDocument(tens_one=np.random.randn(10), tens_two=np.random.randn(10)) for i in range(10)])
store.index(data)
result = store.find(np.random.random((10,)), search_field='tens_one', limit=5)
result = store.find(np.random.random((10,)), search_field='tens_two', limit=5)

# use query builder
q = store.build_query()
q.find(query=data[0], search_field='tens_one', limit=5)
q.find(query=data[1], search_field='tens_two', limit=5)
q.filter(filter_query={'tens_two': {'$exists': True}})
result = store.execute_query(q.build())

# nested search
class InnerDocument(BaseDocument):
    tens_one: NdArray = Field(dim=10)

class MyDocument(BaseDocument):
    d: InnerDocument
    tens: NdArray = Field(dim=10)

store = HnswDocumentStore[MyDocument](work_dir='path/to/work/dir')
data = DocumentArray[MyDocument]([MyDocument(d=InnerDocument(tens_one=np.random.randn(10)), tens=np.random.randn(10)) for i in range(10)])
store.index(data)

result = store.find(np.random.random((10,)), search_field='tens', limit=5)
result = store.find(np.random.random((10,)), search_field='d__tens_one', limit=5)