Skip to content

Commit 31c2bb9

Browse files
author
Joan Fontanals
authored
feat: add description and example to ID field of BaseDoc (#1737)
Signed-off-by: Joan Fontanals Martinez <[email protected]>
1 parent d2c82d4 commit 31c2bb9

File tree

5 files changed

+32
-35
lines changed

5 files changed

+32
-35
lines changed

‎docarray/base_doc/doc.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,11 @@ class MyDoc(BaseDoc):
7272
https://docs.pydantic.dev/usage/models/) and can be used in a similar way.
7373
"""
7474

75-
id: Optional[ID] = Field(default_factory=lambda: ID(os.urandom(16).hex()))
75+
id: Optional[ID] = Field(
76+
description='The ID of the BaseDoc. This is useful for indexing in vector stores. If not set by user, it will automatically be assigned a random value',
77+
default_factory=lambda: ID(os.urandom(16).hex()),
78+
example=os.urandom(16).hex(),
79+
)
7680

7781
class Config:
7882
json_loads = orjson.loads

‎docarray/index/backends/elastic.py‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,9 @@
6767

6868

6969
class ElasticDocIndex(BaseDocIndex, Generic[TSchema]):
70+
_index_vector_params: Optional[Tuple[str]] = ('dims', 'similarity', 'index')
71+
_index_vector_options: Optional[Tuple[str]] = ('m', 'ef_construction')
72+
7073
def __init__(self, db_config=None, **kwargs):
7174
"""Initialize ElasticDocIndex"""
7275
super().__init__(db_config=db_config, **kwargs)
@@ -82,9 +85,6 @@ def __init__(self, db_config=None, **kwargs):
8285
self._logger.debug('ElasticSearch client has been created')
8386

8487
# ElasticSearh index setup
85-
self._index_vector_params = ('dims', 'similarity', 'index')
86-
self._index_vector_options = ('m', 'ef_construction')
87-
8888
mappings: Dict[str, Any] = {
8989
'dynamic': True,
9090
'_source': {'enabled': 'true'},
@@ -572,20 +572,23 @@ def _filter_by_parent_id(self, id: str) -> List[str]:
572572
# Helpers #
573573
###############################################
574574

575-
def _create_index_mapping(self, col: '_ColumnInfo') -> Dict[str, Any]:
575+
@classmethod
576+
def _create_index_mapping(cls, col: '_ColumnInfo') -> Dict[str, Any]:
576577
"""Create a new HNSW index for a column, and initialize it."""
577578

578579
index = {'type': col.config['type'] if 'type' in col.config else col.db_type}
579580

580581
if col.db_type == 'dense_vector':
581-
for k in self._index_vector_params:
582-
index[k] = col.config[k]
582+
if cls._index_vector_params is not None:
583+
for k in cls._index_vector_params:
584+
index[k] = col.config[k]
583585
if col.n_dim:
584586
index['dims'] = col.n_dim
585-
index['index_options'] = dict(
586-
(k, col.config[k]) for k in self._index_vector_options
587-
)
588-
index['index_options']['type'] = 'hnsw'
587+
if cls._index_vector_options is not None:
588+
index['index_options'] = dict(
589+
(k, col.config[k]) for k in cls._index_vector_options
590+
)
591+
index['index_options']['type'] = 'hnsw'
589592
return index
590593

591594
def _send_requests(

‎docarray/index/backends/elasticv7.py‎

Lines changed: 5 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
import warnings
22
from dataclasses import dataclass
3-
from typing import Any, Dict, List, Optional, Sequence, TypeVar, Union
3+
from typing import Any, Dict, List, Optional, Sequence, TypeVar, Union, Tuple
44

55
import numpy as np
66
from pydantic import parse_obj_as
77

88
from docarray import BaseDoc
99
from docarray.index import ElasticDocIndex
10-
from docarray.index.abstract import BaseDocIndex, _ColumnInfo
10+
from docarray.index.abstract import BaseDocIndex
1111
from docarray.typing import AnyTensor
1212
from docarray.typing.tensor.ndarray import NdArray
1313
from docarray.utils.find import _FindResult
@@ -17,6 +17,9 @@
1717

1818

1919
class ElasticV7DocIndex(ElasticDocIndex):
20+
_index_vector_params: Optional[Tuple[str]] = ('dims',)
21+
_index_vector_options: Optional[Tuple[str]] = None
22+
2023
def __init__(self, db_config=None, **kwargs):
2124
"""Initialize ElasticV7DocIndex"""
2225
from elasticsearch import __version__ as __es__version__
@@ -130,19 +133,6 @@ def execute_query(self, query: Dict[str, Any], *args, **kwargs) -> Any:
130133
# Helpers #
131134
###############################################
132135

133-
# ElasticSearch helpers
134-
def _create_index_mapping(self, col: '_ColumnInfo') -> Dict[str, Any]:
135-
"""Create a new HNSW index for a column, and initialize it."""
136-
137-
index = col.config.copy()
138-
if 'type' not in index:
139-
index['type'] = col.db_type
140-
141-
if col.db_type == 'dense_vector' and col.n_dim:
142-
index['dims'] = col.n_dim
143-
144-
return index
145-
146136
def _form_search_body(self, query: np.ndarray, limit: int, search_field: str = '') -> Dict[str, Any]: # type: ignore
147137
body = {
148138
'size': limit,

‎docarray/index/backends/hnswlib.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ def __init__(self, db_config=None, **kwargs):
115115
sub_docs_exist = True
116116
if safe_issubclass(col.docarray_type, AnyDocArray):
117117
continue
118-
if not col.config:
118+
if not col.config or 'dim' not in col.config:
119119
# non-tensor type; don't create an index
120120
continue
121121
if not load_existing and (

‎tests/index/base_classes/test_base_doc_store.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ def test_create_columns():
157157
assert index._column_infos['id'].docarray_type == ID
158158
assert index._column_infos['id'].db_type == str
159159
assert index._column_infos['id'].n_dim is None
160-
assert index._column_infos['id'].config == {'hi': 'there'}
160+
assert index._column_infos['id'].config['hi'] == 'there'
161161

162162
assert issubclass(index._column_infos['tens'].docarray_type, AbstractTensor)
163163
assert index._column_infos['tens'].db_type == str
@@ -171,7 +171,7 @@ def test_create_columns():
171171
assert index._column_infos['id'].docarray_type == ID
172172
assert index._column_infos['id'].db_type == str
173173
assert index._column_infos['id'].n_dim is None
174-
assert index._column_infos['id'].config == {'hi': 'there'}
174+
assert index._column_infos['id'].config['hi'] == 'there'
175175

176176
assert issubclass(index._column_infos['tens_one'].docarray_type, AbstractTensor)
177177
assert index._column_infos['tens_one'].db_type == str
@@ -190,7 +190,7 @@ def test_create_columns():
190190
assert index._column_infos['id'].docarray_type == ID
191191
assert index._column_infos['id'].db_type == str
192192
assert index._column_infos['id'].n_dim is None
193-
assert index._column_infos['id'].config == {'hi': 'there'}
193+
assert index._column_infos['id'].config['hi'] == 'there'
194194

195195
assert issubclass(index._column_infos['d__tens'].docarray_type, AbstractTensor)
196196
assert index._column_infos['d__tens'].db_type == str
@@ -214,7 +214,7 @@ def test_create_columns():
214214
assert index._subindices['d']._column_infos['id'].docarray_type == ID
215215
assert index._subindices['d']._column_infos['id'].db_type == str
216216
assert index._subindices['d']._column_infos['id'].n_dim is None
217-
assert index._subindices['d']._column_infos['id'].config == {'hi': 'there'}
217+
assert index._subindices['d']._column_infos['id'].config['hi'] == 'there'
218218

219219
assert issubclass(
220220
index._subindices['d']._column_infos['tens'].docarray_type, AbstractTensor
@@ -262,10 +262,10 @@ def test_create_columns():
262262
assert (
263263
index._subindices['d_root']._subindices['d']._column_infos['id'].n_dim is None
264264
)
265-
assert index._subindices['d_root']._subindices['d']._column_infos['id'].config == {
266-
'hi': 'there'
267-
}
268-
265+
assert (
266+
index._subindices['d_root']._subindices['d']._column_infos['id'].config['hi']
267+
== 'there'
268+
)
269269
assert issubclass(
270270
index._subindices['d_root']
271271
._subindices['d']

0 commit comments

Comments
 (0)