Skip to content

ElasticV7DocIndex

docarray.index.backends.elasticv7.ElasticV7DocIndex

Bases: ElasticDocIndex

Source code in docarray/index/backends/elasticv7.py 
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
class ElasticV7DocIndex(ElasticDocIndex):
    def __init__(self, db_config=None, **kwargs):
        """Initialize ElasticV7DocIndex"""
        from elasticsearch import __version__ as __es__version__

        if __es__version__[0] > 7:
            raise ImportError(
                'ElasticV7DocIndex requires the elasticsearch library to be version 7.10.1'
            )

        super().__init__(db_config, **kwargs)

    ###############################################
    # Inner classes for query builder and configs #
    ###############################################

    class QueryBuilder(ElasticDocIndex.QueryBuilder):
        def build(self, *args, **kwargs) -> Any:
            """Build the elastic search v7 query object."""
            if (
                'script_score' in self._query['query']
                and 'bool' in self._query['query']
                and len(self._query['query']['bool']) > 0
            ):
                self._query['query']['script_score']['query'] = {}
                self._query['query']['script_score']['query']['bool'] = self._query[
                    'query'
                ]['bool']
                del self._query['query']['bool']

            return self._query

        def find(
            self,
            query: Union[AnyTensor, BaseDoc],
            search_field: str = 'embedding',
            limit: int = 10,
            num_candidates: Optional[int] = None,
        ):
            """
            Find k-nearest neighbors of the query.

            :param query: query vector for KNN/ANN search. Has single axis.
            :param search_field: name of the field to search on
            :param limit: maximum number of documents to return per query
            :return: self
            """
            if num_candidates:
                warnings.warn('`num_candidates` is not supported in ElasticV7DocIndex')

            if isinstance(query, BaseDoc):
                query_vec = BaseDocIndex._get_values_by_column([query], search_field)[0]
            else:
                query_vec = query
            query_vec_np = BaseDocIndex._to_numpy(self._outer_instance, query_vec)
            self._query['size'] = limit
            self._query['query'][
                'script_score'
            ] = self._outer_instance._form_search_body(
                query_vec_np, limit, search_field
            )[
                'query'
            ][
                'script_score'
            ]

            return self

    @dataclass
    class DBConfig(ElasticDocIndex.DBConfig):
        """Dataclass that contains all "static" configurations of ElasticDocIndex."""

        hosts: Union[str, List[str], None] = 'http://localhost:9200'  # type: ignore

    @dataclass
    class RuntimeConfig(ElasticDocIndex.RuntimeConfig):
        """Dataclass that contains all "dynamic" configurations of ElasticDocIndex."""

        def dense_vector_config(self):
            return {'dims': 128}

    ###############################################
    # Implementation of abstract methods          #
    ###############################################

    def execute_query(self, query: Dict[str, Any], *args, **kwargs) -> Any:
        """
        Execute a query on the ElasticDocIndex.

        Can take two kinds of inputs:

        1. A native query of the underlying database. This is meant as a passthrough so that you
        can enjoy any functionality that is not available through the Document index API.
        2. The output of this Document index' `QueryBuilder.build()` method.

        :param query: the query to execute
        :return: the result of the query
        """
        if args or kwargs:
            raise ValueError(
                f'args and kwargs not supported for `execute_query` on {type(self)}'
            )

        resp = self._client.search(index=self.index_name, body=query)
        docs, scores = self._format_response(resp)

        return _FindResult(documents=docs, scores=parse_obj_as(NdArray, scores))

    ###############################################
    # Helpers                                     #
    ###############################################

    # ElasticSearch helpers
    def _create_index_mapping(self, col: '_ColumnInfo') -> Dict[str, Any]:
        """Create a new HNSW index for a column, and initialize it."""

        index = col.config.copy()
        if 'type' not in index:
            index['type'] = col.db_type

        if col.db_type == 'dense_vector' and col.n_dim:
            index['dims'] = col.n_dim

        return index

    def _form_search_body(self, query: np.ndarray, limit: int, search_field: str = '') -> Dict[str, Any]:  # type: ignore
        body = {
            'size': limit,
            'query': {
                'script_score': {
                    'query': {'match_all': {}},
                    'script': {
                        'source': f'cosineSimilarity(params.query_vector, \'{search_field}\') + 1.0',
                        'params': {'query_vector': query},
                    },
                }
            },
        }
        return body

    ###############################################
    # API Wrappers                                #
    ###############################################

    def _client_put_mapping(self, mappings: Dict[str, Any]):
        self._client.indices.put_mapping(index=self.index_name, body=mappings)

    def _client_create(self, mappings: Dict[str, Any]):
        body = {'mappings': mappings}
        self._client.indices.create(index=self.index_name, body=body)

    def _client_put_settings(self, settings: Dict[str, Any]):
        self._client.indices.put_settings(index=self.index_name, body=settings)

    def _client_mget(self, ids: Sequence[str]):
        return self._client.mget(index=self.index_name, body={'ids': ids})

    def _client_search(self, **kwargs):
        return self._client.search(index=self.index_name, body=kwargs)

    def _client_msearch(self, request: List[Dict[str, Any]]):
        return self._client.msearch(index=self.index_name, body=request)

DBConfig dataclass

Bases: ElasticDocIndex.DBConfig

Dataclass that contains all "static" configurations of ElasticDocIndex.

Source code in docarray/index/backends/elasticv7.py 
87
88
89
90
91
@dataclass
class DBConfig(ElasticDocIndex.DBConfig):
    """Dataclass that contains all "static" configurations of ElasticDocIndex."""

    hosts: Union[str, List[str], None] = 'http://localhost:9200'  # type: ignore

QueryBuilder

Bases: ElasticDocIndex.QueryBuilder

