Skip to content

Commit

Permalink
feat: support list of strings for Weaviate
Browse files Browse the repository at this point in the history
  • Loading branch information
JoanFM committed Feb 8, 2024
1 parent 5037635 commit af219af
Showing 1 changed file with 23 additions and 1 deletion.
24 changes: 23 additions & 1 deletion docarray/index/backends/weaviate.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
TypeVar,
Union,
cast,
get_origin,
get_args,
)

import numpy as np
Expand All @@ -43,7 +45,6 @@
TSchema = TypeVar('TSchema', bound=BaseDoc)
T = TypeVar('T', bound='WeaviateDocumentIndex')


DEFAULT_BATCH_CONFIG = {
"batch_size": 20,
"dynamic": False,
Expand Down Expand Up @@ -210,6 +211,8 @@ def _create_schema(self) -> None:
self.bytes_columns.append(column_name)
if column_info.db_type == 'number[]':
self.nonembedding_array_columns.append(column_name)
if column_info.db_type == 'text[]':
self.nonembedding_array_columns.append(column_name)

Check warning on line 215 in docarray/index/backends/weaviate.py

View check run for this annotation

Codecov / codecov/patch

docarray/index/backends/weaviate.py#L215

Added line #L215 was not covered by tests
prop = {
"name": column_name
if column_name != 'id'
Expand Down Expand Up @@ -253,6 +256,8 @@ class DBConfig(BaseDocIndex.DBConfig):
'number': {},
'boolean': {},
'number[]': {},
'int[]': {},
'text[]': {},
'blob': {},
}
)
Expand Down Expand Up @@ -717,6 +722,23 @@ def python_type_to_db_type(self, python_type: Type) -> Any:
bytes: 'blob',
}

if get_origin(python_type) == list:
py_weaviate_list_type_map = {

Check warning on line 726 in docarray/index/backends/weaviate.py

View check run for this annotation

Codecov / codecov/patch

docarray/index/backends/weaviate.py#L726

Added line #L726 was not covered by tests
int: 'int[]',
float: 'number[]',
str: 'text[]',
}

container_type = None
args = get_args(python_type)
if args:
container_type = args[0]
if (

Check warning on line 736 in docarray/index/backends/weaviate.py

View check run for this annotation

Codecov / codecov/patch

docarray/index/backends/weaviate.py#L732-L736

Added lines #L732 - L736 were not covered by tests
container_type is not None
and container_type in py_weaviate_list_type_map
):
return py_weaviate_list_type_map[container_type]

Check warning on line 740 in docarray/index/backends/weaviate.py

View check run for this annotation

Codecov / codecov/patch

docarray/index/backends/weaviate.py#L740

Added line #L740 was not covered by tests

for py_type, weaviate_type in py_weaviate_type_map.items():
if safe_issubclass(python_type, py_type):
return weaviate_type
Expand Down

0 comments on commit af219af

Please sign in to comment.