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]