Source code in docarray/index/backends/elasticv7.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
class QueryBuilder(ElasticDocIndex.QueryBuilder):
    def build(self, *args, **kwargs) -> Any:
        """Build the elastic search v7 query object."""
        if (
            'script_score' in self._query['query']
            and 'bool' in self._query['query']
            and len(self._query['query']['bool']) > 0
        ):
            self._query['query']['script_score']['query'] = {}
            self._query['query']['script_score']['query']['bool'] = self._query[
                'query'
            ]['bool']
            del self._query['query']['bool']

        return self._query

    def find(
        self,
        query: Union[AnyTensor, BaseDoc],
        search_field: str = 'embedding',
        limit: int = 10,
        num_candidates: Optional[int] = None,
    ):
        """
        Find k-nearest neighbors of the query.

        :param query: query vector for KNN/ANN search. Has single axis.
        :param search_field: name of the field to search on
        :param limit: maximum number of documents to return per query
        :return: self
        """
        if num_candidates:
            warnings.warn('`num_candidates` is not supported in ElasticV7DocIndex')

        if isinstance(query, BaseDoc):
            query_vec = BaseDocIndex._get_values_by_column([query], search_field)[0]
        else:
            query_vec = query
        query_vec_np = BaseDocIndex._to_numpy(self._outer_instance, query_vec)
        self._query['size'] = limit
        self._query['query'][
            'script_score'
        ] = self._outer_instance._form_search_body(
            query_vec_np, limit, search_field
        )[
            'query'
        ][
            'script_score'
        ]

        return self

build(*args, **kwargs)

Build the elastic search v7 query object.

Source code in docarray/index/backends/elasticv7.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def build(self, *args, **kwargs) -> Any:
    """Build the elastic search v7 query object."""
    if (
        'script_score' in self._query['query']
        and 'bool' in self._query['query']
        and len(self._query['query']['bool']) > 0
    ):
        self._query['query']['script_score']['query'] = {}
        self._query['query']['script_score']['query']['bool'] = self._query[
            'query'
        ]['bool']
        del self._query['query']['bool']

    return self._query

find(query, search_field='embedding', limit=10, num_candidates=None)

Find k-nearest neighbors of the query.

Parameters:

Name Type Description Default
query Union[AnyTensor, BaseDoc]

query vector for KNN/ANN search. Has single axis.

 required
search_field str

name of the field to search on

 'embedding'
limit int

maximum number of documents to return per query

 10

Returns:

Type Description

self
Source code in docarray/index/backends/elasticv7.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def find(
    self,
    query: Union[AnyTensor, BaseDoc],
    search_field: str = 'embedding',
    limit: int = 10,
    num_candidates: Optional[int] = None,
):
    """
    Find k-nearest neighbors of the query.

    :param query: query vector for KNN/ANN search. Has single axis.
    :param search_field: name of the field to search on
    :param limit: maximum number of documents to return per query
    :return: self
    """
    if num_candidates:
        warnings.warn('`num_candidates` is not supported in ElasticV7DocIndex')

    if isinstance(query, BaseDoc):
        query_vec = BaseDocIndex._get_values_by_column([query], search_field)[0]
    else:
        query_vec = query
    query_vec_np = BaseDocIndex._to_numpy(self._outer_instance, query_vec)
    self._query['size'] = limit
    self._query['query'][
        'script_score'
    ] = self._outer_instance._form_search_body(
        query_vec_np, limit, search_field
    )[
        'query'
    ][
        'script_score'
    ]

    return self

RuntimeConfig dataclass

Bases: ElasticDocIndex.RuntimeConfig

Dataclass that contains all "dynamic" configurations of ElasticDocIndex.

Source code in docarray/index/backends/elasticv7.py 
93
94
95
96
97
98
@dataclass
class RuntimeConfig(ElasticDocIndex.RuntimeConfig):
    """Dataclass that contains all "dynamic" configurations of ElasticDocIndex."""

    def dense_vector_config(self):
        return {'dims': 128}

__init__(db_config=None, **kwargs)

Initialize ElasticV7DocIndex

Source code in docarray/index/backends/elasticv7.py 
20
21
22
23
24
25
26
27
28
29
def __init__(self, db_config=None, **kwargs):
    """Initialize ElasticV7DocIndex"""
    from elasticsearch import __version__ as __es__version__

    if __es__version__[0] > 7:
        raise ImportError(
            'ElasticV7DocIndex requires the elasticsearch library to be version 7.10.1'
        )

    super().__init__(db_config, **kwargs)

execute_query(query, *args, **kwargs)

Execute a query on the ElasticDocIndex.

Can take two kinds of inputs:

  1. A native query of the underlying database. This is meant as a passthrough so that you can enjoy any functionality that is not available through the Document index API.
  2. The output of this Document index' QueryBuilder.build() method.

Parameters:

Name Type Description Default
query Dict[str, Any]

the query to execute

 required

Returns:

Type Description
Any

the result of the query
Source code in docarray/index/backends/elasticv7.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def execute_query(self, query: Dict[str, Any], *args, **kwargs) -> Any:
    """
    Execute a query on the ElasticDocIndex.

    Can take two kinds of inputs:

    1. A native query of the underlying database. This is meant as a passthrough so that you
    can enjoy any functionality that is not available through the Document index API.
    2. The output of this Document index' `QueryBuilder.build()` method.

    :param query: the query to execute
    :return: the result of the query
    """
    if args or kwargs:
        raise ValueError(
            f'args and kwargs not supported for `execute_query` on {type(self)}'
        )

    resp = self._client.search(index=self.index_name, body=query)
    docs, scores = self._format_response(resp)

    return _FindResult(documents=docs, scores=parse_obj_as(NdArray, scores))