Skip to content
This repository was archived by the owner on Sep 9, 2026. It is now read-only.

Commit 69dc861

Browse files
author
Joan Fontanals
authored
refactor: implementation of InMemoryExactNNIndex follows DBConfig way (#1649)
Signed-off-by: Joan Fontanals Martinez <joan.martinez@jina.ai> Signed-off-by: Joan Fontanals <joan.martinez@jina.ai>
1 parent eae4495 commit 69dc861

2 files changed

Lines changed: 39 additions & 23 deletions

File tree

‎docarray/index/backends/in_memory.py‎

Lines changed: 20 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -48,38 +48,38 @@ class InMemoryExactNNIndex(BaseDocIndex, Generic[TSchema]):
4848
def __init__(
4949
self,
5050
docs: Optional[DocList] = None,
51-
index_file_path: Optional[str] = None,
51+
db_config=None,
5252
**kwargs,
5353
):
5454
"""Initialize InMemoryExactNNIndex"""
55-
if 'db_config' in kwargs:
56-
kwargs.pop('db_config')
57-
super().__init__(db_config=None, **kwargs)
55+
super().__init__(db_config=db_config, **kwargs)
5856
self._runtime_config = self.RuntimeConfig()
57+
self._db_config = cast(InMemoryExactNNIndex.DBConfig, self._db_config)
58+
self._index_file_path = self._db_config.index_file_path
5959

60-
if docs and index_file_path:
60+
if docs and self._index_file_path:
6161
raise ValueError(
6262
'Initialize `InMemoryExactNNIndex` with either `docs` or '
6363
'`index_file_path`, not both. Provide `docs` for a fresh index, or '
6464
'`index_file_path` to use an existing file.'
6565
)
6666

67-
if index_file_path:
68-
if os.path.exists(index_file_path):
67+
if self._index_file_path:
68+
if os.path.exists(self._index_file_path):
6969
self._logger.info(
70-
f'Loading index from a binary file: {index_file_path}'
70+
f'Loading index from a binary file: {self._index_file_path}'
7171
)
7272
self._docs = DocList.__class_getitem__(
7373
cast(Type[BaseDoc], self._schema)
74-
).load_binary(file=index_file_path)
74+
).load_binary(file=self._index_file_path)
7575

7676
data_by_columns = self._get_col_value_dict(self._docs)
7777
self._update_subindex_data(self._docs)
7878
self._index_subindex(data_by_columns)
7979

8080
else:
8181
self._logger.warning(
82-
f'Index file does not exist: {index_file_path}. '
82+
f'Index file does not exist: {self._index_file_path}. '
8383
f'Initializing empty InMemoryExactNNIndex.'
8484
)
8585
self._docs = DocList.__class_getitem__(
@@ -137,6 +137,7 @@ def build(self, *args, **kwargs) -> Any:
137137
class DBConfig(BaseDocIndex.DBConfig):
138138
"""Dataclass that contains all "static" configurations of InMemoryExactNNIndex."""
139139

140+
index_file_path: Optional[str] = None
140141
default_column_config: Dict[Type, Dict[str, Any]] = field(
141142
default_factory=lambda: defaultdict(
142143
dict,
@@ -442,9 +443,16 @@ def __contains__(self, item: BaseDoc):
442443
f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'"
443444
)
444445

445-
def persist(self, file: str = 'in_memory_index.bin') -> None:
446+
def persist(self, file: Optional[str] = None) -> None:
446447
"""Persist InMemoryExactNNIndex into a binary file."""
447-
self._docs.save_binary(file=file)
448+
DEFAULT_INDEX_FILE_PATH = 'in_memory_index.bin'
449+
file_to_save = self._index_file_path or file
450+
if file_to_save is None:
451+
self._logger.warning(
452+
f'persisting index to {DEFAULT_INDEX_FILE_PATH} because no `index_file_path` has been used inside DBConfig and no `file` has been passed as argument'
453+
)
454+
file_to_save = file_to_save or DEFAULT_INDEX_FILE_PATH
455+
self._docs.save_binary(file=file_to_save)
448456

449457
def _get_root_doc_id(self, id: str, root: str, sub: str) -> str:
450458
"""Get the root_id given the id of a subindex Document and the root and subindex name

‎docs/user_guide/storing/index_in_memory.md‎

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -29,14 +29,20 @@ docs = DocList[MyDoc](MyDoc() for _ in range(10))
2929
doc_index = InMemoryExactNNIndex[MyDoc]()
3030
doc_index.index(docs)
3131

32-
# or in one step:
32+
# or in one step, create with inserted docs.
3333
doc_index = InMemoryExactNNIndex[MyDoc](docs)
3434
```
3535

36-
Additionally, you can preserve your index as a binary file and instantiate a new one using this file:
36+
Alternatively, you can pass an `index_file_path` argument to make sure that the index can be restored if persisted from that specific file.
3737
```python
3838
# Save your existing index as a binary file
39-
doc_index.persist('docs.bin')
39+
docs = DocList[MyDoc](MyDoc() for _ in range(10))
40+
41+
doc_index = InMemoryExactNNIndex[MyDoc](index_file_path='docs.bin')
42+
doc_index.index(docs)
43+
44+
# or in one step:
45+
doc_index.persist()
4046

4147
# Initialize a new document index using the saved binary file
4248
new_doc_index = InMemoryExactNNIndex[MyDoc](index_file_path='docs.bin')
@@ -46,21 +52,23 @@ new_doc_index = InMemoryExactNNIndex[MyDoc](index_file_path='docs.bin')
4652

4753
This section lays out the configurations and options that are specific to [InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex].
4854

49-
The `DBConfig` of [InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex] contains only one entry:
50-
the default mapping from Python types to column configurations.
55+
The `DBConfig` of [InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex] contains two entries:
56+
`index_file_path` and `default_column_mapping`, the default mapping from Python types to column configurations.
5157

5258
You can see in the [section below](#field-wise-configurations) how to override configurations for specific fields.
5359
If you want to set configurations globally, i.e. for all vector fields in your Documents, you can do that using `DBConfig` or passing it at `__init__`::
5460

5561
```python
5662
from collections import defaultdict
5763
from docarray.typing import AbstractTensor
58-
new_doc_index = InMemoryExactNNIndex[MyDoc](default_column_config=defaultdict(
59-
dict,
60-
{
61-
AbstractTensor: {'space': 'cosine_sim'},
62-
},
63-
))
64+
new_doc_index = InMemoryExactNNIndex[MyDoc](
65+
default_column_config=defaultdict(
66+
dict,
67+
{
68+
AbstractTensor: {'space': 'cosine_sim'},
69+
},
70+
)
71+
)
6472
```
6573

6674
This will set the default configuration for all vector fields to the one specified in the example above.

0 commit comments

Comments
 (0)