import copy import logging from abc import ABC, abstractmethod from dataclasses import dataclass, field, replace from typing import ( TYPE_CHECKING, Any, Dict, Generator, Generic, Iterable, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union, cast, ) import numpy as np from pydantic.error_wrappers import ValidationError from typing_inspect import get_args, is_optional_type, is_union_type from docarray import BaseDoc, DocList from docarray.array.any_array import AnyDocArray from docarray.typing import ID, AnyTensor from docarray.typing.tensor.abstract_tensor import AbstractTensor from docarray.utils._internal._typing import is_tensor_union, safe_issubclass from docarray.utils._internal.misc import import_library from docarray.utils._internal.pydantic import is_pydantic_v2 from docarray.utils.find import ( FindResult, FindResultBatched, SubindexFindResult, _FindResult, _FindResultBatched, ) if TYPE_CHECKING: import tensorflow as tf # type: ignore import torch from pydantic.fields import ModelField from docarray.typing import TensorFlowTensor else: tf = import_library('tensorflow', raise_error=False) if tf is not None: from docarray.typing import TensorFlowTensor torch = import_library('torch', raise_error=False) TSchema = TypeVar('TSchema', bound=BaseDoc) def _raise_not_composable(name): def _inner(self, *args, **kwargs): raise NotImplementedError( f'`{name}` is not usable through the query builder of this Document index ({type(self)}). ' f'But you can call `{type(self)}.{name}()` directly.' ) return _inner def _raise_not_supported(name): def _inner(self, *args, **kwargs): raise NotImplementedError( f'`{name}` is not usable through the query builder of this Document index ({type(self)}). ' ) return _inner @dataclass class _ColumnInfo: docarray_type: Type db_type: Any n_dim: Optional[int] config: Dict[str, Any] class BaseDocIndex(ABC, Generic[TSchema]): """Abstract class for all Document Stores""" # the BaseDoc that defines the schema of the store # for subclasses this is filled automatically _schema: Optional[Type[BaseDoc]] = None def __init__(self, db_config=None, subindex: bool = False, **kwargs): if self._schema is None: raise ValueError( 'A DocumentIndex must be typed with a Document type.' 'To do so, use the syntax: DocumentIndex[DocumentType]' ) if subindex: class _NewSchema(self._schema): # type: ignore parent_id: Optional[ID] = None self._ori_schema = self._schema self._schema = cast(Type[BaseDoc], _NewSchema) self._logger = logging.getLogger('docarray') self._db_config = db_config or self.DBConfig(**kwargs) if not isinstance(self._db_config, self.DBConfig): raise ValueError(f'db_config must be of type {self.DBConfig}') self._logger.info('DB config created') self._runtime_config = self.RuntimeConfig() self._logger.info('Runtime config created') self._column_infos: Dict[str, _ColumnInfo] = self._create_column_infos( self._schema ) self._is_subindex = subindex self._subindices: Dict[str, BaseDocIndex] = {} self._init_subindex() ############################################### # Inner classes for query builder and configs # # Subclasses must subclass & implement these # ############################################### class QueryBuilder(ABC): @abstractmethod def build(self, *args, **kwargs) -> Any: """Build the DB specific query object. The DB specific implementation can leverage self._queries to do so. The output of this should be able to be passed to execute_query(). """ ... # TODO support subindex in QueryBuilder # the methods below need to be implemented by subclasses # If, in your subclass, one of these is not usable in a query builder, but # can be called directly on the DocumentIndex, use `_raise_not_composable`. # If the method is not supported _at all_, use `_raise_not_supported`. find = abstractmethod(lambda *args, **kwargs: ...) filter = abstractmethod(lambda *args, **kwargs: ...) text_search = abstractmethod(lambda *args, **kwargs: ...) find_batched = abstractmethod(lambda *args, **kwargs: ...) filter_batched = abstractmethod(lambda *args, **kwargs: ...) text_search_batched = abstractmethod(lambda *args, **kwargs: ...) @dataclass class DBConfig(ABC): index_name: Optional[str] = None # default configurations for every column type # a dictionary from a column type (DB specific) to a dictionary # of default configurations for that type # These configs are used if no configs are specified in the `Field(...)` # of a field in the Document schema (`cls._schema`) # Example: `default_column_config['VARCHAR'] = {'length': 255}` default_column_config: Dict[Type, Dict[str, Any]] = field(default_factory=dict) @dataclass class RuntimeConfig(ABC): pass @property def index_name(self): """Return the name of the index in the database.""" ... ##################################### # Abstract methods # # Subclasses must implement these # ##################################### @abstractmethod def python_type_to_db_type(self, python_type: Type) -> Any: """Map python type to database type. Takes any python type and returns the corresponding database column type. :param python_type: a python type. :return: the corresponding database column type, or None if ``python_type`` is not supported. """ ... @abstractmethod def _index(self, column_to_data: Dict[str, Generator[Any, None, None]]): """index a document into the store""" # `column_to_data` is a dictionary from column name to a generator # that yields the data for that column. # If you want to work directly on documents, you can implement index() instead # If you implement index(), _index() only needs a dummy implementation. ... @abstractmethod def num_docs(self) -> int: """Return the number of indexed documents""" ... @property def _is_index_empty(self) -> bool: """ Check if index is empty by comparing the number of documents to zero. :return: True if the index is empty, False otherwise. """ return self.num_docs() == 0 @abstractmethod def _del_items(self, doc_ids: Sequence[str]): """Delete Documents from the index. :param doc_ids: ids to delete from the Document Store """ ... @abstractmethod def _get_items( self, doc_ids: Sequence[str] ) -> Union[Sequence[TSchema], Sequence[Dict[str, Any]]]: """Get Documents from the index, by `id`. If no document is found, a KeyError is raised. :param doc_ids: ids to get from the Document index :return: Sequence of Documents, sorted corresponding to the order of `doc_ids`. Duplicate `doc_ids` can be omitted in the output. """ ... @abstractmethod def execute_query(self, query: Any, *args, **kwargs) -> Any: """ Execute a query on the database. 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 :param args: positional arguments to pass to the query :param kwargs: keyword arguments to pass to the query :return: the result of the query """ ... @abstractmethod def _doc_exists(self, doc_id: str) -> bool: """ Checks if a given document exists in the index. :param doc_id: The id of a document to check. :return: True if the document exists in the index, False otherwise. """ ... @abstractmethod def _find( self, query: np.ndarray, limit: int, search_field: str = '', ) -> _FindResult: """Find documents in the index :param query: query vector for KNN/ANN search. Has single axis. :param limit: maximum number of documents to return per query :param search_field: name of the field to search on :return: a named tuple containing `documents` and `scores` """ # NOTE: in standard implementations, # `search_field` is equal to the column name to search on ... @abstractmethod def _find_batched( self, queries: np.ndarray, limit: int, search_field: str = '', ) -> _FindResultBatched: """Find documents in the index :param queries: query vectors for KNN/ANN search. Has shape (batch_size, vector_dim) :param limit: maximum number of documents to return :param search_field: name of the field to search on :return: a named tuple containing `documents` and `scores` """ ... @abstractmethod def _filter( self, filter_query: Any, limit: int, ) -> Union[DocList, List[Dict]]: """Find documents in the index based on a filter query :param filter_query: the DB specific filter query to execute :param limit: maximum number of documents to return :return: a DocList containing the documents that match the filter query """ ... @abstractmethod def _filter_batched( self, filter_queries: Any, limit: int, ) -> Union[List[DocList], List[List[Dict]]]: """Find documents in the index based on multiple filter queries. Each query is considered individually, and results are returned per query. :param filter_queries: the DB specific filter queries to execute :param limit: maximum number of documents to return per query :return: List of DocLists containing the documents that match the filter queries """ ... @abstractmethod def _text_search( self, query: str, limit: int, search_field: str = '', ) -> _FindResult: """Find documents in the index based on a text search query :param query: The text to search for :param limit: maximum number of documents to return :param search_field: name of the field to search on :return: a named tuple containing `documents` and `scores` """ # NOTE: in standard implementations, # `search_field` is equal to the column name to search on ... @abstractmethod def _text_search_batched( self, queries: Sequence[str], limit: int, search_field: str = '', ) -> _FindResultBatched: """Find documents in the index based on a text search query :param queries: The texts to search for :param limit: maximum number of documents to return per query :param search_field: name of the field to search on :return: a named tuple containing `documents` and `scores` """ # NOTE: in standard implementations, # `search_field` is equal to the column name to search on ... #################################################### # Optional overrides # # Subclasses may or may not need to change these # #################################################### def __getitem__( self, key: Union[str, Sequence[str]] ) -> Union[TSchema, DocList[TSchema]]: """Get one or multiple Documents into the index, by `id`. If no document is found, a KeyError is raised. :param key: id or ids to get from the Document index """ # normalize input if isinstance(key, str): return_singleton = True key = [key] else: return_singleton = False # retrieve data doc_sequence = self._get_items(key) # check data if len(doc_sequence) == 0: raise KeyError(f'No document with id {key} found') # retrieve nested data for field_name, type_, _ in self._flatten_schema( cast(Type[BaseDoc], self._schema) ): if safe_issubclass(type_, AnyDocArray) and isinstance( doc_sequence[0], Dict ): for doc in doc_sequence: self._get_subindex_doclist(doc, field_name) # type: ignore # cast output if isinstance(doc_sequence, DocList): out_docs: DocList[TSchema] = doc_sequence elif isinstance(doc_sequence[0], Dict): out_docs = self._dict_list_to_docarray(doc_sequence) # type: ignore else: docs_cls = DocList.__class_getitem__(cast(Type[BaseDoc], self._schema)) out_docs = docs_cls(doc_sequence) return out_docs[0] if return_singleton else out_docs def __delitem__(self, key: Union[str, Sequence[str]]): """Delete one or multiple Documents from the index, by `id`. If no document is found, a KeyError is raised. :param key: id or ids to delete from the Document index """ self._logger.info(f'Deleting documents with id(s) {key} from the index') if isinstance(key, str): key = [key] # delete nested data for field_name, type_, _ in self._flatten_schema( cast(Type[BaseDoc], self._schema) ): if safe_issubclass(type_, AnyDocArray): for doc_id in key: nested_docs_id = self._subindices[field_name]._filter_by_parent_id( doc_id ) if nested_docs_id: del self._subindices[field_name][nested_docs_id] # delete data self._del_items(key) def __contains__(self, item: BaseDoc) -> bool: """ Checks if a given document exists in the index. :param item: The document to check. It must be an instance of BaseDoc or its subclass. :return: True if the document exists in the index, False otherwise. """ if safe_issubclass(type(item), BaseDoc): return self._doc_exists(str(item.id)) else: raise TypeError( f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'" ) def configure(self, runtime_config=None, **kwargs): """ Configure the DocumentIndex. You can either pass a config object to `config` or pass individual config parameters as keyword arguments. If a configuration object is passed, it will replace the current configuration. If keyword arguments are passed, they will update the current configuration. :param runtime_config: the configuration to apply :param kwargs: individual configuration parameters """ if runtime_config is None: self._runtime_config = replace(self._runtime_config, **kwargs) else: if not isinstance(runtime_config, self.RuntimeConfig): raise ValueError(f'runtime_config must be of type {self.RuntimeConfig}') self._runtime_config = runtime_config def index(self, docs: Union[BaseDoc, Sequence[BaseDoc]], **kwargs): """index Documents into the index. !!! note Passing a sequence of Documents that is not a DocList (such as a List of Docs) comes at a performance penalty. This is because the Index needs to check compatibility between itself and the data. With a DocList as input this is a single check; for other inputs compatibility needs to be checked for every Document individually. :param docs: Documents to index. """ n_docs = 1 if isinstance(docs, BaseDoc) else len(docs) self._logger.debug(f'Indexing {n_docs} documents') docs_validated = self._validate_docs(docs) self._update_subindex_data(docs_validated) data_by_columns = self._get_col_value_dict(docs_validated) self._index(data_by_columns, **kwargs) def find( self, query: Union[AnyTensor, BaseDoc], search_field: str = '', limit: int = 10, **kwargs, ) -> FindResult: """Find documents in the index using nearest neighbor search. :param query: query vector for KNN/ANN search. Can be either a tensor-like (np.array, torch.Tensor, etc.) with a single axis, or a Document :param search_field: name of the field to search on. Documents in the index are retrieved based on this similarity of this field to the query. :param limit: maximum number of documents to return :return: a named tuple containing `documents` and `scores` """ self._logger.debug(f'Executing `find` for search field {search_field}') self._validate_search_field(search_field) if isinstance(query, BaseDoc): query_vec = self._get_values_by_column([query], search_field)[0] else: query_vec = query query_vec_np = self._to_numpy(query_vec) docs, scores = self._find( query_vec_np, search_field=search_field, limit=limit, **kwargs ) if isinstance(docs, List) and not isinstance(docs, DocList): docs = self._dict_list_to_docarray(docs) return FindResult(documents=docs, scores=scores) def find_subindex( self, query: Union[AnyTensor, BaseDoc], subindex: str = '', search_field: str = '', limit: int = 10, **kwargs, ) -> SubindexFindResult: """Find documents in subindex level. :param query: query vector for KNN/ANN search. Can be either a tensor-like (np.array, torch.Tensor, etc.) with a single axis, or a Document :param subindex: name of the subindex to search on :param search_field: name of the field to search on :param limit: maximum number of documents to return :return: a named tuple containing root docs, subindex docs and scores """ self._logger.debug(f'Executing `find_subindex` for search field {search_field}') sub_docs, scores = self._find_subdocs( query, subindex=subindex, search_field=search_field, limit=limit, **kwargs ) fields = subindex.split('__') root_ids = [ self._get_root_doc_id(doc.id, fields[0], '__'.join(fields[1:])) for doc in sub_docs ] root_docs = DocList[self._schema]() # type: ignore for id in root_ids: root_docs.append(self[id]) return SubindexFindResult( root_documents=root_docs, sub_documents=sub_docs, scores=scores # type: ignore ) def find_batched( self, queries: Union[AnyTensor, DocList], search_field: str = '', limit: int = 10, **kwargs, ) -> FindResultBatched: """Find documents in the index using nearest neighbor search. :param queries: query vector for KNN/ANN search. Can be either a tensor-like (np.array, torch.Tensor, etc.) with a, or a DocList. If a tensor-like is passed, it should have shape (batch_size, vector_dim) :param search_field: name of the field to search on. Documents in the index are retrieved based on this similarity of this field to the query. :param limit: maximum number of documents to return per query :return: a named tuple containing `documents` and `scores` """ self._logger.debug(f'Executing `find_batched` for search field {search_field}') if search_field: if '__' in search_field: fields = search_field.split('__') if safe_issubclass(self._schema._get_field_annotation(fields[0]), AnyDocArray): # type: ignore return self._subindices[fields[0]].find_batched( queries, search_field='__'.join(fields[1:]), limit=limit, **kwargs, ) self._validate_search_field(search_field) if isinstance(queries, Sequence): query_vec_list = self._get_values_by_column(queries, search_field) query_vec_np = np.stack( tuple(self._to_numpy(query_vec) for query_vec in query_vec_list) ) else: query_vec_np = self._to_numpy(queries) da_list, scores = self._find_batched( query_vec_np, search_field=search_field, limit=limit, **kwargs ) if ( len(da_list) > 0 and isinstance(da_list[0], List) and not isinstance(da_list[0], DocList) ): da_list = [self._dict_list_to_docarray(docs) for docs in da_list] return FindResultBatched(documents=da_list, scores=scores) # type: ignore def filter( self, filter_query: Any, limit: int = 10, **kwargs, ) -> DocList: """Find documents in the index based on a filter query :param filter_query: the DB specific filter query to execute :param limit: maximum number of documents to return :return: a DocList containing the documents that match the filter query """ self._logger.debug(f'Executing `filter` for the query {filter_query}') docs = self._filter(filter_query, limit=limit, **kwargs) if isinstance(docs, List) and not isinstance(docs, DocList): docs = self._dict_list_to_docarray(docs) return docs def filter_subindex( self, filter_query: Any, subindex: str, limit: int = 10, **kwargs, ) -> DocList: """Find documents in subindex level based on a filter query :param filter_query: the DB specific filter query to execute :param subindex: name of the subindex to search on :param limit: maximum number of documents to return :return: a DocList containing the subindex level documents that match the filter query """ self._logger.debug( f'Executing `filter` for the query {filter_query} in subindex {subindex}' ) if '__' in subindex: fields = subindex.split('__') return self._subindices[fields[0]].filter_subindex( filter_query, '__'.join(fields[1:]), limit=limit, **kwargs ) else: return self._subindices[subindex].filter( filter_query, limit=limit, **kwargs ) def filter_batched( self, filter_queries: Any, limit: int = 10, **kwargs, ) -> List[DocList]: """Find documents in the index based on multiple filter queries. :param filter_queries: the DB specific filter query to execute :param limit: maximum number of documents to return :return: a DocList containing the documents that match the filter query """ self._logger.debug( f'Executing `filter_batched` for the queries {filter_queries}' ) da_list = self._filter_batched(filter_queries, limit=limit, **kwargs) if len(da_list) > 0 and isinstance(da_list[0], List): da_list = [self._dict_list_to_docarray(docs) for docs in da_list] return da_list # type: ignore def text_search( self, query: Union[str, BaseDoc], search_field: str = '', limit: int = 10, **kwargs, ) -> FindResult: """Find documents in the index based on a text search query. :param query: The text to search for :param search_field: name of the field to search on :param limit: maximum number of documents to return :return: a named tuple containing `documents` and `scores` """ self._logger.debug(f'Executing `text_search` for search field {search_field}') self._validate_search_field(search_field) if isinstance(query, BaseDoc): query_text = self._get_values_by_column([query], search_field)[0] else: query_text = query docs, scores = self._text_search( query_text, search_field=search_field, limit=limit, **kwargs ) if isinstance(docs, List) and not isinstance(docs, DocList): docs = self._dict_list_to_docarray(docs) return FindResult(documents=docs, scores=scores) def text_search_batched( self, queries: Union[Sequence[str], Sequence[BaseDoc]], search_field: str = '', limit: int = 10, **kwargs, ) -> FindResultBatched: """Find documents in the index based on a text search query. :param queries: The texts to search for :param search_field: name of the field to search on :param limit: maximum number of documents to return :return: a named tuple containing `documents` and `scores` """ self._logger.debug( f'Executing `text_search_batched` for search field {search_field}' ) self._validate_search_field(search_field) if isinstance(queries[0], BaseDoc): query_docs: Sequence[BaseDoc] = cast(Sequence[BaseDoc], queries) query_texts: Sequence[str] = self._get_values_by_column( query_docs, search_field ) else: query_texts = cast(Sequence[str], queries) da_list, scores = self._text_search_batched( query_texts, search_field=search_field, limit=limit, **kwargs ) if len(da_list) > 0 and isinstance(da_list[0], List): docs = [self._dict_list_to_docarray(docs) for docs in da_list] return FindResultBatched(documents=docs, scores=scores) da_list_ = cast(List[DocList], da_list) return FindResultBatched(documents=da_list_, scores=scores) def _filter_by_parent_id(self, id: str) -> Optional[List[str]]: """Filter the ids of the subindex documents given id of root document. :param id: the root document id to filter by :return: a list of ids of the subindex documents """ return None ########################################################## # Helper methods # # These might be useful in your subclass implementation # ########################################################## @staticmethod def _get_values_by_column(docs: Sequence[BaseDoc], col_name: str) -> List[Any]: """Get the value of a column of a document. :param docs: The DocList to get the values from :param col_name: The name of the column, e.g. 'text' or 'image__tensor' :return: The value of the column of `doc` """ leaf_vals = [] for doc in docs: if '__' in col_name: fields = col_name.split('__') leaf_doc: BaseDoc = doc for f in fields[:-1]: leaf_doc = getattr(leaf_doc, f) leaf_vals.append(getattr(leaf_doc, fields[-1])) else: leaf_vals.append(getattr(doc, col_name)) return leaf_vals @staticmethod def _transpose_col_value_dict( col_value_dict: Mapping[str, Iterable[Any]] ) -> Generator[Dict[str, Any], None, None]: """'Transpose' the output of `_get_col_value_dict()`: Yield rows of columns, where each row represent one Document. Since a generator is returned, this process comes at negligible cost. :param docs: The DocList to get the values from :return: The `docs` flattened out as rows. Each row is a dictionary mapping from column name to value """ return (dict(zip(col_value_dict, row)) for row in zip(*col_value_dict.values())) def _get_col_value_dict( self, docs: Union[BaseDoc, Sequence[BaseDoc]] ) -> Dict[str, Generator[Any, None, None]]: """ Get all data from a (sequence of) document(s), flattened out by column. This can be seen as the transposed representation of `_get_rows()`. :param docs: The document(s) to get the data from :return: A dictionary mapping column names to a generator of values """ if isinstance(docs, BaseDoc): docs_seq: Sequence[BaseDoc] = [docs] else: docs_seq = docs def _col_gen(col_name: str): return ( self._to_numpy( self._get_values_by_column([doc], col_name)[0], allow_passthrough=True, ) for doc in docs_seq ) return {col_name: _col_gen(col_name) for col_name in self._column_infos} def _update_subindex_data( self, docs: DocList[BaseDoc], ): """ Add `parent_id` to all sublevel documents. :param docs: The document(s) to update the `parent_id` for """ for field_name, type_, _ in self._flatten_schema( cast(Type[BaseDoc], self._schema) ): if safe_issubclass(type_, AnyDocArray): for doc in docs: _list = getattr(doc, field_name) for i, nested_doc in enumerate(_list): nested_doc = self._subindices[field_name]._schema( # type: ignore **nested_doc.__dict__ ) nested_doc.parent_id = doc.id _list[i] = nested_doc ################################################## # Behind-the-scenes magic # # Subclasses should not need to implement these # ################################################## def __class_getitem__(cls, item: Type[TSchema]): if not isinstance(item, type): # do nothing # enables use in static contexts with type vars, e.g. as type annotation return Generic.__class_getitem__.__func__(cls, item) if not safe_issubclass(item, BaseDoc): raise ValueError( f'{cls.__name__}[item] `item` should be a Document not a {item} ' ) class _DocumentIndexTyped(cls): # type: ignore _schema: Type[TSchema] = item _DocumentIndexTyped.__name__ = f'{cls.__name__}[{item.__name__}]' _DocumentIndexTyped.__qualname__ = f'{cls.__qualname__}[{item.__name__}]' return _DocumentIndexTyped def build_query(self) -> QueryBuilder: """ Build a query for this DocumentIndex. :return: a new `QueryBuilder` object for this DocumentIndex """ return self.QueryBuilder() # type: ignore @classmethod def _flatten_schema( cls, schema: Type[BaseDoc], name_prefix: str = '' ) -> List[Tuple[str, Type, 'ModelField']]: """Flatten the schema of a Document into a list of column names and types. Nested Documents are handled in a recursive manner by adding `'__'` as a prefix to the column name. :param schema: The schema to flatten :param name_prefix: prefix to append to the column names. Used for recursive calls to handle nesting. :return: A list of column names, types, and fields """ names_types_fields: List[Tuple[str, Type, 'ModelField']] = [] for field_name, field_ in schema._docarray_fields().items(): t_ = schema._get_field_annotation(field_name) inner_prefix = name_prefix + field_name + '__' if is_union_type(t_): union_args = get_args(t_) if is_tensor_union(t_): names_types_fields.append( (name_prefix + field_name, AbstractTensor, field_) ) elif len(union_args) == 2 and type(None) in union_args: # simple "Optional" type, treat as special case: # treat as if it was a single non-optional type for t_arg in union_args: if t_arg is not type(None): if safe_issubclass(t_arg, BaseDoc): names_types_fields.extend( cls._flatten_schema(t_arg, name_prefix=inner_prefix) ) else: names_types_fields.append( (name_prefix + field_name, t_arg, field_) ) else: raise ValueError( f'Union type {t_} is not supported. Only Union of subclasses of AbstractTensor or Union[type, None] are supported.' ) elif safe_issubclass(t_, BaseDoc): names_types_fields.extend( cls._flatten_schema(t_, name_prefix=inner_prefix) ) elif safe_issubclass(t_, AbstractTensor): names_types_fields.append( (name_prefix + field_name, AbstractTensor, field_) ) else: names_types_fields.append((name_prefix + field_name, t_, field_)) return names_types_fields def _create_column_infos(self, schema: Type[BaseDoc]) -> Dict[str, _ColumnInfo]: """Collects information about every column that is implied by a given schema. :param schema: The schema (subclass of BaseDoc) to analyze and parse columns from :returns: A dictionary mapping from column names to column information. """ column_infos: Dict[str, _ColumnInfo] = dict() for field_name, type_, field_ in self._flatten_schema(schema): # Union types are handle in _flatten_schema if safe_issubclass(type_, AnyDocArray): column_infos[field_name] = _ColumnInfo( docarray_type=type_, db_type=None, config=dict(), n_dim=None ) else: column_infos[field_name] = self._create_single_column(field_, type_) return column_infos def _create_single_column(self, field: 'ModelField', type_: Type) -> _ColumnInfo: custom_config = ( field.json_schema_extra if is_pydantic_v2 else field.field_info.extra ) if custom_config is None: custom_config = dict() if 'col_type' in custom_config.keys(): db_type = custom_config['col_type'] custom_config.pop('col_type') if db_type not in self._db_config.default_column_config.keys(): raise ValueError( f'The given col_type is not a valid db type: {db_type}' ) else: db_type = self.python_type_to_db_type(type_) config = self._db_config.default_column_config[db_type].copy() config.update(custom_config) # parse n_dim from parametrized tensor type field_type = field.annotation if is_pydantic_v2 else field.type_ if ( hasattr(field_type, '__docarray_target_shape__') and field_type.__docarray_target_shape__ ): if len(field_type.__docarray_target_shape__) == 1: n_dim = field_type.__docarray_target_shape__[0] else: n_dim = field_type.__docarray_target_shape__ else: n_dim = None return _ColumnInfo( docarray_type=type_, db_type=db_type, config=config, n_dim=n_dim ) def _init_subindex( self, ): """Initialize subindices if any column is subclass of AnyDocArray.""" for col_name, col in self._column_infos.items(): if safe_issubclass(col.docarray_type, AnyDocArray): sub_db_config = copy.deepcopy(self._db_config) sub_db_config.index_name = f'{self.index_name}__{col_name}' self._subindices[col_name] = self.__class__[col.docarray_type.doc_type]( # type: ignore db_config=sub_db_config, subindex=True ) def _validate_docs( self, docs: Union[BaseDoc, Sequence[BaseDoc]] ) -> DocList[BaseDoc]: """Validates Document against the schema of the Document Index. For validation to pass, the schema of `docs` and the schema of the Document Index need to evaluate to the same flattened columns. If Validation fails, a ValueError is raised. :param docs: Document to evaluate. If this is a DocList, validation is performed using its `doc_type` (parametrization), without having to check ever Document in `docs`. If this check fails, or if `docs` is not a DocList, evaluation is performed for every Document in `docs`. :return: A DocList containing the Documents in `docs` """ if isinstance(docs, BaseDoc): docs = [docs] if isinstance(docs, DocList): # validation shortcut for DocList; only look at the schema reference_schema_flat = self._flatten_schema( cast(Type[BaseDoc], self._schema) ) reference_names = [name for (name, _, _) in reference_schema_flat] reference_types = [t_ for (_, t_, _) in reference_schema_flat] try: input_schema_flat = self._flatten_schema(docs.doc_type) except ValueError: pass else: input_names = [name for (name, _, _) in input_schema_flat] input_types = [t_ for (_, t_, _) in input_schema_flat] # this could be relaxed in the future, # see schema translation ideas in the design doc names_compatible = reference_names == input_names types_compatible = all( (safe_issubclass(t2, t1)) for (t1, t2) in zip(reference_types, input_types) ) if names_compatible and types_compatible: return docs out_docs = [] for i in range(len(docs)): # validate the data try: out_docs.append( cast(Type[BaseDoc], self._schema).parse_obj(dict(docs[i])) ) except (ValueError, ValidationError) as e: raise ValueError( 'The schema of the input Documents is not compatible with the schema of the Document Index.' ' Ensure that the field names of your data match the field names of the Document Index schema,' ' and that the types of your data match the types of the Document Index schema.' f'original error {e}' ) return DocList[BaseDoc].construct(out_docs) def _validate_search_field(self, search_field: Union[str, None]) -> bool: """ Validate if the given `search_field` corresponds to one of the columns that was parsed from the schema. Some backends, like weaviate, don't use search fields, so the function returns True if `search_field` is empty or None. :param search_field: search field to validate. :return: True if the field exists, False otherwise. """ if not search_field or search_field in self._column_infos.keys(): if not search_field: self._logger.info('Empty search field was passed') return True else: valid_search_fields = ', '.join(self._column_infos.keys()) raise ValueError( f'{search_field} is not a valid search field. Valid search fields are: {valid_search_fields}' ) def _to_numpy(self, val: Any, allow_passthrough=False) -> Any: """ Converts a value to a numpy array, if possible. :param val: The value to convert :param allow_passthrough: If True, the value is returned as-is if it is not convertible to a numpy array. If False, a `ValueError` is raised if the value is not convertible to a numpy array. :return: The value as a numpy array, or as-is if `allow_passthrough` is True and the value is not convertible """ if isinstance(val, np.ndarray): return val if tf is not None and isinstance(val, TensorFlowTensor): return val.unwrap().numpy() if isinstance(val, (list, tuple)): return np.array(val) if torch is not None and isinstance(val, torch.Tensor): return val.detach().numpy() if tf is not None and isinstance(val, tf.Tensor): return val.numpy() if allow_passthrough: return val raise ValueError(f'Unsupported input type for {type(self)}: {type(val)}') def _convert_dict_to_doc( self, doc_dict: Dict[str, Any], schema: Type[BaseDoc], inner=False ) -> BaseDoc: """ Convert a dict to a Document object. :param doc_dict: A dict that contains all the flattened fields of a Document, the field names are the keys and follow the pattern {field_name} or {field_name}__{nested_name} :param schema: The schema of the Document object :return: A Document object """ for field_name, _ in schema._docarray_fields().items(): t_ = schema._get_field_annotation(field_name) if not is_union_type(t_) and safe_issubclass(t_, AnyDocArray): self._get_subindex_doclist(doc_dict, field_name) if is_optional_type(t_): for t_arg in get_args(t_): if t_arg is not type(None): t_ = t_arg if not is_union_type(t_) and safe_issubclass(t_, BaseDoc): inner_dict = {} fields = [ key for key in doc_dict.keys() if key.startswith(f'{field_name}__') ] for key in fields: nested_name = key[len(f'{field_name}__') :] inner_dict[nested_name] = doc_dict.pop(key) doc_dict[field_name] = self._convert_dict_to_doc( inner_dict, t_, inner=True ) if self._is_subindex and not inner: doc_dict.pop('parent_id', None) schema_cls = cast(Type[BaseDoc], self._ori_schema) else: schema_cls = cast(Type[BaseDoc], schema) doc = schema_cls(**doc_dict) return doc def _dict_list_to_docarray(self, dict_list: Sequence[Dict[str, Any]]) -> DocList: """Convert a list of docs in dict type to a DocList of the schema type.""" doc_list = [self._convert_dict_to_doc(doc_dict, self._schema) for doc_dict in dict_list] # type: ignore if self._is_subindex: docs_cls = DocList.__class_getitem__(cast(Type[BaseDoc], self._ori_schema)) else: docs_cls = DocList.__class_getitem__(cast(Type[BaseDoc], self._schema)) return docs_cls(doc_list) def __len__(self) -> int: return self.num_docs() def _index_subindex(self, column_to_data: Dict[str, Generator[Any, None, None]]): """Index subindex documents in the corresponding subindex. :param column_to_data: A dictionary from column name to a generator """ for col_name, col in self._column_infos.items(): if safe_issubclass(col.docarray_type, AnyDocArray): docs = [ doc for doc_list in column_to_data[col_name] for doc in doc_list ] self._subindices[col_name].index(docs) column_to_data.pop(col_name, None) def _get_subindex_doclist(self, doc: Dict[str, Any], field_name: str): """Get subindex Documents from the index and assign them to `field_name`. :param doc: a dictionary mapping from column name to value :param field_name: field name of the subindex Documents """ if field_name not in doc.keys(): parent_id = doc['id'] nested_docs_id = self._subindices[field_name]._filter_by_parent_id( parent_id ) if nested_docs_id: doc[field_name] = self._subindices[field_name].__getitem__( nested_docs_id ) def _find_subdocs( self, query: Union[AnyTensor, BaseDoc], subindex: str = '', search_field: str = '', limit: int = 10, **kwargs, ) -> FindResult: """Find documents in the subindex and return subindex docs and scores.""" fields = subindex.split('__') if not subindex or not safe_issubclass( self._schema._get_field_annotation(fields[0]), AnyDocArray # type: ignore ): raise ValueError(f'subindex {subindex} is not valid') if len(fields) == 1: return self._subindices[fields[0]].find( query, search_field=search_field, limit=limit, **kwargs ) return self._subindices[fields[0]]._find_subdocs( query, subindex='___'.join(fields[1:]), search_field=search_field, limit=limit, **kwargs, ) def _get_root_doc_id(self, id: str, root: str, sub: str) -> str: """Get the root_id given the id of a subindex Document and the root and subindex name :param id: id of the subindex Document :param root: root index name :param sub: subindex name :return: the root_id of the Document """ subindex = self._subindices[root] if not sub: sub_doc = subindex._get_items([id]) parent_id = ( sub_doc[0]['parent_id'] if isinstance(sub_doc[0], dict) else sub_doc[0].parent_id ) return parent_id else: fields = sub.split('__') cur_root_id = subindex._get_root_doc_id( id, fields[0], '__'.join(fields[1:]) ) return self._get_root_doc_id(cur_root_id, root, '') def subindex_contains(self, item: BaseDoc) -> bool: """Checks if a given BaseDoc item is contained in the index or any of its subindices. :param item: the given BaseDoc :return: if the given BaseDoc item is contained in the index/subindices """ if self._is_index_empty: return False if safe_issubclass(type(item), BaseDoc): return self.__contains__(item) or any( index.subindex_contains(item) for index in self._subindices.values() ) else: raise TypeError( f"item must be an instance of BaseDoc or its subclass, not '{type(item).__name__}'" )