-
-
- {% block content %}{{ body }}{% endblock %}
-
-
-
-
-
-
-
+
-The data structure for unstructured data
+The data structure for multimodal data
-
-
+
+
+
-
-
+# Load image tensor from URL
+doc.image_tensor = doc.image_url.load()
-Our problem is given an image from `/left` and find its most-similar image in `/right` (without looking at the filename of course).
-### Load images
+# Compute embedding with any model of your choice
+def clip_image_encoder(image_tensor: TorchTensor) -> TorchTensor: # dummy function
+ return torch.rand(512)
-First load images and preprocess them with standard computer vision techniques:
-```python
-from docarray import DocumentArray, Document
+doc.embedding = clip_image_encoder(doc.image_tensor)
-left_da = DocumentArray.from_files('left/*.jpg')
+print(doc.embedding.shape) # torch.Size([512])
```
-To get a feeling of the data you will handle, plot them in one sprite image:
+### Compose nested Documents
+
+Of course, you can compose Documents into a nested structure:
```python
-left_da.plot_image_sprites()
+from docarray import BaseDoc
+from docarray.documents import ImageDoc, TextDoc
+import numpy as np
+
+
+class MultiModalDocument(BaseDoc):
+ image_doc: ImageDoc
+ text_doc: TextDoc
+
+
+doc = MultiModalDocument(
+ image_doc=ImageDoc(tensor=np.zeros((3, 224, 224))), text_doc=TextDoc(text='hi!')
+)
```
-
+
+## Send
+
+DocArray facilitates the **transmission of your data** in a manner inherently compatible with machine learning.
+
+This includes native support for **Protobuf and gRPC**, along with **HTTP** and serialization to JSON, JSONSchema, Base64, and Bytes.
+
+This feature proves beneficial for several scenarios:
+
+- :cloud: You are **serving a model**, perhaps through frameworks like **[Jina](https://github.com/jina-ai/jina/)** or **[FastAPI](https://github.com/tiangolo/fastapi/)**
+- :spider_web: You are **distributing your model** across multiple machines and need an efficient means of transmitting your data between nodes
+- :gear: You are architecting a **microservice** environment and require a method for data transmission between microservices
+
+> :bulb: **Are you familiar with FastAPI?** You'll be delighted to learn
+> that DocArray maintains full compatibility with FastAPI!
+> Plus, we have a [dedicated section](#coming-from-fastapi) specifically for you!
+
+When it comes to data transmission, serialization is a crucial step. Let's delve into how DocArray streamlines this process:
-Fun is fun, but recall our goal is to match left images against right images and so far we have only handled the left. Let's repeat the same procedure for the right:
```python
-right_da = (DocumentArray.from_files('right/*.jpg')
- .apply(preproc)
- .embed(model, device='cuda'))
+from docarray import BaseDoc
+from docarray.typing import ImageTorchTensor
+import torch
+
+
+# model your data
+class MyDocument(BaseDoc):
+ description: str
+ image: ImageTorchTensor[3, 224, 224]
+
+
+# create a Document
+doc = MyDocument(
+ description="This is a description",
+ image=torch.zeros((3, 224, 224)),
+)
+
+# serialize it!
+proto = doc.to_protobuf()
+bytes_ = doc.to_bytes()
+json = doc.json()
+
+# deserialize it!
+doc_2 = MyDocument.from_protobuf(proto)
+doc_4 = MyDocument.from_bytes(bytes_)
+doc_5 = MyDocument.parse_raw(json)
```
-### Match nearest neighbours
+Of course, serialization is not all you need. So check out how DocArray integrates with **[Jina](https://github.com/jina-ai/jina/)** and **[FastAPI](https://github.com/tiangolo/fastapi/)**.
+
+## Store
+
+After modeling and possibly distributing your data, you'll typically want to **store it** somewhere. That's where DocArray steps in!
-We can now match the left to the right and take the top-9 results.
+**Document Stores** provide a seamless way to, as the name suggests, store your Documents. Be it locally or remotely, you can do it all through the same user interface:
+
+- :cd: **On disk**, as a file in your local filesystem
+- :bucket: On **[AWS S3](https://aws.amazon.com/de/s3/)**
+- :cloud: On **[Jina AI Cloud](https://cloud.jina.ai/)**
+
+The Document Store interface lets you push and pull Documents to and from multiple data sources, all with the same user interface.
+
+For example, let's see how that works with on-disk storage:
```python
-left_da.match(right_da, limit=9)
+from docarray import BaseDoc, DocList
+
+
+class SimpleDoc(BaseDoc):
+ text: str
+
+
+docs = DocList[SimpleDoc]([SimpleDoc(text=f'doc {i}') for i in range(8)])
+docs.push('file://simple_docs')
+
+docs_pull = DocList[SimpleDoc].pull('file://simple_docs')
```
-Let's inspect what's inside `left_da` matches now:
+## Retrieve
+
+**Document Indexes** let you index your Documents in a **vector database** for efficient similarity-based retrieval.
+
+This is useful for:
+
+- :left_speech_bubble: Augmenting **LLMs and Chatbots** with domain knowledge ([Retrieval Augmented Generation](https://arxiv.org/abs/2005.11401))
+- :mag: **Neural search** applications
+- :bulb: **Recommender systems**
+
+Currently, Document Indexes support **[Weaviate](https://weaviate.io/)**, **[Qdrant](https://qdrant.tech/)**, **[ElasticSearch](https://www.elastic.co/)**, **[Redis](https://redis.io/)**, **[Mongo Atlas](https://www.mongodb.com/)**, and **[HNSWLib](https://github.com/nmslib/hnswlib)**, with more to come!
+
+The Document Index interface lets you index and retrieve Documents from multiple vector databases, all with the same user interface.
+
+It supports ANN vector search, text search, filtering, and hybrid search.
```python
-for d in left_da:
- for m in d.matches:
- print(d.uri, m.uri, m.scores['cosine'].value)
+from docarray import DocList, BaseDoc
+from docarray.index import HnswDocumentIndex
+import numpy as np
+
+from docarray.typing import ImageUrl, ImageTensor, NdArray
+
+
+class ImageDoc(BaseDoc):
+ url: ImageUrl
+ tensor: ImageTensor
+ embedding: NdArray[128]
+
+
+# create some data
+dl = DocList[ImageDoc](
+ [
+ ImageDoc(
+ url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg",
+ tensor=np.zeros((3, 224, 224)),
+ embedding=np.random.random((128,)),
+ )
+ for _ in range(100)
+ ]
+)
+
+# create a Document Index
+index = HnswDocumentIndex[ImageDoc](work_dir='/tmp/test_index')
+
+
+# index your data
+index.index(dl)
+
+# find similar Documents
+query = dl[0]
+results, scores = index.find(query, limit=10, search_field='embedding')
```
-```text
-left/02262.jpg right/03459.jpg 0.21102
-left/02262.jpg right/02964.jpg 0.13871843
-left/02262.jpg right/02103.jpg 0.18265384
-left/02262.jpg right/04520.jpg 0.16477376
-...
+---
+
+## Learn DocArray
+
+Depending on your background and use case, there are different ways for you to understand DocArray.
+
+### Coming from DocArray <=0.21
+
+Click for more details
+ +Let's take a closer look at how you can represent your data with DocArray: -Here the feature embedding is done by simple [feature hashing](https://en.wikipedia.org/wiki/Feature_hashing) and distance metric is [Jaccard distance](https://en.wikipedia.org/wiki/Jaccard_index). You got better embedding? Of course you do! Looking forward to seeing your results. +```python +from docarray import BaseDoc +from docarray.typing import TorchTensor, ImageUrl +from typing import Optional +import torch + + +# Define your data model +class MyDocument(BaseDoc): + description: str + image_url: ImageUrl # could also be VideoUrl, AudioUrl, etc. + image_tensor: Optional[ + TorchTensor[1704, 2272, 3] + ] = None # could also be NdArray or TensorflowTensor + embedding: Optional[TorchTensor] = None +``` -### A complete workflow of visual search +So not only can you define the types of your data, you can even **specify the shape of your tensors!** -Let's use DocArray and [Totally Looks Like](https://sites.google.com/view/totally-looks-like-dataset) dataset to build simple meme image search. The dataset contains 6016 image-pairs stored in `/left` and `/right`. Images that shares the same filename are perceptually similar. For example, +```python +# Create a document +doc = MyDocument( + description="This is a photo of a mountain", + image_url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg", +) -| left/00018.jpg | -right/00018.jpg | -left/00131.jpg | -right/00131.jpg | -
|---|---|---|---|
 |
-  |
-  |
-  |
-
-
-
-
-
+
+
+### Coming from Pydantic
+
+Click to expand
+ +If you are using DocArray version 0.30.0 or lower, you will be familiar with its [dataclass API](https://docarray.jina.ai/fundamentals/dataclass/). + +_DocArray >=0.30 is that idea, taken seriously._ Every document is created through a dataclass-like interface, +courtesy of [Pydantic](https://pydantic-docs.helpmanual.io/usage/models/). + +This gives the following advantages: +- **Flexibility:** No need to conform to a fixed set of fields -- your data defines the schema +- **Multimodality:** At their core, documents are just dictionaries. This makes it easy to create and send them from any language, not just Python. + +You may also be familiar with our old Document Stores for vector DB integration. +They are now called **Document Indexes** and offer the following improvements (see [here](#store) for the new API): + +- **Hybrid search:** You can now combine vector search with text search, and even filter by arbitrary fields +- **Production-ready:** The new Document Indexes are a much thinner wrapper around the various vector DB libraries, making them more robust and easier to maintain +- **Increased flexibility:** We strive to support any configuration or setting that you could perform through the DB's first-party client + +For now, Document Indexes support **[Weaviate](https://weaviate.io/)**, **[Qdrant](https://qdrant.tech/)**, **[ElasticSearch](https://www.elastic.co/)**, **[Redis](https://redis.io/)**, **[Mongo Atlas](https://www.mongodb.com/)**, Exact Nearest Neighbour search and **[HNSWLib](https://github.com/nmslib/hnswlib)**, with more to come. + +
+
+
+### Coming from PyTorch
+
+Click to expand
+ +If you come from Pydantic, you can see DocArray documents as juiced up Pydantic models, and DocArray as a collection of goodies around them. + +More specifically, we set out to **make Pydantic fit for the ML world** - not by replacing it, but by building on top of it! + +This means you get the following benefits: + +- **ML-focused types**: Tensor, TorchTensor, Embedding, ..., including **tensor shape validation** +- Full compatibility with **FastAPI** +- **DocList** and **DocVec** generalize the idea of a model to a _sequence_ or _batch_ of models. Perfect for **use in ML models** and other batch processing tasks. +- **Types that are alive**: ImageUrl can `.load()` a URL to image tensor, TextUrl can load and tokenize text documents, etc. +- Cloud-ready: Serialization to **Protobuf** for use with microservices and **gRPC** +- **Pre-built multimodal documents** for different data modalities: Image, Text, 3DMesh, Video, Audio and more. Note that all of these are valid Pydantic models! +- **Document Stores** and **Document Indexes** let you store your data and retrieve it using **vector search** + +The most obvious advantage here is **first-class support for ML centric data**, such as `{Torch, TF, ...}Tensor`, `Embedding`, etc. + +This includes handy features such as validating the shape of a tensor: + +```python +from docarray import BaseDoc +from docarray.typing import TorchTensor +import torch + + +class MyDoc(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + +doc = MyDoc(tensor=torch.zeros(3, 224, 224)) # works +doc = MyDoc(tensor=torch.zeros(224, 224, 3)) # works by reshaping + +try: + doc = MyDoc(tensor=torch.zeros(224)) # fails validation +except Exception as e: + print(e) + # tensor + # Cannot reshape tensor of shape (224,) to shape (3, 224, 224) (type=value_error) + + +class Image(BaseDoc): + tensor: TorchTensor[3, 'x', 'x'] + + +Image(tensor=torch.zeros(3, 224, 224)) # works + +try: + Image( + tensor=torch.zeros(3, 64, 128) + ) # fails validation because second dimension does not match third +except Exception as e: + print() + + +try: + Image( + tensor=torch.zeros(4, 224, 224) + ) # fails validation because of the first dimension +except Exception as e: + print(e) + # Tensor shape mismatch. Expected(3, 'x', 'x'), got(4, 224, 224)(type=value_error) + +try: + Image( + tensor=torch.zeros(3, 64) + ) # fails validation because it does not have enough dimensions +except Exception as e: + print(e) + # Tensor shape mismatch. Expected (3, 'x', 'x'), got (3, 64) (type=value_error) ``` -Or shorten the loop as one-liner: +
+
-### Quantitative evaluation
-Serious as you are, visual inspection is surely not enough. Let's calculate the recall@K. First we construct the groundtruth matches:
+### Coming from TensorFlow
+
+Click to expand
+ +If you come from PyTorch, you can see DocArray mainly as a way of _organizing your data as it flows through your model_. + +It offers you several advantages: + +- Express **tensor shapes in type hints** +- **Group tensors that belong to the same object**, e.g. an audio track and an image +- **Go directly to deployment**, by re-using your data model as a [FastAPI](https://fastapi.tiangolo.com/) or [Jina](https://github.com/jina-ai/jina) API schema +- Connect model components between **microservices**, using Protobuf and gRPC + +DocArray can be used directly inside ML models to handle and represent multimodaldata. +This allows you to reason about your data using DocArray's abstractions deep inside of `nn.Module`, +and provides a FastAPI-compatible schema that eases the transition between model training and model serving. + +To see the effect of this, let's first observe a vanilla PyTorch implementation of a tri-modal ML model: ```python -print(left_da['@m', ('uri', 'scores__cosine__value')]) +import torch +from torch import nn + + +def encoder(x): + return torch.rand(512) + + +class MyMultiModalModel(nn.Module): + def __init__(self): + super().__init__() + self.audio_encoder = encoder() + self.image_encoder = encoder() + self.text_encoder = encoder() + + def forward(self, text_1, text_2, image_1, image_2, audio_1, audio_2): + embedding_text_1 = self.text_encoder(text_1) + embedding_text_2 = self.text_encoder(text_2) + + embedding_image_1 = self.image_encoder(image_1) + embedding_image_2 = self.image_encoder(image_2) + + embedding_audio_1 = self.image_encoder(audio_1) + embedding_audio_2 = self.image_encoder(audio_2) + + return ( + embedding_text_1, + embedding_text_2, + embedding_image_1, + embedding_image_2, + embedding_audio_1, + embedding_audio_2, + ) ``` -Better see it. +Not very easy on the eyes if you ask us. And even worse, if you need to add one more modality you have to touch every part of your code base, changing the `forward()` return type and making a whole lot of changes downstream from that. + +So, now let's see what the same code looks like with DocArray: ```python -(DocumentArray(left_da[8].matches, copy=True) - .apply(lambda d: d.set_image_blob_channel_axis(0, -1) - .set_image_blob_inv_normalization()) - .plot_image_sprites('result.png')) +from docarray import DocList, BaseDoc +from docarray.documents import ImageDoc, TextDoc, AudioDoc +from docarray.typing import TorchTensor +from torch import nn +import torch + + +def encoder(x): + return torch.rand(512) + + +class Podcast(BaseDoc): + text: TextDoc + image: ImageDoc + audio: AudioDoc + + +class PairPodcast(BaseDoc): + left: Podcast + right: Podcast + + +class MyPodcastModel(nn.Module): + def __init__(self): + super().__init__() + self.audio_encoder = encoder() + self.image_encoder = encoder() + self.text_encoder = encoder() + + def forward_podcast(self, docs: DocList[Podcast]) -> DocList[Podcast]: + docs.audio.embedding = self.audio_encoder(docs.audio.tensor) + docs.text.embedding = self.text_encoder(docs.text.tensor) + docs.image.embedding = self.image_encoder(docs.image.tensor) + + return docs + + def forward(self, docs: DocList[PairPodcast]) -> DocList[PairPodcast]: + docs.left = self.forward_podcast(docs.left) + docs.right = self.forward_podcast(docs.right) + + return docs ``` -
-
-
-
+
+
+### Coming from FastAPI
+
+Click to expand
+ +Like the [PyTorch approach](#coming-from-pytorch), you can also use DocArray with TensorFlow to handle and represent multimodal data inside your ML model. + +First off, to use DocArray with TensorFlow we first need to install it as follows: -```python -groundtruth = DocumentArray( - Document(uri=d.uri, matches=[Document(uri=d.uri.replace('left', 'right'))]) for d in left_da) ``` +pip install tensorflow==2.12.0 +pip install protobuf==3.19.0 +``` + +Compared to using DocArray with PyTorch, there is one main difference when using it with TensorFlow: +While DocArray's `TorchTensor` is a subclass of `torch.Tensor`, this is not the case for the `TensorFlowTensor`: Due to some technical limitations of `tf.Tensor`, DocArray's `TensorFlowTensor` is not a subclass of `tf.Tensor` but rather stores a `tf.Tensor` in its `.tensor` attribute. -Here we create a new DocumentArray with real matches by simply replacing the filename, e.g. `left/00001.jpg` to `right/00001.jpg`. That's all we need: if the predicted match has the identical `uri` as the groundtruth match, then it is correct. +How does this affect you? Whenever you want to access the tensor data to, let's say, do operations with it or hand it to your ML model, instead of handing over your `TensorFlowTensor` instance, you need to access its `.tensor` attribute. -Now let's check recall rate from 1 to 5 over the full dataset: +This would look like the following: ```python -for k in range(1, 6): - print(f'recall@{k}', - left_da.evaluate( - groundtruth, - hash_fn=lambda d: d.uri, - metric='recall_at_k', - k=k, - max_rel=1)) +from typing import Optional + +from docarray import DocList, BaseDoc + +import tensorflow as tf + + +class Podcast(BaseDoc): + audio_tensor: Optional[AudioTensorFlowTensor] = None + embedding: Optional[AudioTensorFlowTensor] = None + + +class MyPodcastModel(tf.keras.Model): + def __init__(self): + super().__init__() + self.audio_encoder = AudioEncoder() + + def call(self, inputs: DocList[Podcast]) -> DocList[Podcast]: + inputs.audio_tensor.embedding = self.audio_encoder( + inputs.audio_tensor.tensor + ) # access audio_tensor's .tensor attribute + return inputs ``` -```text -recall@1 0.02726063829787234 -recall@2 0.03873005319148936 -recall@3 0.04670877659574468 -recall@4 0.052194148936170214 -recall@5 0.0573470744680851 +
+
+### Coming from Jina
-### Save results
+Click to expand
+ +Documents are Pydantic Models (with a twist), and as such they are fully compatible with FastAPI! + +But why should you use them, and not the Pydantic models you already know and love? +Good question! + +- Because of the ML-first features, types and validations, [here](#coming-from-pydantic) +- Because DocArray can act as an [ORM for vector databases](#coming-from-a-vector-database), similar to what SQLModel does for SQL databases + +And to seal the deal, let us show you how easily documents slot into your FastAPI app: + +```python +import numpy as np +from fastapi import FastAPI +from docarray.base_doc import DocArrayResponse +from docarray import BaseDoc +from docarray.documents import ImageDoc +from docarray.typing import NdArray, ImageTensor + + +class InputDoc(BaseDoc): + img: ImageDoc + text: str + + +class OutputDoc(BaseDoc): + embedding_clip: NdArray + embedding_bert: NdArray + + +app = FastAPI() + + +def model_img(img: ImageTensor) -> NdArray: + return np.zeros((100, 1)) + + +def model_text(text: str) -> NdArray: + return np.zeros((100, 1)) + + +@app.post("/embed/", response_model=OutputDoc, response_class=DocArrayResponse) +async def create_item(doc: InputDoc) -> OutputDoc: + doc = OutputDoc( + embedding_clip=model_img(doc.img.tensor), embedding_bert=model_text(doc.text) + ) + return doc + + +input_doc = InputDoc(text='', img=ImageDoc(tensor=np.random.random((3, 224, 224)))) + +async with AsyncClient(app=app, base_url="http://test") as ac: + response = await ac.post("/embed/", data=input_doc.json()) ``` -More metrics can be used such as `precision_at_k`, `ndcg_at_k`, `hit_at_k`. +Just like a vanilla Pydantic model! -If you think a pretrained ResNet50 is good enough, let me tell you with [Finetuner](https://github.com/jina-ai/finetuner) one could do much better in just 10 extra lines of code. [Here is how](https://finetuner.jina.ai/get-started/totally-looks-like/). +
+
+
+### Coming from a vector database
+
+Click to expand
-You can save a DocumentArray to binary, JSON, dict, dataframe, CSV or Protobuf message with/without compression. In its simplest form, +Jina has adopted docarray as their library for representing and serializing Documents. + +Jina allows to serve models and services that are built with DocArray allowing you to serve and scale these applications +making full use of DocArray's serialization capabilites. ```python -left_da.save('left_da.bin') +import numpy as np +from jina import Deployment, Executor, requests +from docarray import BaseDoc, DocList +from docarray.documents import ImageDoc +from docarray.typing import NdArray, ImageTensor + + +class InputDoc(BaseDoc): + img: ImageDoc + text: str + + +class OutputDoc(BaseDoc): + embedding_clip: NdArray + embedding_bert: NdArray + + +def model_img(img: ImageTensor) -> NdArray: + return np.zeros((100, 1)) + + +def model_text(text: str) -> NdArray: + return np.zeros((100, 1)) + + +class MyEmbeddingExecutor(Executor): + @requests(on='/embed') + def encode(self, docs: DocList[InputDoc], **kwargs) -> DocList[OutputDoc]: + ret = DocList[OutputDoc]() + for doc in docs: + output = OutputDoc( + embedding_clip=model_img(doc.img.tensor), + embedding_bert=model_text(doc.text), + ) + ret.append(output) + return ret + + +with Deployment( + protocols=['grpc', 'http'], ports=[12345, 12346], uses=MyEmbeddingExecutor +) as dep: + resp = dep.post( + on='/embed', + inputs=DocList[InputDoc]( + [InputDoc(text='', img=ImageDoc(tensor=np.random.random((3, 224, 224))))] + ), + return_type=DocList[OutputDoc], + ) + print(resp) ``` -To reuse it, do `left_da = DocumentArray.load('left_da.bin')`. +
+
+
+
+### Coming from Langchain
+
+Click to expand
-If you want to transfer a DoucmentArray from one machine to another or share it with your colleagues, you can do: +If you came across DocArray as a universal vector database client, you can best think of it as **a new kind of ORM for vector databases**. +DocArray's job is to take multimodal, nested and domain-specific data and to map it to a vector database, +store it there, and thus make it searchable: ```python -left_da.push(token='my_shared_da') +from docarray import DocList, BaseDoc +from docarray.index import HnswDocumentIndex +import numpy as np + +from docarray.typing import ImageUrl, ImageTensor, NdArray + + +class ImageDoc(BaseDoc): + url: ImageUrl + tensor: ImageTensor + embedding: NdArray[128] + + +# create some data +dl = DocList[ImageDoc]( + [ + ImageDoc( + url="https://upload.wikimedia.org/wikipedia/commons/2/2f/Alpamayo.jpg", + tensor=np.zeros((3, 224, 224)), + embedding=np.random.random((128,)), + ) + for _ in range(100) + ] +) + +# create a Document Index +index = HnswDocumentIndex[ImageDoc](work_dir='/tmp/test_index2') + + +# index your data +index.index(dl) + +# find similar Documents +query = dl[0] +results, scores = index.find(query, limit=10, search_field='embedding') ``` -Now anyone who knows the token `my_shared_da` can pull and work on it. +Currently, DocArray supports the following vector databases: + +- [Weaviate](https://www.weaviate.io/) +- [Qdrant](https://qdrant.tech/) +- [Elasticsearch](https://www.elastic.co/elasticsearch/) v8 and v7 +- [Redis](https://redis.io/) +- [Milvus](https://milvus.io) +- ExactNNMemorySearch as a local alternative with exact kNN search. +- [HNSWlib](https://github.com/nmslib/hnswlib) as a local-first ANN alternative +- [Mongo Atlas](https://www.mongodb.com/) + +An integration of [OpenSearch](https://opensearch.org/) is currently in progress. + +Of course this is only one of the things that DocArray can do, so we encourage you to check out the rest of this readme! + +
+
-- Use [Discussions](https://github.com/jina-ai/docarray/discussions) to talk about your use cases, questions, and
- support queries.
-- Join our [Slack community](https://slack.jina.ai) and chat with other community members about ideas.
-- Join our [Engineering All Hands](https://youtube.com/playlist?list=PL3UBBWOUVhFYRUa_gpYYKBqEAkO4sxmne) meet-up to discuss your use case and learn Jina's new features.
- - **When?** The second Tuesday of every month
- - **Where?**
- Zoom ([see our public events calendar](https://calendar.google.com/calendar/embed?src=c_1t5ogfp2d45v8fit981j08mcm4%40group.calendar.google.com&ctz=Europe%2FBerlin)/[.ical](https://calendar.google.com/calendar/ical/c_1t5ogfp2d45v8fit981j08mcm4%40group.calendar.google.com/public/basic.ics))
- and [live stream on YouTube](https://youtube.com/c/jina-ai)
-- Subscribe to the latest video tutorials on our [YouTube channel](https://youtube.com/c/jina-ai)
-## Join Us
+## See also
-DocArray is backed by [Jina AI](https://jina.ai) and licensed under [Apache-2.0](./LICENSE). [We are actively hiring](https://jobs.jina.ai) AI engineers, solution engineers to build the next neural search ecosystem in opensource.
+- [Documentation](https://docs.docarray.org)
+- [DocArray<=0.21 documentation](https://docarray.jina.ai/)
+- [Join our Discord server](https://discord.gg/WaMp6PVPgR)
+- [Donation to Linux Foundation AI&Data blog post](https://jina.ai/news/donate-docarray-lf-for-inclusive-standard-multimodal-data-model/)
+- [Roadmap](https://github.com/docarray/docarray/issues/1714)
-
+> DocArray is a trademark of LF AI Projects, LLC
+>
diff --git a/SECURITY.md b/SECURITY.md
new file mode 100644
index 00000000000..6816dd5cd3a
--- /dev/null
+++ b/SECURITY.md
@@ -0,0 +1,16 @@
+# Security Policy
+
+Security is very important for DocArray and its community.
+
+## Supported Versions
+
+The latest versions of DocArray are supported.
+
+
+## Reporting a Vulnerability
+
+If you think you've found a vulnerability, even if you are not sure about it, please report it right away by sending an email to security@jina.ai. Please try to be as explicit as possible, describing all the steps and example code to reproduce the security issue.
+
+Our team will review it thoroughly and get back to you.
+
+Please refrain from publicly discussing a potential security vulnerability on Github issues or in Slack community. It's better to discuss it privately and try to find a solution first, to limit the potential impact as much as possible.
diff --git a/docarray/__init__.py b/docarray/__init__.py
index ca79cbc3a54..20b08ba1735 100644
--- a/docarray/__init__.py
+++ b/docarray/__init__.py
@@ -1,4 +1,99 @@
-__version__ = '0.1.3'
+# Licensed to the LF AI & Data foundation under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+__version__ = '0.40.2'
-from .document import Document
-from .array import DocumentArray
+import logging
+
+from docarray.array import DocList, DocVec
+from docarray.base_doc.doc import BaseDoc
+from docarray.utils._internal.misc import _get_path_from_docarray_root_level
+from docarray.utils._internal.pydantic import is_pydantic_v2
+
+
+def unpickle_doclist(doc_type, b):
+ return DocList[doc_type].from_bytes(b, protocol="protobuf")
+
+
+def unpickle_docvec(doc_type, tensor_type, b):
+ return DocVec[doc_type].from_bytes(b, protocol="protobuf", tensor_type=tensor_type)
+
+
+if is_pydantic_v2:
+ # Register the pickle functions
+ def register_serializers():
+ import copyreg
+ from functools import partial
+
+ unpickle_doc_fn = partial(BaseDoc.from_bytes, protocol="protobuf")
+
+ def pickle_doc(doc):
+ b = doc.to_bytes(protocol='protobuf')
+ return unpickle_doc_fn, (doc.__class__, b)
+
+ # Register BaseDoc serialization
+ copyreg.pickle(BaseDoc, pickle_doc)
+
+ # For DocList, we need to hook into __reduce__ since it's a generic
+
+ def pickle_doclist(doc_list):
+ b = doc_list.to_bytes(protocol='protobuf')
+ doc_type = doc_list.doc_type
+ return unpickle_doclist, (doc_type, b)
+
+ # Replace DocList.__reduce__ with a method that returns the correct format
+ def doclist_reduce(self):
+ return pickle_doclist(self)
+
+ DocList.__reduce__ = doclist_reduce
+
+ # For DocVec, we need to hook into __reduce__ since it's a generic
+
+ def pickle_docvec(doc_vec):
+ b = doc_vec.to_bytes(protocol='protobuf')
+ doc_type = doc_vec.doc_type
+ tensor_type = doc_vec.tensor_type
+ return unpickle_docvec, (doc_type, tensor_type, b)
+
+ # Replace DocList.__reduce__ with a method that returns the correct format
+ def docvec_reduce(self):
+ return pickle_docvec(self)
+
+ DocVec.__reduce__ = docvec_reduce
+
+ register_serializers()
+
+__all__ = ['BaseDoc', 'DocList', 'DocVec']
+
+logger = logging.getLogger('docarray')
+
+handler = logging.StreamHandler()
+formatter = logging.Formatter("%(levelname)s - %(name)s - %(message)s")
+handler.setFormatter(formatter)
+logger.addHandler(handler)
+
+
+def __getattr__(name: str):
+ if name in ['Document', 'DocumentArray']:
+ raise ImportError(
+ f'Cannot import name \'{name}\' from \'{_get_path_from_docarray_root_level(__file__)}\'.\n'
+ f'The object named \'{name}\' does not exist anymore in this version of docarray.\n'
+ f'If you still want to use \'{name}\' please downgrade to version <=0.21.0 '
+ f'with: `pip install -U docarray==0.21.0`.'
+ )
+ else:
+ raise ImportError(
+ f'cannot import name \'{name}\' from \'{_get_path_from_docarray_root_level(__file__)}\''
+ )
diff --git a/docarray/array/__init__.py b/docarray/array/__init__.py
index f3d8e8486f7..8fc423c13f0 100644
--- a/docarray/array/__init__.py
+++ b/docarray/array/__init__.py
@@ -1 +1,20 @@
-from .document import DocumentArray
+# Licensed to the LF AI & Data foundation under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from docarray.array.any_array import AnyDocArray
+from docarray.array.doc_list.doc_list import DocList
+from docarray.array.doc_vec.doc_vec import DocVec
+
+__all__ = ['DocList', 'DocVec', 'AnyDocArray']
diff --git a/docarray/array/any_array.py b/docarray/array/any_array.py
new file mode 100644
index 00000000000..0c29e54ae82
--- /dev/null
+++ b/docarray/array/any_array.py
@@ -0,0 +1,353 @@
+import sys
+import random
+from abc import abstractmethod
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Dict,
+ Generator,
+ Generic,
+ Iterable,
+ List,
+ MutableSequence,
+ Sequence,
+ Type,
+ TypeVar,
+ Union,
+ cast,
+ overload,
+)
+
+import numpy as np
+
+from docarray.base_doc.doc import BaseDocWithoutId
+from docarray.display.document_array_summary import DocArraySummary
+from docarray.exceptions.exceptions import UnusableObjectError
+from docarray.typing.abstract_type import AbstractType
+from docarray.utils._internal._typing import change_cls_name, safe_issubclass
+from docarray.utils._internal.pydantic import is_pydantic_v2
+
+if TYPE_CHECKING:
+ from docarray.proto import DocListProto, NodeProto
+ from docarray.typing.tensor.abstract_tensor import AbstractTensor
+
+if sys.version_info >= (3, 12):
+ from types import GenericAlias
+
+T = TypeVar('T', bound='AnyDocArray')
+T_doc = TypeVar('T_doc', bound=BaseDocWithoutId)
+IndexIterType = Union[slice, Iterable[int], Iterable[bool], None]
+
+UNUSABLE_ERROR_MSG = (
+ 'This {cls} instance is in an unusable state. \n'
+ 'The most common cause of this is converting a DocVec to a DocList. '
+ 'After you call `doc_vec.to_doc_list()`, `doc_vec` cannot be used anymore. '
+ 'Instead, you should do `doc_list = doc_vec.to_doc_list()` and only use `doc_list`.'
+)
+
+
+class AnyDocArray(Sequence[T_doc], Generic[T_doc], AbstractType):
+ doc_type: Type[BaseDocWithoutId]
+ __typed_da__: Dict[Type['AnyDocArray'], Dict[Type[BaseDocWithoutId], Type]] = {}
+
+ def __repr__(self):
+ return f'<{self.__class__.__name__} (length={len(self)})>'
+
+ @classmethod
+ def __class_getitem__(cls, item: Union[Type[BaseDocWithoutId], TypeVar, str]):
+ if not isinstance(item, type):
+ if sys.version_info < (3, 12):
+ return Generic.__class_getitem__.__func__(cls, item) # type: ignore
+ # this do nothing that checking that item is valid type var or str
+ # Keep the approach in #1147 to be compatible with lower versions of Python.
+ else:
+ return GenericAlias(cls, item) # type: ignore
+ if not safe_issubclass(item, BaseDocWithoutId):
+ raise ValueError(
+ f'{cls.__name__}[item] item should be a Document not a {item} '
+ )
+
+ if cls not in cls.__typed_da__:
+ cls.__typed_da__[cls] = {}
+
+ if item not in cls.__typed_da__[cls]:
+ # Promote to global scope so multiprocessing can pickle it
+ global _DocArrayTyped
+
+ if not is_pydantic_v2:
+
+ class _DocArrayTyped(cls): # type: ignore
+ doc_type: Type[BaseDocWithoutId] = cast(
+ Type[BaseDocWithoutId], item
+ )
+
+ else:
+
+ class _DocArrayTyped(cls, Generic[T_doc]): # type: ignore
+ doc_type: Type[BaseDocWithoutId] = cast(
+ Type[BaseDocWithoutId], item
+ )
+
+ for field in _DocArrayTyped.doc_type._docarray_fields().keys():
+
+ def _property_generator(val: str):
+ def _getter(self):
+ if getattr(self, '_is_unusable', False):
+ raise UnusableObjectError(
+ UNUSABLE_ERROR_MSG.format(cls=cls.__name__)
+ )
+ return self._get_data_column(val)
+
+ def _setter(self, value):
+ if getattr(self, '_is_unusable', False):
+ raise UnusableObjectError(
+ UNUSABLE_ERROR_MSG.format(cls=cls.__name__)
+ )
+ self._set_data_column(val, value)
+
+ # need docstring for the property
+ return property(fget=_getter, fset=_setter)
+
+ setattr(_DocArrayTyped, field, _property_generator(field))
+ # this generates property on the fly based on the schema of the item
+
+ # # The global scope and qualname need to refer to this class a unique name.
+ # # Otherwise, creating another _DocArrayTyped will overwrite this one.
+ if not is_pydantic_v2:
+ change_cls_name(
+ _DocArrayTyped, f'{cls.__name__}[{item.__name__}]', globals()
+ )
+
+ cls.__typed_da__[cls][item] = _DocArrayTyped
+ else:
+ change_cls_name(_DocArrayTyped, f'{cls.__name__}', globals())
+ if sys.version_info < (3, 12):
+ cls.__typed_da__[cls][item] = Generic.__class_getitem__.__func__(
+ _DocArrayTyped, item
+ ) # type: ignore
+ # this do nothing that checking that item is valid type var or str
+ # Keep the approach in #1147 to be compatible with lower versions of Python.
+ else:
+ cls.__typed_da__[cls][item] = GenericAlias(_DocArrayTyped, item) # type: ignore
+ return cls.__typed_da__[cls][item]
+
+ @overload
+ def __getitem__(self: T, item: int) -> T_doc:
+ ...
+
+ @overload
+ def __getitem__(self: T, item: IndexIterType) -> T:
+ ...
+
+ @abstractmethod
+ def __getitem__(self, item: Union[int, IndexIterType]) -> Union[T_doc, T]:
+ ...
+
+ def __getattr__(self, item: str):
+ # Needs to be explicitly defined here for the purpose to disable PyCharm's complaints
+ # about not detected properties: https://youtrack.jetbrains.com/issue/PY-47991
+ return super().__getattribute__(item)
+
+ @abstractmethod
+ def _get_data_column(
+ self: T,
+ field: str,
+ ) -> Union[MutableSequence, T, 'AbstractTensor', None]:
+ """Return all values of the fields from all docs this array contains
+
+ :param field: name of the fields to extract
+ :return: Returns a list of the field value for each document
+ in the array like container
+ """
+ ...
+
+ @abstractmethod
+ def _set_data_column(
+ self: T,
+ field: str,
+ values: Union[List, T, 'AbstractTensor'],
+ ):
+ """Set all Documents in this [`DocList`][docarray.array.doc_list.doc_list.DocList] using the passed values
+
+ :param field: name of the fields to extract
+ :values: the values to set at the DocList level
+ """
+ ...
+
+ @classmethod
+ @abstractmethod
+ def from_protobuf(cls: Type[T], pb_msg: 'DocListProto') -> T:
+ """create a Document from a protobuf message"""
+ ...
+
+ @abstractmethod
+ def to_protobuf(self) -> 'DocListProto':
+ """Convert DocList into a Protobuf message"""
+ ...
+
+ def _to_node_protobuf(self) -> 'NodeProto':
+ """Convert a [`DocList`][docarray.array.doc_list.doc_list.DocList] into a NodeProto
+ protobuf message.
+ This function should be called when a DocList is nested into
+ another Document that need to be converted into a protobuf.
+
+ :return: the nested item protobuf message
+ """
+ from docarray.proto import NodeProto
+
+ return NodeProto(doc_array=self.to_protobuf())
+
+ @abstractmethod
+ def traverse_flat(
+ self: 'AnyDocArray',
+ access_path: str,
+ ) -> Union[List[Any], 'AbstractTensor']:
+ """
+ Return a List of the accessed objects when applying the `access_path`. If this
+ results in a nested list or list of [`DocList`s][docarray.array.doc_list.doc_list.DocList], the list will be flattened
+ on the first level. The access path is a string that consists of attribute
+ names, concatenated and `"__"`-separated. It describes the path from the first
+ level to an arbitrary one, e.g. `'content__image__url'`.
+
+
+ ```python
+ from docarray import BaseDoc, DocList, Text
+
+
+ class Author(BaseDoc):
+ name: str
+
+
+ class Book(BaseDoc):
+ author: Author
+ content: Text
+
+
+ docs = DocList[Book](
+ Book(author=Author(name='Jenny'), content=Text(text=f'book_{i}'))
+ for i in range(10) # noqa: E501
+ )
+
+ books = docs.traverse_flat(access_path='content') # list of 10 Text objs
+
+ authors = docs.traverse_flat(access_path='author__name') # list of 10 strings
+ ```
+
+ If the resulting list is a nested list, it will be flattened:
+
+ ```python
+ from docarray import BaseDoc, DocList
+
+
+ class Chapter(BaseDoc):
+ content: str
+
+
+ class Book(BaseDoc):
+ chapters: DocList[Chapter]
+
+
+ docs = DocList[Book](
+ Book(chapters=DocList[Chapter]([Chapter(content='some_content') for _ in range(3)]))
+ for _ in range(10)
+ )
+
+ chapters = docs.traverse_flat(access_path='chapters') # list of 30 strings
+ ```
+
+ If your [`DocList`][docarray.array.doc_list.doc_list.DocList] is in doc_vec mode, and you want to access a field of
+ type `AnyTensor`, the doc_vec tensor will be returned instead of a list:
+
+ ```python
+ class Image(BaseDoc):
+ tensor: TorchTensor[3, 224, 224]
+
+
+ batch = DocList[Image](
+ [
+ Image(
+ tensor=torch.zeros(3, 224, 224),
+ )
+ for _ in range(2)
+ ]
+ )
+
+ batch_stacked = batch.stack()
+ tensors = batch_stacked.traverse_flat(
+ access_path='tensor'
+ ) # tensor of shape (2, 3, 224, 224)
+ ```
+
+ :param access_path: a string that represents the access path ("__"-separated).
+ :return: list of the accessed objects, flattened if nested.
+ """
+ ...
+
+ @staticmethod
+ def _traverse(node: Any, access_path: str):
+ if access_path:
+ curr_attr, _, path_attrs = access_path.partition('__')
+
+ from docarray.array import DocList
+
+ if isinstance(node, (DocList, list)):
+ for n in node:
+ x = getattr(n, curr_attr)
+ yield from AnyDocArray._traverse(x, path_attrs)
+ else:
+ x = getattr(node, curr_attr)
+ yield from AnyDocArray._traverse(x, path_attrs)
+ else:
+ yield node
+
+ @staticmethod
+ def _flatten_one_level(sequence: List[Any]) -> List[Any]:
+ from docarray import DocList
+
+ if len(sequence) == 0 or not isinstance(sequence[0], (list, DocList)):
+ return sequence
+ else:
+ return [item for sublist in sequence for item in sublist]
+
+ def summary(self):
+ """
+ Print a summary of this [`DocList`][docarray.array.doc_list.doc_list.DocList] object and a summary of the schema of its
+ Document type.
+ """
+ DocArraySummary(self).summary()
+
+ def _batch(
+ self: T,
+ batch_size: int,
+ shuffle: bool = False,
+ show_progress: bool = False,
+ ) -> Generator[T, None, None]:
+ """
+ Creates a `Generator` that yields [`DocList`][docarray.array.doc_list.doc_list.DocList] of size `batch_size`.
+ Note, that the last batch might be smaller than `batch_size`.
+
+ :param batch_size: Size of each generated batch.
+ :param shuffle: If set, shuffle the Documents before dividing into minibatches.
+ :param show_progress: if set, show a progress bar when batching documents.
+ :yield: a Generator of [`DocList`][docarray.array.doc_list.doc_list.DocList], each in the length of `batch_size`
+ """
+ from rich.progress import track
+
+ if not (isinstance(batch_size, int) and batch_size > 0):
+ raise ValueError(
+ f'`batch_size` should be a positive integer, received: {batch_size}'
+ )
+
+ N = len(self)
+ indices = list(range(N))
+ n_batches = int(np.ceil(N / batch_size))
+
+ if shuffle:
+ random.shuffle(indices)
+
+ for i in track(
+ range(n_batches),
+ description='Batching documents',
+ disable=not show_progress,
+ ):
+ yield self[indices[i * batch_size : (i + 1) * batch_size]]
diff --git a/docarray/array/chunk.py b/docarray/array/chunk.py
deleted file mode 100644
index ea445702f62..00000000000
--- a/docarray/array/chunk.py
+++ /dev/null
@@ -1,81 +0,0 @@
-import itertools
-from typing import (
- TYPE_CHECKING,
- Generator,
- Iterator,
- Sequence,
-)
-
-from .document import DocumentArray
-
-if TYPE_CHECKING:
- from ..document import Document
-
-
-class ChunkArray(DocumentArray):
- """
- :class:`ChunkArray` inherits from :class:`DocumentArray`.
- It's a subset of Documents.
-
- :param docs: Set of sub-documents (i.e chunks) of `reference_doc`
- :param reference_doc: Reference :class:`Document` for the sub-documents
- """
-
- def __init__(self, docs, reference_doc: 'Document'):
- """
- Set constructor method.
-
- :param doc_views: protobuf representation of the chunks
- :param reference_doc: parent document
- """
- self._ref_doc = reference_doc
- super().__init__(docs)
- if (
- isinstance(
- docs, (DocumentArray, Sequence, Generator, Iterator, itertools.chain)
- )
- and self._ref_doc is not None
- ):
- for d in docs:
- d.parent_id = self._ref_doc.id
- d.granularity = self._ref_doc.granularity + 1
-
- def append(self, document: 'Document'):
- """Add a sub-document (i.e chunk) to the current Document.
-
- :param document: Sub-document to be appended
-
- .. note::
- Comparing to :attr:`DocumentArray.append()`, this method adds more safeguard to
- make sure the added chunk is legit.
- """
- document.parent_id = self._ref_doc.id
- document.granularity = self._ref_doc.granularity + 1
- super().append(document)
-
- @property
- def reference_doc(self) -> 'Document':
- """
- Get the document that :class:`ChunkArray` belongs to.
-
- :return: reference doc
- """
- return self._ref_doc
-
- @property
- def granularity(self) -> int:
- """
- Get granularity of all document in this array.
-
- :return: granularity
- """
- return self._ref_doc.granularity + 1
-
- @property
- def adjacency(self) -> int:
- """
- Get adjacency of all document in this array.
-
- :return: adjacency
- """
- return self._ref_doc.adjacency
diff --git a/docarray/array/doc_list/__init__.py b/docarray/array/doc_list/__init__.py
new file mode 100644
index 00000000000..74f8f7582cd
--- /dev/null
+++ b/docarray/array/doc_list/__init__.py
@@ -0,0 +1,15 @@
+# Licensed to the LF AI & Data foundation under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/docarray/array/doc_list/doc_list.py b/docarray/array/doc_list/doc_list.py
new file mode 100644
index 00000000000..49236199153
--- /dev/null
+++ b/docarray/array/doc_list/doc_list.py
@@ -0,0 +1,373 @@
+import io
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Iterable,
+ List,
+ MutableSequence,
+ Optional,
+ Sequence,
+ Type,
+ TypeVar,
+ Union,
+ cast,
+ overload,
+ Callable,
+)
+
+from pydantic import parse_obj_as
+from typing_extensions import SupportsIndex
+from typing_inspect import is_typevar, is_union_type
+
+from docarray.array.any_array import AnyDocArray
+from docarray.array.doc_list.io import IOMixinDocList
+from docarray.array.doc_list.pushpull import PushPullMixin
+from docarray.array.list_advance_indexing import IndexIterType, ListAdvancedIndexing
+from docarray.base_doc import AnyDoc
+from docarray.base_doc.doc import BaseDocWithoutId
+from docarray.typing import NdArray
+from docarray.utils._internal.pydantic import is_pydantic_v2
+
+if is_pydantic_v2:
+ from pydantic_core import core_schema
+
+from docarray.utils._internal._typing import safe_issubclass
+
+if TYPE_CHECKING:
+
+ from docarray.array.doc_vec.doc_vec import DocVec
+ from docarray.proto import DocListProto
+ from docarray.typing import TorchTensor
+ from docarray.typing.tensor.abstract_tensor import AbstractTensor
+
+T = TypeVar('T', bound='DocList')
+T_doc = TypeVar('T_doc', bound=BaseDocWithoutId)
+
+
+class DocList(
+ ListAdvancedIndexing[T_doc], PushPullMixin, IOMixinDocList, AnyDocArray[T_doc]
+):
+ """
+ DocList is a container of Documents.
+
+ A DocList is a list of Documents of any schema. However, many
+ DocList features are only available if these Documents are
+ homogeneous and follow the same schema. To precise this schema you can use
+ the `DocList[MyDocument]` syntax where MyDocument is a Document class
+ (i.e. schema). This creates a DocList that can only contains Documents of
+ the type `MyDocument`.
+
+
+ ```python
+ from docarray import BaseDoc, DocList
+ from docarray.typing import NdArray, ImageUrl
+ from typing import Optional
+
+
+ class Image(BaseDoc):
+ tensor: Optional[NdArray[100]] = None
+ url: ImageUrl
+
+
+ docs = DocList[Image](
+ Image(url='http://url.com/foo.png') for _ in range(10)
+ ) # noqa: E510
+
+
+ # If your DocList is homogeneous (i.e. follows the same schema), you can access
+ # fields at the DocList level (for example `docs.tensor` or `docs.url`).
+
+ print(docs.url)
+ # [ImageUrl('http://url.com/foo.png', host_type='domain'), ...]
+
+
+ # You can also set fields, with `docs.tensor = np.random.random([10, 100])`:
+
+
+ import numpy as np
+
+ docs.tensor = np.random.random([10, 100])
+
+ print(docs.tensor)
+ # [NdArray([0.11299577, 0.47206767, 0.481723 , 0.34754724, 0.15016037,
+ # 0.88861321, 0.88317666, 0.93845579, 0.60486676, ... ]), ...]
+
+
+ # You can index into a DocList like a numpy doc_list or torch tensor:
+
+ docs[0] # index by position
+ docs[0:5:2] # index by slice
+ docs[[0, 2, 3]] # index by list of indices
+ docs[True, False, True, True, ...] # index by boolean mask
+
+
+ # You can delete items from a DocList like a Python List
+
+ del docs[0] # remove first element from DocList
+ del docs[0:5] # remove elements for 0 to 5 from DocList
+ ```
+
+ !!! note
+ If the DocList is homogeneous and its schema contains nested BaseDoc
+ (i.e, BaseDoc inside a BaseDoc) where the nested Document is `Optional`, calling
+ `docs.nested_doc` will return a List of the nested BaseDoc instead of DocList.
+ This is because the nested field could be None and therefore could not fit into
+ a DocList.
+
+ :param docs: iterable of Document
+
+ """
+
+ doc_type: Type[BaseDocWithoutId] = AnyDoc
+
+ def __init__(
+ self,
+ docs: Optional[Iterable[T_doc]] = None,
+ validate_input_docs: bool = True,
+ ):
+ if validate_input_docs:
+ docs = self._validate_docs(docs) if docs else []
+ else:
+ docs = docs if docs else []
+ super().__init__(docs)
+
+ @classmethod
+ def construct(
+ cls: Type[T],
+ docs: Sequence[T_doc],
+ ) -> T:
+ """
+ Create a `DocList` without validation any data. The data must come from a
+ trusted source
+ :param docs: a Sequence (list) of Document with the same schema
+ :return: a `DocList` object
+ """
+ return cls(docs, False)
+
+ def __eq__(self, other: Any) -> bool:
+ if self.__len__() != other.__len__():
+ return False
+ for doc_self, doc_other in zip(self, other):
+ if doc_self != doc_other:
+ return False
+ return True
+
+ def _validate_docs(self, docs: Iterable[T_doc]) -> Iterable[T_doc]:
+ """
+ Validate if an Iterable of Document are compatible with this `DocList`
+ """
+ for doc in docs:
+ yield self._validate_one_doc(doc)
+
+ def _validate_one_doc(self, doc: T_doc) -> T_doc:
+ """Validate if a Document is compatible with this `DocList`"""
+ if not safe_issubclass(self.doc_type, AnyDoc) and not isinstance(
+ doc, self.doc_type
+ ):
+ raise ValueError(f'{doc} is not a {self.doc_type}')
+ return doc
+
+ def __bytes__(self) -> bytes:
+ with io.BytesIO() as bf:
+ self._write_bytes(bf=bf)
+ return bf.getvalue()
+
+ def append(self, doc: T_doc):
+ """
+ Append a Document to the `DocList`. The Document must be from the same class
+ as the `.doc_type` of this `DocList` otherwise it will fail.
+ :param doc: A Document
+ """
+ return super().append(self._validate_one_doc(doc))
+
+ def extend(self, docs: Iterable[T_doc]):
+ """
+ Extend a `DocList` with an Iterable of Document. The Documents must be from
+ the same class as the `.doc_type` of this `DocList` otherwise it will
+ fail.
+ :param docs: Iterable of Documents
+ """
+ it: Iterable[T_doc] = list()
+ if self is docs:
+ # see https://github.com/docarray/docarray/issues/1489
+ it = list(docs)
+ else:
+ it = self._validate_docs(docs)
+
+ return super().extend(it)
+
+ def insert(self, i: SupportsIndex, doc: T_doc):
+ """
+ Insert a Document to the `DocList`. The Document must be from the same
+ class as the doc_type of this `DocList` otherwise it will fail.
+ :param i: index to insert
+ :param doc: A Document
+ """
+ super().insert(i, self._validate_one_doc(doc))
+
+ def _get_data_column(
+ self: T,
+ field: str,
+ ) -> Union[MutableSequence, T, 'TorchTensor', 'NdArray']:
+ """Return all v @classmethod
+ def __class_getitem__(cls, item: Union[Type[BaseDoc], TypeVar, str]):alues of the fields from all docs this doc_list contains
+ @classmethod
+ def __class_getitem__(cls, item: Union[Type[BaseDoc], TypeVar, str]):
+ :param field: name of the fields to extract
+ :return: Returns a list of the field value for each document
+ in the doc_list like container
+ """
+ field_type = self.__class__.doc_type._get_field_annotation(field)
+ field_info = self.__class__.doc_type._docarray_fields()[field]
+ is_field_required = (
+ field_info.is_required() if is_pydantic_v2 else field_info.required
+ )
+
+ if (
+ not is_union_type(field_type)
+ and is_field_required
+ and isinstance(field_type, type)
+ and safe_issubclass(field_type, BaseDocWithoutId)
+ ):
+ # calling __class_getitem__ ourselves is a hack otherwise mypy complain
+ # most likely a bug in mypy though
+ # bug reported here https://github.com/python/mypy/issues/14111
+ return DocList.__class_getitem__(field_type)(
+ (getattr(doc, field) for doc in self),
+ )
+ else:
+ return [getattr(doc, field) for doc in self]
+
+ def _set_data_column(
+ self: T,
+ field: str,
+ values: Union[List, T, 'AbstractTensor'],
+ ):
+ """Set all Documents in this `DocList` using the passed values
+
+ :param field: name of the fields to set
+ :values: the values to set at the `DocList` level
+ """
+ ...
+
+ for doc, value in zip(self, values):
+ setattr(doc, field, value)
+
+ def to_doc_vec(
+ self,
+ tensor_type: Type['AbstractTensor'] = NdArray,
+ ) -> 'DocVec':
+ """
+ Convert the `DocList` into a `DocVec`. `Self` cannot be used
+ afterward
+ :param tensor_type: Tensor Class used to wrap the doc_vec tensors. This is useful
+ if the BaseDoc has some undefined tensor type like AnyTensor or Union of NdArray and TorchTensor
+ :return: A `DocVec` of the same document type as self
+ """
+ from docarray.array.doc_vec.doc_vec import DocVec
+
+ return DocVec.__class_getitem__(self.doc_type)(self, tensor_type=tensor_type)
+
+ @classmethod
+ def _docarray_validate(
+ cls: Type[T],
+ value: Union[T, Iterable[BaseDocWithoutId]],
+ ):
+ from docarray.array.doc_vec.doc_vec import DocVec
+
+ if isinstance(value, cls):
+ return value
+ elif isinstance(value, DocVec):
+ if (
+ safe_issubclass(value.doc_type, cls.doc_type)
+ or value.doc_type == cls.doc_type
+ ):
+ return cast(T, value.to_doc_list())
+ else:
+ raise ValueError(
+ f'DocList[value.doc_type] is not compatible with {cls}'
+ )
+ elif isinstance(value, cls):
+ return cls(value)
+ elif isinstance(value, Iterable):
+ docs = []
+ for doc in value:
+ docs.append(parse_obj_as(cls.doc_type, doc))
+ return cls(docs)
+ else:
+ raise TypeError(f'Expecting an Iterable of {cls.doc_type}')
+
+ def traverse_flat(
+ self: 'DocList',
+ access_path: str,
+ ) -> List[Any]:
+ nodes = list(AnyDocArray._traverse(node=self, access_path=access_path))
+ flattened = AnyDocArray._flatten_one_level(nodes)
+
+ return flattened
+
+ @classmethod
+ def from_protobuf(cls: Type[T], pb_msg: 'DocListProto') -> T:
+ """create a Document from a protobuf message
+ :param pb_msg: The protobuf message from where to construct the `DocList`
+ """
+ return super().from_protobuf(pb_msg)
+
+ @classmethod
+ def _get_proto_class(cls: Type[T]):
+ from docarray.proto import DocListProto
+
+ return DocListProto
+
+ @overload
+ def __getitem__(self, item: SupportsIndex) -> T_doc:
+ ...
+
+ @overload
+ def __getitem__(self: T, item: IndexIterType) -> T:
+ ...
+
+ def __getitem__(self, item):
+ return super().__getitem__(item)
+
+ @classmethod
+ def __class_getitem__(cls, item: Union[Type[BaseDocWithoutId], TypeVar, str]):
+ if cls.doc_type != AnyDoc:
+ raise TypeError(f'{cls} object is not subscriptable')
+
+ if isinstance(item, type) and safe_issubclass(item, BaseDocWithoutId):
+ return AnyDocArray.__class_getitem__.__func__(cls, item) # type: ignore
+ if (
+ isinstance(item, object)
+ and not is_typevar(item)
+ and not isinstance(item, str)
+ and item is not Any
+ ):
+ raise TypeError('Expecting a type, got object instead')
+
+ return super().__class_getitem__(item)
+
+ def __repr__(self):
+ return AnyDocArray.__repr__(self) # type: ignore
+
+ if is_pydantic_v2:
+
+ @classmethod
+ def __get_pydantic_core_schema__(
+ cls, source: Any, handler: Callable[[Any], core_schema.CoreSchema]
+ ) -> core_schema.CoreSchema:
+ instance_schema = core_schema.is_instance_schema(cls)
+ args = getattr(source, '__args__', None)
+ if args:
+ sequence_t_schema = handler(Sequence[args[0]])
+ else:
+ sequence_t_schema = handler(Sequence)
+
+ def validate_fn(v, info):
+ # input has already been validated
+ return cls(v, validate_input_docs=False)
+
+ non_instance_schema = core_schema.with_info_after_validator_function(
+ validate_fn, sequence_t_schema
+ )
+ return core_schema.union_schema([instance_schema, non_instance_schema])
diff --git a/docarray/array/doc_list/io.py b/docarray/array/doc_list/io.py
new file mode 100644
index 00000000000..3acb66bf6e8
--- /dev/null
+++ b/docarray/array/doc_list/io.py
@@ -0,0 +1,860 @@
+import base64
+import csv
+import io
+import os
+import pathlib
+import pickle
+from abc import abstractmethod
+from contextlib import nullcontext
+from io import StringIO, TextIOWrapper
+from itertools import compress
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ BinaryIO,
+ ContextManager,
+ Dict,
+ Generator,
+ Iterable,
+ Iterator,
+ List,
+ Optional,
+ Tuple,
+ Type,
+ TypeVar,
+ Union,
+ cast,
+)
+
+import orjson
+
+from docarray.base_doc import AnyDoc, BaseDoc
+from docarray.base_doc.io.json import orjson_dumps
+from docarray.helper import (
+ _access_path_dict_to_nested_dict,
+ _all_access_paths_valid,
+ _dict_to_access_paths,
+)
+from docarray.utils._internal.compress import _decompress_bytes, _get_compress_ctx
+from docarray.utils._internal.misc import import_library, ProtocolType
+
+if TYPE_CHECKING:
+ import pandas as pd
+
+ from docarray.array.doc_vec.doc_vec import DocVec
+ from docarray.array.doc_vec.io import IOMixinDocVec
+ from docarray.proto import DocListProto
+ from docarray.typing.tensor.abstract_tensor import AbstractTensor
+
+T = TypeVar('T', bound='IOMixinDocList')
+T_doc = TypeVar('T_doc', bound=BaseDoc)
+
+ARRAY_PROTOCOLS = {'protobuf-array', 'pickle-array', 'json-array'}
+SINGLE_PROTOCOLS = {'pickle', 'protobuf', 'json'}
+ALLOWED_PROTOCOLS = ARRAY_PROTOCOLS.union(SINGLE_PROTOCOLS)
+ALLOWED_COMPRESSIONS = {'lz4', 'bz2', 'lzma', 'zlib', 'gzip'}
+
+
+def _protocol_and_compress_from_file_path(
+ file_path: Union[pathlib.Path, str],
+ default_protocol: Optional[ProtocolType] = None,
+ default_compress: Optional[str] = None,
+) -> Tuple[Optional[ProtocolType], Optional[str]]:
+ """Extract protocol and compression algorithm from a string, use defaults if not found.
+ :param file_path: path of a file.
+ :param default_protocol: default serialization protocol used in case not found.
+ :param default_compress: default compression method used in case not found.
+ Examples:
+ >>> _protocol_and_compress_from_file_path('./docarray_fashion_mnist.protobuf.gzip')
+ ('protobuf', 'gzip')
+ >>> _protocol_and_compress_from_file_path('/Documents/docarray_fashion_mnist.protobuf')
+ ('protobuf', None)
+ >>> _protocol_and_compress_from_file_path('/Documents/docarray_fashion_mnist.gzip')
+ (None, gzip)
+ """
+
+ protocol = default_protocol
+ compress = default_compress
+
+ file_extensions = [e.replace('.', '') for e in pathlib.Path(file_path).suffixes]
+ for extension in file_extensions:
+ if extension in ALLOWED_PROTOCOLS:
+ protocol = cast(ProtocolType, extension)
+ elif extension in ALLOWED_COMPRESSIONS:
+ compress = extension
+
+ return protocol, compress
+
+
+class _LazyRequestReader:
+ def __init__(self, r):
+ self._data = r.iter_content(chunk_size=1024 * 1024)
+ self.content = b''
+
+ def __getitem__(self, item: slice):
+ while len(self.content) < item.stop:
+ try:
+ self.content += next(self._data)
+ except StopIteration:
+ return self.content[item.start : -1 : item.step]
+ return self.content[item]
+
+
+class IOMixinDocList(Iterable[T_doc]):
+ doc_type: Type[T_doc]
+
+ @abstractmethod
+ def __len__(self):
+ ...
+
+ @abstractmethod
+ def __init__(
+ self,
+ docs: Optional[Iterable[BaseDoc]] = None,
+ ):
+ ...
+
+ @classmethod
+ def from_protobuf(cls: Type[T], pb_msg: 'DocListProto') -> T:
+ """create a Document from a protobuf message
+ :param pb_msg: The protobuf message from where to construct the DocList
+ """
+ return cls(cls.doc_type.from_protobuf(doc_proto) for doc_proto in pb_msg.docs)
+
+ def to_protobuf(self) -> 'DocListProto':
+ """Convert `DocList` into a Protobuf message"""
+ from docarray.proto import DocListProto
+
+ da_proto = DocListProto()
+ for doc in self:
+ da_proto.docs.append(doc.to_protobuf())
+
+ return da_proto
+
+ @classmethod
+ def from_bytes(
+ cls: Type[T],
+ data: bytes,
+ protocol: ProtocolType = 'protobuf-array',
+ compress: Optional[str] = None,
+ show_progress: bool = False,
+ ) -> T:
+ """Deserialize bytes into a `DocList`.
+
+ :param data: Bytes from which to deserialize
+ :param protocol: protocol that was used to serialize
+ :param compress: compression algorithm that was used to serialize between `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
+ :param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
+ :return: the deserialized `DocList`
+ """
+ return cls._load_binary_all(
+ file_ctx=nullcontext(data),
+ protocol=protocol,
+ compress=compress,
+ show_progress=show_progress,
+ )
+
+ def _write_bytes(
+ self,
+ bf: BinaryIO,
+ protocol: ProtocolType = 'protobuf-array',
+ compress: Optional[str] = None,
+ show_progress: bool = False,
+ ) -> None:
+ if protocol in ARRAY_PROTOCOLS:
+ compress_ctx = _get_compress_ctx(compress)
+ else:
+ # delegate the compression to per-doc compression
+ compress_ctx = None
+
+ fc: ContextManager
+ if compress_ctx is None:
+ # if compress do not support streaming then postpone the compress
+ # into the for-loop
+ f, fc = bf, nullcontext()
+ else:
+ f = compress_ctx(bf)
+ fc = f
+ compress = None
+
+ with fc:
+ if protocol == 'protobuf-array':
+ f.write(self.to_protobuf().SerializePartialToString())
+ elif protocol == 'pickle-array':
+ f.write(pickle.dumps(self))
+ elif protocol == 'json-array':
+ f.write(self.to_json().encode())
+ elif protocol in SINGLE_PROTOCOLS:
+ f.write(
+ b''.join(
+ self._to_binary_stream(
+ protocol=protocol,
+ compress=compress,
+ show_progress=show_progress,
+ )
+ )
+ )
+ else:
+ raise ValueError(
+ f'protocol={protocol} is not supported. Can be only {ALLOWED_PROTOCOLS}.'
+ )
+
+ def _to_binary_stream(
+ self,
+ protocol: ProtocolType = 'protobuf',
+ compress: Optional[str] = None,
+ show_progress: bool = False,
+ ) -> Iterator[bytes]:
+ from rich import filesize
+
+ if show_progress:
+ from docarray.utils._internal.progress_bar import _get_progressbar
+
+ pbar, t = _get_progressbar(
+ 'Serializing', disable=not show_progress, total=len(self)
+ )
+ else:
+ from contextlib import nullcontext
+
+ pbar = nullcontext()
+
+ yield self._stream_header
+
+ with pbar:
+ if show_progress:
+ _total_size = 0
+ pbar.start_task(t)
+ for doc in self:
+ doc_bytes = doc.to_bytes(protocol=protocol, compress=compress)
+ len_doc_as_bytes = len(doc_bytes).to_bytes(4, 'big', signed=False)
+ all_bytes = len_doc_as_bytes + doc_bytes
+
+ yield all_bytes
+
+ if show_progress:
+ _total_size += len(all_bytes)
+ pbar.update(
+ t,
+ advance=1,
+ total_size=str(filesize.decimal(_total_size)),
+ )
+
+ def to_bytes(
+ self,
+ protocol: ProtocolType = 'protobuf-array',
+ compress: Optional[str] = None,
+ file_ctx: Optional[BinaryIO] = None,
+ show_progress: bool = False,
+ ) -> Optional[bytes]:
+ """Serialize itself into `bytes`.
+
+ For more Pythonic code, please use ``bytes(...)``.
+
+ :param protocol: protocol to use. It can be 'pickle-array', 'protobuf-array', 'pickle' or 'protobuf'
+ :param compress: compress algorithm to use between : `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
+ :param file_ctx: File or filename or serialized bytes where the data is stored.
+ :param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
+ :return: the binary serialization in bytes or None if file_ctx is passed where to store
+ """
+ with file_ctx or io.BytesIO() as bf:
+ self._write_bytes(
+ bf=bf,
+ protocol=protocol,
+ compress=compress,
+ show_progress=show_progress,
+ )
+ if isinstance(bf, io.BytesIO):
+ return bf.getvalue()
+
+ return None
+
+ @classmethod
+ def from_base64(
+ cls: Type[T],
+ data: str,
+ protocol: ProtocolType = 'protobuf-array',
+ compress: Optional[str] = None,
+ show_progress: bool = False,
+ ) -> T:
+ """Deserialize base64 strings into a `DocList`.
+
+ :param data: Base64 string to deserialize
+ :param protocol: protocol that was used to serialize
+ :param compress: compress algorithm that was used to serialize between `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
+ :param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
+ :return: the deserialized `DocList`
+ """
+ return cls._load_binary_all(
+ file_ctx=nullcontext(base64.b64decode(data)),
+ protocol=protocol,
+ compress=compress,
+ show_progress=show_progress,
+ )
+
+ def to_base64(
+ self,
+ protocol: ProtocolType = 'protobuf-array',
+ compress: Optional[str] = None,
+ show_progress: bool = False,
+ ) -> str:
+ """Serialize itself into base64 encoded string.
+
+ :param protocol: protocol to use. It can be 'pickle-array', 'protobuf-array', 'pickle' or 'protobuf'
+ :param compress: compress algorithm to use between `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
+ :param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
+ :return: the binary serialization in bytes or None if file_ctx is passed where to store
+ """
+ with io.BytesIO() as bf:
+ self._write_bytes(
+ bf=bf,
+ compress=compress,
+ protocol=protocol,
+ show_progress=show_progress,
+ )
+ return base64.b64encode(bf.getvalue()).decode('utf-8')
+
+ @classmethod
+ def from_json(
+ cls: Type[T],
+ file: Union[str, bytes, bytearray],
+ ) -> T:
+ """Deserialize JSON strings or bytes into a `DocList`.
+
+ :param file: JSON object from where to deserialize a `DocList`
+ :return: the deserialized `DocList`
+ """
+ json_docs = orjson.loads(file)
+ return cls([cls.doc_type(**v) for v in json_docs])
+
+ def to_json(self) -> str:
+ """Convert the object into JSON bytes. Can be loaded via `.from_json`.
+ :return: JSON serialization of `DocList`
+ """
+ return orjson_dumps(self).decode('UTF-8')
+
+ @classmethod
+ def from_csv(
+ cls: Type['T'],
+ file_path: str,
+ encoding: str = 'utf-8',
+ dialect: Union[str, csv.Dialect] = 'excel',
+ ) -> 'T':
+ """
+ Load a DocList from a csv file following the schema defined in the
+ [`.doc_type`][docarray.DocList] attribute.
+ Every row of the csv file will be mapped to one document in the doc_list.
+ The column names (defined in the first row) have to match the field names
+ of the Document type.
+ For nested fields use "__"-separated access paths, such as `'image__url'`.
+
+ List-like fields (including field of type DocList) are not supported.
+
+ :param file_path: path to csv file to load DocList from.
+ :param encoding: encoding used to read the csv file. Defaults to 'utf-8'.
+ :param dialect: defines separator and how to handle whitespaces etc.
+ Can be a [`csv.Dialect`](https://docs.python.org/3/library/csv.html#csv.Dialect)
+ instance or one string of:
+ `'excel'` (for comma separated values),
+ `'excel-tab'` (for tab separated values),
+ `'unix'` (for csv file generated on UNIX systems).
+
+ :return: `DocList` object
+ """
+ if cls.doc_type == AnyDoc or cls.doc_type == BaseDoc:
+ raise TypeError(
+ 'There is no document schema defined. '
+ f'Please specify the {cls}\'s Document type using `{cls}[MyDoc]`.'
+ )
+
+ if file_path.startswith('http'):
+ import urllib.request
+
+ with urllib.request.urlopen(file_path) as f:
+ file = StringIO(f.read().decode(encoding))
+ return cls._from_csv_file(file, dialect)
+ else:
+ with open(file_path, 'r', encoding=encoding) as fp:
+ return cls._from_csv_file(fp, dialect)
+
+ @classmethod
+ def _from_csv_file(
+ cls: Type['T'],
+ file: Union[StringIO, TextIOWrapper],
+ dialect: Union[str, csv.Dialect],
+ ) -> 'T':
+ rows = csv.DictReader(file, dialect=dialect)
+
+ doc_type = cls.doc_type
+ docs = []
+
+ field_names: List[str] = (
+ [] if rows.fieldnames is None else [str(f) for f in rows.fieldnames]
+ )
+ if field_names is None or len(field_names) == 0:
+ raise TypeError("No field names are given.")
+
+ valid_paths = _all_access_paths_valid(
+ doc_type=doc_type, access_paths=field_names
+ )
+ if not all(valid_paths):
+ raise ValueError(
+ f'Column names do not match the schema of the DocList\'s '
+ f'document type ({cls.doc_type.__name__}): '
+ f'{list(compress(field_names, [not v for v in valid_paths]))}'
+ )
+
+ for access_path2val in rows:
+ doc_dict: Dict[Any, Any] = _access_path_dict_to_nested_dict(access_path2val)
+ docs.append(doc_type.parse_obj(doc_dict))
+
+ return cls(docs)
+
+ def to_csv(
+ self, file_path: str, dialect: Union[str, csv.Dialect] = 'excel'
+ ) -> None:
+ """
+ Save a `DocList` to a csv file.
+ The field names will be stored in the first row. Each row corresponds to the
+ information of one Document.
+ Columns for nested fields will be named after the "__"-seperated access paths,
+ such as `'image__url'` for `image.url`.
+
+ :param file_path: path to a csv file.
+ :param dialect: defines separator and how to handle whitespaces etc.
+ Can be a [`csv.Dialect`](https://docs.python.org/3/library/csv.html#csv.Dialect)
+ instance or one string of:
+ `'excel'` (for comma separated values),
+ `'excel-tab'` (for tab separated values),
+ `'unix'` (for csv file generated on UNIX systems).
+
+ """
+ if self.doc_type == AnyDoc or self.doc_type == BaseDoc:
+ raise TypeError(
+ f'{type(self)} must be homogeneous to be converted to a csv.'
+ 'There is no document schema defined. '
+ f'Please specify the {type(self)}\'s Document type using `{type(self)}[MyDoc]`.'
+ )
+ fields = self.doc_type._get_access_paths()
+
+ with open(file_path, 'w') as csv_file:
+ writer = csv.DictWriter(csv_file, fieldnames=fields, dialect=dialect)
+ writer.writeheader()
+
+ for doc in self:
+ doc_dict = _dict_to_access_paths(doc.dict())
+ writer.writerow(doc_dict)
+
+ @classmethod
+ def from_dataframe(cls: Type['T'], df: 'pd.DataFrame') -> 'T':
+ """
+ Load a `DocList` from a `pandas.DataFrame` following the schema
+ defined in the [`.doc_type`][docarray.DocList] attribute.
+ Every row of the dataframe will be mapped to one Document in the doc_list.
+ The column names of the dataframe have to match the field names of the
+ Document type.
+ For nested fields use "__"-separated access paths as column names,
+ such as `'image__url'`.
+
+ List-like fields (including field of type DocList) are not supported.
+
+ ---
+
+ ```python
+ import pandas as pd
+
+ from docarray import BaseDoc, DocList
+
+
+ class Person(BaseDoc):
+ name: str
+ follower: int
+
+
+ df = pd.DataFrame(
+ data=[['Maria', 12345], ['Jake', 54321]], columns=['name', 'follower']
+ )
+
+ docs = DocList[Person].from_dataframe(df)
+
+ assert docs.name == ['Maria', 'Jake']
+ assert docs.follower == [12345, 54321]
+ ```
+
+ ---
+
+ :param df: `pandas.DataFrame` to extract Document's information from
+ :return: `DocList` where each Document contains the information of one
+ corresponding row of the `pandas.DataFrame`.
+ """
+ from docarray import DocList
+
+ if cls.doc_type == AnyDoc or cls.doc_type == BaseDoc:
+ raise TypeError(
+ 'There is no document schema defined. '
+ f'Please specify the {cls}\'s Document type using `{cls}[MyDoc]`.'
+ )
+
+ doc_type = cls.doc_type
+ docs = DocList.__class_getitem__(doc_type)()
+ field_names = df.columns.tolist()
+
+ if field_names is None or len(field_names) == 0:
+ raise TypeError("No field names are given.")
+
+ valid_paths = _all_access_paths_valid(
+ doc_type=doc_type, access_paths=field_names
+ )
+ if not all(valid_paths):
+ raise ValueError(
+ f'Column names do not match the schema of the DocList\'s '
+ f'document type ({cls.doc_type.__name__}): '
+ f'{list(compress(field_names, [not v for v in valid_paths]))}'
+ )
+
+ for row in df.itertuples():
+ access_path2val = row._asdict()
+ access_path2val.pop('index', None)
+ doc_dict = _access_path_dict_to_nested_dict(access_path2val)
+ docs.append(doc_type.parse_obj(doc_dict))
+
+ return docs
+
+ def to_dataframe(self) -> 'pd.DataFrame':
+ """
+ Save a DocList to a `pandas.DataFrame`.
+ The field names will be stored as column names. Each row of the dataframe corresponds
+ to the information of one Document.
+ Columns for nested fields will be named after the "__"-seperated access paths,
+ such as `'image__url'` for `image.url`.
+
+ :return: `pandas.DataFrame`
+ """
+ if TYPE_CHECKING:
+ import pandas as pd
+ else:
+ pd = import_library('pandas', raise_error=True)
+
+ if self.doc_type == AnyDoc:
+ raise TypeError(
+ 'DocList must be homogeneous to be converted to a DataFrame.'
+ 'There is no document schema defined. '
+ 'Please specify the DocList\'s Document type using `DocList[MyDoc]`.'
+ )
+
+ fields = self.doc_type._get_access_paths()
+ df = pd.DataFrame(columns=fields)
+
+ for doc in self:
+ doc_dict = _dict_to_access_paths(doc.dict())
+ doc_dict = {k: [v] for k, v in doc_dict.items()}
+ df = pd.concat([df, pd.DataFrame.from_dict(doc_dict)], ignore_index=True)
+
+ return df
+
+ # Methods to load from/to files in different formats
+ @property
+ def _stream_header(self) -> bytes:
+ # Binary format for streaming case
+
+ # V2 DocList streaming serialization format
+ # | 1 byte | 8 bytes | 4 bytes | variable(DocArray >=0.30) | 4 bytes | variable(DocArray >=0.30) ...
+
+ # 1 byte (uint8)
+ version_byte = b'\x02'
+ # 8 bytes (uint64)
+ num_docs_as_bytes = len(self).to_bytes(8, 'big', signed=False)
+ return version_byte + num_docs_as_bytes
+
+ @classmethod
+ @abstractmethod
+ def _get_proto_class(cls: Type[T]):
+ ...
+
+ @classmethod
+ def _load_binary_all(
+ cls: Type[T],
+ file_ctx: Union[ContextManager[io.BufferedReader], ContextManager[bytes]],
+ protocol: Optional[ProtocolType],
+ compress: Optional[str],
+ show_progress: bool,
+ tensor_type: Optional[Type['AbstractTensor']] = None,
+ ):
+ """Read a `DocList` object from a binary file
+ :param protocol: protocol to use. It can be 'pickle-array', 'protobuf-array', 'pickle' or 'protobuf'
+ :param compress: compress algorithm to use between `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
+ :param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
+ :param tensor_type: only relevant for DocVec; tensor_type of the DocVec
+ :return: a `DocList`
+ """
+ with file_ctx as fp:
+ if isinstance(fp, bytes):
+ d = fp
+ else:
+ d = fp.read()
+
+ if protocol is not None and protocol in (
+ 'pickle-array',
+ 'protobuf-array',
+ 'json-array',
+ ):
+ if _get_compress_ctx(algorithm=compress) is not None:
+ d = _decompress_bytes(d, algorithm=compress)
+ compress = None
+
+ if protocol is not None and protocol == 'protobuf-array':
+ proto = cls._get_proto_class()()
+ proto.ParseFromString(d)
+
+ if tensor_type is not None:
+ cls_ = cast('IOMixinDocVec', cls)
+ return cls_.from_protobuf(proto, tensor_type=tensor_type)
+ else:
+ return cls.from_protobuf(proto)
+ elif protocol is not None and protocol == 'pickle-array':
+ return pickle.loads(d)
+
+ elif protocol is not None and protocol == 'json-array':
+ if tensor_type is not None:
+ cls_ = cast('IOMixinDocVec', cls)
+ return cls_.from_json(d, tensor_type=tensor_type)
+ else:
+ return cls.from_json(d)
+
+ # Binary format for streaming case
+ else:
+ from rich import filesize
+
+ from docarray.utils._internal.progress_bar import _get_progressbar
+
+ # 1 byte (uint8)
+ version_num = int.from_bytes(d[0:1], 'big', signed=False)
+ if version_num != 2:
+ raise ValueError(
+ f'Unsupported version number {version_num} in binary format, expected 2'
+ )
+
+ # 8 bytes (uint64)
+ num_docs = int.from_bytes(d[1:9], 'big', signed=False)
+
+ pbar, t = _get_progressbar(
+ 'Deserializing', disable=not show_progress, total=num_docs
+ )
+
+ # this 9 is version + num_docs bytes used
+ start_pos = 9
+ docs = []
+ with pbar:
+ _total_size = 0
+ pbar.start_task(t)
+
+ for _ in range(num_docs):
+ # 4 bytes (uint32)
+ len_current_doc_in_bytes = int.from_bytes(
+ d[start_pos : start_pos + 4], 'big', signed=False
+ )
+ start_doc_pos = start_pos + 4
+ end_doc_pos = start_doc_pos + len_current_doc_in_bytes
+ start_pos = end_doc_pos
+
+ # variable length bytes doc
+ load_protocol: ProtocolType = protocol or cast(
+ ProtocolType, 'protobuf'
+ )
+ doc = cls.doc_type.from_bytes(
+ d[start_doc_pos:end_doc_pos],
+ protocol=load_protocol,
+ compress=compress,
+ )
+ docs.append(doc)
+ _total_size += len_current_doc_in_bytes
+ pbar.update(
+ t, advance=1, total_size=str(filesize.decimal(_total_size))
+ )
+ if tensor_type is not None:
+ cls__ = cast(Type['DocVec'], cls)
+ # mypy doesn't realize that cls_ is callable
+ return cls__(docs, tensor_type=tensor_type) # type: ignore
+ return cls(docs)
+
+ @classmethod
+ def _load_binary_stream(
+ cls: Type[T],
+ file_ctx: ContextManager[io.BufferedReader],
+ protocol: ProtocolType = 'protobuf',
+ compress: Optional[str] = None,
+ show_progress: bool = False,
+ ) -> Generator['T_doc', None, None]:
+ """Yield `Document` objects from a binary file
+
+ :param protocol: protocol to use. It can be 'pickle' or 'protobuf'
+ :param compress: compress algorithm to use between `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
+ :param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
+ :return: a generator of `Document` objects
+ """
+
+ from rich import filesize
+
+ with file_ctx as f:
+ version_numdocs_lendoc0 = f.read(9)
+ # 1 byte (uint8)
+ version_num = int.from_bytes(
+ version_numdocs_lendoc0[0:1], 'big', signed=False
+ )
+ if version_num != 2:
+ raise ValueError(
+ f'Unsupported version number {version_num} in binary format, expected 2'
+ )
+
+ # 8 bytes (uint64)
+ num_docs = int.from_bytes(version_numdocs_lendoc0[1:9], 'big', signed=False)
+
+ if show_progress:
+ from docarray.utils._internal.progress_bar import _get_progressbar
+
+ pbar, t = _get_progressbar(
+ 'Deserializing', disable=not show_progress, total=num_docs
+ )
+ else:
+ from contextlib import nullcontext
+
+ pbar = nullcontext()
+
+ with pbar:
+ if show_progress:
+ _total_size = 0
+ pbar.start_task(t)
+ for _ in range(num_docs):
+ # 4 bytes (uint32)
+ len_current_doc_in_bytes = int.from_bytes(
+ f.read(4), 'big', signed=False
+ )
+ load_protocol: ProtocolType = protocol
+ yield cls.doc_type.from_bytes(
+ f.read(len_current_doc_in_bytes),
+ protocol=load_protocol,
+ compress=compress,
+ )
+ if show_progress:
+ _total_size += len_current_doc_in_bytes
+ pbar.update(
+ t, advance=1, total_size=str(filesize.decimal(_total_size))
+ )
+
+ @staticmethod
+ def _get_file_context(
+ file: Union[str, bytes, pathlib.Path, io.BufferedReader, _LazyRequestReader],
+ protocol: ProtocolType,
+ compress: Optional[str] = None,
+ ) -> Tuple[
+ Union[nullcontext, io.BufferedReader], Optional[ProtocolType], Optional[str]
+ ]:
+ load_protocol: Optional[ProtocolType] = protocol
+ load_compress: Optional[str] = compress
+ file_ctx: Union[nullcontext, io.BufferedReader]
+ if isinstance(file, (io.BufferedReader, _LazyRequestReader, bytes)):
+ file_ctx = nullcontext(file)
+ # by checking path existence we allow file to be of type Path, LocalPath, PurePath and str
+ elif isinstance(file, (str, pathlib.Path)) and os.path.exists(file):
+ load_protocol, load_compress = _protocol_and_compress_from_file_path(
+ file, protocol, compress
+ )
+ file_ctx = open(file, 'rb')
+ else:
+ raise FileNotFoundError(f'cannot find file {file}')
+ return file_ctx, load_protocol, load_compress
+
+ @classmethod
+ def load_binary(
+ cls: Type[T],
+ file: Union[str, bytes, pathlib.Path, io.BufferedReader, _LazyRequestReader],
+ protocol: ProtocolType = 'protobuf-array',
+ compress: Optional[str] = None,
+ show_progress: bool = False,
+ streaming: bool = False,
+ ) -> Union[T, Generator['T_doc', None, None]]:
+ """Load doc_list elements from a compressed binary file.
+
+ In case protocol is pickle the `Documents` are streamed from disk to save memory usage
+
+ !!! note
+ If `file` is `str` it can specify `protocol` and `compress` as file extensions.
+ This functionality assumes `file=file_name.$protocol.$compress` where `$protocol` and `$compress` refer to a
+ string interpolation of the respective `protocol` and `compress` methods.
+ For example if `file=my_docarray.protobuf.lz4` then the binary data will be loaded assuming `protocol=protobuf`
+ and `compress=lz4`.
+
+ :param file: File or filename or serialized bytes where the data is stored.
+ :param protocol: protocol to use. It can be 'pickle-array', 'protobuf-array', 'pickle' or 'protobuf'
+ :param compress: compress algorithm to use between `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
+ :param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
+ :param streaming: if `True` returns a generator over `Document` objects.
+
+ :return: a `DocList` object
+
+ """
+ file_ctx, load_protocol, load_compress = cls._get_file_context(
+ file, protocol, compress
+ )
+ if streaming:
+ if load_protocol not in SINGLE_PROTOCOLS:
+ raise ValueError(
+ f'`streaming` is only available when using {" or ".join(map(lambda x: f"`{x}`", SINGLE_PROTOCOLS))} as protocol, '
+ f'got {load_protocol}'
+ )
+ else:
+ return cls._load_binary_stream(
+ file_ctx,
+ protocol=load_protocol,
+ compress=load_compress,
+ show_progress=show_progress,
+ )
+ else:
+ return cls._load_binary_all(
+ file_ctx, load_protocol, load_compress, show_progress
+ )
+
+ def save_binary(
+ self,
+ file: Union[str, pathlib.Path],
+ protocol: ProtocolType = 'protobuf-array',
+ compress: Optional[str] = None,
+ show_progress: bool = False,
+ ) -> None:
+ """Save DocList into a binary file.
+
+ It will use the protocol to pick how to save the DocList.
+ If used `picke-doc_list` and `protobuf-array` the DocList will be stored
+ and compressed at complete level using `pickle` or `protobuf`.
+ When using `protobuf` or `pickle` as protocol each Document in DocList
+ will be stored individually and this would make it available for streaming.
+
+ !!! note
+ If `file` is `str` it can specify `protocol` and `compress` as file extensions.
+ This functionality assumes `file=file_name.$protocol.$compress` where `$protocol` and `$compress` refer to a
+ string interpolation of the respective `protocol` and `compress` methods.
+ For example if `file=my_docarray.protobuf.lz4` then the binary data will be created using `protocol=protobuf`
+ and `compress=lz4`.
+
+ :param file: File or filename to which the data is saved.
+ :param protocol: protocol to use. It can be 'pickle-array', 'protobuf-array', 'pickle' or 'protobuf'
+ :param compress: compress algorithm to use between `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
+ :param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
+ """
+ if isinstance(file, io.BufferedWriter):
+ file_ctx = nullcontext(file)
+ else:
+ _protocol, _compress = _protocol_and_compress_from_file_path(file)
+
+ if _protocol is not None:
+ protocol = _protocol
+ if _compress is not None:
+ compress = _compress
+
+ file_ctx = open(file, 'wb')
+
+ self.to_bytes(
+ protocol=protocol,
+ compress=compress,
+ file_ctx=file_ctx,
+ show_progress=show_progress,
+ )
diff --git a/docarray/array/doc_list/pushpull.py b/docarray/array/doc_list/pushpull.py
new file mode 100644
index 00000000000..5784610633b
--- /dev/null
+++ b/docarray/array/doc_list/pushpull.py
@@ -0,0 +1,166 @@
+import logging
+from abc import abstractmethod
+from typing import (
+ TYPE_CHECKING,
+ Dict,
+ Iterable,
+ Iterator,
+ Tuple,
+ Type,
+ TypeVar,
+ cast,
+)
+
+from typing_extensions import Literal
+from typing_inspect import get_args
+
+PUSH_PULL_PROTOCOL = Literal['s3', 'file']
+SUPPORTED_PUSH_PULL_PROTOCOLS = get_args(PUSH_PULL_PROTOCOL)
+
+if TYPE_CHECKING: # pragma: no cover
+ from docarray import BaseDoc, DocList
+ from docarray.store.abstract_doc_store import AbstractDocStore
+
+
+SelfPushPullMixin = TypeVar('SelfPushPullMixin', bound='PushPullMixin')
+
+
+class PushPullMixin(Iterable['BaseDoc']):
+ """Mixin class for push/pull functionality."""
+
+ __backends__: Dict[str, Type['AbstractDocStore']] = {}
+ doc_type: Type['BaseDoc']
+
+ @abstractmethod
+ def __len__(self) -> int:
+ ...
+
+ @staticmethod
+ def resolve_url(url: str) -> Tuple[PUSH_PULL_PROTOCOL, str]:
+ """Resolve the URL to the correct protocol and name.
+ :param url: url to resolve
+ """
+ protocol, name = url.split('://', 2)
+ if protocol in SUPPORTED_PUSH_PULL_PROTOCOLS:
+ protocol = cast(PUSH_PULL_PROTOCOL, protocol)
+ return protocol, name
+ else:
+ raise ValueError(f'Unsupported protocol {protocol}')
+
+ @classmethod
+ def get_pushpull_backend(
+ cls: Type[SelfPushPullMixin], protocol: PUSH_PULL_PROTOCOL
+ ) -> Type['AbstractDocStore']:
+ """
+ Get the backend for the given protocol.
+
+ :param protocol: the protocol to use, e.g. 'file', 's3'
+ :return: the backend class
+ """
+ if protocol in cls.__backends__:
+ return cls.__backends__[protocol]
+
+ if protocol == 'file':
+ from docarray.store.file import FileDocStore
+
+ cls.__backends__[protocol] = FileDocStore
+ logging.debug('Loaded Local Filesystem backend')
+ elif protocol == 's3':
+ from docarray.store.s3 import S3DocStore
+
+ cls.__backends__[protocol] = S3DocStore
+ logging.debug('Loaded S3 backend')
+ else:
+ raise NotImplementedError(f'protocol {protocol} not supported')
+
+ return cls.__backends__[protocol]
+
+ def push(
+ self,
+ url: str,
+ show_progress: bool = False,
+ **kwargs,
+ ) -> Dict:
+ """Push this `DocList` object to the specified url.
+
+ :param url: url specifying the protocol and save name of the `DocList`. Should be of the form ``protocol://namespace/name``. e.g. ``s3://bucket/path/to/namespace/name``, ``file:///path/to/folder/name``
+ :param show_progress: If true, a progress bar will be displayed.
+ """
+ logging.info(f'Pushing {len(self)} docs to {url}')
+ protocol, name = self.__class__.resolve_url(url)
+ return self.__class__.get_pushpull_backend(protocol).push(
+ self, name, show_progress # type: ignore
+ )
+
+ @classmethod
+ def push_stream(
+ cls: Type[SelfPushPullMixin],
+ docs: Iterator['BaseDoc'],
+ url: str,
+ show_progress: bool = False,
+ ) -> Dict:
+ """Push a stream of documents to the specified url.
+
+ :param docs: a stream of documents
+ :param url: url specifying the protocol and save name of the `DocList`. Should be of the form ``protocol://namespace/name``. e.g. ``s3://bucket/path/to/namespace/name``, ``file:///path/to/folder/name``
+ :param show_progress: If true, a progress bar will be displayed.
+ """
+ logging.info(f'Pushing stream to {url}')
+ protocol, name = cls.resolve_url(url)
+ return cls.get_pushpull_backend(protocol).push_stream(docs, name, show_progress)
+
+ @classmethod
+ def pull(
+ cls: Type[SelfPushPullMixin],
+ url: str,
+ show_progress: bool = False,
+ local_cache: bool = True,
+ ) -> 'DocList':
+ """Pull a `DocList` from the specified url.
+
+ :param url: url specifying the protocol and save name of the `DocList`. Should be of the form ``protocol://namespace/name``. e.g. ``s3://bucket/path/to/namespace/name``, ``file:///path/to/folder/name``
+ :param show_progress: if true, display a progress bar.
+ :param local_cache: store the downloaded `DocList` to local folder
+ :return: a `DocList` object
+ """
+ from docarray.base_doc import AnyDoc
+
+ if cls.doc_type == AnyDoc:
+ raise TypeError(
+ 'There is no document schema defined. '
+ 'Please specify the `DocList`\'s Document type using `DocList[MyDoc]`.'
+ )
+
+ logging.info(f'Pulling {url}')
+ protocol, name = cls.resolve_url(url)
+ return cls.get_pushpull_backend(protocol).pull(
+ cls, name, show_progress, local_cache # type: ignore
+ )
+
+ @classmethod
+ def pull_stream(
+ cls: Type[SelfPushPullMixin],
+ url: str,
+ show_progress: bool = False,
+ local_cache: bool = False,
+ ) -> Iterator['BaseDoc']:
+ """Pull a stream of Documents from the specified url.
+
+ :param url: url specifying the protocol and save name of the `DocList`. Should be of the form ``protocol://namespace/name``. e.g. ``s3://bucket/path/to/namespace/name``, ``file:///path/to/folder/name``
+ :param show_progress: if true, display a progress bar.
+ :param local_cache: store the downloaded `DocList` to local folder
+ :return: Iterator of Documents
+ """
+ from docarray.base_doc import AnyDoc
+
+ if cls.doc_type == AnyDoc:
+ raise TypeError(
+ 'There is no document schema defined. '
+ 'Please specify the `DocList`\'s Document type using `DocList[MyDoc]`.'
+ )
+
+ logging.info(f'Pulling Document stream from {url}')
+ protocol, name = cls.resolve_url(url)
+ return cls.get_pushpull_backend(protocol).pull_stream(
+ cls, name, show_progress, local_cache # type: ignore
+ )
diff --git a/docarray/array/doc_vec/__init__.py b/docarray/array/doc_vec/__init__.py
new file mode 100644
index 00000000000..74f8f7582cd
--- /dev/null
+++ b/docarray/array/doc_vec/__init__.py
@@ -0,0 +1,15 @@
+# Licensed to the LF AI & Data foundation under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/docarray/array/doc_vec/column_storage.py b/docarray/array/doc_vec/column_storage.py
new file mode 100644
index 00000000000..ea5da4291fa
--- /dev/null
+++ b/docarray/array/doc_vec/column_storage.py
@@ -0,0 +1,225 @@
+from collections import ChainMap
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Dict,
+ ItemsView,
+ Iterable,
+ MutableMapping,
+ NamedTuple,
+ Optional,
+ Type,
+ TypeVar,
+ Union,
+ ValuesView,
+)
+
+from docarray.array.list_advance_indexing import ListAdvancedIndexing
+from docarray.typing import NdArray
+from docarray.typing.tensor.abstract_tensor import AbstractTensor
+
+if TYPE_CHECKING:
+ from docarray.array.doc_vec.doc_vec import DocVec
+
+IndexIterType = Union[slice, Iterable[int], Iterable[bool], None]
+
+
+T = TypeVar('T', bound='ColumnStorage')
+
+
+class ColumnsJsonCompatible(NamedTuple):
+ tensor_columns: Dict[str, Any]
+ doc_columns: Dict[str, Any]
+ docs_vec_columns: Dict[str, Any]
+ any_columns: Dict[str, Any]
+
+
+class ColumnStorage:
+ """
+ ColumnStorage is a container to store the columns of the
+ :class:`~docarray.array.doc_vec.DocVec`.
+
+ :param tensor_columns: a Dict of AbstractTensor
+ :param doc_columns: a Dict of :class:`~docarray.array.doc_vec.DocVec`
+ :param docs_vec_columns: a Dict of List of :class:`~docarray.array.doc_vec.DocVec`
+ :param any_columns: a Dict of List
+ :param tensor_type: Class used to wrap the doc_vec tensors
+ """
+
+ def __init__(
+ self,
+ tensor_columns: Dict[str, Optional[AbstractTensor]],
+ doc_columns: Dict[str, Optional['DocVec']],
+ docs_vec_columns: Dict[str, Optional[ListAdvancedIndexing['DocVec']]],
+ any_columns: Dict[str, ListAdvancedIndexing],
+ tensor_type: Type[AbstractTensor] = NdArray,
+ ):
+ self.tensor_columns = tensor_columns
+ self.doc_columns = doc_columns
+ self.docs_vec_columns = docs_vec_columns
+ self.any_columns = any_columns
+
+ self.tensor_type = tensor_type
+
+ self.columns = ChainMap( # type: ignore
+ self.tensor_columns, # type: ignore
+ self.doc_columns, # type: ignore
+ self.docs_vec_columns, # type: ignore
+ self.any_columns, # type: ignore
+ ) # type: ignore
+
+ def __len__(self) -> int:
+ return len(self.any_columns['id']) # TODO what if ID are None ?
+
+ def __getitem__(self: T, item: IndexIterType) -> T:
+ if isinstance(item, tuple):
+ item = list(item)
+ tensor_columns = {
+ key: col[item] if col is not None else None
+ for key, col in self.tensor_columns.items()
+ }
+ doc_columns = {
+ key: col[item] if col is not None else None
+ for key, col in self.doc_columns.items()
+ }
+ docs_vec_columns = {
+ key: col[item] if col is not None else None
+ for key, col in self.docs_vec_columns.items()
+ }
+ any_columns = {
+ key: col[item] if col is not None else None
+ for key, col in self.any_columns.items()
+ }
+
+ return self.__class__(
+ tensor_columns,
+ doc_columns,
+ docs_vec_columns,
+ any_columns,
+ self.tensor_type,
+ )
+
+ def columns_json_compatible(self) -> ColumnsJsonCompatible:
+ tens_cols = {
+ key: value._docarray_to_json_compatible() if value is not None else value
+ for key, value in self.tensor_columns.items()
+ }
+ doc_cols = {
+ key: value._docarray_to_json_compatible() if value is not None else value
+ for key, value in self.doc_columns.items()
+ }
+ doc_vec_cols = {
+ key: [vec._docarray_to_json_compatible() for vec in value]
+ if value is not None
+ else value
+ for key, value in self.docs_vec_columns.items()
+ }
+ return ColumnsJsonCompatible(
+ tens_cols, doc_cols, doc_vec_cols, self.any_columns
+ )
+
+ def __eq__(self, other: Any) -> bool:
+ if not isinstance(other, ColumnStorage):
+ return False
+ if self.tensor_type != other.tensor_type:
+ return False
+ for col_map_self, col_map_other in zip(self.columns.maps, other.columns.maps):
+ if col_map_self.keys() != col_map_other.keys():
+ return False
+ for key_self in col_map_self.keys():
+ if key_self == 'id':
+ continue
+
+ val1, val2 = col_map_self[key_self], col_map_other[key_self]
+ if isinstance(val1, AbstractTensor):
+ values_are_equal = val1.get_comp_backend().equal(val1, val2)
+ elif isinstance(val2, AbstractTensor):
+ values_are_equal = val2.get_comp_backend().equal(val1, val2)
+ else:
+ values_are_equal = val1 == val2
+ if not values_are_equal:
+ return False
+ return True
+
+
+class ColumnStorageView(dict, MutableMapping[str, Any]):
+ index: int
+ storage: ColumnStorage
+
+ def __init__(self, index: int, storage: ColumnStorage):
+ super().__init__()
+ self.index = index
+ self.storage = storage
+
+ def __getitem__(self, name: str) -> Any:
+ if name in self.storage.tensor_columns.keys():
+ tensor = self.storage.tensor_columns[name]
+ if tensor is None:
+ return None
+ if tensor.get_comp_backend().n_dim(tensor) == 1:
+ # to ensure consistensy between numpy and pytorch
+ # we wrap the scalr in a tensor of ndim = 1
+ # otherwise numpy pass by value whereas torch by reference
+ col = self.storage.tensor_columns[name]
+
+ if col is not None:
+ return col[self.index : self.index + 1]
+ else:
+ return None
+
+ col = self.storage.columns[name]
+
+ if col is None:
+ return None
+ return col[self.index]
+
+ def __reduce__(self):
+ # implementing __reduce__ to solve a pickle issue when subclassing dict
+ # see here: https://stackoverflow.com/questions/21144845/how-can-i-unpickle-a-subclass-of-dict-that-validates-with-setitem-in-pytho
+ return (ColumnStorageView, (self.index, self.storage))
+
+ def __setitem__(self, name, value) -> None:
+ if self.storage.columns[name] is None:
+ raise ValueError(
+ f'Cannot set an item to a None column. This mean that '
+ f'the DocVec that encapsulate this doc has the field '
+ f'{name} set to None. If you want to modify that you need to do it at the'
+ f'DocVec level. `docs.field = np.zeros(10)`'
+ )
+
+ self.storage.columns[name][self.index] = value
+
+ def __delitem__(self, key):
+ raise RuntimeError('Cannot delete an item from a StorageView')
+
+ def __iter__(self):
+ return self.storage.columns.keys()
+
+ def __len__(self):
+ return len(self.storage.columns)
+
+ def _local_dict(self):
+ """The storage.columns dictionary with every value at position self.index"""
+
+ return {key: self[key] for key in self.storage.columns.keys()}
+
+ def keys(self):
+ return self.storage.columns.keys()
+
+ # type ignore because return type dict_values is private and we cannot use it.
+ # context: https://github.com/python/typing/discussions/1033
+ def values(self) -> ValuesView: # type: ignore
+ return ValuesView(self._local_dict())
+
+ # type ignore because return type dict_items is private and we cannot use it.
+ # context: https://github.com/python/typing/discussions/1033
+ def items(self) -> ItemsView: # type: ignore
+ return ItemsView(self._local_dict())
+
+ def to_dict(self) -> Dict[str, Any]:
+ """
+ Return a dictionary with the same keys as the storage.columns
+ and the values at position self.index.
+ Warning: modification on the dict will not be reflected on the storage.
+ """
+ return {key: self[key] for key in self.storage.columns.keys()}
diff --git a/docarray/array/doc_vec/doc_vec.py b/docarray/array/doc_vec/doc_vec.py
new file mode 100644
index 00000000000..0cc462f173d
--- /dev/null
+++ b/docarray/array/doc_vec/doc_vec.py
@@ -0,0 +1,675 @@
+from collections import ChainMap
+from typing import (
+ Any,
+ Dict,
+ Iterable,
+ List,
+ MutableSequence,
+ Optional,
+ Sequence,
+ Tuple,
+ Type,
+ TypeVar,
+ Union,
+ cast,
+ no_type_check,
+ overload,
+)
+
+from pydantic import parse_obj_as
+from typing_inspect import typingGenericAlias
+
+from docarray.array.any_array import AnyDocArray
+from docarray.array.doc_list.doc_list import DocList
+from docarray.array.doc_vec.column_storage import ColumnStorage, ColumnStorageView
+from docarray.array.doc_vec.io import IOMixinDocVec
+from docarray.array.list_advance_indexing import ListAdvancedIndexing
+from docarray.base_doc import AnyDoc, BaseDoc
+from docarray.typing import NdArray
+from docarray.typing.tensor.abstract_tensor import AbstractTensor
+from docarray.utils._internal.pydantic import is_pydantic_v2
+
+if is_pydantic_v2:
+ from pydantic import GetCoreSchemaHandler
+ from pydantic_core import core_schema
+
+from docarray.utils._internal._typing import is_tensor_union, safe_issubclass
+from docarray.utils._internal.misc import (
+ is_jax_available,
+ is_tf_available,
+ is_torch_available,
+)
+
+torch_available = is_torch_available()
+if torch_available:
+ from docarray.typing import TorchTensor
+else:
+ TorchTensor = None # type: ignore
+
+tf_available = is_tf_available()
+if tf_available:
+ import tensorflow as tf # type: ignore
+
+ from docarray.typing import TensorFlowTensor # noqa: F401
+else:
+ TensorFlowTensor = None # type: ignore
+
+jnp_available = is_jax_available()
+if jnp_available:
+ import jax.numpy as jnp # type: ignore
+
+ from docarray.typing import JaxArray # noqa: F401
+else:
+ JaxArray = None # type: ignore
+
+T_doc = TypeVar('T_doc', bound=BaseDoc)
+T = TypeVar('T', bound='DocVec')
+T_io_mixin = TypeVar('T_io_mixin', bound='IOMixinDocVec')
+
+IndexIterType = Union[slice, Iterable[int], Iterable[bool], None]
+
+
+# type ignore because from_protobuf has a different signature
+class DocVec(IOMixinDocVec, AnyDocArray[T_doc]): # type: ignore
+ """
+ DocVec is a container of Documents appropriates to perform
+ computation that require batches of data (ex: matrix multiplication, distance
+ calculation, deep learning forward pass)
+
+ A DocVec has a similar interface as [`DocList`][docarray.array.DocList]
+ but with an underlying implementation that is column based instead of row based.
+ Each field of the schema of the `DocVec` (the `.doc_type` which is a
+ [`BaseDoc`][docarray.BaseDoc]) will be stored in a column.
+
+ If the field is a tensor, the data from all Documents will be stored as a single
+ (torch/np/tf) tensor.
+
+ If the tensor field is `AnyTensor` or a Union of tensor types, the
+ `.tensor_type` will be used to determine the type of the column.
+
+ If the field is another [`BaseDoc`][docarray.BaseDoc] the column will be another
+ `DocVec` that follows the schema of the nested Document.
+
+ If the field is a [`DocList`][docarray.DocList] or `DocVec` then the column will
+ be a list of `DocVec`.
+
+ For any other type the column is a Python list.
+
+ Every `Document` inside a `DocVec` is a view into the data columns stored at the
+ `DocVec` level. The `BaseDoc` does not hold any data itself. The behavior of
+ this Document "view" is similar to the behavior of `view = tensor[i]` in
+ numpy/PyTorch.
+
+ !!! note
+ DocVec supports optional fields. Nevertheless if a field is optional it needs to
+ be homogeneous. This means that if the first document has a None value all of the
+ other documents should have a None value as well.
+ !!! note
+ If one field is Optional the column will be stored
+ * as None if the first doc is as the field as None
+ * as a normal column otherwise that cannot contain None value
+
+ :param docs: a homogeneous sequence of `BaseDoc`
+ :param tensor_type: Tensor Class used to wrap the doc_vec tensors. This is useful
+ if the BaseDoc of this DocVec has some undefined tensor type like
+ AnyTensor or Union of NdArray and TorchTensor
+ """
+
+ doc_type: Type[T_doc] = BaseDoc # type: ignore
+
+ def __init__(
+ self: T,
+ docs: Sequence[T_doc],
+ tensor_type: Type['AbstractTensor'] = NdArray,
+ ):
+
+ if (
+ not hasattr(self, 'doc_type')
+ or self.doc_type == AnyDoc
+ or self.doc_type == BaseDoc
+ ):
+ raise TypeError(
+ f'{self.__class__.__name__} does not precise a doc_type. You probably should do'
+ f'docs = DocVec[MyDoc](docs) instead of DocVec(docs)'
+ )
+ self.tensor_type = tensor_type
+ self._is_unusable = False
+
+ tensor_columns: Dict[str, Optional[AbstractTensor]] = dict()
+ doc_columns: Dict[str, Optional['DocVec']] = dict()
+ docs_vec_columns: Dict[str, Optional[ListAdvancedIndexing['DocVec']]] = dict()
+ any_columns: Dict[str, ListAdvancedIndexing] = dict()
+
+ if len(docs) == 0:
+ raise ValueError(f'docs {docs}: should not be empty')
+ docs = (
+ docs
+ if isinstance(docs, DocList)
+ else DocList.__class_getitem__(self.doc_type)(docs)
+ )
+
+ for field_name, field in self.doc_type._docarray_fields().items():
+ # here we iterate over the field of the docs schema, and we collect the data
+ # from each document and put them in the corresponding column
+ field_type: Type = self.doc_type._get_field_annotation(field_name)
+
+ field_info = self.doc_type._docarray_fields()[field_name]
+ is_field_required = (
+ field_info.is_required() if is_pydantic_v2 else field_info.required
+ )
+
+ first_doc_is_none = getattr(docs[0], field_name) is None
+
+ def _verify_optional_field_of_docs(docs):
+
+ if is_field_required:
+ if first_doc_is_none:
+ raise ValueError(
+ f'Field {field_name} is None for {docs[0]} even though it is required'
+ )
+
+ if first_doc_is_none:
+ for i, doc in enumerate(docs):
+ if getattr(doc, field_name) is not None:
+ raise ValueError(
+ f'Field {field_name} is put to None for the first doc. This mean that '
+ f'all of the other docs should have this field set to None as well. '
+ f'This is not the case for {doc} at index {i}'
+ )
+
+ def _check_doc_field_not_none(field_name, doc):
+ if getattr(doc, field_name) is None:
+ raise ValueError(
+ f'Field {field_name} is None for {doc} even though it is not None for the first doc'
+ )
+
+ if is_tensor_union(field_type):
+ field_type = tensor_type
+ # all generic tensor types such as AnyTensor, ImageTensor, etc. are subclasses of AbstractTensor.
+ # Perform check only if the field_type is not an alias and is a subclass of AbstractTensor
+ elif not isinstance(field_type, typingGenericAlias) and safe_issubclass(
+ field_type, AbstractTensor
+ ):
+ # check if the tensor associated with the field_name in the document is a subclass of the tensor_type
+ # e.g. if the field_type is AnyTensor but the type(docs[0][field_name]) is ImageTensor,
+ # then we change the field_type to ImageTensor, since AnyTensor is a union of all the tensor types
+ # and does not override any methods of specific tensor types
+ tensor = getattr(docs[0], field_name)
+ if safe_issubclass(tensor.__class__, tensor_type):
+ field_type = tensor_type
+
+ if isinstance(field_type, type) or safe_issubclass(field_type, AnyDocArray):
+ if tf_available and safe_issubclass(field_type, TensorFlowTensor):
+ # tf.Tensor does not allow item assignment, therefore the
+ # optimized way
+ # of initializing an empty array and assigning values to it
+ # iteratively
+ # does not work here, therefore handle separately.
+
+ if first_doc_is_none:
+ _verify_optional_field_of_docs(docs)
+ tensor_columns[field_name] = None
+ else:
+ tf_stack = []
+ for i, doc in enumerate(docs):
+ val = getattr(doc, field_name)
+ _check_doc_field_not_none(field_name, doc)
+ tf_stack.append(val.tensor)
+
+ stacked: tf.Tensor = tf.stack(tf_stack)
+ tensor_columns[field_name] = TensorFlowTensor(stacked)
+ elif jnp_available and safe_issubclass(field_type, JaxArray):
+ if first_doc_is_none:
+ _verify_optional_field_of_docs(docs)
+ tensor_columns[field_name] = None
+ else:
+ tf_stack = []
+ for i, doc in enumerate(docs):
+ val = getattr(doc, field_name)
+ _check_doc_field_not_none(field_name, doc)
+ tf_stack.append(val.tensor)
+
+ jax_stacked: jnp.ndarray = jnp.stack(tf_stack)
+ tensor_columns[field_name] = JaxArray(jax_stacked)
+
+ elif safe_issubclass(field_type, AbstractTensor):
+ if first_doc_is_none:
+ _verify_optional_field_of_docs(docs)
+ tensor_columns[field_name] = None
+ else:
+ tensor = getattr(docs[0], field_name)
+ column_shape = (
+ (len(docs), *tensor.shape)
+ if tensor is not None
+ else (len(docs),)
+ )
+ tensor_columns[field_name] = field_type._docarray_from_native(
+ field_type.get_comp_backend().empty(
+ column_shape,
+ dtype=tensor.dtype
+ if hasattr(tensor, 'dtype')
+ else None,
+ device=tensor.device
+ if hasattr(tensor, 'device')
+ else None,
+ )
+ )
+
+ for i, doc in enumerate(docs):
+ _check_doc_field_not_none(field_name, doc)
+ val = getattr(doc, field_name)
+ cast(AbstractTensor, tensor_columns[field_name])[i] = val
+
+ elif safe_issubclass(field_type, BaseDoc):
+ if first_doc_is_none:
+ _verify_optional_field_of_docs(docs)
+ doc_columns[field_name] = None
+ else:
+ if is_field_required:
+ doc_columns[field_name] = getattr(
+ docs, field_name
+ ).to_doc_vec(tensor_type=self.tensor_type)
+ else:
+ doc_columns[field_name] = DocList.__class_getitem__(
+ field_type
+ )(getattr(docs, field_name)).to_doc_vec(
+ tensor_type=self.tensor_type
+ )
+
+ elif safe_issubclass(field_type, AnyDocArray):
+ if first_doc_is_none:
+ _verify_optional_field_of_docs(docs)
+ docs_vec_columns[field_name] = None
+ else:
+ docs_list = list()
+ for doc in docs:
+ docs_nested = getattr(doc, field_name)
+ _check_doc_field_not_none(field_name, doc)
+ if isinstance(docs_nested, DocList):
+ docs_nested = docs_nested.to_doc_vec(
+ tensor_type=self.tensor_type
+ )
+ docs_list.append(docs_nested)
+ docs_vec_columns[field_name] = ListAdvancedIndexing(docs_list)
+ else:
+ any_columns[field_name] = ListAdvancedIndexing(
+ getattr(docs, field_name)
+ )
+ else:
+ any_columns[field_name] = ListAdvancedIndexing(
+ getattr(docs, field_name)
+ )
+
+ self._storage = ColumnStorage(
+ tensor_columns,
+ doc_columns,
+ docs_vec_columns,
+ any_columns,
+ tensor_type,
+ )
+
+ @classmethod
+ def from_columns_storage(cls: Type[T], storage: ColumnStorage) -> T:
+ """
+ Create a DocVec directly from a storage object
+ :param storage: the underlying storage.
+ :return: a DocVec
+ """
+ docs = cls.__new__(cls)
+ docs.tensor_type = storage.tensor_type
+ docs._storage = storage
+ return docs
+
+ @classmethod
+ def _docarray_validate(
+ cls: Type[T],
+ value: Union[T, Iterable[T_doc]],
+ ) -> T:
+ if isinstance(value, cls):
+ return value
+ elif isinstance(value, DocList):
+ if (
+ safe_issubclass(value.doc_type, cls.doc_type)
+ or value.doc_type == cls.doc_type
+ ):
+ return cast(T, value.to_doc_vec())
+ else:
+ raise ValueError(f'DocVec[value.doc_type] is not compatible with {cls}')
+ elif not is_pydantic_v2 and isinstance(
+ value, DocList.__class_getitem__(cls.doc_type)
+ ):
+ return cast(T, value.to_doc_vec())
+ elif isinstance(value, Sequence):
+ return cls(value)
+ elif isinstance(value, Iterable):
+ return cls(list(value))
+ else:
+ raise TypeError(f'Expecting an Iterable of {cls.doc_type}')
+
+ def to(self: T, device: str) -> T:
+ """Move all tensors of this DocVec to the given device
+
+ :param device: the device to move the data to
+ """
+ for field, col_tens in self._storage.tensor_columns.items():
+ if col_tens is not None:
+ self._storage.tensor_columns[
+ field
+ ] = col_tens.get_comp_backend().to_device(col_tens, device)
+
+ for field, col_doc in self._storage.doc_columns.items():
+ if col_doc is not None:
+ self._storage.doc_columns[field] = col_doc.to(device)
+ for _, col_da in self._storage.docs_vec_columns.items():
+ if col_da is not None:
+ for docs in col_da:
+ docs.to(device)
+
+ return self
+
+ ################################################
+ # Accessing data : Indexing / Getitem related #
+ ################################################
+
+ @overload
+ def __getitem__(self: T, item: int) -> T_doc:
+ ...
+
+ @overload
+ def __getitem__(self: T, item: IndexIterType) -> T:
+ ...
+
+ def __getitem__(self: T, item: Union[int, IndexIterType]) -> Union[T_doc, T]:
+ if item is None:
+ return self # PyTorch behaviour
+ # multiple docs case
+ if isinstance(item, (slice, Iterable)):
+ return self.__class__.from_columns_storage(self._storage[item])
+ # single doc case
+ return self.doc_type.from_view(ColumnStorageView(item, self._storage))
+
+ def _get_data_column(
+ self: T,
+ field: str,
+ ) -> Union[MutableSequence, 'DocVec', AbstractTensor, None]:
+ """Return one column of the data
+
+ :param field: name of the fields to extract
+ :return: Returns a list of the field value for each document
+ in the array like container
+ """
+ if field in self._storage.any_columns.keys():
+ return self._storage.any_columns[field]
+ elif field in self._storage.docs_vec_columns.keys():
+ return self._storage.docs_vec_columns[field]
+ elif field in self._storage.columns.keys():
+ return self._storage.columns[field]
+ else:
+ raise ValueError(f'{field} does not exist in {self}')
+
+ ####################################
+ # Updating data : Setitem related #
+ ####################################
+
+ @overload
+ def __setitem__(self: T, key: int, value: T_doc):
+ ...
+
+ @overload
+ def __setitem__(self: T, key: IndexIterType, value: T):
+ ...
+
+ @no_type_check
+ def __setitem__(self: T, key, value):
+ # single doc case
+ if not isinstance(key, (slice, Iterable)):
+ if not isinstance(value, self.doc_type):
+ raise ValueError(f'{value} is not a {self.doc_type}')
+
+ for field, value in value.dict().items():
+ self._storage.columns[field][key] = value # todo we might want to
+ # define a safety mechanism in someone put a wrong value
+ else:
+ # multiple docs case
+ self._set_data_and_columns(key, value)
+
+ def _set_data_and_columns(
+ self: T,
+ index_item: Union[Tuple, Iterable, slice],
+ value: Union[T, DocList[T_doc]],
+ ) -> None:
+ """Delegates the setting to the data and the columns.
+
+ :param index_item: the key used as index. Needs to be a valid index for both
+ DocList (data) and column types (torch/tensorflow/numpy tensors)
+ :value: the value to set at the `key` location
+ """
+ if isinstance(index_item, tuple):
+ index_item = list(index_item)
+
+ # set data and prepare columns
+ processed_value: T
+ if isinstance(value, DocList):
+ if not safe_issubclass(value.doc_type, self.doc_type):
+ raise TypeError(
+ f'{value} schema : {value.doc_type} is not compatible with '
+ f'this DocVec schema : {self.doc_type}'
+ )
+ processed_value = cast(
+ T, value.to_doc_vec(tensor_type=self.tensor_type)
+ ) # we need to copy data here
+
+ elif isinstance(value, DocVec):
+ if not safe_issubclass(value.doc_type, self.doc_type):
+ raise TypeError(
+ f'{value} schema : {value.doc_type} is not compatible with '
+ f'this DocVec schema : {self.doc_type}'
+ )
+ processed_value = value
+ else:
+ raise TypeError(f'Can not set a DocVec with {type(value)}')
+
+ for field, col in self._storage.columns.items():
+ col[index_item] = processed_value._storage.columns[field]
+
+ def _set_data_column(
+ self: T,
+ field: str,
+ values: Union[
+ Sequence[DocList[T_doc]],
+ Sequence[Any],
+ T,
+ DocList,
+ AbstractTensor,
+ None,
+ ],
+ ) -> None:
+ """Set all Documents in this DocList using the passed values
+
+ :param field: name of the fields to set
+ :values: the values to set at the DocList level
+ """
+ if values is None:
+ if field in self._storage.tensor_columns.keys():
+ self._storage.tensor_columns[field] = values
+ elif field in self._storage.doc_columns.keys():
+ self._storage.doc_columns[field] = values
+ elif field in self._storage.docs_vec_columns.keys():
+ self._storage.docs_vec_columns[field] = values
+ elif field in self._storage.any_columns.keys():
+ raise ValueError(
+ f'column {field} cannot be set to None, try to pass '
+ f'a list of None instead'
+ )
+ else:
+ raise ValueError(f'{field} does not exist in {self}')
+
+ else:
+ if len(values) != len(self._storage):
+ raise ValueError(
+ f'{values} has not the right length, expected '
+ f'{len(self._storage)} , got {len(values)}'
+ )
+ if field in self._storage.tensor_columns.keys():
+
+ col = self._storage.tensor_columns[field]
+ if col is not None:
+ validation_class = col.__unparametrizedcls__ or col.__class__
+ else:
+ validation_class = self.doc_type._get_field_annotation(field)
+
+ # TODO shape check should be handle by the tensor validation
+
+ values = parse_obj_as(validation_class, values)
+ self._storage.tensor_columns[field] = values
+
+ elif field in self._storage.doc_columns.keys():
+ values_ = parse_obj_as(
+ DocVec.__class_getitem__(
+ self.doc_type._get_field_annotation(field)
+ ),
+ values,
+ )
+ self._storage.doc_columns[field] = values_
+
+ elif field in self._storage.docs_vec_columns.keys():
+ values_ = cast(Sequence[DocList[T_doc]], values)
+ # TODO here we should actually check if this is correct
+ self._storage.docs_vec_columns[field] = values_
+ elif field in self._storage.any_columns.keys():
+ # TODO here we should actually check if this is correct
+ values_ = cast(Sequence, values)
+ self._storage.any_columns[field] = values_
+ else:
+ raise KeyError(f'{field} is not a valid field for this DocList')
+
+ ####################
+ # Deleting data #
+ ####################
+
+ def __delitem__(self, key: Union[int, IndexIterType]) -> None:
+ raise NotImplementedError(
+ f'{self.__class__.__name__} does not implement '
+ f'__del_item__. You are trying to delete an element'
+ f'from {self.__class__.__name__} which is not '
+ f'designed for this operation. Please `unstack`'
+ f' before doing the deletion'
+ )
+
+ ####################
+ # Sequence related #
+ ####################
+ def __iter__(self):
+ for i in range(len(self)):
+ yield self[i]
+
+ def __len__(self):
+ return len(self._storage)
+
+ def __eq__(self, other: Any) -> bool:
+ if not isinstance(other, DocVec):
+ return False
+ if self.doc_type != other.doc_type:
+ return False
+ if self.tensor_type != other.tensor_type:
+ return False
+ if self._storage != other._storage:
+ return False
+ return True
+
+ ####################
+ # IO related #
+ ####################
+
+ @classmethod
+ def _get_proto_class(cls: Type[T]):
+ from docarray.proto import DocVecProto
+
+ return DocVecProto
+
+ def _docarray_to_json_compatible(self) -> Dict[str, Dict[str, Any]]:
+ tup = self._storage.columns_json_compatible()
+ return tup._asdict()
+
+ def to_doc_list(self: T) -> DocList[T_doc]:
+ """Convert DocVec into a DocList.
+
+ Note this destroys the arguments and returns a new DocList
+ """
+
+ unstacked_doc_column: Dict[str, Optional[DocList]] = dict()
+ unstacked_da_column: Dict[str, Optional[List[DocList]]] = dict()
+ unstacked_tensor_column: Dict[str, Optional[List[AbstractTensor]]] = dict()
+ unstacked_any_column = self._storage.any_columns
+
+ for field, doc_col in self._storage.doc_columns.items():
+ unstacked_doc_column[field] = doc_col.to_doc_list() if doc_col else None
+
+ for field, da_col in self._storage.docs_vec_columns.items():
+ unstacked_da_column[field] = (
+ [docs.to_doc_list() for docs in da_col] if da_col else None
+ )
+
+ for field, tensor_col in list(self._storage.tensor_columns.items()):
+ # list is needed here otherwise we cannot delete the column
+ if tensor_col is not None:
+ tensors = list()
+ for tensor in tensor_col:
+ tensor_copy = tensor.get_comp_backend().copy(tensor)
+ tensors.append(tensor_copy)
+
+ unstacked_tensor_column[field] = tensors
+ del self._storage.tensor_columns[field]
+
+ unstacked_column = ChainMap( # type: ignore
+ unstacked_any_column, # type: ignore
+ unstacked_tensor_column, # type: ignore
+ unstacked_da_column, # type: ignore
+ unstacked_doc_column, # type: ignore
+ ) # type: ignore
+
+ docs = []
+
+ for i in range(len(self)):
+ data = {field: col[i] for field, col in unstacked_column.items()}
+ docs.append(self.doc_type.construct(**data))
+
+ del self._storage
+
+ doc_type = self.doc_type
+
+ # Setting _is_unusable will raise an Exception if someone interacts with this instance from hereon out.
+ # I don't like relying on this state, but we can't override the getattr/setattr directly:
+ # https://stackoverflow.com/questions/10376604/overriding-special-methods-on-an-instance
+ self._is_unusable = True
+
+ return DocList.__class_getitem__(doc_type).construct(docs)
+
+ def traverse_flat(
+ self,
+ access_path: str,
+ ) -> Union[List[Any], 'TorchTensor', 'NdArray']:
+ nodes = list(AnyDocArray._traverse(node=self, access_path=access_path))
+ flattened = AnyDocArray._flatten_one_level(nodes)
+
+ cls_to_check = (NdArray, TorchTensor) if TorchTensor is not None else (NdArray,)
+
+ if len(flattened) == 1 and isinstance(flattened[0], cls_to_check):
+ return flattened[0]
+ else:
+ return flattened
+
+ @classmethod
+ def __class_getitem__(cls, item: Union[Type[BaseDoc], TypeVar, str]):
+ # call implementation in AnyDocArray
+ return super(IOMixinDocVec, cls).__class_getitem__(item)
+
+ if is_pydantic_v2:
+
+ @classmethod
+ def __get_pydantic_core_schema__(
+ cls, _source_type: Any, _handler: GetCoreSchemaHandler
+ ) -> core_schema.CoreSchema:
+ return core_schema.general_plain_validator_function(
+ cls.validate,
+ )
diff --git a/docarray/array/doc_vec/io.py b/docarray/array/doc_vec/io.py
new file mode 100644
index 00000000000..dd7213252fa
--- /dev/null
+++ b/docarray/array/doc_vec/io.py
@@ -0,0 +1,507 @@
+import base64
+import io
+import pathlib
+from abc import abstractmethod
+from contextlib import nullcontext
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Dict,
+ Generator,
+ Optional,
+ Type,
+ TypeVar,
+ Union,
+ cast,
+)
+
+import numpy as np
+import orjson
+from pydantic import parse_obj_as
+
+from docarray.array.doc_list.io import (
+ SINGLE_PROTOCOLS,
+ IOMixinDocList,
+ _LazyRequestReader,
+)
+from docarray.array.doc_vec.column_storage import ColumnStorage
+from docarray.array.list_advance_indexing import ListAdvancedIndexing
+from docarray.base_doc import BaseDoc
+from docarray.base_doc.mixins.io import _type_to_protobuf
+from docarray.typing import NdArray
+from docarray.typing.tensor.abstract_tensor import AbstractTensor
+from docarray.utils._internal.pydantic import is_pydantic_v2
+from docarray.utils._internal.misc import ProtocolType
+
+if TYPE_CHECKING:
+ import csv
+
+ import pandas as pd
+
+ from docarray.array.doc_vec.doc_vec import DocVec
+ from docarray.proto import (
+ DocVecProto,
+ ListOfDocArrayProto,
+ ListOfDocVecProto,
+ NdArrayProto,
+ )
+
+
+T = TypeVar('T', bound='IOMixinDocVec')
+T_doc = TypeVar('T_doc', bound=BaseDoc)
+
+NONE_NDARRAY_PROTO_SHAPE = (0,)
+NONE_NDARRAY_PROTO_DTYPE = 'None'
+
+
+def _none_ndarray_proto() -> 'NdArrayProto':
+ from docarray.proto import NdArrayProto
+
+ zeros_arr = parse_obj_as(NdArray, np.zeros(NONE_NDARRAY_PROTO_SHAPE))
+ nd_proto = NdArrayProto()
+ nd_proto.dense.buffer = zeros_arr.tobytes()
+ nd_proto.dense.ClearField('shape')
+ nd_proto.dense.shape.extend(list(zeros_arr.shape))
+ nd_proto.dense.dtype = NONE_NDARRAY_PROTO_DTYPE
+
+ return nd_proto
+
+
+def _none_docvec_proto() -> 'DocVecProto':
+ from docarray.proto import DocVecProto
+
+ return DocVecProto()
+
+
+def _none_list_of_docvec_proto() -> 'ListOfDocArrayProto':
+ from docarray.proto import ListOfDocVecProto
+
+ return ListOfDocVecProto()
+
+
+def _is_none_ndarray_proto(proto: 'NdArrayProto') -> bool:
+ return (
+ proto.dense.shape == list(NONE_NDARRAY_PROTO_SHAPE)
+ and proto.dense.dtype == NONE_NDARRAY_PROTO_DTYPE
+ )
+
+
+def _is_none_docvec_proto(proto: 'DocVecProto') -> bool:
+ return (
+ proto.tensor_columns == {}
+ and proto.doc_columns == {}
+ and proto.docs_vec_columns == {}
+ and proto.any_columns == {}
+ )
+
+
+def _is_none_list_of_docvec_proto(proto: 'ListOfDocVecProto') -> bool:
+ from docarray.proto import ListOfDocVecProto
+
+ return isinstance(proto, ListOfDocVecProto) and len(proto.data) == 0
+
+
+class IOMixinDocVec(IOMixinDocList):
+ @classmethod
+ @abstractmethod
+ def from_columns_storage(cls: Type[T], storage: ColumnStorage) -> T:
+ ...
+
+ @classmethod
+ @abstractmethod
+ def __class_getitem__(cls, item: Union[Type[BaseDoc], TypeVar, str]):
+ ...
+
+ @classmethod
+ def from_json(
+ cls: Type[T],
+ file: Union[str, bytes, bytearray],
+ tensor_type: Type[AbstractTensor] = NdArray,
+ ) -> T:
+ """Deserialize JSON strings or bytes into a `DocList`.
+
+ :param file: JSON object from where to deserialize a `DocList`
+ :param tensor_type: the tensor type to use for the tensor columns.
+ Could be NdArray, TorchTensor, or TensorFlowTensor. Defaults to NdArray.
+ All tensors of the output DocVec will be of this type.
+ :return: the deserialized `DocList`
+ """
+ json_columns = orjson.loads(file)
+ return cls._from_json_col_dict(json_columns, tensor_type=tensor_type)
+
+ @classmethod
+ def _from_json_col_dict(
+ cls: Type[T],
+ json_columns: Dict[str, Any],
+ tensor_type: Type[AbstractTensor] = NdArray,
+ ) -> T:
+ tensor_cols = json_columns['tensor_columns']
+ doc_cols = json_columns['doc_columns']
+ docs_vec_cols = json_columns['docs_vec_columns']
+ any_cols = json_columns['any_columns']
+
+ for key, col in tensor_cols.items():
+ if col is not None:
+ tensor_cols[key] = parse_obj_as(tensor_type, col)
+ else:
+ tensor_cols[key] = None
+
+ for key, col in doc_cols.items():
+ if col is not None:
+ col_doc_type = cls.doc_type._get_field_annotation(key)
+ doc_cols[key] = cls.__class_getitem__(col_doc_type)._from_json_col_dict(
+ col, tensor_type=tensor_type
+ )
+ else:
+ doc_cols[key] = None
+
+ for key, col in docs_vec_cols.items():
+ if col is not None:
+ col_doc_type = cls.doc_type._get_field_annotation(key).doc_type
+ col_ = ListAdvancedIndexing(
+ cls.__class_getitem__(col_doc_type)._from_json_col_dict(
+ vec, tensor_type=tensor_type
+ )
+ for vec in col
+ )
+ docs_vec_cols[key] = col_
+ else:
+ docs_vec_cols[key] = None
+
+ for key, col in any_cols.items():
+ if col is not None:
+ col_type = cls.doc_type._get_field_annotation(key)
+
+ field_required = (
+ cls.doc_type._docarray_fields()[key].is_required()
+ if is_pydantic_v2
+ else cls.doc_type._docarray_fields()[key].required
+ )
+
+ col_type = col_type if field_required else Optional[col_type]
+ col_ = ListAdvancedIndexing(parse_obj_as(col_type, val) for val in col)
+ any_cols[key] = col_
+ else:
+ any_cols[key] = None
+
+ return cls.from_columns_storage(
+ ColumnStorage(
+ tensor_cols, doc_cols, docs_vec_cols, any_cols, tensor_type=tensor_type
+ )
+ )
+
+ @classmethod
+ def from_protobuf(
+ cls: Type[T], pb_msg: 'DocVecProto', tensor_type: Type[AbstractTensor] = NdArray
+ ) -> T:
+ """create a DocVec from a protobuf message
+ :param pb_msg: the protobuf message to deserialize
+ :param tensor_type: the tensor type to use for the tensor columns.
+ Could be NdArray, TorchTensor, or TensorFlowTensor. Defaults to NdArray.
+ All tensors of the output DocVec will be of this type.
+ :return: The deserialized DocVec
+ """
+ tensor_columns: Dict[str, Optional[AbstractTensor]] = {}
+ doc_columns: Dict[str, Optional['DocVec']] = {}
+ docs_vec_columns: Dict[str, Optional[ListAdvancedIndexing['DocVec']]] = {}
+ any_columns: Dict[str, ListAdvancedIndexing] = {}
+
+ for tens_col_name, tens_col_proto in pb_msg.tensor_columns.items():
+ if _is_none_ndarray_proto(tens_col_proto):
+ # handle values that were None before serialization
+ tensor_columns[tens_col_name] = None
+ else:
+ tensor_columns[tens_col_name] = tensor_type.from_protobuf(
+ tens_col_proto
+ )
+
+ for doc_col_name, doc_col_proto in pb_msg.doc_columns.items():
+ if _is_none_docvec_proto(doc_col_proto):
+ # handle values that were None before serialization
+ doc_columns[doc_col_name] = None
+ else:
+ col_doc_type: Type = cls.doc_type._get_field_annotation(doc_col_name)
+ doc_columns[doc_col_name] = cls.__class_getitem__(
+ col_doc_type
+ ).from_protobuf(doc_col_proto, tensor_type=tensor_type)
+
+ for docs_vec_col_name, docs_vec_col_proto in pb_msg.docs_vec_columns.items():
+ vec_list: Optional[ListAdvancedIndexing]
+ if _is_none_list_of_docvec_proto(docs_vec_col_proto):
+ # handle values that were None before serialization
+ vec_list = None
+ else:
+ vec_list = ListAdvancedIndexing()
+ for doc_list_proto in docs_vec_col_proto.data:
+ col_doc_type = cls.doc_type._get_field_annotation(
+ docs_vec_col_name
+ ).doc_type
+ vec_list.append(
+ cls.__class_getitem__(col_doc_type).from_protobuf(
+ doc_list_proto, tensor_type=tensor_type
+ )
+ )
+ docs_vec_columns[docs_vec_col_name] = vec_list
+
+ for any_col_name, any_col_proto in pb_msg.any_columns.items():
+ any_column: ListAdvancedIndexing = ListAdvancedIndexing()
+ for node_proto in any_col_proto.data:
+ content = cls.doc_type._get_content_from_node_proto(
+ node_proto, any_col_name
+ )
+ any_column.append(content)
+ any_columns[any_col_name] = any_column
+
+ storage = ColumnStorage(
+ tensor_columns=tensor_columns,
+ doc_columns=doc_columns,
+ docs_vec_columns=docs_vec_columns,
+ any_columns=any_columns,
+ tensor_type=tensor_type,
+ )
+
+ return cls.from_columns_storage(storage)
+
+ def to_protobuf(self) -> 'DocVecProto':
+ """Convert DocVec into a Protobuf message"""
+ from docarray.proto import (
+ DocVecProto,
+ ListOfAnyProto,
+ ListOfDocArrayProto,
+ ListOfDocVecProto,
+ NdArrayProto,
+ )
+
+ self_ = cast('DocVec', self)
+
+ doc_columns_proto: Dict[str, DocVecProto] = dict()
+ tensor_columns_proto: Dict[str, NdArrayProto] = dict()
+ da_columns_proto: Dict[str, ListOfDocArrayProto] = dict()
+ any_columns_proto: Dict[str, ListOfAnyProto] = dict()
+
+ for field, col_doc in self_._storage.doc_columns.items():
+ if col_doc is None:
+ # put dummy empty DocVecProto for serialization
+ doc_columns_proto[field] = _none_docvec_proto()
+ else:
+ doc_columns_proto[field] = col_doc.to_protobuf()
+ for field, col_tens in self_._storage.tensor_columns.items():
+ if col_tens is None:
+ # put dummy empty NdArrayProto for serialization
+ tensor_columns_proto[field] = _none_ndarray_proto()
+ else:
+ tensor_columns_proto[field] = (
+ col_tens.to_protobuf() if col_tens is not None else None
+ )
+ for field, col_da in self_._storage.docs_vec_columns.items():
+ list_proto = ListOfDocVecProto()
+ if col_da:
+ for docs in col_da:
+ list_proto.data.append(docs.to_protobuf())
+ else:
+ # put dummy empty ListOfDocVecProto for serialization
+ list_proto = _none_list_of_docvec_proto()
+ da_columns_proto[field] = list_proto
+ for field, col_any in self_._storage.any_columns.items():
+ list_proto = ListOfAnyProto()
+ for data in col_any:
+ list_proto.data.append(_type_to_protobuf(data))
+ any_columns_proto[field] = list_proto
+
+ return DocVecProto(
+ doc_columns=doc_columns_proto,
+ tensor_columns=tensor_columns_proto,
+ docs_vec_columns=da_columns_proto,
+ any_columns=any_columns_proto,
+ )
+
+ def to_csv(
+ self, file_path: str, dialect: Union[str, 'csv.Dialect'] = 'excel'
+ ) -> None:
+ """
+ DocVec does not support `.to_csv()`. This is because CSV is a row-based format
+ while DocVec has a column-based data layout.
+ To overcome this, do: `doc_vec.to_doc_list().to_csv(...)`.
+ """
+ raise NotImplementedError(
+ f'{type(self)} does not support `.to_csv()`. This is because CSV is a row-based format'
+ f'while {type(self)} has a column-based data layout. '
+ f'To overcome this, do: `doc_vec.to_doc_list().to_csv(...)`.'
+ )
+
+ @classmethod
+ def from_csv(
+ cls: Type['T'],
+ file_path: str,
+ encoding: str = 'utf-8',
+ dialect: Union[str, 'csv.Dialect'] = 'excel',
+ ) -> 'T':
+ """
+ DocVec does not support `.from_csv()`. This is because CSV is a row-based format
+ while DocVec has a column-based data layout.
+ To overcome this, do: `DocList[MyDoc].from_csv(...).to_doc_vec()`.
+ """
+ raise NotImplementedError(
+ f'{cls} does not support `.from_csv()`. This is because CSV is a row-based format while'
+ f'{cls} has a column-based data layout. '
+ f'To overcome this, do: `DocList[MyDoc].from_csv(...).to_doc_vec()`.'
+ )
+
+ @classmethod
+ def from_base64(
+ cls: Type[T],
+ data: str,
+ protocol: ProtocolType = 'protobuf-array',
+ compress: Optional[str] = None,
+ show_progress: bool = False,
+ tensor_type: Type['AbstractTensor'] = NdArray,
+ ) -> T:
+ """Deserialize base64 strings into a `DocVec`.
+
+ :param data: Base64 string to deserialize
+ :param protocol: protocol that was used to serialize
+ :param compress: compress algorithm that was used to serialize between `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
+ :param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
+ :param tensor_type: the tensor type of the resulting DocVEc
+ :return: the deserialized `DocVec`
+ """
+ return cls._load_binary_all(
+ file_ctx=nullcontext(base64.b64decode(data)),
+ protocol=protocol,
+ compress=compress,
+ show_progress=show_progress,
+ tensor_type=tensor_type,
+ )
+
+ @classmethod
+ def from_bytes(
+ cls: Type[T],
+ data: bytes,
+ protocol: ProtocolType = 'protobuf-array',
+ compress: Optional[str] = None,
+ show_progress: bool = False,
+ tensor_type: Type['AbstractTensor'] = NdArray,
+ ) -> T:
+ """Deserialize bytes into a `DocList`.
+
+ :param data: Bytes from which to deserialize
+ :param protocol: protocol that was used to serialize
+ :param compress: compression algorithm that was used to serialize between `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
+ :param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
+ :param tensor_type: the tensor type of the resulting DocVec
+ :return: the deserialized `DocVec`
+ """
+ return cls._load_binary_all(
+ file_ctx=nullcontext(data),
+ protocol=protocol,
+ compress=compress,
+ show_progress=show_progress,
+ tensor_type=tensor_type,
+ )
+
+ @classmethod
+ def from_dataframe(
+ cls: Type['T'],
+ df: 'pd.DataFrame',
+ tensor_type: Type['AbstractTensor'] = NdArray,
+ ) -> 'T':
+ """
+ Load a `DocVec` from a `pandas.DataFrame` following the schema
+ defined in the [`.doc_type`][docarray.DocVec] attribute.
+ Every row of the dataframe will be mapped to one Document in the doc_vec.
+ The column names of the dataframe have to match the field names of the
+ Document type.
+ For nested fields use "__"-separated access paths as column names,
+ such as `'image__url'`.
+
+ List-like fields (including field of type DocList) are not supported.
+
+ ---
+
+ ```python
+ import pandas as pd
+
+ from docarray import BaseDoc, DocVec
+
+
+ class Person(BaseDoc):
+ name: str
+ follower: int
+
+
+ df = pd.DataFrame(
+ data=[['Maria', 12345], ['Jake', 54321]], columns=['name', 'follower']
+ )
+
+ docs = DocVec[Person].from_dataframe(df)
+
+ assert docs.name == ['Maria', 'Jake']
+ assert docs.follower == [12345, 54321]
+ ```
+
+ ---
+
+ :param df: `pandas.DataFrame` to extract Document's information from
+ :param tensor_type: the tensor type of the resulting DocVec
+ :return: `DocList` where each Document contains the information of one
+ corresponding row of the `pandas.DataFrame`.
+ """
+ # type ignore could be avoided by simply putting this implementation in the DocVec class
+ # but leaving it here for code separation
+ return cls(super().from_dataframe(df), tensor_type=tensor_type) # type: ignore
+
+ @classmethod
+ def load_binary(
+ cls: Type[T],
+ file: Union[str, bytes, pathlib.Path, io.BufferedReader, _LazyRequestReader],
+ protocol: ProtocolType = 'protobuf-array',
+ compress: Optional[str] = None,
+ show_progress: bool = False,
+ streaming: bool = False,
+ tensor_type: Type['AbstractTensor'] = NdArray,
+ ) -> Union[T, Generator['T_doc', None, None]]:
+ """Load doc_vec elements from a compressed binary file.
+
+ In case protocol is pickle the `Documents` are streamed from disk to save memory usage
+
+ !!! note
+ If `file` is `str` it can specify `protocol` and `compress` as file extensions.
+ This functionality assumes `file=file_name.$protocol.$compress` where `$protocol` and `$compress` refer to a
+ string interpolation of the respective `protocol` and `compress` methods.
+ For example if `file=my_docarray.protobuf.lz4` then the binary data will be loaded assuming `protocol=protobuf`
+ and `compress=lz4`.
+
+ :param file: File or filename or serialized bytes where the data is stored.
+ :param protocol: protocol to use. It can be 'pickle-array', 'protobuf-array', 'pickle' or 'protobuf'
+ :param compress: compress algorithm to use between `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
+ :param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
+ :param streaming: if `True` returns a generator over `Document` objects.
+ :param tensor_type: the tensor type of the resulting DocVEc
+
+ :return: a `DocVec` object
+
+ """
+ file_ctx, load_protocol, load_compress = cls._get_file_context(
+ file, protocol, compress
+ )
+ if streaming:
+ if load_protocol not in SINGLE_PROTOCOLS:
+ raise ValueError(
+ f'`streaming` is only available when using {" or ".join(map(lambda x: f"`{x}`", SINGLE_PROTOCOLS))} as protocol, '
+ f'got {load_protocol}'
+ )
+ else:
+ return cls._load_binary_stream(
+ file_ctx,
+ protocol=load_protocol,
+ compress=load_compress,
+ show_progress=show_progress,
+ )
+ else:
+ return cls._load_binary_all(
+ file_ctx,
+ load_protocol,
+ load_compress,
+ show_progress,
+ tensor_type=tensor_type,
+ )
diff --git a/docarray/array/document.py b/docarray/array/document.py
deleted file mode 100644
index be70a655bb9..00000000000
--- a/docarray/array/document.py
+++ /dev/null
@@ -1,369 +0,0 @@
-import itertools
-from typing import (
- Optional,
- TYPE_CHECKING,
- Generator,
- Iterator,
- Dict,
- Union,
- MutableSequence,
- Sequence,
- Iterable,
- overload,
- Any,
- List,
-)
-
-import numpy as np
-
-from .mixins import AllMixins
-from .. import Document
-from ..helper import typename
-
-if TYPE_CHECKING:
- from ..types import (
- DocumentArraySourceType,
- DocumentArrayIndexType,
- DocumentArraySingletonIndexType,
- DocumentArrayMultipleIndexType,
- DocumentArrayMultipleAttributeType,
- DocumentArraySingleAttributeType,
- )
-
-
-class DocumentArray(AllMixins, MutableSequence[Document]):
- def __init__(
- self, docs: Optional['DocumentArraySourceType'] = None, copy: bool = False
- ):
- super().__init__()
- self._data = []
- if docs is None:
- return
- elif isinstance(
- docs, (DocumentArray, Sequence, Generator, Iterator, itertools.chain)
- ):
- if copy:
- self._data = [Document(d, copy=True) for d in docs]
- self._rebuild_id2offset()
- elif isinstance(docs, DocumentArray):
- self._data = docs._data
- self._id_to_index = docs._id_to_index
- else:
- self._data = list(docs)
- self._rebuild_id2offset()
- else:
- if isinstance(docs, Document):
- if copy:
- self.append(Document(docs, copy=True))
- else:
- self.append(docs)
-
- @property
- def _id2offset(self) -> Dict[str, int]:
- """Return the `_id_to_index` map
-
- :return: a Python dict.
- """
- if not hasattr(self, '_id_to_index'):
- self._rebuild_id2offset()
- return self._id_to_index
-
- def _rebuild_id2offset(self) -> None:
- """Update the id_to_index map by enumerating all Documents in self._data.
-
- Very costy! Only use this function when self._data is dramtically changed.
- """
-
- self._id_to_index = {
- d.id: i for i, d in enumerate(self._data)
- } # type: Dict[str, int]
-
- def insert(self, index: int, value: 'Document'):
- """Insert `doc` at `index`.
-
- :param index: Position of the insertion.
- :param value: The doc needs to be inserted.
- """
- self._data.insert(index, value)
- self._id2offset[value.id] = index
-
- def __eq__(self, other):
- return (
- type(self) is type(other)
- and type(self._data) is type(other._data)
- and self._data == other._data
- )
-
- def __len__(self):
- return len(self._data)
-
- def __iter__(self) -> Iterator['Document']:
- yield from self._data
-
- def __contains__(self, x: Union[str, 'Document']):
- if isinstance(x, str):
- return x in self._id2offset
- elif isinstance(x, Document):
- return x.id in self._id2offset
- else:
- return False
-
- @overload
- def __getitem__(self, index: 'DocumentArraySingletonIndexType') -> 'Document':
- ...
-
- @overload
- def __getitem__(self, index: 'DocumentArrayMultipleIndexType') -> 'DocumentArray':
- ...
-
- @overload
- def __getitem__(self, index: 'DocumentArraySingleAttributeType') -> List[Any]:
- ...
-
- @overload
- def __getitem__(
- self, index: 'DocumentArrayMultipleAttributeType'
- ) -> List[List[Any]]:
- ...
-
- def __getitem__(
- self, index: 'DocumentArrayIndexType'
- ) -> Union['Document', 'DocumentArray']:
- if isinstance(index, (int, np.generic)):
- return self._data[int(index)]
- elif isinstance(index, str):
- if index.startswith('@'):
- return self.traverse_flat(index[1:])
- else:
- return self._data[self._id2offset[index]]
- elif isinstance(index, slice):
- return DocumentArray(self._data[index])
- elif index is Ellipsis:
- return self.flatten()
- elif isinstance(index, Sequence):
- if (
- isinstance(index, tuple)
- and len(index) == 2
- and isinstance(index[0], (slice, Sequence))
- ):
- _docs = self[index[0]]
- _attrs = index[1]
- if isinstance(_attrs, str):
- _attrs = (index[1],)
-
- return _docs._get_attributes(*_attrs)
- elif isinstance(index[0], bool):
- return DocumentArray(itertools.compress(self._data, index))
- elif isinstance(index[0], int):
- return DocumentArray(self._data[t] for t in index)
- elif isinstance(index[0], str):
- return DocumentArray(self._data[self._id2offset[t]] for t in index)
- elif isinstance(index, np.ndarray):
- index = index.squeeze()
- if index.ndim == 1:
- return self[index.tolist()]
- else:
- raise IndexError(
- f'When using np.ndarray as index, its `ndim` must =1. However, receiving ndim={index.ndim}'
- )
- raise IndexError(f'Unsupported index type {typename(index)}: {index}')
-
- @overload
- def __setitem__(
- self,
- index: 'DocumentArrayMultipleAttributeType',
- value: List[List['Any']],
- ):
- ...
-
- @overload
- def __setitem__(
- self,
- index: 'DocumentArraySingleAttributeType',
- value: List['Any'],
- ):
- ...
-
- @overload
- def __setitem__(
- self,
- index: 'DocumentArraySingletonIndexType',
- value: 'Document',
- ):
- ...
-
- @overload
- def __setitem__(
- self,
- index: 'DocumentArrayMultipleIndexType',
- value: Sequence['Document'],
- ):
- ...
-
- def __setitem__(
- self,
- index: 'DocumentArrayIndexType',
- value: Union['Document', Sequence['Document']],
- ):
- if isinstance(index, (int, np.generic)):
- index = int(index)
- self._data[index] = value
- self._id2offset[value.id] = index
- elif isinstance(index, str):
- if index.startswith('@'):
- for _d, _v in zip(self.traverse_flat(index[1:]), value):
- _d._data = _v._data
- self._rebuild_id2offset()
- else:
- old_idx = self._id2offset.pop(index)
- self._data[old_idx] = value
- self._id2offset[value.id] = old_idx
- elif isinstance(index, slice):
- self._data[index] = value
- self._rebuild_id2offset()
- elif index is Ellipsis:
- for _d, _v in zip(self.flatten(), value):
- _d._data = _v._data
- self._rebuild_id2offset()
- elif isinstance(index, Sequence):
- if (
- isinstance(index, tuple)
- and len(index) == 2
- and isinstance(index[0], (slice, Sequence))
- ):
- _docs = self[index[0]]
- _attrs = index[1]
-
- if isinstance(_attrs, str):
- # a -> [a]
- # [a, a] -> [a, a]
- _attrs = (index[1],)
- if isinstance(value, (list, tuple)) and not any(
- isinstance(el, (tuple, list)) for el in value
- ):
- # [x] -> [[x]]
- # [[x], [y]] -> [[x], [y]]
- value = (value,)
- if not isinstance(value, (list, tuple)):
- # x -> [x]
- value = (value,)
-
- for _a, _v in zip(_attrs, value):
- if _a == 'blob':
- _docs.blobs = _v
- elif _a == 'embedding':
- _docs.embeddings = _v
- else:
- for _d, _vv in zip(_docs, _v):
- setattr(_d, _a, _vv)
- elif isinstance(index[0], bool):
- if len(index) != len(self._data):
- raise IndexError(
- f'Boolean mask index is required to have the same length as {len(self._data)}, '
- f'but receiving {len(index)}'
- )
- _selected = itertools.compress(self._data, index)
- for _idx, _val in zip(_selected, value):
- self[_idx.id] = _val
- elif isinstance(index[0], (int, str)):
- if not isinstance(value, Sequence) or len(index) != len(value):
- raise ValueError(
- f'Number of elements for assigning must be '
- f'the same as the index length: {len(index)}'
- )
- if isinstance(value, Document):
- for si in index:
- self[si] = value
- else:
- for si, _val in zip(index, value):
- self[si] = _val
- elif isinstance(index, np.ndarray):
- index = index.squeeze()
- if index.ndim == 1:
- self[index.tolist()] = value
- else:
- raise IndexError(
- f'When using np.ndarray as index, its `ndim` must =1. However, receiving ndim={index.ndim}'
- )
- else:
- raise IndexError(f'Unsupported index type {typename(index)}: {index}')
-
- def __delitem__(self, index: 'DocumentArrayIndexType'):
- if isinstance(index, (int, np.generic)):
- index = int(index)
- self._id2offset.pop(self._data[index].id)
- del self._data[index]
- elif isinstance(index, str):
- if index.startswith('@'):
- raise NotImplementedError(
- 'Delete elements along traversal paths is not implemented'
- )
- else:
- del self._data[self._id2offset[index]]
- self._id2offset.pop(index)
- elif isinstance(index, slice):
- del self._data[index]
- self._rebuild_id2offset()
- elif index is Ellipsis:
- self._data.clear()
- self._id2offset.clear()
- elif isinstance(index, Sequence):
- if (
- isinstance(index, tuple)
- and len(index) == 2
- and isinstance(index[0], (slice, Sequence))
- ):
- _docs = self[index[0]]
- _attrs = index[1]
- if isinstance(_attrs, str):
- _attrs = (index[1],)
- for _d in _docs:
- _d.pop(*_attrs)
- elif isinstance(index[0], bool):
- self._data = list(
- itertools.compress(self._data, (not _i for _i in index))
- )
- self._rebuild_id2offset()
- elif isinstance(index[0], int):
- for t in sorted(index, reverse=True):
- del self[t]
- elif isinstance(index[0], str):
- for t in index:
- del self[t]
- elif isinstance(index, np.ndarray):
- index = index.squeeze()
- if index.ndim == 1:
- del self[index.tolist()]
- else:
- raise IndexError(
- f'When using np.ndarray as index, its `ndim` must =1. However, receiving ndim={index.ndim}'
- )
- else:
- raise IndexError(f'Unsupported index type {typename(index)}: {index}')
-
- def clear(self):
- """Clear the data of :class:`DocumentArray`"""
- self._data.clear()
- self._id2offset.clear()
-
- def __bool__(self):
- """To simulate ```l = []; if l: ...```
-
- :return: returns true if the length of the array is larger than 0
- """
- return len(self) > 0
-
- def __repr__(self):
- return f'<{self.__class__.__name__} (length={len(self)}) at {id(self)}>'
-
- def __add__(self, other: 'Document'):
- v = type(self)()
- for doc in self:
- v.append(doc)
- for doc in other:
- v.append(doc)
- return v
-
- def extend(self, values: Iterable['Document']) -> None:
- self._data.extend(values)
- self._rebuild_id2offset()
diff --git a/docarray/array/list_advance_indexing.py b/docarray/array/list_advance_indexing.py
new file mode 100644
index 00000000000..c3d80ad2f6c
--- /dev/null
+++ b/docarray/array/list_advance_indexing.py
@@ -0,0 +1,223 @@
+from typing import (
+ Any,
+ Iterable,
+ List,
+ Sequence,
+ TypeVar,
+ Union,
+ cast,
+ no_type_check,
+ overload,
+)
+
+import numpy as np
+from typing_extensions import SupportsIndex
+
+from docarray.utils._internal.misc import (
+ is_jax_available,
+ is_tf_available,
+ is_torch_available,
+)
+
+torch_available = is_torch_available()
+if torch_available:
+ import torch
+tf_available = is_tf_available()
+if tf_available:
+ import tensorflow as tf # type: ignore
+
+ from docarray.typing.tensor.tensorflow_tensor import TensorFlowTensor
+jax_available = is_jax_available()
+if jax_available:
+ import jax.numpy as jnp
+
+ from docarray.typing.tensor.jaxarray import JaxArray
+
+T_item = TypeVar('T_item')
+T = TypeVar('T', bound='ListAdvancedIndexing')
+
+IndexIterType = Union[slice, Iterable[int], Iterable[bool], None]
+
+
+def _is_np_int(item: Any) -> bool:
+ dtype = getattr(item, 'dtype', None)
+ ndim = getattr(item, 'ndim', None)
+ if dtype is not None and ndim is not None:
+ try:
+ return ndim == 0 and np.issubdtype(dtype, np.integer)
+ except TypeError:
+ return False
+ return False # this is unreachable, but mypy wants it
+
+
+class ListAdvancedIndexing(List[T_item]):
+ """
+ A list wrapper that implements custom indexing
+
+ You can index into a ListAdvanceIndex like a numpy array or torch tensor:
+
+ ---
+
+ ```python
+ docs[0] # index by position
+ docs[0:5:2] # index by slice
+ docs[[0, 2, 3]] # index by list of indices
+ docs[True, False, True, True, ...] # index by boolean mask
+ ```
+
+ ---
+
+ """
+
+ @staticmethod
+ def _normalize_index_item(
+ item: Any,
+ ) -> Union[int, slice, Iterable[int], Iterable[bool], None]:
+ # basic index types
+ if item is None or isinstance(item, (int, slice, tuple, list)):
+ return item
+
+ # numpy index types
+ if _is_np_int(item):
+ return item.item()
+
+ index_has_getitem = hasattr(item, '__getitem__')
+ is_valid_bulk_index = index_has_getitem and isinstance(item, Iterable)
+ if not is_valid_bulk_index:
+ raise ValueError(f'Invalid index type {type(item)}')
+
+ if isinstance(item, np.ndarray) and (
+ item.dtype == np.bool_ or np.issubdtype(item.dtype, np.integer)
+ ):
+ return item.tolist()
+
+ # torch index types
+ if torch_available:
+
+ allowed_torch_dtypes = [
+ torch.bool,
+ torch.int64,
+ ]
+ if isinstance(item, torch.Tensor) and (item.dtype in allowed_torch_dtypes):
+ return item.tolist()
+
+ if tf_available:
+ if isinstance(item, tf.Tensor):
+ return item.numpy().tolist()
+ if isinstance(item, TensorFlowTensor):
+ return item.tensor.numpy().tolist()
+
+ if jax_available:
+ if isinstance(item, jnp.ndarray):
+ return item.__array__().tolist()
+ if isinstance(item, JaxArray):
+ return item.tensor.__array__().tolist()
+
+ return item
+
+ def _get_from_indices(self: T, item: Iterable[int]) -> T:
+ results = []
+ for ix in item:
+ results.append(self[ix])
+ return self.__class__(results)
+
+ def _set_by_indices(self: T, item: Iterable[int], value: Iterable[T_item]):
+ for ix, doc_to_set in zip(item, value):
+ try:
+ self[ix] = doc_to_set
+ except KeyError:
+ raise IndexError(f'Index {ix} is out of range')
+
+ def _get_from_mask(self: T, item: Iterable[bool]) -> T:
+ return self.__class__(
+ [doc for doc, mask_value in zip(self, item) if mask_value]
+ )
+
+ def _set_by_mask(self: T, item: Iterable[bool], value: Sequence[T_item]):
+ i_value = 0
+ for i, mask_value in zip(range(len(self)), item):
+ if mask_value:
+ self[i] = value[i_value]
+ i_value += 1
+
+ def _del_from_mask(self: T, item: Iterable[bool]) -> None:
+ idx_to_delete = [i for i, val in enumerate(item) if val]
+ self._del_from_indices(idx_to_delete)
+
+ def _del_from_indices(self: T, item: Iterable[int]) -> None:
+ for ix in sorted(item, reverse=True):
+ # reversed is needed here otherwise some the indices are not up to date after
+ # each delete
+ del self[ix]
+
+ def __delitem__(self, key: Union[SupportsIndex, IndexIterType]) -> None:
+ item = self._normalize_index_item(key)
+
+ if item is None:
+ return
+ elif isinstance(item, (int, slice)):
+ super().__delitem__(item)
+ else:
+ head = item[0] # type: ignore
+ if isinstance(head, bool):
+ item_ = cast(Iterable[bool], item)
+ return self._del_from_mask(item_)
+ elif isinstance(head, int):
+ return self._del_from_indices(item)
+ else:
+ raise TypeError(f'Invalid type {type(head)} for indexing')
+
+ @overload
+ def __getitem__(self: T, item: SupportsIndex) -> T_item:
+ ...
+
+ @overload
+ def __getitem__(self: T, item: IndexIterType) -> T:
+ ...
+
+ def __getitem__(self, item):
+ item = self._normalize_index_item(item)
+
+ if type(item) == slice:
+ return self.__class__(super().__getitem__(item))
+
+ if isinstance(item, int):
+ return super().__getitem__(item)
+
+ if item is None:
+ return self
+
+ # _normalize_index_item() guarantees the line below is correct
+ head = item[0] # type: ignore
+ if isinstance(head, bool):
+ return self._get_from_mask(item)
+ elif isinstance(head, int):
+ return self._get_from_indices(item)
+ else:
+ raise TypeError(f'Invalid type {type(head)} for indexing')
+
+ @overload
+ def __setitem__(self: T, key: SupportsIndex, value: T_item) -> None:
+ ...
+
+ @overload
+ def __setitem__(self: T, key: IndexIterType, value: Iterable[T_item]):
+ ...
+
+ @no_type_check
+ def __setitem__(self: T, key, value):
+ key_norm = self._normalize_index_item(key)
+
+ if isinstance(key_norm, int):
+ super().__setitem__(key_norm, value)
+ elif isinstance(key_norm, slice):
+ super().__setitem__(key_norm, value)
+ else:
+ # _normalize_index_item() guarantees the line below is correct
+ head = key_norm[0]
+ if isinstance(head, bool):
+ return self._set_by_mask(key_norm, value)
+ elif isinstance(head, int):
+ return self._set_by_indices(key_norm, value)
+ else:
+ raise TypeError(f'Invalid type {type(head)} for indexing')
diff --git a/docarray/array/match.py b/docarray/array/match.py
deleted file mode 100644
index 2b33828f3fb..00000000000
--- a/docarray/array/match.py
+++ /dev/null
@@ -1,63 +0,0 @@
-import itertools
-from typing import (
- TYPE_CHECKING,
- Generator,
- Iterator,
- Sequence,
-)
-
-from .. import DocumentArray
-
-if TYPE_CHECKING:
- from ..document import Document
-
-
-class MatchArray(DocumentArray):
- """
- :class:`MatchArray` inherits from :class:`DocumentArray`.
- It's a subset of Documents that represents the matches
-
- :param docs: Set of matches of the `reference_doc`
- :param reference_doc: Reference :class:`Document` for the sub-documents
- """
-
- def __init__(self, docs, reference_doc: 'Document'):
- self._ref_doc = reference_doc
- super().__init__(docs)
- if (
- isinstance(
- docs, (DocumentArray, Sequence, Generator, Iterator, itertools.chain)
- )
- and self._ref_doc is not None
- ):
- for d in docs:
- d.adjacency = self._ref_doc.adjacency + 1
-
- def append(self, document: 'Document'):
- """Add a matched document to the current Document.
-
- :param document: Sub-document to be added
- """
- document.adjacency = self._ref_doc.adjacency + 1
- super().append(document)
-
- @property
- def reference_doc(self) -> 'Document':
- """Get the document that this :class:`MatchArray` referring to.
- :return: the document the match refers to
- """
- return self._ref_doc
-
- @property
- def granularity(self) -> int:
- """Get granularity of all document in this array.
- :return: the granularity of the documents of which these are match
- """
- return self._ref_doc.granularity
-
- @property
- def adjacency(self) -> int:
- """Get the adjacency of all document in this array.
- :return: the adjacency of the array of matches
- """
- return self._ref_doc.adjacency + 1
diff --git a/docarray/array/mixins/__init__.py b/docarray/array/mixins/__init__.py
deleted file mode 100644
index b031b2e4769..00000000000
--- a/docarray/array/mixins/__init__.py
+++ /dev/null
@@ -1,50 +0,0 @@
-from abc import ABC
-
-from .content import ContentPropertyMixin
-from .embed import EmbedMixin
-from .empty import EmptyMixin
-from .evaluation import EvaluationMixin
-from .getattr import GetAttributeMixin
-from .group import GroupMixin
-from .io.binary import BinaryIOMixin
-from .io.common import CommonIOMixin
-from .io.csv import CsvIOMixin
-from .io.dataframe import DataframeIOMixin
-from .io.from_gen import FromGeneratorMixin
-from .io.json import JsonIOMixin
-from .io.pushpull import PushPullMixin
-from .match import MatchMixin
-from .parallel import ParallelMixin
-from .plot import PlotMixin
-from .reduce import ReduceMixin
-from .sample import SampleMixin
-from .text import TextToolsMixin
-from .traverse import TraverseMixin
-
-
-class AllMixins(
- GetAttributeMixin,
- ContentPropertyMixin,
- GroupMixin,
- EmptyMixin,
- CsvIOMixin,
- JsonIOMixin,
- BinaryIOMixin,
- CommonIOMixin,
- EmbedMixin,
- PushPullMixin,
- FromGeneratorMixin,
- MatchMixin,
- TraverseMixin,
- PlotMixin,
- SampleMixin,
- TextToolsMixin,
- EvaluationMixin,
- ReduceMixin,
- ParallelMixin,
- DataframeIOMixin,
- ABC,
-):
- """All plugins that can be used in :class:`DocumentArray`. """
-
- ...
diff --git a/docarray/array/mixins/content.py b/docarray/array/mixins/content.py
deleted file mode 100644
index d2692758b49..00000000000
--- a/docarray/array/mixins/content.py
+++ /dev/null
@@ -1,160 +0,0 @@
-from typing import List, Sequence, TYPE_CHECKING, Optional, Union
-
-from ...math.ndarray import ravel, unravel
-
-if TYPE_CHECKING:
- from ...types import ArrayType, DocumentContentType
-
-
-class ContentPropertyMixin:
- """Helpers that provide faster getter & setter for :attr:`.content` by using protobuf directly. """
-
- def _check_length(self, target_len: int):
- if target_len != len(self):
- raise ValueError(
- f'Length must match {self!r}, but {target_len} != {len(self)}'
- )
-
- @property
- def embeddings(self) -> Optional['ArrayType']:
- """Return a :class:`ArrayType` stacking all the `embedding` attributes as rows.
-
- :return: a :class:`ArrayType` of embedding
- """
- if self:
- return unravel(self, 'embedding')
-
- @embeddings.setter
- def embeddings(self, value: 'ArrayType'):
- """Set the :attr:`.embedding` of the Documents.
-
- To remove all embeddings of all Documents:
- .. highlight:: python
- .. code-block:: python
-
- da.embeddings = None
-
- :param value: The embedding matrix to set
- """
-
- if value is None:
- for d in self:
- d.embedding = None
- else:
- emb_shape0 = _get_len(value)
- self._check_length(emb_shape0)
- ravel(value, self, 'embedding')
-
- @property
- def blobs(self) -> Optional['ArrayType']:
- """Return a :class:`ArrayType` stacking all :attr:`.blob`.
-
- The `blob` attributes are stacked together along a newly created first
- dimension (as if you would stack using ``np.stack(X, axis=0)``).
-
- .. warning:: This operation assumes all blobs have the same shape and dtype.
- All dtype and shape values are assumed to be equal to the values of the
- first element in the DocumentArray
-
- :return: a :class:`ArrayType` of blobs
- """
- if self and self[0].content_type == 'blob':
- if self:
- return unravel(self, 'blob')
-
- @blobs.setter
- def blobs(self, value: 'ArrayType'):
- """Set :attr:`.blob` of the Documents. To clear all :attr:`blob`, set it to ``None``.
-
- :param value: The blob array to set. The first axis is the "row" axis.
- """
-
- if value is None:
- for d in self:
- d.blob = None
- else:
- blobs_shape0 = _get_len(value)
- self._check_length(blobs_shape0)
-
- ravel(value, self, 'blob')
-
- @property
- def texts(self) -> Optional[List[str]]:
- """Get :attr:`.text` of all Documents
-
- :return: a list of texts
- """
- if self and self[0].content_type == 'text':
- if self:
- return [d.text for d in self]
-
- @texts.setter
- def texts(self, value: Sequence[str]):
- """Set :attr:`.text` for all Documents. To clear all :attr:`text`, set it to ``None``.
-
- :param value: A sequence of texts to set, should be the same length as the
- number of Documents
- """
- if value is None:
- for d in self:
- d.text = None
- else:
- self._check_length(len(value))
-
- for doc, text in zip(self, value):
- doc.text = text
-
- @property
- def buffers(self) -> Optional[List[bytes]]:
- """Get the buffer attribute of all Documents.
-
- :return: a list of buffers
- """
- if self and self[0].content_type == 'buffer':
- if self:
- return [d.buffer for d in self]
-
- @buffers.setter
- def buffers(self, value: List[bytes]):
- """Set the buffer attribute for all Documents. To clear all :attr:`buffer`, set it to ``None``.
-
- :param value: A sequence of buffer to set, should be the same length as the
- number of Documents
- """
-
- if value is None:
- for d in self:
- d.buffer = None
- else:
- self._check_length(len(value))
-
- for doc, buffer in zip(self, value):
- doc.buffer = buffer
-
- @property
- def contents(self) -> Optional[Union[Sequence['DocumentContentType'], 'ArrayType']]:
- """Get the :attr:`.content` of all Documents.
-
- :return: a list of texts, buffers or :class:`ArrayType`
- """
- if self:
- content_type = self[0].content_type
- if content_type:
- return getattr(self, f'{self[0].content_type}s')
-
- @contents.setter
- def contents(
- self, value: Sequence[Union[Sequence['DocumentContentType'], 'ArrayType']]
- ):
- """Set the :attr:`.content` of all Documents.
-
- :param value: a list of texts, buffers or :class:`ArrayType`
- """
- if self:
- content_type = self[0].content_type
- if content_type:
- setattr(self, f'{self[0].content_type}s', value)
-
-
-def _get_len(value):
- return len(value) if isinstance(value, (list, tuple)) else value.shape[0]
diff --git a/docarray/array/mixins/embed.py b/docarray/array/mixins/embed.py
deleted file mode 100644
index 93a354d5637..00000000000
--- a/docarray/array/mixins/embed.py
+++ /dev/null
@@ -1,150 +0,0 @@
-import warnings
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
- from ...types import T, AnyDNN
-
-
-class EmbedMixin:
- """Helper functions for embedding with a model"""
-
- def embed(
- self: 'T',
- embed_model: 'AnyDNN',
- device: str = 'cpu',
- batch_size: int = 256,
- to_numpy: bool = False,
- ) -> 'T':
- """Fill :attr:`.embedding` of Documents inplace by using `embed_model`
-
- :param embed_model: the embedding model written in Keras/Pytorch/Paddle
- :param device: the computational device for `embed_model`, can be either
- `cpu` or `cuda`.
- :param batch_size: number of Documents in a batch for embedding
- :param to_numpy: if to store embeddings back to Document in ``numpy.ndarray`` or original framework format.
- :return: itself after modified.
- """
-
- fm = get_framework(embed_model)
- getattr(self, f'_set_embeddings_{fm}')(
- embed_model, device, batch_size, to_numpy
- )
- return self
-
- def _set_embeddings_keras(
- self: 'T',
- embed_model: 'AnyDNN',
- device: str = 'cpu',
- batch_size: int = 256,
- to_numpy: bool = False,
- ):
- import tensorflow as tf
-
- device = tf.device('/GPU:0') if device == 'cuda' else tf.device('/CPU:0')
- with device:
- for b in self.batch(batch_size):
- r = embed_model(b.blobs, training=False)
- b.embeddings = r.numpy() if to_numpy else r
-
- def _set_embeddings_torch(
- self: 'T',
- embed_model: 'AnyDNN',
- device: str = 'cpu',
- batch_size: int = 256,
- to_numpy: bool = False,
- ):
- import torch
-
- embed_model = embed_model.to(device)
- is_training_before = embed_model.training
- embed_model.eval()
- with torch.inference_mode():
- for b in self.batch(batch_size):
- batch_inputs = torch.tensor(b.blobs, device=device)
- r = embed_model(batch_inputs).cpu().detach()
- b.embeddings = r.numpy() if to_numpy else r
- if is_training_before:
- embed_model.train()
-
- def _set_embeddings_paddle(
- self: 'T',
- embed_model,
- device: str = 'cpu',
- batch_size: int = 256,
- to_numpy: bool = False,
- ):
- import paddle
-
- is_training_before = embed_model.training
- embed_model.to(device=device)
- embed_model.eval()
- for b in self.batch(batch_size):
- batch_inputs = paddle.to_tensor(b.blobs, place=device)
- r = embed_model(batch_inputs)
- b.embeddings = r.numpy() if to_numpy else r
- if is_training_before:
- embed_model.train()
-
- def _set_embeddings_onnx(
- self: 'T',
- embed_model,
- device: str = 'cpu',
- batch_size: int = 256,
- *args,
- **kwargs,
- ):
- # embed_model is always an onnx.InferenceSession
- if device != 'cpu':
- import onnxruntime as ort
-
- support_device = ort.get_device()
- if device.lower().strip() != support_device.lower().strip():
- warnings.warn(
- f'Your installed `onnxruntime` supports `{support_device}`, but you give {device}'
- )
-
- for b in self.batch(batch_size):
- b.embeddings = embed_model.run(
- None, {embed_model.get_inputs()[0].name: b.blobs}
- )[0]
-
-
-def get_framework(dnn_model) -> str:
- """Return the framework that powers a DNN model.
-
- .. note::
- This is not a solid implementation. It is based on ``__module__`` name,
- the key idea is to tell ``dnn_model`` without actually importing the
- framework.
-
- :param dnn_model: a DNN model
- :return: `keras`, `torch`, `paddle` or ValueError
-
- """
- import importlib.util
-
- if importlib.util.find_spec('torch'):
- import torch
-
- if isinstance(dnn_model, torch.nn.Module):
- return 'torch'
-
- if importlib.util.find_spec('paddle'):
- import paddle
-
- if isinstance(dnn_model, paddle.nn.Layer):
- return 'paddle'
-
- if importlib.util.find_spec('tensorflow'):
- from tensorflow import keras
-
- if isinstance(dnn_model, keras.layers.Layer):
- return 'keras'
-
- if importlib.util.find_spec('onnx'):
- from onnxruntime import InferenceSession
-
- if isinstance(dnn_model, InferenceSession):
- return 'onnx'
-
- raise ValueError(f'can not determine the backend of {dnn_model!r}')
diff --git a/docarray/array/mixins/empty.py b/docarray/array/mixins/empty.py
deleted file mode 100644
index 9959562fdbb..00000000000
--- a/docarray/array/mixins/empty.py
+++ /dev/null
@@ -1,20 +0,0 @@
-from typing import Type, TYPE_CHECKING
-
-from ... import Document
-
-if TYPE_CHECKING:
- from ...types import T
-
-
-class EmptyMixin:
- """Helper functions for building arrays with empty Document."""
-
- @classmethod
- def empty(cls: Type['T'], size: int = 0) -> 'T':
- """Create a :class:`DocumentArray` object with :attr:`size` empty
- :class:`Document` objects.
-
- :param size: the number of empty Documents in this container
- :return: a :class:`DocumentArray` object
- """
- return cls(Document() for _ in range(size))
diff --git a/docarray/array/mixins/evaluation.py b/docarray/array/mixins/evaluation.py
deleted file mode 100644
index 82e9704ebea..00000000000
--- a/docarray/array/mixins/evaluation.py
+++ /dev/null
@@ -1,86 +0,0 @@
-import warnings
-from typing import Optional, Union, TYPE_CHECKING, Callable
-
-import numpy as np
-
-from ...score import NamedScore
-
-if TYPE_CHECKING:
- from ... import Document, DocumentArray
-
-
-class EvaluationMixin:
- """A mixin that provides ranking evaluation functionality to DocumentArrayLike objects"""
-
- def evaluate(
- self,
- other: 'DocumentArray',
- metric: Union[str, Callable[..., float]],
- hash_fn: Optional[Callable[['Document'], str]] = None,
- metric_name: Optional[str] = None,
- strict: bool = True,
- **kwargs,
- ) -> Optional[float]:
- """Compute ranking evaluation metrics for a given `DocumentArray` when compared with a groundtruth.
-
- This implementation expects to provide a `groundtruth` DocumentArray that is structurally identical to `self`. It is based
- on comparing the `matches` of `documents` inside the `DocumentArray.
-
- This method will fill the `evaluations` field of Documents inside this `DocumentArray` and will return the average of the computations
-
- :param other: The groundtruth DocumentArray` that the `DocumentArray` compares to.
- :param metric: The name of the metric, or multiple metrics to be computed
- :param hash_fn: The function used for identifying the uniqueness of Documents. If not given, then ``Document.id`` is used.
- :param metric_name: If provided, the results of the metrics computation will be stored in the `evaluations` field of each Document. If not provided, the name will be computed based on the metrics name.
- :param strict: If set, then left and right sides are required to be fully aligned: on the length, and on the semantic of length. These are preventing
- you to evaluate on irrelevant matches accidentally.
- :param kwargs: Additional keyword arguments to be passed to `metric_fn`
- :return: The average evaluation computed or a list of them if multiple metrics are required
- """
- if strict:
- self._check_length(len(other))
-
- if hash_fn is None:
- hash_fn = lambda d: d.id
-
- if callable(metric):
- metric_fn = metric
- elif isinstance(metric, str):
- from ...math import evaluation
-
- metric_fn = getattr(evaluation, metric)
-
- metric_name = metric_name or metric_fn.__name__
- results = []
- for d, gd in zip(self, other):
- if strict and hash_fn(d) != hash_fn(gd):
- raise ValueError(
- f'Document {d} from the left-hand side and '
- f'{gd} from the right-hand are not hashed to the same value. '
- f'This means your left and right DocumentArray may not be aligned; or it means your '
- f'`hash_fn` is badly designed.'
- )
- if not d.matches or not gd.matches:
- raise ValueError(
- f'Document {d!r} or {gd!r} has no matches, please check your Document'
- )
-
- desired = {hash_fn(m) for m in gd.matches}
- if len(desired) != len(gd.matches):
- warnings.warn(
- f'{hash_fn!r} may not be valid, as it maps multiple Documents into the same hash. '
- f'Evaluation results may be affected'
- )
-
- binary_relevance = [1 if hash_fn(m) in desired else 0 for m in d.matches]
-
- if 'max_rel' not in kwargs:
- kwargs['max_rel'] = len(gd.matches)
-
- r = metric_fn(binary_relevance, **kwargs)
- d.evaluations[metric_name] = NamedScore(
- value=r, op_name=str(metric_fn), ref_id=d.id
- )
- results.append(r)
- if results:
- return float(np.mean(results))
diff --git a/docarray/array/mixins/getattr.py b/docarray/array/mixins/getattr.py
deleted file mode 100644
index 598e20ab6a0..00000000000
--- a/docarray/array/mixins/getattr.py
+++ /dev/null
@@ -1,51 +0,0 @@
-from typing import List
-
-
-class GetAttributeMixin:
- """Helpers that provide attributes getter in bulk """
-
- def _get_attributes(self, *fields: str) -> List:
- """Return all nonempty values of the fields from all docs this array contains
-
- :param fields: Variable length argument with the name of the fields to extract
- :return: Returns a list of the values for these fields.
- When `fields` has multiple values, then it returns a list of list.
- """
- e_index, b_index = None, None
- fields = list(fields)
- if 'embedding' in fields:
- e_index = fields.index('embedding')
- if 'blob' in fields:
- b_index = fields.index('blob')
- fields.remove('blob')
-
- if 'embedding' in fields:
- fields.remove('embedding')
- if 'blob' in fields:
- fields.remove('blob')
-
- if fields:
- contents = [doc._get_attributes(*fields) for doc in self]
- if len(fields) > 1:
- contents = list(map(list, zip(*contents)))
- if b_index is None and e_index is None:
- return contents
-
- if len(fields) == 1:
- contents = [contents]
- if b_index is not None:
- contents.insert(b_index, self.blobs)
- if e_index is not None:
- contents.insert(e_index, self.embeddings)
- return contents
-
- if b_index is not None and e_index is None:
- return self.blobs
- if b_index is None and e_index is not None:
- return self.embeddings
- if b_index is not None and e_index is not None:
- return (
- [self.embeddings, self.blobs]
- if b_index > e_index
- else [self.blobs, self.embeddings]
- )
diff --git a/docarray/array/mixins/group.py b/docarray/array/mixins/group.py
deleted file mode 100644
index 17cb179ab33..00000000000
--- a/docarray/array/mixins/group.py
+++ /dev/null
@@ -1,66 +0,0 @@
-import random
-from collections import defaultdict
-from typing import Dict, Any, TYPE_CHECKING, Generator
-from ...helper import dunder_get
-import numpy as np
-
-if TYPE_CHECKING:
- from ... import DocumentArray
-
-
-class GroupMixin:
- """These helpers yield groups of :class:`DocumentArray` from
- a source :class:`DocumentArray`."""
-
- def split_by_tag(self, tag: str) -> Dict[Any, 'DocumentArray']:
- """Split the `DocumentArray` into multiple DocumentArray according to the tag value of each `Document`.
-
- :param tag: the tag name to split stored in tags.
- :return: a dict where Documents with the same value on `tag` are grouped together, their orders
- are preserved from the original :class:`DocumentArray`.
-
- .. note::
- If the :attr:`tags` of :class:`Document` do not contains the specified :attr:`tag`,
- return an empty dict.
- """
- from ... import DocumentArray
-
- rv = defaultdict(DocumentArray)
- for doc in self:
- if '__' in tag:
- value = dunder_get(doc.tags, tag)
- elif tag in doc.tags:
- value = doc.tags[tag]
- else:
- continue
- rv[value].append(doc)
- return dict(rv)
-
- def batch(
- self,
- batch_size: int,
- shuffle: bool = False,
- ) -> Generator['DocumentArray', None, None]:
- """
- Creates a `Generator` that yields `DocumentArray` of size `batch_size` until `docs` is fully traversed along
- the `traversal_path`. The None `docs` are filtered out and optionally the `docs` can be filtered by checking for
- the existence of a `Document` attribute.
- Note, that the last batch might be smaller than `batch_size`.
-
- :param batch_size: Size of each generated batch (except the last one, which might be smaller, default: 32)
- :param shuffle: If set, shuffle the Documents before dividing into minibatches.
- :yield: a Generator of `DocumentArray`, each in the length of `batch_size`
- """
-
- if not (isinstance(batch_size, int) and batch_size > 0):
- raise ValueError('`batch_size` should be a positive integer')
-
- N = len(self)
- ix = list(range(N))
- n_batches = int(np.ceil(N / batch_size))
-
- if shuffle:
- random.shuffle(ix)
-
- for i in range(n_batches):
- yield self[ix[i * batch_size : (i + 1) * batch_size]]
diff --git a/docarray/array/mixins/io/binary.py b/docarray/array/mixins/io/binary.py
deleted file mode 100644
index 044c24a6cc1..00000000000
--- a/docarray/array/mixins/io/binary.py
+++ /dev/null
@@ -1,177 +0,0 @@
-import io
-import os.path
-import pickle
-from contextlib import nullcontext
-from typing import Union, BinaryIO, TYPE_CHECKING, Type, Optional
-
-from ....helper import random_uuid, __windows__, get_compress_ctx, decompress_bytes
-
-if TYPE_CHECKING:
- from ....types import T
- from ....proto.docarray_pb2 import DocumentArrayProto
-
-
-class BinaryIOMixin:
- """Save/load an array to a binary file. """
-
- @classmethod
- def load_binary(
- cls: Type['T'],
- file: Union[str, BinaryIO, bytes],
- protocol: str = 'pickle-array',
- compress: Optional[str] = None,
- _show_progress: bool = False,
- ) -> 'T':
- """Load array elements from a LZ4-compressed binary file.
-
- :param file: File or filename or serialized bytes where the data is stored.
- :param protocol: protocol to use
- :param compress: compress algorithm to use
- :param _show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
-
- :return: a DocumentArray object
- """
-
- if isinstance(file, io.BufferedReader):
- file_ctx = nullcontext(file)
- elif isinstance(file, bytes):
- file_ctx = nullcontext(file)
- elif os.path.exists(file):
- file_ctx = open(file, 'rb')
- else:
- raise ValueError(f'unsupported input {file!r}')
-
- from .... import Document
-
- with file_ctx as fp:
- d = fp.read() if hasattr(fp, 'read') else fp
- if get_compress_ctx(algorithm=compress) is not None:
- d = decompress_bytes(d, algorithm=compress)
- compress = None
-
- if protocol == 'protobuf-array':
- from ....proto.docarray_pb2 import DocumentArrayProto
-
- dap = DocumentArrayProto()
- dap.ParseFromString(d)
-
- return cls.from_protobuf(dap)
- elif protocol == 'pickle-array':
- return pickle.loads(d)
- else:
- _len = len(random_uuid().bytes)
- _binary_delimiter = d[:_len] # first get delimiter
- if _show_progress:
- from rich.progress import track as _track
-
- track = lambda x: _track(x, description='Deserializing')
- else:
- track = lambda x: x
- return cls(
- Document.from_bytes(od, protocol=protocol, compress=compress)
- for od in track(d[_len:].split(_binary_delimiter))
- )
-
- @classmethod
- def from_bytes(
- cls: Type['T'],
- data: bytes,
- protocol: str = 'pickle-array',
- compress: Optional[str] = None,
- _show_progress: bool = False,
- ) -> 'T':
- return cls.load_binary(
- data, protocol=protocol, compress=compress, _show_progress=_show_progress
- )
-
- def save_binary(
- self,
- file: Union[str, BinaryIO],
- protocol: str = 'pickle-array',
- compress: Optional[str] = None,
- ) -> None:
- """Save array elements into a binary file.
-
- Comparing to :meth:`save_json`, it is faster and the file is smaller, but not human-readable.
-
- .. note::
- To get a binary presentation in memory, use ``bytes(...)``.
-
- :param protocol: protocol to use
- :param compress: compress algorithm to use
- :param file: File or filename to which the data is saved.
- """
- if isinstance(file, io.BufferedWriter):
- file_ctx = nullcontext(file)
- else:
- if __windows__:
- file_ctx = open(file, 'wb', newline='')
- else:
- file_ctx = open(file, 'wb')
-
- self.to_bytes(protocol=protocol, compress=compress, _file_ctx=file_ctx)
-
- def to_bytes(
- self,
- protocol: str = 'pickle-array',
- compress: Optional[str] = None,
- _file_ctx: Optional[BinaryIO] = None,
- _show_progress: bool = False,
- ) -> bytes:
- """Serialize itself into bytes.
-
- For more Pythonic code, please use ``bytes(...)``.
-
- :param _file_ctx: File or filename or serialized bytes where the data is stored.
- :param protocol: protocol to use
- :param compress: compress algorithm to use
- :param _show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
- :return: the binary serialization in bytes
- """
-
- _binary_delimiter = random_uuid().bytes
- compress_ctx = get_compress_ctx(compress, mode='wb')
- with (_file_ctx or io.BytesIO()) as bf:
- if compress_ctx is None:
- # if compress do not support streaming then postpone the compress
- # into the for-loop
- f, fc = bf, nullcontext()
- else:
- f = compress_ctx(bf)
- fc = f
- compress = None
- with fc:
- if protocol == 'protobuf-array':
- f.write(self.to_protobuf().SerializePartialToString())
- elif protocol == 'pickle-array':
- f.write(pickle.dumps(self))
- else:
- if _show_progress:
- from rich.progress import track as _track
-
- track = lambda x: _track(x, description='Serializing')
- else:
- track = lambda x: x
-
- for d in track(self):
- f.write(_binary_delimiter)
- f.write(d.to_bytes(protocol=protocol, compress=compress))
- if not _file_ctx:
- return bf.getvalue()
-
- def to_protobuf(self) -> 'DocumentArrayProto':
- from ....proto.docarray_pb2 import DocumentArrayProto
-
- dap = DocumentArrayProto()
- for d in self:
- dap.docs.append(d.to_protobuf())
- return dap
-
- @classmethod
- def from_protobuf(cls: Type['T'], pb_msg: 'DocumentArrayProto') -> 'T':
- from .... import Document
-
- return cls(Document.from_protobuf(od) for od in pb_msg.docs)
-
- def __bytes__(self):
- return self.to_bytes()
diff --git a/docarray/array/mixins/io/common.py b/docarray/array/mixins/io/common.py
deleted file mode 100644
index 17b04e6e80d..00000000000
--- a/docarray/array/mixins/io/common.py
+++ /dev/null
@@ -1,49 +0,0 @@
-from typing import Union, TextIO, BinaryIO, TYPE_CHECKING, Type
-
-if TYPE_CHECKING:
- from ....types import T
-
-
-class CommonIOMixin:
- """The common IO helper function for arrays. """
-
- def save(
- self, file: Union[str, TextIO, BinaryIO], file_format: str = 'binary'
- ) -> None:
- """Save array elements into a JSON, a binary file or a CSV file.
-
- :param file: File or filename to which the data is saved.
- :param file_format: `json` or `binary` or `csv`. JSON and CSV files are human-readable,
- but binary format gives much smaller size and faster save/load speed. Note that, CSV file has very limited
- compatability, complex DocumentArray with nested structure can not be restored from a CSV file.
- """
- if file_format == 'json':
- self.save_json(file)
- elif file_format == 'binary':
- self.save_binary(file)
- elif file_format == 'csv':
- self.save_csv(file)
- else:
- raise ValueError('`format` must be one of [`json`, `binary`, `csv`]')
-
- @classmethod
- def load(
- cls: Type['T'], file: Union[str, TextIO, BinaryIO], file_format: str = 'binary'
- ) -> 'T':
- """Load array elements from a JSON or a binary file, or a CSV file.
-
- :param file: File or filename to which the data is saved.
- :param file_format: `json` or `binary` or `csv`. JSON and CSV files are human-readable,
- but binary format gives much smaller size and faster save/load speed. CSV file has very limited compatability,
- complex DocumentArray with nested structure can not be restored from a CSV file.
-
- :return: the loaded DocumentArray object
- """
- if file_format == 'json':
- return cls.load_json(file)
- elif file_format == 'binary':
- return cls.load_binary(file)
- elif file_format == 'csv':
- return cls.load_csv(file)
- else:
- raise ValueError('`format` must be one of [`json`, `binary`, `csv`]')
diff --git a/docarray/array/mixins/io/csv.py b/docarray/array/mixins/io/csv.py
deleted file mode 100644
index c3406223b46..00000000000
--- a/docarray/array/mixins/io/csv.py
+++ /dev/null
@@ -1,104 +0,0 @@
-import csv
-from contextlib import nullcontext
-from typing import Union, TextIO, Optional, Dict, TYPE_CHECKING, Type, Sequence
-
-import numpy as np
-
-if TYPE_CHECKING:
- from ....types import T
-
-
-class CsvIOMixin:
- """CSV IO helper.
-
- can be applied to DA & DAM
- """
-
- def save_embeddings_csv(self, file: Union[str, TextIO], **kwargs) -> None:
- """Save embeddings to a CSV file
-
- This function utilizes :meth:`numpy.savetxt` internal.
-
- :param file: File or filename to which the data is saved.
- :param kwargs: extra kwargs will be passed to :meth:`numpy.savetxt`.
- """
- if hasattr(file, 'write'):
- file_ctx = nullcontext(file)
- else:
- file_ctx = open(file, 'w')
- np.savetxt(file_ctx, self.embeddings, **kwargs)
-
- def save_csv(
- self,
- file: Union[str, TextIO],
- flatten_tags: bool = True,
- exclude_fields: Optional[Sequence[str]] = None,
- dialect: Union[str, 'csv.Dialect'] = 'excel',
- with_header: bool = True,
- ) -> None:
- """Save array elements into a CSV file.
-
- :param file: File or filename to which the data is saved.
- :param flatten_tags: if set, then all fields in ``Document.tags`` will be flattened into ``tag__fieldname`` and
- stored as separated columns. It is useful when ``tags`` contain a lot of information.
- :param exclude_fields: if set, those fields wont show up in the output CSV
- :param dialect: define a set of parameters specific to a particular CSV dialect. could be a string that represents
- predefined dialects in your system, or could be a :class:`csv.Dialect` class that groups specific formatting
- parameters together.
- """
- if hasattr(file, 'write'):
- file_ctx = nullcontext(file)
- else:
- file_ctx = open(file, 'w')
-
- with file_ctx as fp:
- if flatten_tags and self[0].tags:
- keys = list(self[0].non_empty_fields) + list(
- f'tag__{k}' for k in self[0].tags
- )
- keys.remove('tags')
- else:
- flatten_tags = False
- keys = list(self[0].non_empty_fields)
-
- if exclude_fields:
- for k in exclude_fields:
- if k in keys:
- keys.remove(k)
-
- writer = csv.DictWriter(fp, fieldnames=keys, dialect=dialect)
-
- if with_header:
- writer.writeheader()
-
- from .... import Document
-
- for d in self:
- _d = d
- if exclude_fields:
- _d = Document(d, copy=True)
- _d.pop(*exclude_fields)
-
- pd = _d.to_dict()
- if flatten_tags:
- t = pd.pop('tags')
- pd.update({f'tag__{k}': v for k, v in t.items()})
- writer.writerow(pd)
-
- @classmethod
- def load_csv(
- cls: Type['T'],
- file: Union[str, TextIO],
- field_resolver: Optional[Dict[str, str]] = None,
- ) -> 'T':
- """Load array elements from a binary file.
-
- :param file: File or filename to which the data is saved.
- :param field_resolver: a map from field names defined in JSON, dict to the field
- names defined in Document.
- :return: a DocumentArray object
- """
-
- from ....document.generators import from_csv
-
- return cls(from_csv(file, field_resolver=field_resolver))
diff --git a/docarray/array/mixins/io/dataframe.py b/docarray/array/mixins/io/dataframe.py
deleted file mode 100644
index 65d2fd7637e..00000000000
--- a/docarray/array/mixins/io/dataframe.py
+++ /dev/null
@@ -1,43 +0,0 @@
-from typing import TYPE_CHECKING, Type
-
-if TYPE_CHECKING:
- from pandas import DataFrame
- from ....types import T
-
-
-class DataframeIOMixin:
- """Save/load from :class:`pandas.dataframe`
-
- .. note::
- These functions require you to install `pandas`
-
- """
-
- def to_dataframe(self, **kwargs) -> 'DataFrame':
- """Export itself to a :class:`pandas.DataFrame` object.
-
- :param kwargs: the extra kwargs will be passed to :meth:`pandas.DataFrame.from_dict`.
- :return: a :class:`pandas.DataFrame` object
- """
- from pandas import DataFrame
-
- return DataFrame.from_dict(self.to_list(), **kwargs)
-
- @classmethod
- def from_dataframe(cls: Type['T'], df: 'DataFrame') -> 'T':
- """Import a :class:`DocumentArray` from a :class:`pandas.DataFrame` object.
-
- :param df: a :class:`pandas.DataFrame` object.
- :return: a :class:`DocumentArray` object
- """
- da = cls()
- from .... import Document
-
- for m in df.to_dict(orient='records'):
- # drop nan
- da.append(
- Document(
- {k: v for k, v in m.items() if (not isinstance(v, float) or v == v)}
- )
- )
- return da
diff --git a/docarray/array/mixins/io/from_gen.py b/docarray/array/mixins/io/from_gen.py
deleted file mode 100644
index e038b526ffb..00000000000
--- a/docarray/array/mixins/io/from_gen.py
+++ /dev/null
@@ -1,229 +0,0 @@
-from typing import (
- Type,
- TYPE_CHECKING,
- Optional,
- overload,
- Union,
- List,
- TextIO,
- Dict,
- Iterable,
-)
-
-if TYPE_CHECKING:
- from ....types import T
- import numpy as np
- import csv
-
-
-class FromGeneratorMixin:
- """Provide helper functions filling a :class:`DocumentArray`-like object with a generator."""
-
- @classmethod
- def _from_generator(cls: Type['T'], meth: str, *args, **kwargs) -> 'T':
- from ....document import generators
-
- from_fn = getattr(generators, meth)
- da_like = cls()
- da_like.extend(from_fn(*args, **kwargs))
- return da_like
-
- @classmethod
- @overload
- def from_ndarray(
- cls: Type['T'],
- array: 'np.ndarray',
- axis: int = 0,
- size: Optional[int] = None,
- shuffle: bool = False,
- ) -> 'T':
- """Build from a numpy array.
-
- :param array: the numpy ndarray data source
- :param axis: iterate over that axis
- :param size: the maximum number of the sub arrays
- :param shuffle: shuffle the numpy data source beforehand
- """
- ...
-
- @classmethod
- def from_ndarray(cls: Type['T'], *args, **kwargs) -> 'T':
- """
- # noqa: DAR101
- # noqa: DAR102
- # noqa: DAR201
- """
- return cls._from_generator('from_ndarray', *args, **kwargs)
-
- @classmethod
- @overload
- def from_files(
- cls: Type['T'],
- patterns: Union[str, List[str]],
- recursive: bool = True,
- size: Optional[int] = None,
- sampling_rate: Optional[float] = None,
- read_mode: Optional[str] = None,
- to_dataturi: bool = False,
- ) -> 'T':
- """Build from a list of file path or the content of the files.
-
- :param patterns: The pattern may contain simple shell-style wildcards, e.g. '\*.py', '[\*.zip, \*.gz]'
- :param recursive: If recursive is true, the pattern '**' will match any files
- and zero or more directories and subdirectories
- :param size: the maximum number of the files
- :param sampling_rate: the sampling rate between [0, 1]
- :param read_mode: specifies the mode in which the file is opened.
- 'r' for reading in text mode, 'rb' for reading in binary mode.
- If `read_mode` is None, will iterate over filenames.
- :param to_dataturi: if set, then the Document.uri will be filled with DataURI instead of the plan URI
- """
- ...
-
- @classmethod
- def from_files(cls: Type['T'], *args, **kwargs) -> 'T':
- """
- # noqa: DAR101
- # noqa: DAR102
- # noqa: DAR201
- """
- return cls._from_generator('from_files', *args, **kwargs)
-
- @classmethod
- @overload
- def from_csv(
- cls: Type['T'],
- file: Union[str, TextIO],
- field_resolver: Optional[Dict[str, str]] = None,
- size: Optional[int] = None,
- sampling_rate: Optional[float] = None,
- dialect: Union[str, 'csv.Dialect'] = 'excel',
- ) -> 'T':
- """Build from CSV.
-
- :param file: file paths or file handler
- :param field_resolver: a map from field names defined in JSON, dict to the field
- names defined in Document.
- :param size: the maximum number of the documents
- :param sampling_rate: the sampling rate between [0, 1]
- :param dialect: define a set of parameters specific to a particular CSV dialect. could be a string that represents
- predefined dialects in your system, or could be a :class:`csv.Dialect` class that groups specific formatting
- parameters together. If you don't know the dialect and the default one does not work for you,
- you can try set it to ``auto``.
- """
- ...
-
- @classmethod
- def from_csv(cls: Type['T'], *args, **kwargs) -> 'T':
- """
- # noqa: DAR101
- # noqa: DAR102
- # noqa: DAR201
- """
- return cls._from_generator('from_csv', *args, **kwargs)
-
- @classmethod
- @overload
- def from_huggingface_datasets(
- cls: Type['T'],
- dataset_path: str,
- field_resolver: Optional[Dict[str, str]] = None,
- size: Optional[int] = None,
- sampling_rate: Optional[float] = None,
- filter_fields: bool = False,
- **datasets_kwargs,
- ) -> 'T':
- """Build from Hugging Face Datasets. Yields documents.
-
- This function helps to load datasets from Hugging Face Datasets Hub
- (https://huggingface.co/datasets) in Jina. Additional parameters can be
- passed to the ``datasets`` library using keyword arguments. The ``load_dataset``
- method from ``datasets`` library is used to load the datasets.
-
- :param dataset_path: a valid dataset path for Hugging Face Datasets library.
- :param field_resolver: a map from field names defined in ``document`` (JSON, dict) to the field
- names defined in Protobuf. This is only used when the given ``document`` is
- a JSON string or a Python dict.
- :param size: the maximum number of the documents
- :param sampling_rate: the sampling rate between [0, 1]
- :param filter_fields: specifies whether to filter the dataset with the fields
- given in ```field_resolver`` argument.
- :param **datasets_kwargs: additional arguments for ``load_dataset`` method
- from Datasets library. More details at
- https://huggingface.co/docs/datasets/package_reference/loading_methods.html#datasets.load_dataset
- """
- ...
-
- @classmethod
- def from_huggingface_datasets(cls: Type['T'], *args, **kwargs) -> 'T':
- """
- # noqa: DAR101
- # noqa: DAR102
- # noqa: DAR201
- """
- return cls._from_generator('from_huggingface_datasets', *args, **kwargs)
-
- @classmethod
- @overload
- def from_ndjson(
- cls: Type['T'],
- fp: Iterable[str],
- field_resolver: Optional[Dict[str, str]] = None,
- size: Optional[int] = None,
- sampling_rate: Optional[float] = None,
- ) -> 'T':
- """Build from line separated JSON. Yields documents.
-
- :param fp: file paths
- :param field_resolver: a map from field names defined in ``document`` (JSON, dict) to the field
- names defined in Protobuf. This is only used when the given ``document`` is
- a JSON string or a Python dict.
- :param size: the maximum number of the documents
- :param sampling_rate: the sampling rate between [0, 1]
- """
- ...
-
- @classmethod
- def from_ndjson(cls: Type['T'], *args, **kwargs) -> 'T':
- """
- # noqa: DAR101
- # noqa: DAR102
- # noqa: DAR201
- """
- return cls._from_generator('from_ndjson', *args, **kwargs)
-
- @classmethod
- @overload
- def from_lines(
- cls: 'T',
- lines: Optional[Iterable[str]] = None,
- filepath: Optional[str] = None,
- read_mode: str = 'r',
- line_format: str = 'json',
- field_resolver: Optional[Dict[str, str]] = None,
- size: Optional[int] = None,
- sampling_rate: Optional[float] = None,
- ) -> 'T':
- """Build from lines, json and csv. Yields documents or strings.
-
- :param lines: a list of strings, each is considered as a document
- :param filepath: a text file that each line contains a document
- :param read_mode: specifies the mode in which the file
- is opened. 'r' for reading in text mode, 'rb' for reading in binary
- :param line_format: the format of each line ``json`` or ``csv``
- :param field_resolver: a map from field names defined in ``document`` (JSON, dict) to the field
- names defined in Protobuf. This is only used when the given ``document`` is
- a JSON string or a Python dict.
- :param size: the maximum number of the documents
- :param sampling_rate: the sampling rate between [0, 1]
- """
- ...
-
- @classmethod
- def from_lines(cls: Type['T'], *args, **kwargs) -> 'T':
- """
- # noqa: DAR101
- # noqa: DAR102
- # noqa: DAR201
- """
- return cls._from_generator('from_lines', *args, **kwargs)
diff --git a/docarray/array/mixins/io/json.py b/docarray/array/mixins/io/json.py
deleted file mode 100644
index 4b962632cde..00000000000
--- a/docarray/array/mixins/io/json.py
+++ /dev/null
@@ -1,78 +0,0 @@
-import json
-import os.path
-from contextlib import nullcontext
-from typing import Union, TextIO, TYPE_CHECKING, Type, List
-
-if TYPE_CHECKING:
- from ....types import T
-
-
-class JsonIOMixin:
- """Save/load a array into a JSON file."""
-
- def save_json(self, file: Union[str, TextIO]) -> None:
- """Save array elements into a JSON file.
-
- Comparing to :meth:`save_binary`, it is human-readable but slower to save/load and the file size larger.
-
- :param file: File or filename to which the data is saved.
- """
- if hasattr(file, 'write'):
- file_ctx = nullcontext(file)
- else:
- file_ctx = open(file, 'w')
-
- with file_ctx as fp:
- for d in self:
- json.dump(d.to_dict(), fp)
- fp.write('\n')
-
- @classmethod
- def load_json(cls: Type['T'], file: Union[str, TextIO]) -> 'T':
- """Load array elements from a JSON file.
-
- :param file: File or filename or a JSON string to which the data is saved.
-
- :return: a DocumentArrayLike object
- """
-
- from .... import Document
-
- constructor = Document.from_json
- if hasattr(file, 'read'):
- file_ctx = nullcontext(file)
- elif os.path.exists(file):
- file_ctx = open(file)
- else:
- file_ctx = nullcontext(json.loads(file))
- constructor = Document.from_dict
-
- with file_ctx as fp:
- return cls(constructor(v) for v in fp)
-
- @classmethod
- def from_json(cls: Type['T'], file: Union[str, TextIO]) -> 'T':
- return cls.load_json(file)
-
- @classmethod
- def from_list(cls: Type['T'], values: List) -> 'T':
- from .... import Document
-
- return cls(Document.from_dict(v) for v in values)
-
- def to_list(self, strict: bool = True) -> List:
- """Convert the object into a Python list.
-
- .. note::
- Array like object such as :class:`numpy.ndarray` will be converted to Python list.
-
- :return: a Python list
- """
- return [d.to_dict(strict=strict) for d in self]
-
- def to_json(self) -> str:
- """Convert the object into a JSON string. Can be loaded via :meth:`.load_json`.
-
- :return: a Python list
- """
- return json.dumps(self.to_list())
diff --git a/docarray/array/mixins/io/pushpull.py b/docarray/array/mixins/io/pushpull.py
deleted file mode 100644
index 77fbed2cf95..00000000000
--- a/docarray/array/mixins/io/pushpull.py
+++ /dev/null
@@ -1,155 +0,0 @@
-import io
-from contextlib import nullcontext
-from typing import Type, TYPE_CHECKING, Optional
-
-from ....helper import get_request_header
-
-if TYPE_CHECKING:
- from ....types import T
-
-
-class PushPullMixin:
- """Transmitting :class:`DocumentArray` via Jina Cloud Service"""
-
- _service_url = 'https://apihubble.jina.ai/v2/rpc/da.'
- _max_bytes = 4 * 1024 * 1024 * 1024
-
- def push(self, token: str, show_progress: bool = False) -> None:
- """Push this DocumentArray object to Jina Cloud which can be later retrieved via :meth:`.push`
-
- .. note::
- - Push with the same ``token`` will override the existing content.
- - Kinda like a public clipboard where everyone can override anyone's content.
- So to make your content survive longer, you may want to use longer & more complicated token.
- - The lifetime of the content is not promised atm, could be a day, could be a week. Do not use it for
- persistence. Only use this full temporary transmission/storage/clipboard.
-
- :param token: a key that later can be used for retrieve this :class:`DocumentArray`.
- :param show_progress: if to show a progress bar on pulling
- """
- import requests
-
- dict_data = self._get_dict_data(token, show_progress)
-
- progress = _get_progressbar(show_progress)
- task_id = progress.add_task('upload', start=False) if show_progress else None
-
- class BufferReader(io.BytesIO):
- def __init__(self, buf=b'', p_bar=None, task_id=None):
- super().__init__(buf)
- self._len = len(buf)
- self._p_bar = p_bar
- self._task_id = task_id
- if show_progress:
- progress.update(task_id, total=self._len)
- progress.start_task(task_id)
-
- def __len__(self):
- return self._len
-
- def read(self, n=-1):
- chunk = io.BytesIO.read(self, n)
- if self._p_bar:
- self._p_bar.update(self._task_id, advance=len(chunk))
- return chunk
-
- (data, ctype) = requests.packages.urllib3.filepost.encode_multipart_formdata(
- dict_data
- )
-
- headers = {'Content-Type': ctype, **get_request_header()}
-
- with progress as p_bar:
- body = BufferReader(data, p_bar, task_id)
- requests.post(self._service_url + 'push', data=body, headers=headers)
-
- @classmethod
- def pull(
- cls: Type['T'],
- token: str,
- show_progress: bool = False,
- ) -> 'T':
- """Pulling a :class:`DocumentArray` from Jina Cloud Service to local.
-
- :param token: the upload token set during :meth:`.push`
- :param show_progress: if to show a progress bar on pulling
- :return: a :class:`DocumentArray` object
- """
- import requests
-
- url = f'{cls._service_url}pull?token={token}'
- response = requests.get(url)
-
- progress = _get_progressbar(show_progress)
-
- url = response.json()['data']['download']
-
- with requests.get(
- url,
- stream=True,
- headers=get_request_header(),
- ) as r, progress:
- r.raise_for_status()
- if show_progress:
- task_id = progress.add_task('download', start=False)
- progress.update(task_id, total=int(r.headers['Content-length']))
- with io.BytesIO() as f:
- chunk_size = 8192
- if show_progress:
- progress.start_task(task_id)
- for chunk in r.iter_content(chunk_size=chunk_size):
- f.write(chunk)
- if show_progress:
- progress.update(task_id, advance=len(chunk))
-
- if show_progress:
- progress.stop()
- return cls.from_bytes(
- f.getvalue(),
- protocol='protobuf',
- compress='gzip',
- _show_progress=show_progress,
- )
-
- def _get_dict_data(self, token, show_progress):
- _serialized = self.to_bytes(
- protocol='protobuf', compress='gzip', _show_progress=show_progress
- )
- if len(_serialized) > self._max_bytes:
- raise ValueError(
- f'DocumentArray is too big. '
- f'Size of the serialization {len(_serialized)} is larger than {self._max_bytes}.'
- )
-
- return {
- 'file': (
- 'DocumentArray',
- _serialized,
- ),
- 'token': token,
- }
-
-
-def _get_progressbar(show_progress):
- if show_progress:
- from rich.progress import (
- BarColumn,
- DownloadColumn,
- Progress,
- TimeRemainingColumn,
- TransferSpeedColumn,
- )
-
- return Progress(
- BarColumn(bar_width=None),
- "[progress.percentage]{task.percentage:>3.1f}%",
- "•",
- DownloadColumn(),
- "•",
- TransferSpeedColumn(),
- "•",
- TimeRemainingColumn(),
- transient=True,
- )
- else:
- return nullcontext()
diff --git a/docarray/array/mixins/match.py b/docarray/array/mixins/match.py
deleted file mode 100644
index 930454e18e5..00000000000
--- a/docarray/array/mixins/match.py
+++ /dev/null
@@ -1,234 +0,0 @@
-from typing import Optional, Union, Callable, Tuple, TYPE_CHECKING
-
-import numpy as np
-
-from ...math.helper import top_k, minmax_normalize, update_rows_x_mat_best
-from ...score import NamedScore
-
-if TYPE_CHECKING:
- from ...types import Document, ArrayType
- from ... import DocumentArray
-
-
-class MatchMixin:
- """ A mixin that provides match functionality to DocumentArrays """
-
- def match(
- self,
- darray: 'DocumentArray',
- metric: Union[
- str, Callable[['ArrayType', 'ArrayType'], 'np.ndarray']
- ] = 'cosine',
- limit: Optional[Union[int, float]] = 20,
- normalization: Optional[Tuple[float, float]] = None,
- metric_name: Optional[str] = None,
- batch_size: Optional[int] = None,
- exclude_self: bool = False,
- only_id: bool = False,
- use_scipy: bool = False,
- device: str = 'cpu',
- num_worker: Optional[int] = 1,
- **kwargs,
- ) -> None:
- """Compute embedding based nearest neighbour in `another` for each Document in `self`,
- and store results in `matches`.
- .. note::
- 'cosine', 'euclidean', 'sqeuclidean' are supported natively without extra dependency.
- You can use other distance metric provided by ``scipy``, such as `braycurtis`, `canberra`, `chebyshev`,
- `cityblock`, `correlation`, `cosine`, `dice`, `euclidean`, `hamming`, `jaccard`, `jensenshannon`,
- `kulsinski`, `mahalanobis`, `matching`, `minkowski`, `rogerstanimoto`, `russellrao`, `seuclidean`,
- `sokalmichener`, `sokalsneath`, `sqeuclidean`, `wminkowski`, `yule`.
- To use scipy metric, please set ``use_scipy=True``.
- - To make all matches values in [0, 1], use ``dA.match(dB, normalization=(0, 1))``
- - To invert the distance as score and make all values in range [0, 1],
- use ``dA.match(dB, normalization=(1, 0))``. Note, how ``normalization`` differs from the previous.
- - If a custom metric distance is provided. Make sure that it returns scores as distances and not similarity, meaning the smaller the better.
- :param darray: the other DocumentArray to match against
- :param metric: the distance metric
- :param limit: the maximum number of matches, when not given defaults to 20.
- :param normalization: a tuple [a, b] to be used with min-max normalization,
- the min distance will be rescaled to `a`, the max distance will be rescaled to `b`
- all values will be rescaled into range `[a, b]`.
- :param metric_name: if provided, then match result will be marked with this string.
- :param batch_size: if provided, then ``darray`` is loaded in batches, where each of them is at most ``batch_size``
- elements. When `darray` is big, this can significantly speedup the computation.
- :param exclude_self: if set, Documents in ``darray`` with same ``id`` as the left-hand values will not be
- considered as matches.
- :param only_id: if set, then returning matches will only contain ``id``
- :param use_scipy: if set, use ``scipy`` as the computation backend. Note, ``scipy`` does not support distance
- on sparse matrix.
- :param device: the computational device for ``.match()``, can be either `cpu` or `cuda`.
- :param num_worker: the number of parallel workers. If not given, then the number of CPUs in the system will be used.
-
- .. note::
- This argument is only effective when ``batch_size`` is set.
-
- :param kwargs: other kwargs.
- """
- if limit is not None:
- if limit <= 0:
- raise ValueError(f'`limit` must be larger than 0, receiving {limit}')
- else:
- limit = int(limit)
-
- if batch_size is not None:
- if batch_size <= 0:
- raise ValueError(
- f'`batch_size` must be larger than 0, receiving {batch_size}'
- )
- else:
- batch_size = int(batch_size)
-
- lhv = self
- rhv = darray
-
- if not (lhv and rhv):
- return
-
- if callable(metric):
- cdist = metric
- elif isinstance(metric, str):
- if use_scipy:
- from scipy.spatial.distance import cdist as cdist
- else:
- from ...math.distance import cdist as _cdist
-
- cdist = lambda *x: _cdist(*x, device=device)
- else:
- raise TypeError(
- f'metric must be either string or a 2-arity function, received: {metric!r}'
- )
-
- metric_name = metric_name or (metric.__name__ if callable(metric) else metric)
- _limit = len(rhv) if limit is None else (limit + (1 if exclude_self else 0))
-
- if batch_size:
- dist, idx = lhv._match_online(
- rhv, cdist, _limit, normalization, metric_name, batch_size, num_worker
- )
- else:
- dist, idx = lhv._match(rhv, cdist, _limit, normalization, metric_name)
-
- from ... import Document
-
- for _q, _ids, _dists in zip(lhv, idx, dist):
- _q.matches.clear()
- num_matches = 0
- for _id, _dist in zip(_ids, _dists):
- # Note, when match self with other, or both of them share the same Document
- # we might have recursive matches .
- # checkout https://github.com/jina-ai/jina/issues/3034
- if only_id:
- d = Document(id=rhv[_id].id)
- else:
- d = rhv[int(_id)] # type: Document
-
- if d.id in lhv:
- d = Document(
- d, copy=True
- ) # to prevent self-reference and override on matches
- d.pop('matches')
- if not (d.id == _q.id and exclude_self):
- d.scores[metric_name] = NamedScore(value=_dist, ref_id=_q.id)
- _q.matches.append(d)
- num_matches += 1
- if num_matches >= (limit or _limit):
- break
-
- def _match(self, darray, cdist, limit, normalization, metric_name):
- """
- Computes the matches between self and `darray` loading `darray` into main memory.
- :param darray: the other DocumentArray or to match against
- :param cdist: the distance metric
- :param limit: the maximum number of matches, when not given
- all Documents in `darray` are considered as matches
- :param normalization: a tuple [a, b] to be used with min-max normalization,
- the min distance will be rescaled to `a`, the max distance will be rescaled to `b`
- all values will be rescaled into range `[a, b]`.
- :param metric_name: if provided, then match result will be marked with this string.
- :return: distances and indices
- """
-
- x_mat = self.embeddings
- y_mat = darray.embeddings
-
- dists = cdist(x_mat, y_mat, metric_name)
- dist, idx = top_k(dists, min(limit, len(darray)), descending=False)
- if isinstance(normalization, (tuple, list)) and normalization is not None:
- # normalization bound uses original distance not the top-k trimmed distance
- min_d = np.min(dists, axis=-1, keepdims=True)
- max_d = np.max(dists, axis=-1, keepdims=True)
- dist = minmax_normalize(dist, normalization, (min_d, max_d))
-
- return dist, idx
-
- def _match_online(
- self,
- darray,
- cdist,
- limit,
- normalization,
- metric_name,
- batch_size,
- num_worker,
- ):
- """
- Computes the matches between self and `darray` loading `darray` into main memory in chunks of size `batch_size`.
-
- :param darray: the other DocumentArray or to match against
- :param cdist: the distance metric
- :param limit: the maximum number of matches, when not given
- all Documents in `another` are considered as matches
- :param normalization: a tuple [a, b] to be used with min-max normalization,
- the min distance will be rescaled to `a`, the max distance will be rescaled to `b`
- all values will be rescaled into range `[a, b]`.
- :param batch_size: length of the chunks loaded into memory from darray.
- :param metric_name: if provided, then match result will be marked with this string.
- :param num_worker: the number of parallel workers. If not given, then the number of CPUs in the system will be used.
- :return: distances and indices
- """
-
- x_mat = self.embeddings
- n_x = x_mat.shape[0]
-
- idx = 0
- top_dists = np.inf * np.ones((n_x, limit))
- top_inds = np.zeros((n_x, limit), dtype=int)
-
- def _get_dist(da: 'DocumentArray'):
- y_batch = da.embeddings
-
- distances = cdist(x_mat, y_batch, metric_name)
- dists, inds = top_k(distances, limit, descending=False)
-
- if isinstance(normalization, (tuple, list)) and normalization is not None:
- dists = minmax_normalize(dists, normalization)
-
- return dists, inds, y_batch.shape[0]
-
- if num_worker is None or num_worker > 1:
- # notice that all most all computations (regardless the framework) are conducted in C
- # hence there is no worry on Python GIL and the backend can be safely put to `thread` to
- # save unnecessary data passing. This in fact gives a huge boost on the performance.
- _gen = darray.map_batch(
- _get_dist,
- batch_size=batch_size,
- backend='thread',
- num_worker=num_worker,
- )
- else:
- _gen = (_get_dist(b) for b in darray.batch(batch_size=batch_size))
-
- for (dists, inds, _bs) in _gen:
- inds += idx
- idx += _bs
- top_dists, top_inds = update_rows_x_mat_best(
- top_dists, top_inds, dists, inds, limit
- )
-
- # sort final the final `top_dists` and `top_inds` per row
- permutation = np.argsort(top_dists, axis=1)
- dist = np.take_along_axis(top_dists, permutation, axis=1)
- idx = np.take_along_axis(top_inds, permutation, axis=1)
-
- return dist, idx
diff --git a/docarray/array/mixins/parallel.py b/docarray/array/mixins/parallel.py
deleted file mode 100644
index eaaf5d66815..00000000000
--- a/docarray/array/mixins/parallel.py
+++ /dev/null
@@ -1,200 +0,0 @@
-import sys
-from types import LambdaType
-from typing import Callable, TYPE_CHECKING, Generator, Optional, overload, TypeVar
-
-if TYPE_CHECKING:
- from ...types import T
- from ... import Document, DocumentArray
-
-T_DA = TypeVar('T_DA')
-
-
-class ParallelMixin:
- """Helper functions that provide parallel map to :class:`DocumentArray`"""
-
- @overload
- def apply(
- self: 'T',
- func: Callable[['Document'], 'Document'],
- backend: str = 'thread',
- num_worker: Optional[int] = None,
- ) -> 'T':
- """Apply each element in itself with ``func``, return itself after modified.
-
- :param func: a function that takes :class:`Document` as input and outputs :class:`Document`.
- :param backend: if to use multi-`process` or multi-`thread` as the parallelization backend. In general, if your
- ``func`` is IO-bound then perhaps `thread` is good enough. If your ``func`` is CPU-bound then you may use `process`.
- In practice, you should try yourselves to figure out the best value. However, if you wish to modify the elements
- in-place, regardless of IO/CPU-bound, you should always use `thread` backend.
-
- .. warning::
- When using `process` backend, you should not expect ``func`` modify elements in-place. This is because
- the multiprocessing backing pass the variable via pickle and work in another process. The passed object
- and the original object do **not** share the same memory.
-
- :param num_worker: the number of parallel workers. If not given, then the number of CPUs in the system will be used.
-
- """
- ...
-
- def apply(self: 'T', *args, **kwargs) -> 'T':
- """
- # noqa: DAR102
- # noqa: DAR101
- # noqa: DAR201
- :return: a new :class:`DocumentArray`
- """
- from ... import DocumentArray
-
- new_da = DocumentArray(self.map(*args, **kwargs))
- self.clear()
- self.extend(new_da)
- return self
-
- def map(
- self,
- func: Callable[['Document'], 'T'],
- backend: str = 'thread',
- num_worker: Optional[int] = None,
- ) -> Generator['T', None, None]:
- """Return an iterator that applies function to every **element** of iterable in parallel, yielding the results.
-
- .. seealso::
- - To process on a batch of elements, please use :meth:`.map_batch`;
- - To return a :class:`DocumentArray`, please use :meth:`.apply`.
-
- :param func: a function that takes :class:`Document` as input and outputs anything. You can either modify elements
- in-place (only with `thread` backend) or work later on return elements.
- :param backend: if to use multi-`process` or multi-`thread` as the parallelization backend. In general, if your
- ``func`` is IO-bound then perhaps `thread` is good enough. If your ``func`` is CPU-bound then you may use `process`.
- In practice, you should try yourselves to figure out the best value. However, if you wish to modify the elements
- in-place, regardless of IO/CPU-bound, you should always use `thread` backend.
-
- .. warning::
- When using `process` backend, you should not expect ``func`` modify elements in-place. This is because
- the multiprocessing backing pass the variable via pickle and work in another process. The passed object
- and the original object do **not** share the same memory.
-
- :param num_worker: the number of parallel workers. If not given, then the number of CPUs in the system will be used.
- :yield: anything return from ``func``
- """
- if _is_lambda_or_local_function(func) and backend == 'process':
- func = _globalize_lambda_function(func)
- with _get_pool(backend, num_worker) as p:
- for x in p.imap(func, self):
- yield x
-
- @overload
- def apply_batch(
- self: 'T',
- func: Callable[['DocumentArray'], 'DocumentArray'],
- batch_size: int,
- backend: str = 'thread',
- num_worker: Optional[int] = None,
- shuffle: bool = False,
- ) -> 'T':
- """Apply each element in itself with ``func``, return itself after modified.
-
- :param func: a function that takes :class:`Document` as input and outputs :class:`Document`.
- :param backend: if to use multi-`process` or multi-`thread` as the parallelization backend. In general, if your
- ``func`` is IO-bound then perhaps `thread` is good enough. If your ``func`` is CPU-bound then you may use `process`.
- In practice, you should try yourselves to figure out the best value. However, if you wish to modify the elements
- in-place, regardless of IO/CPU-bound, you should always use `thread` backend.
-
- .. warning::
- When using `process` backend, you should not expect ``func`` modify elements in-place. This is because
- the multiprocessing backing pass the variable via pickle and work in another process. The passed object
- and the original object do **not** share the same memory.
-
- :param num_worker: the number of parallel workers. If not given, then the number of CPUs in the system will be used.
- :param batch_size: Size of each generated batch (except the last one, which might be smaller, default: 32)
- :param shuffle: If set, shuffle the Documents before dividing into minibatches.
- """
- ...
-
- def apply_batch(self: 'T', *args, **kwargs) -> 'T':
- """
- # noqa: DAR102
- # noqa: DAR101
- # noqa: DAR201
- :return: a new :class:`DocumentArray`
- """
- from ... import DocumentArray
-
- new_da = DocumentArray()
- for _b in self.map_batch(*args, **kwargs):
- new_da.extend(_b)
- self.clear()
- self.extend(new_da)
- return self
-
- def map_batch(
- self: 'T_DA',
- func: Callable[['DocumentArray'], 'T'],
- batch_size: int,
- backend: str = 'thread',
- num_worker: Optional[int] = None,
- shuffle: bool = False,
- ) -> Generator['T', None, None]:
- """Return an iterator that applies function to every **minibatch** of iterable in parallel, yielding the results.
- Each element in the returned iterator is :class:`DocumentArray`.
-
- .. seealso::
- - To process single element, please use :meth:`.map`;
- - To return :class:`DocumentArray`, please use :meth:`.apply_batch`.
-
- :param batch_size: Size of each generated batch (except the last one, which might be smaller, default: 32)
- :param shuffle: If set, shuffle the Documents before dividing into minibatches.
- :param func: a function that takes :class:`DocumentArray` as input and outputs anything. You can either modify elements
- in-place (only with `thread` backend) or work later on return elements.
- :param backend: if to use multi-`process` or multi-`thread` as the parallelization backend. In general, if your
- ``func`` is IO-bound then perhaps `thread` is good enough. If your ``func`` is CPU-bound then you may use `process`.
- In practice, you should try yourselves to figure out the best value. However, if you wish to modify the elements
- in-place, regardless of IO/CPU-bound, you should always use `thread` backend.
-
- .. warning::
- When using `process` backend, you should not expect ``func`` modify elements in-place. This is because
- the multiprocessing backing pass the variable via pickle and work in another process. The passed object
- and the original object do **not** share the same memory.
-
- :param num_worker: the number of parallel workers. If not given, then the number of CPUs in the system will be used.
- :yield: anything return from ``func``
- """
-
- if _is_lambda_or_local_function(func) and backend == 'process':
- func = _globalize_lambda_function(func)
- with _get_pool(backend, num_worker) as p:
- for x in p.imap(func, self.batch(batch_size=batch_size, shuffle=shuffle)):
- yield x
-
-
-def _get_pool(backend, num_worker):
- if backend == 'thread':
- from multiprocessing.pool import ThreadPool as Pool
-
- return Pool(processes=num_worker)
- elif backend == 'process':
- from multiprocessing.pool import Pool
-
- return Pool(processes=num_worker)
- else:
- raise ValueError(
- f'`backend` must be either `process` or `thread`, receiving {backend}'
- )
-
-
-def _is_lambda_or_local_function(func):
- return (isinstance(func, LambdaType) and func.__name__ == 'Click to expand
+ +With DocArray, you can connect external data to LLMs through Langchain. DocArray gives you the freedom to establish +flexible document schemas and choose from different backends for document storage. +After creating your document index, you can connect it to your Langchain app using [DocArrayRetriever](https://python.langchain.com/docs/modules/data_connection/retrievers/integrations/docarray_retriever). +Install Langchain via: +```shell +pip install langchain +``` + +1. Define a schema and create documents: ```python -left_da = DocumentArray.pull(token='my_shared_da') +from docarray import BaseDoc, DocList +from docarray.typing import NdArray +from langchain.embeddings.openai import OpenAIEmbeddings + +embeddings = OpenAIEmbeddings() + + +# Define a document schema +class MovieDoc(BaseDoc): + title: str + description: str + year: int + embedding: NdArray[1536] + + +movies = [ + {"title": "#1 title", "description": "#1 description", "year": 1999}, + {"title": "#2 title", "description": "#2 description", "year": 2001}, +] + +# Embed `description` and create documents +docs = DocList[MovieDoc]( + MovieDoc(embedding=embeddings.embed_query(movie["description"]), **movie) + for movie in movies +) +``` + +2. Initialize a document index using any supported backend: +```python +from docarray.index import ( + InMemoryExactNNIndex, + HnswDocumentIndex, + WeaviateDocumentIndex, + QdrantDocumentIndex, + ElasticDocIndex, + RedisDocumentIndex, + MongoDBAtlasDocumentIndex, +) + +# Select a suitable backend and initialize it with data +db = InMemoryExactNNIndex[MovieDoc](docs) ``` -Intrigued? That's only scratching the surface of what DocArray is capable of. [Read our docs to learn more](https://docarray.jina.ai). +3. Finally, initialize a retriever and integrate it into your chain! +```python +from langchain.chat_models import ChatOpenAI +from langchain.chains import ConversationalRetrievalChain +from langchain.retrievers import DocArrayRetriever + + +# Create a retriever +retriever = DocArrayRetriever( + index=db, + embeddings=embeddings, + search_field="embedding", + content_field="description", +) + +# Use the retriever in your chain +model = ChatOpenAI() +qa = ConversationalRetrievalChain.from_llm(model, retriever=retriever) +``` +Alternatively, you can use built-in vector stores. Langchain supports two vector stores: [DocArrayInMemorySearch](https://python.langchain.com/docs/modules/data_connection/vectorstores/integrations/docarray_in_memory) and [DocArrayHnswSearch](https://python.langchain.com/docs/modules/data_connection/vectorstores/integrations/docarray_hnsw). +Both are user-friendly and are best suited to small to medium-sized datasets. - -## Support +
-
-
-{%- endif %}
-
-
-
-
-
-
-{%- endblock %}
\ No newline at end of file
diff --git a/docs/_templates/sidebar/brand.html b/docs/_templates/sidebar/brand.html
deleted file mode 100644
index 1531f31b10f..00000000000
--- a/docs/_templates/sidebar/brand.html
+++ /dev/null
@@ -1,41 +0,0 @@
-
- {% block brand_content %}
- {%- if logo_url %}
-
-
-
-
-
-
-
-
-
-
-
-
-
- 
-
- {%- endif %}
- {%- if theme_light_logo and theme_dark_logo %}
-
-  }})
-  }})
-
- {%- endif %}
- {% if not theme_sidebar_hide_name %}
-
- {%- endif %}
- {% endblock brand_content %}
-
-
- Star
- {% if versions %}
-
- {% endif %}
-
diff --git a/docs/_templates/sidebar/navigation.html b/docs/_templates/sidebar/navigation.html
deleted file mode 100644
index 7dd7c6134b0..00000000000
--- a/docs/_templates/sidebar/navigation.html
+++ /dev/null
@@ -1,23 +0,0 @@
-
- {{ furo_navigation_tree }}
-
\ No newline at end of file
diff --git a/docs/_versions.json b/docs/_versions.json
new file mode 100644
index 00000000000..f318a2796a0
--- /dev/null
+++ b/docs/_versions.json
@@ -0,0 +1 @@
+[{"version": "v0.40.1"}, {"version": "v0.40.0"}, {"version": "v0.39.1"}, {"version": "v0.39.0"}, {"version": "v0.38.0"}, {"version": "v0.37.1"}, {"version": "v0.37.0"}, {"version": "v0.36.0"}, {"version": "v0.35.0"}, {"version": "v0.34.0"}, {"version": "v0.33.0"}, {"version": "v0.32.1"}, {"version": "v0.32.0"}, {"version": "v0.31.1"}, {"version": "v0.31.0"}, {"version": "v0.30.0"}, {"version": "v0.21.0"}, {"version": "v0.20.1"}, {"version": "v0.20.0"}, {"version": "v0.19.0"}, {"version": "v0.18.1"}, {"version": "v0.18.0"}, {"version": "v0.17.0"}, {"version": "v0.16.5"}, {"version": "v0.16.4"}, {"version": "v0.16.3"}, {"version": "v0.16.2"}, {"version": "v0.16.1"}, {"version": "v0.16.0"}, {"version": "v0.15.4"}, {"version": "v0.15.3"}, {"version": "v0.15.2"}, {"version": "v0.15.1"}, {"version": "v0.15.0"}, {"version": "v0.14.11"}, {"version": "v0.14.10"}, {"version": "v0.14.9"}, {"version": "v0.14.8"}, {"version": "v0.14.7"}, {"version": "v0.14.6"}, {"version": "v0.14.5"}, {"version": "v0.14.4"}, {"version": "v0.14.3"}, {"version": "v0.14.2"}, {"version": "v0.14.1"}, {"version": "v0.14.0"}, {"version": "v0.13.33"}, {"version": "v0.13.0"}, {"version": "v0.12.9"}, {"version": "v0.12.0"}, {"version": "v0.11.3"}, {"version": "v0.11.2"}, {"version": "v0.11.1"}, {"version": "v0.11.0"}, {"version": "v0.10.5"}, {"version": "v0.10.4"}, {"version": "v0.10.3"}, {"version": "v0.10.2"}, {"version": "v0.10.1"}, {"version": "v0.10.0"}]
\ No newline at end of file
diff --git a/docs/assets/docarray-colorful.svg b/docs/assets/docarray-colorful.svg
new file mode 100644
index 00000000000..ed803d09d56
--- /dev/null
+++ b/docs/assets/docarray-colorful.svg
@@ -0,0 +1,16 @@
+
+
\ No newline at end of file
diff --git a/docs/_static/docarray-dark.svg b/docs/assets/docarray-dark.svg
similarity index 100%
rename from docs/_static/docarray-dark.svg
rename to docs/assets/docarray-dark.svg
diff --git a/docs/_static/docarray-light.svg b/docs/assets/docarray-light.svg
similarity index 100%
rename from docs/_static/docarray-light.svg
rename to docs/assets/docarray-light.svg
diff --git a/docs/_static/favicon.png b/docs/assets/favicon.png
similarity index 100%
rename from docs/_static/favicon.png
rename to docs/assets/favicon.png
diff --git a/docs/_static/logo-light.svg b/docs/assets/logo-dark.svg
similarity index 100%
rename from docs/_static/logo-light.svg
rename to docs/assets/logo-dark.svg
diff --git a/docs/_static/logo-dark.svg b/docs/assets/logo-light.svg
similarity index 100%
rename from docs/_static/logo-dark.svg
rename to docs/assets/logo-light.svg
diff --git a/docs/conf.py b/docs/conf.py
deleted file mode 100644
index bf986e8d96f..00000000000
--- a/docs/conf.py
+++ /dev/null
@@ -1,227 +0,0 @@
-import os
-import re
-import sys
-from os import path
-
-sys.path.insert(0, path.abspath('..'))
-
-project = 'DocArray'
-slug = re.sub(r'\W+', '-', project.lower())
-author = 'Jina AI'
-copyright = 'Jina AI Limited. All rights reserved.'
-source_suffix = ['.rst', '.md']
-master_doc = 'index'
-language = 'en'
-repo_dir = '../'
-
-try:
- if 'DOCARRAY_VERSION' not in os.environ:
- pkg_name = 'docarray'
- libinfo_py = path.join(repo_dir, pkg_name, '__init__.py')
- libinfo_content = open(libinfo_py, 'r').readlines()
- version_line = [
- l.strip() for l in libinfo_content if l.startswith('__version__')
- ][0]
- exec(version_line)
- else:
- __version__ = os.environ['DOCARRAY_VERSION']
-except FileNotFoundError:
- __version__ = '0.0.0'
-
-version = __version__
-release = __version__
-
-templates_path = ['_templates']
-exclude_patterns = [
- '_build',
- 'Thumbs.db',
- '.DS_Store',
- 'tests',
- 'page_templates',
- '.github',
-]
-pygments_style = 'rainbow_dash'
-html_theme = 'furo'
-
-base_url = '/'
-html_baseurl = 'https://docarray.jina.ai'
-sitemap_url_scheme = '{link}'
-sitemap_locales = [None]
-sitemap_filename = "sitemap.xml"
-
-html_theme_options = {
- 'light_logo': 'logo-light.svg',
- 'dark_logo': 'logo-dark.svg',
- "sidebar_hide_name": True,
- "light_css_variables": {
- "color-brand-primary": "#009191",
- "color-brand-content": "#009191",
- },
- "dark_css_variables": {
- "color-brand-primary": "#FBCB67",
- "color-brand-content": "#FBCB67",
- },
-
- # PLEASE DO NOT DELETE the empty line between `start-announce` and `end-announce`
- # PLEASE DO NOT DELETE `start-announce`/ `end-announce` it is used for our dev bot to inject announcement from GH
-
- # start-announce
-
- # end-announce
-}
-
-html_static_path = ['_static']
-html_extra_path = ['html_extra']
-html_css_files = [
- 'main.css',
- 'https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0-beta2/css/all.min.css',
-]
-html_js_files = ['https://cdn.jsdelivr.net/npm/vue@2/dist/vue.min.js', 'docbot.js']
-htmlhelp_basename = slug
-html_show_sourcelink = False
-html_favicon = '_static/favicon.png'
-
-latex_documents = [(master_doc, f'{slug}.tex', project, author, 'manual')]
-man_pages = [(master_doc, slug, project, [author], 1)]
-texinfo_documents = [
- (master_doc, slug, project, author, slug, project, 'Miscellaneous')
-]
-epub_title = project
-epub_exclude_files = ['search.html']
-
-# -- Extension configuration -------------------------------------------------
-
-extensions = [
- 'sphinx.ext.autodoc',
- 'sphinx_autodoc_typehints',
- 'sphinx.ext.viewcode',
- 'sphinx.ext.coverage',
- 'sphinxcontrib.apidoc',
- 'sphinxarg.ext',
- 'sphinx_markdown_tables',
- 'sphinx_copybutton',
- 'sphinx_sitemap',
- 'sphinx.ext.intersphinx',
- 'sphinxext.opengraph',
- 'notfound.extension',
- 'myst_parser',
- 'sphinx_design',
- 'sphinx_inline_tabs',
-]
-
-myst_enable_extensions = ['colon_fence', 'substitution']
-
-# -- Custom 404 page
-
-# sphinx-notfound-page
-# https://github.com/readthedocs/sphinx-notfound-page
-notfound_context = {
- 'title': 'Page Not Found',
- 'body': '''
-Ecosystem
--
-
-
-
-
 }})
-  }})
- Jina
- -
-
 }})
-  }})
- Jina Hub
- -
-
 }})
-  }})
- Finetuner
- -
-
 }})
-  }})
- DocArray
-
Page Not Found
-Oops, we couldn't find that page.
-You can try "asking our docs" on the right corner of the page to find answer.
-Otherwise, please create a Github issue and one of our team will respond.
- -''', -} -notfound_no_urls_prefix = True - -apidoc_module_dir = repo_dir -apidoc_output_dir = 'api' -apidoc_excluded_paths = ['tests', 'legacy', 'hub', 'toy*', 'setup.py'] -apidoc_separate_modules = True -apidoc_extra_args = ['-t', 'template/'] -autodoc_member_order = 'bysource' -autodoc_mock_imports = ['argparse', 'numpy', 'np', 'tensorflow', 'torch', 'scipy'] -autoclass_content = 'both' -set_type_checking_flag = False -html_last_updated_fmt = '' -nitpicky = True -nitpick_ignore = [('py:class', 'type')] -linkcheck_ignore = [ - # Avoid link check on local uri - 'http://0.0.0.0:*', - 'pods/encode.yml', - 'https://github.com/jina-ai/docarray/commit/*', - '.github/*', - 'extra-requirements.txt', - 'fastentrypoints.py' '../../101', - '../../102', - 'http://www.twinsun.com/tz/tz-link.htm', # Broken link from pytz library - 'https://urllib3.readthedocs.io/en/latest/contrib.html#google-app-engine', # Broken link from urllib3 library - 'https://linuxize.com/post/how-to-add-swap-space-on-ubuntu-20-04/', - # This link works but gets 403 error on linkcheck -] -linkcheck_timeout = 20 -linkcheck_retries = 2 -linkcheck_anchors = False - -ogp_site_url = 'https://docarray.jina.ai/' -ogp_image = 'https://docarray.jina.ai/_static/banner.png' -ogp_use_first_image = True -ogp_description_length = 300 -ogp_type = 'website' -ogp_site_name = f'DocArray {os.environ.get("SPHINX_MULTIVERSION_VERSION", version)} Documentation' - -ogp_custom_meta_tags = [ - '', - '', - '', - '', - '', - ''' - - - - - - ''', -] - - -def add_server_address(app): - # This makes variable `server_address` available to docbot.js - server_address = app.config['server_address'] - js_text = "var server_address = '%s';" % server_address - app.add_js_file(None, body=js_text) - - -def setup(app): - from sphinx.domains.python import PyField - from sphinx.util.docfields import Field - from sphinx.locale import _ - - app.add_object_type( - 'confval', - 'confval', - objname='configuration value', - indextemplate='pair: %s; configuration value', - doc_field_types=[ - PyField( - 'type', - label=_('Type'), - has_arg=False, - names=('type',), - bodyrolename='class', - ), - Field( - 'default', - label=_('Default'), - has_arg=False, - names=('default',), - ), - ], - ) - # app.add_config_value( - # name='server_address', - # default=os.getenv('JINA_DOCSBOT_SERVER', 'https://docsbot.jina.ai'), - # rebuild='', - # ) - # app.connect('builder-inited', add_server_address) \ No newline at end of file diff --git a/docs/data_types/3d_mesh/3d_mesh.md b/docs/data_types/3d_mesh/3d_mesh.md new file mode 100644 index 00000000000..4727f12cb78 --- /dev/null +++ b/docs/data_types/3d_mesh/3d_mesh.md @@ -0,0 +1,2697 @@ +# 🧬 3D Mesh + +DocArray supports many different modalities including `3D Mesh`. +This section will show you how to load and handle 3D data using DocArray. + +A 3D mesh is the structural build of a 3D model consisting of polygons. Most 3D meshes are created via professional software packages, such as commercial suites like [Unity](https://unity.com/), or the open-source [Blender](https://www.blender.org/). + +!!! note + This feature requires `trimesh`. You can install all necessary dependencies via: + + ```cmd + pip install "docarray[mesh]" + ``` + +## Vertices and Faces representation + +A 3D mesh can be represented by its vertices and faces: + +- **Vertices** are points in a 3D space, represented as a tensor of shape `(n_points, 3)`. +- **Faces** are triangular surfaces that are defined by three points in 3D space, corresponding to the three vertices of a triangle. They can be represented as a tensor of shape `(n_faces, 3)`. Each number in that tensor refers to an index of a vertex in the tensor of vertices. + +### Load vertices and faces + +First, let's define our class `MyMesh3D`, which extends [`BaseDoc`][docarray.base_doc.doc.BaseDoc] and provides attributes to store our 3D data: + +- The `mesh_url` attribute of type [`Mesh3DUrl`][docarray.typing.url.url_3d.mesh_url.Mesh3DUrl]. +- The optional `tensors` attribute, of type [`VerticesAndFaces`][docarray.documents.mesh.vertices_and_faces.VerticesAndFaces] + - The `VerticesAndFaces` class has the attributes `vertices` and `faces`, both of type [`AnyTensor`](../../../../api_references/typing/tensor/tensor). This especially comes in handy later when we want to display our 3D mesh. + +!!! tip + Check out our predefined [`Mesh3D`](#getting-started-predefined-docs) to get started and play around with our 3D features. + +But for now, let's create a `MyMesh3D` instance with a URL to a remote `.obj` file: + +```python +from typing import Optional + +from docarray import BaseDoc +from docarray.documents.mesh.vertices_and_faces import VerticesAndFaces +from docarray.typing import Mesh3DUrl + + +class MyMesh3D(BaseDoc): + mesh_url: Mesh3DUrl + tensors: Optional[VerticesAndFaces] = None + + +doc = MyMesh3D(mesh_url="https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj") +``` + +To load the vertices and faces information, you can call [`.load()`][docarray.typing.url.url_3d.mesh_url.Mesh3DUrl.load] on the [`Mesh3DUrl`][docarray.typing.url.url_3d.mesh_url.Mesh3DUrl] instance. This will return a [`VerticesAndFaces`][docarray.documents.mesh.vertices_and_faces.VerticesAndFaces] object: + +```python +doc.tensors = doc.mesh_url.load() +doc.summary() +``` + +
+
+
+### Display 3D mesh in notebook
+
+You can display your 3D mesh interactively from its URL as well as a [`VerticesAndFaces`][docarray.documents.mesh.vertices_and_faces.VerticesAndFaces] instance, by calling `.display()` on either one. The latter will always display without color, whereas the display from the URL will show with color if this information is included in the file content.
+
+``` { .python }
+doc.mesh_url.display()
+```
+
+
+
+## Point cloud representation
+
+A point cloud is a representation of a 3D mesh. It is made by repeatedly and uniformly sampling points within the surface of the 3D body. Compared to the mesh representation, the point cloud is a fixed size `ndarray` and hence easier for deep learning algorithms to handle.
+
+### Load point cloud
+
+!!! tip
+ Check out our predefined [`PointCloud3D`](#getting-started-predefined-docs) to get started and play around with our 3D features.
+
+In DocArray, loading a point cloud from a [`PointCloud3DUrl`][docarray.typing.url.url_3d.point_cloud_url.PointCloud3DUrl] instance will return a [`PointsAndColors`][docarray.documents.point_cloud.points_and_colors.PointsAndColors] instance. Such an object has a `points` attribute containing the information about the points in 3D space as well as an optional `colors` attribute.
+
+First, let's define our class `MyPointCloud`, which extends [`BaseDoc`][docarray.base_doc.doc.BaseDoc] and provides attributes to store the point cloud information:
+
+```python
+from typing import Optional
+
+from docarray import BaseDoc
+from docarray.documents.point_cloud.points_and_colors import PointsAndColors
+from docarray.typing import PointCloud3DUrl
+
+
+class MyPointCloud(BaseDoc):
+ url: PointCloud3DUrl
+ tensors: Optional[PointsAndColors] = None
+
+
+doc = MyPointCloud(url="https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj")
+```
+
+Next, we can load a point cloud of size `samples` by simply calling [`.load()`][docarray.typing.url.url_3d.point_cloud_url.PointCloud3DUrl.load] on the [`PointCloud3DUrl`][docarray.typing.url.url_3d.point_cloud_url.PointCloud3DUrl] instance:
+
+```python
+doc.tensors = doc.url.load(samples=1000)
+doc.summary()
+```
+
+Output
+ ``` { .text .no-copy } + 📄 MyMesh3D : 9d8c26f ... + ╭─────────────────────┬────────────────────────────────────────────────────────╮ + │ Attribute │ Value │ + ├─────────────────────┼────────────────────────────────────────────────────────┤ + │ mesh_url: Mesh3DUrl │ https://people.sc.fsu.edu/~jburkardt/data/obj/al.o ... │ + │ │ (length: 52) │ + ╰─────────────────────┴────────────────────────────────────────────────────────╯ + └── 🔶 tensors: VerticesAndFaces + └── 📄 VerticesAndFaces : 8cae4c4 ... + ╭───────────────────┬──────────────────────────────────────────────────╮ + │ Attribute │ Value │ + ├───────────────────┼──────────────────────────────────────────────────┤ + │ vertices: NdArray │ NdArray of shape (3980, 3), dtype: float64 │ + │ faces: NdArray │ NdArray of shape (7152, 3), dtype: int64 │ + ╰───────────────────┴──────────────────────────────────────────────────╯ + ``` +
+
+
+### Display 3D point cloud in notebook
+
+You can display your point cloud and interact with it from its URL as well as from a PointsAndColors instance. The first will always display without color, whereas the display from [`PointsAndColors`][docarray.documents.point_cloud.points_and_colors.PointsAndColors] will show with color if `.colors` is not None.
+
+``` { .python}
+doc.url.display()
+```
+
+
+
+## Getting started - Predefined documents
+
+To get started and play around with the 3D modalities, DocArray provides the predefined documents [`Mesh3D`][docarray.documents.mesh.Mesh3D] and [`PointCloud3D`][docarray.documents.point_cloud.PointCloud3D], which includes all of the previously mentioned functionalities.
+
+### `Mesh3D`
+
+The [`Mesh3D`][docarray.documents.mesh.Mesh3D] class provides a [`Mesh3DUrl`][docarray.typing.Mesh3DUrl] field and [`VerticesAndFaces`][docarray.documents.mesh.vertices_and_faces.VerticesAndFaces] field.
+
+``` { .python }
+class Mesh3D(BaseDoc):
+ url: Optional[Mesh3DUrl] = None
+ tensors: Optional[VerticesAndFaces] = None
+ embedding: Optional[AnyEmbedding] = None
+ bytes_: Optional[bytes] = None
+```
+
+### `PointCloud3D`
+
+``` { .python }
+class PointCloud3D(BaseDoc):
+ url: Optional[PointCloud3DUrl] = None
+ tensors: Optional[PointsAndColors] = None
+ embedding: Optional[AnyEmbedding] = None
+ bytes_: Optional[bytes] = None
+```
+
+You can use them directly, extend or compose them:
+
+```python
+from docarray import BaseDoc
+from docarray.documents import Mesh3D, PointCloud3D
+
+
+class My3DObject(BaseDoc):
+ title: str
+ mesh: Mesh3D
+ pc: PointCloud3D
+
+
+obj_file = 'https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj'
+
+doc = My3DObject(
+ title='My first 3D object!',
+ mesh=Mesh3D(url=obj_file),
+ pc=PointCloud3D(url=obj_file),
+)
+
+doc.mesh.tensors = doc.mesh.url.load()
+doc.pc.tensors = doc.pc.url.load(samples=100)
+```
diff --git a/docs/data_types/audio/audio.md b/docs/data_types/audio/audio.md
new file mode 100644
index 00000000000..a676adb4a01
--- /dev/null
+++ b/docs/data_types/audio/audio.md
@@ -0,0 +1,215 @@
+# 🔊 Audio
+
+DocArray supports many different modalities including `Audio`.
+This section will show you how to load and handle audio data using DocArray.
+
+Moreover, you will learn about DocArray's audio-specific types, to represent your audio data ranging from [`AudioUrl`][docarray.typing.url.AudioUrl] to [`AudioBytes`][docarray.typing.bytes.AudioBytes] and [`AudioNdArray`][docarray.typing.tensor.audio.audio_ndarray.AudioNdArray].
+
+!!! note
+ This requires a `pydub` dependency. You can install all necessary dependencies via:
+
+ ```cmd
+ pip install "docarray[audio]"
+ ```
+
+ Additionally, you have to install `ffmpeg` (see more info [here](https://github.com/jiaaro/pydub#getting-ffmpeg-set-up)):
+
+ ```cmd
+ # on Mac with brew:
+ brew install ffmpeg
+ ```
+
+ ```cmd
+ # on Linux with apt-get
+ apt-get install ffmpeg libavcodec-extra
+ ```
+
+## Load audio file
+
+First, let's define a class which extends [`BaseDoc`][docarray.base_doc.doc.BaseDoc] and has a `url` attribute of type [`AudioUrl`][docarray.typing.url.AudioUrl], and an optional `tensor` attribute of type [`AudioTensor`](../../../../api_references/typing/tensor/audio).
+
+!!! tip
+ Check out our predefined [`AudioDoc`](#getting-started-predefined-audiodoc) to get started and play around with our audio features.
+
+Next, you can instantiate an object of that class with a local or remote URL:
+
+```python
+from docarray import BaseDoc
+from docarray.typing import AudioUrl, AudioNdArray
+
+
+class MyAudio(BaseDoc):
+ url: AudioUrl
+ tensor: AudioNdArray = None
+ frame_rate: int = None
+
+
+doc = MyAudio(
+ url='https://github.com/docarray/docarray/blob/main/tests/toydata/hello.mp3?raw=true'
+)
+```
+
+Loading the content of the audio file is as easy as calling [`.load()`][docarray.typing.url.AudioUrl] on the [`AudioUrl`][docarray.typing.url.AudioUrl] instance.
+
+This will return a tuple of:
+
+- An [`AudioNdArray`][docarray.typing.tensor.audio.AudioNdArray] representing the audio file content
+- An integer representing the frame rate (number of signals for a certain period of time)
+
+```python
+doc.tensor, doc.frame_rate = doc.url.load()
+doc.summary()
+```
+
+Output
+ ``` { .text .no-copy } + 📄 MyPointCloud : a63374d ... + ╭──────────────────────────────────┬───────────────────────────────────────────╮ + │ Attribute │ Value │ + ├──────────────────────────────────┼───────────────────────────────────────────┤ + │ url: PointCloud3DUrl │ https://people.sc.fsu.edu/~jburkardt/dat… │ + │ │ ... (length: 52) │ + ╰──────────────────────────────────┴───────────────────────────────────────────╯ + └── 🔶 tensors: PointsAndColors + └── 📄 PointsAndColors : 70ae175 ... + ╭─────────────────┬────────────────────────────────────────────────────╮ + │ Attribute │ Value │ + ├─────────────────┼────────────────────────────────────────────────────┤ + │ points: NdArray │ NdArray of shape (1000, 3), dtype: float64 │ + ╰─────────────────┴────────────────────────────────────────────────────╯ + + ``` +
+
+
+## AudioTensor
+
+DocArray offers several [`AudioTensor`s](../../../../api_references/typing/tensor/audio) to store your data to:
+
+- [`AudioNdArray`][docarray.typing.tensor.audio.audio_ndarray.AudioNdArray]
+- [`AudioTorchTensor`][docarray.typing.AudioTorchTensor]
+- [`AudioTensorFlowTensor`][docarray.typing.AudioTensorFlowTensor]
+
+If you specify the type of your tensor to one of the above, it will be cast to that automatically:
+
+```python hl_lines="7 8 15 16"
+from docarray import BaseDoc
+from docarray.typing import AudioTensorFlowTensor, AudioTorchTensor, AudioUrl
+
+
+class MyAudio(BaseDoc):
+ url: AudioUrl
+ tf_tensor: AudioTensorFlowTensor = None
+ torch_tensor: AudioTorchTensor = None
+
+
+doc = MyAudio(
+ url='https://github.com/docarray/docarray/blob/main/tests/toydata/hello.mp3?raw=true'
+)
+
+doc.tf_tensor, _ = doc.url.load()
+doc.torch_tensor, _ = doc.url.load()
+
+assert isinstance(doc.tf_tensor, AudioTensorFlowTensor)
+assert isinstance(doc.torch_tensor, AudioTorchTensor)
+```
+
+## AudioBytes
+
+Alternatively, you can load your [`AudioUrl`][docarray.typing.url.AudioUrl] instance to [`AudioBytes`][docarray.typing.bytes.AudioBytes], and your [`AudioBytes`][docarray.typing.bytes.AudioBytes] instance to an [`AudioTensor`](../../../../api_references/typing/tensor/audio) of your choice:
+
+```python hl_lines="15 16"
+
+from docarray import BaseDoc
+from docarray.typing import AudioBytes, AudioTensor, AudioUrl
+
+
+class MyAudio(BaseDoc):
+ url: AudioUrl = None
+ bytes_: AudioBytes = None
+ tensor: AudioTensor = None
+
+
+doc = MyAudio(
+ url='https://github.com/docarray/docarray/blob/main/tests/toydata/hello.mp3?raw=true'
+)
+
+doc.bytes_ = doc.url.load_bytes() # type(doc.bytes_) = AudioBytes
+doc.tensor, _ = doc.bytes_.load() # type(doc.tensor) = AudioNdarray
+```
+
+Vice versa, you can also transform an [`AudioTensor`](../../../../api_references/typing/tensor/audio) to [`AudioBytes`][docarray.typing.bytes.AudioBytes]:
+
+```python
+from docarray.typing import AudioBytes
+
+
+bytes_from_tensor = doc.tensor.to_bytes()
+
+assert isinstance(bytes_from_tensor, AudioBytes)
+```
+
+## Save audio to file
+
+You can save your [`AudioTensor`](../../../../api_references/typing/tensor/audio) to an audio file of any format as follows:
+
+``` { .python }
+tensor_reversed = doc.tensor[::-1]
+tensor_reversed.save(
+ file_path='olleh.mp3',
+ format='mp3',
+)
+```
+## Play audio in a notebook
+
+You can play your audio sound in a notebook from its URL or tensor, by calling `.display()` on either one.
+
+Play from `url`:
+``` { .python }
+doc.url.display()
+```
+
+Output
+ ``` { .text .no-copy } + 📄 MyAudio : 2015696 ... + ╭──────────────────────┬───────────────────────────────────────────────────────╮ + │ Attribute │ Value │ + ├──────────────────────┼───────────────────────────────────────────────────────┤ + │ url: AudioUrl │ https://github.com/docarray/docarray/blob/main/tes │ + │ │ ... (length: 90) │ + │ tensor: AudioNdArray │ AudioNdArray of shape (30833,), dtype: float64 │ + │ frame_rate: int │ 44100 │ + ╰──────────────────────┴───────────────────────────────────────────────────────╯ + ``` +| + |
| + |
+
+
+### Access data
+
+After instantiation, each modality can be accessed directly from the `Page` object:
+
+```python
+print(page.main_text)
+print(page.img_url)
+print(page.img_description)
+print(page.img_tensor)
+```
+
+Output
+ ``` { .text .no-copy } + 📄 Page : 8f39674 ... + ╭──────────────────────────────┬───────────────────────────────────────────────╮ + │ Attribute │ Value │ + ├──────────────────────────────┼───────────────────────────────────────────────┤ + │ main_text: str │ Hello world │ + │ img_url: ImageUrl │ https://github.com/docarray/docarray/blob/ma… │ + │ │ ... (length: 90) │ + │ img_description: str │ This is DocArray │ + │ img_tensor: ImageTorchTensor │ ImageTorchTensor of shape (320, 320, 3), │ + │ │ dtype: torch.uint8 │ + ╰──────────────────────────────┴───────────────────────────────────────────────╯ + ``` +
+
+
+### Nested data
+
+If the data you want to model requires a more complex structure, nesting your attributes may be a good solution.
+
+For this example, let's try to define a schema to represent a newspaper. The newspaper should consist of a cover page,
+any number of following pages, and some metadata. Further, each page contains a main text and can contain an image
+and an image description.
+
+To implement this you can add a `Newspaper` class to the previous implementation. The newspaper has a required
+`cover_page` attribute of type `Page` as well as a `pages` attribute, which is a `DocList` of `Page`s.
+
+```python
+from docarray import BaseDoc, DocList
+from docarray.typing import ImageTorchTensor, ImageUrl
+
+
+class Page(BaseDoc):
+ main_text: str
+ img_url: ImageUrl = None
+ img_description: str = None
+ img_tensor: ImageTorchTensor = None
+
+
+class Newspaper(BaseDoc):
+ cover: Page
+ pages: DocList[Page] = None
+ metadata: dict = None
+```
+
+You can instantiate this more complex `Newspaper` object the same way as before:
+
+```python
+cover_page = Page(
+ main_text='DocArray Daily',
+ img_url='https://github.com/docarray/docarray/blob/main/docs/assets/favicon.png',
+)
+
+pages = DocList[Page](
+ [
+ Page(
+ main_text='Hello world',
+ img_url='https://github.com/docarray/docarray/blob/main/docs/assets/favicon.png',
+ img_description='This is the image of an apple',
+ ),
+ Page(main_text='Second page'),
+ Page(main_text='Third page'),
+ ]
+)
+
+docarray_daily = Newspaper(
+ cover=cover_page,
+ pages=pages,
+ metadata={'author': 'DocArray and friends', 'issue': '0.30.0'},
+)
+
+docarray_daily.summary()
+```
+
+Output
+ ``` { .text .no-copy } + Hello world + https://github.com/docarray/docarray/blob/main/docs/assets/favicon.png?raw=true + This is DocArray + ImageTorchTensor([[[0, 0, 0], + [0, 0, 0], + [0, 0, 0], + ..., + [0, 0, 0]]]) + ``` +
+
diff --git a/docs/data_types/table/table.md b/docs/data_types/table/table.md
new file mode 100644
index 00000000000..3ebc90013d1
--- /dev/null
+++ b/docs/data_types/table/table.md
@@ -0,0 +1,234 @@
+# 📊 Table
+
+DocArray supports many different modalities including tabular data.
+This section will show you how to load and handle tabular data using DocArray.
+
+## Load CSV table
+
+A common way to store tabular data is via `CSV` (comma-separated values) files.
+You can load such data from a given `CSV` file into a [`DocList`][docarray.DocList].
+
+Let's take a look at the following example file, which includes data about books and their authors and year of publication:
+
+```text
+title,author,year
+Harry Potter and the Philosopher's Stone,J. K. Rowling,1997
+Klara and the sun,Kazuo Ishiguro,2020
+A little life,Hanya Yanagihara,2015
+```
+
+First, define the Document schema describing the data:
+
+```python
+from docarray import BaseDoc
+
+
+class Book(BaseDoc):
+ title: str
+ author: str
+ year: int
+```
+Next, load the content of the CSV file to a [`DocList`][docarray.DocList] instance of `Book`s via [`.from_csv()`][docarray.array.doc_list.io.IOMixinDocList.from_csv]:
+
+```python
+from docarray import DocList
+
+
+docs = DocList[Book].from_csv(
+ file_path='https://github.com/docarray/docarray/blob/main/tests/toydata/books.csv?raw=true'
+)
+docs.summary()
+```
+
+Output
+ ``` { .text .no-copy } + 📄 Newspaper : 63189f7 ... + ╭────────────────┬─────────────────────────────────────────────────────────────╮ + │ Attribute │ Value │ + ├────────────────┼─────────────────────────────────────────────────────────────┤ + │ metadata: dict │ {'author': 'DocArray and friends', 'issue': '0.0.3 ... } │ + │ │ (length: 2) │ + ╰────────────────┴─────────────────────────────────────────────────────────────╯ + ├── 🔶 cover: Page + │ └── 📄 Page : ca164e3 ... + │ ╭───────────────────┬──────────────────────────────────────────────────╮ + │ │ Attribute │ Value │ + │ ├───────────────────┼──────────────────────────────────────────────────┤ + │ │ main_text: str │ DocArray Daily │ + │ │ img_url: ImageUrl │ https://github.com/docarray/docarray/blob/main/… │ + │ │ │ ... (length: 81) │ + │ ╰───────────────────┴──────────────────────────────────────────────────╯ + └── 💠 pages: DocList[Page] + ├── 📄 Page : 64ed19c ... + │ ╭──────────────────────┬───────────────────────────────────────────────╮ + │ │ Attribute │ Value │ + │ ├──────────────────────┼───────────────────────────────────────────────┤ + │ │ main_text: str │ Hello world │ + │ │ img_url: ImageUrl │ https://github.com/docarray/docarray/blob/ma… │ + │ │ │ ... (length: 81) │ + │ │ img_description: str │ DocArray logoooo │ + │ ╰──────────────────────┴───────────────────────────────────────────────╯ + ├── 📄 Page : 4bd7e45 ... + │ ╭─────────────────────┬────────────────╮ + │ │ Attribute │ Value │ + │ ├─────────────────────┼────────────────┤ + │ │ main_text: str │ Second page │ + │ ╰─────────────────────┴────────────────╯ + └── ... 1 more Page documents + ``` +
+
+
+The resulting [`DocList`][docarray.DocList] object contains three `Book`s, since each row of the CSV file corresponds to one book and is assigned to one `Book` instance.
+
+## Save to CSV file
+
+Vice versa, you can also store your [`DocList`][docarray.DocList] data in a `.csv` file using [`.to_csv()`][docarray.array.doc_list.io.IOMixinDocList.to_csv]:
+
+``` { .python }
+docs.to_csv(file_path='/path/to/my_file.csv')
+```
+
+Tabular data is often not the best choice to represent nested Documents. Hence, nested Documents will be stored flattened and can be accessed by their `'__'`-separated access paths.
+
+Let's take a look at an example. We now want to store not only the book data but moreover book review data. To do so, we define a `BookReview` class that has a nested `book` attribute as well as the non-nested attributes `n_ratings` and `stars`:
+
+```python
+class BookReview(BaseDoc):
+ book: Book
+ n_ratings: int
+ stars: float
+
+
+review_docs = DocList[BookReview](
+ [BookReview(book=book, n_ratings=12345, stars=5) for book in docs]
+)
+review_docs.summary()
+```
+
+Output
+ ``` { .text .no-copy } + ╭────── DocList Summary ──────╮ + │ │ + │ Type DocList[Book] │ + │ Length 3 │ + │ │ + ╰─────────────────────────────╯ + ╭── Document Schema ──╮ + │ │ + │ Book │ + │ ├── title: str │ + │ ├── author: str │ + │ └── year: int │ + │ │ + ╰─────────────────────╯ + ``` +
+
+
+As expected all nested attributes will be stored by their access path:
+
+``` { .python }
+review_docs.to_csv(file_path='/path/to/nested_documents.csv')
+```
+
+``` { .text .no-copy hl_lines="1" }
+id,book__id,book__title,book__author,book__year,n_ratings,stars
+d6363aa3b78b4f4244fb976570a84ff7,8cd85fea52b3a3bc582cf56c9d612cbb,Harry Potter and the Philosopher's Stone,J. K. Rowling,1997,12345,5.0
+5b53fff67e6b6cede5870f2ee09edb05,87b369b93593967226c525cf226e3325,Klara and the sun,Kazuo Ishiguro,2020,12345,5.0
+addca0475756fc12cdec8faf8fb10d71,03194cec1b75927c2259b3c0fff1ab6f,A little life,Hanya Yanagihara,2015,12345,5.0
+```
+
+## Handle TSV tables
+
+Not only can you load and save comma-separated values (`CSV`) data, but also tab-separated values (`TSV`),
+by adjusting the `dialect` parameter in [`.from_csv()`][docarray.array.doc_list.io.IOMixinDocList.from_csv]
+and [`.to_csv()`][docarray.array.doc_list.io.IOMixinDocList.to_csv].
+
+The dialect defaults to `'excel'`, which refers to comma-separated values. For tab-separated values, you can use
+`'excel-tab'`.
+
+Let's take a look at what this would look like with a tab-separated file:
+
+```text
+title author year
+Title1 author1 2020
+Title2 author2 1234
+```
+
+```python
+docs = DocList[Book].from_csv(
+ file_path='https://github.com/docarray/docarray/blob/main/tests/toydata/books.tsv?raw=true',
+ dialect='excel-tab',
+)
+for doc in docs:
+ doc.summary()
+```
+
+Output
+ ``` { .text .no-copy} + ╭───────── DocList Summary ─────────╮ + │ │ + │ Type DocList[BookReview] │ + │ Length 3 │ + │ │ + ╰───────────────────────────────────╯ + ╭──── Document Schema ────╮ + │ │ + │ BookReview │ + │ ├── book: Book │ + │ │ ├── title: str │ + │ │ ├── author: str │ + │ │ └── year: int │ + │ ├── n_ratings: int │ + │ └── stars: float │ + │ │ + ╰─────────────────────────╯ + ``` +
+
+
+Great! All the data is correctly read and stored in `Book` instances.
+
+## Other separators
+
+If your values are separated by yet another separator, you can create your own class that inherits from `csv.Dialect`.
+Within this class, you can define your dialect's behavior by setting the provided [formatting parameters](https://docs.python.org/3/library/csv.html#dialects-and-formatting-parameters).
+
+For instance, let's assume you have a semicolon-separated table:
+
+```text
+first_name;last_name;year
+Jane;Austin;2020
+John;Doe;1234
+```
+
+Now, let's define our `SemicolonSeparator` class. Next to the `delimiter` parameter, we have to set some more formatting parameters such as `doublequote` and `lineterminator`.
+
+```python
+import csv
+
+
+class SemicolonSeparator(csv.Dialect):
+ delimiter = ';'
+ doublequote = True
+ lineterminator = '\r\n'
+ quotechar = '"'
+ quoting = csv.QUOTE_MINIMAL
+```
+
+Finally, you can load your data by setting the `dialect` parameter in [`.from_csv()`][docarray.array.doc_list.io.IOMixinDocList.from_csv] to an instance of your `SemicolonSeparator`.
+
+```python
+docs = DocList[Book].from_csv(
+ file_path='https://github.com/docarray/docarray/blob/main/tests/toydata/books_semicolon_sep.csv?raw=true',
+ dialect=SemicolonSeparator(),
+)
+for doc in docs:
+ doc.summary()
+```
+
+Output
+ ```text + 📄 Book : c1ac9d4 ... + ╭──────────────────────┬───────────────╮ + │ Attribute │ Value │ + ├──────────────────────┼───────────────┤ + │ title: str │ Title1 │ + │ author: str │ author1 │ + │ year: int │ 2020 │ + ╰──────────────────────┴───────────────╯ + 📄 Book : c1ac9d4 ... + ╭──────────────────────┬───────────────╮ + │ Attribute │ Value │ + ├──────────────────────┼───────────────┤ + │ title: str │ Title1 │ + │ author: str │ author1 │ + │ year: int │ 2020 │ + ╰──────────────────────┴───────────────╯ + ``` +
+
diff --git a/docs/data_types/tensor/tensor.md b/docs/data_types/tensor/tensor.md
new file mode 100644
index 00000000000..f16b8a82f9a
--- /dev/null
+++ b/docs/data_types/tensor/tensor.md
@@ -0,0 +1,231 @@
+# 🔢 Tensor
+
+DocArray supports several tensor types that can you can use inside `BaseDoc`.
+
+The main ones are:
+
+- [`NdArray`][docarray.typing.tensor.NdArray] for NumPy tensors
+- [`TorchTensor`][docarray.typing.tensor.TorchTensor] for PyTorch tensors
+- [`TensorFlowTensor`][docarray.typing.tensor.TensorFlowTensor] for TensorFlow tensors
+
+The three of them wrap their respective framework's tensor type.
+
+!!! note
+ [`NdArray`][docarray.typing.tensor.NdArray] and [`TorchTensor`][docarray.typing.tensor.TorchTensor] are a subclass of their native tensor type. This means that they can be used natively in their framework.
+
+!!! warning
+ [`TensorFlowTensor`][docarray.typing.tensor.TensorFlowTensor] stores the pure `tf.Tensor` object inside the `tensor` attribute. This is due to a limitation of the TensorFlow framework that prevents you from subclassing the `tf.Tensor` object.
+
+DocArray also supports [`AnyTensor`][docarray.typing.tensor.AnyTensor], which is the Union of the three previous tensor types.
+This is a generic placeholder to specify that it can work with any tensor type (NumPy, PyTorch, TensorFlow).
+
+## Tensor Shape validation
+
+All three tensor types support shape validation. This means that you can specify the shape of the tensor using type hint syntax: `NdArray[100, 100]`, `TorchTensor[100, 100]`, `TensorFlowTensor[100, 100]`.
+
+Let's take an example:
+
+```python
+from docarray import BaseDoc
+from docarray.typing import NdArray
+
+
+class MyDoc(BaseDoc):
+ tensor: NdArray[100, 100]
+```
+
+If you try to pass a tensor with a different shape, an error will be raised:
+
+```python
+import numpy as np
+
+try:
+ doc = MyDoc(tensor=np.zeros((100, 200)))
+except ValueError as e:
+ print(e)
+```
+
+```bash
+1 validation error for MyDoc
+tensor
+ cannot reshape array of size 20000 into shape (100,100) (type=value_error)
+```
+
+
+Whereas if you just pass a tensor with the correct shape, no error will be raised:
+
+```python
+doc = MyDoc(tensor=np.zeros((100, 100)))
+```
+
+### Axes validation
+
+You can check that the number of axes is correct by specifying `NdArray['x','y']`, `TorchTensor['x','y']`, `TensorFlowTensor['x','y']`.
+
+```python
+from docarray import BaseDoc
+from docarray.typing import NdArray
+
+
+class MyDoc(BaseDoc):
+ tensor: NdArray['x', 'y']
+```
+
+Here you can only pass a tensor with two axes. `np.zeros(10, 12)` will work, but `np.zeros(10, 12, 3)` will raise an error.
+
+### Axis names
+
+You can specify that two axes should have the same dimensions with the syntax `NdArray['x', 'x']`, `TorchTensor['x', 'x']`, `TensorFlowTensor['x', 'x']`.
+
+```python
+from docarray import BaseDoc
+from docarray.typing import NdArray
+
+
+class MyDoc(BaseDoc):
+ tensor: NdArray['x', 'x']
+```
+
+Here you can only pass a tensor with two axes that have the same dimensions. `np.zeros(10, 10)` will work but `np.zeros(10, 12)` will raise an error.
+
+### Arbitrary number of axis
+
+To specify that your shape can have an arbitrary number of axes, use the syntax `NdArray['x', ...]`, or `NdArray[..., 'x']`.
+
+```python
+from docarray import BaseDoc
+from docarray.typing import NdArray
+
+
+class MyDoc(BaseDoc):
+ tensor: NdArray[100, ...]
+```
+
+Here you can only pass a tensor with at least one axis with dimension 100. `np.zeros(100, 10)` will work but `np.zeros(10, 12)` will raise an error.
+
+## Tensor type validation
+
+You don't need to directly instantiate the [`NdArray`][docarray.typing.tensor.NdArray] , [`TorchTensor`][docarray.typing.tensor.TorchTensor], or [`TensorFlowTensor`][docarray.typing.tensor.TensorFlowTensor] by yourself.
+
+Instead, you should use them as type hints on [`BaseDoc`][docarray.base_doc.doc.BaseDoc] fields, where they perform data validation.
+During this process, [`BaseDoc`][docarray.base_doc.doc.BaseDoc] will cast the native tensor type into the respective DocArray tensor type.
+
+Let's look at an example:
+
+```python
+from docarray import BaseDoc
+from docarray.typing import NdArray
+
+import numpy as np
+
+
+class MyDoc(BaseDoc):
+ tensor: NdArray
+
+
+doc = MyDoc(tensor=np.zeros(100))
+
+assert isinstance(doc.tensor, NdArray) # True
+```
+Here you see that the `doc.tensor` is an `NdArray`:
+
+```python
+assert isinstance(doc.tensor, np.ndarray) # True as well
+```
+
+But since it inherits from `np.ndarray`, you can also use it as a normal NumPy array. The same holds for PyTorch and `TorchTensor`.
+
+## Type coercion with different tensor types
+
+DocArray also supports type coercion between different tensor types. This mean that if you pass a different tensor type to a tensor field, it will be converted to the correct tensor type.
+
+For instance, if you define a field of type [`TorchTensor`][docarray.typing.tensor.TorchTensor] and you pass a NumPy array to it, it will be converted to a [`TorchTensor`][docarray.typing.tensor.TorchTensor].
+
+```python
+from docarray import BaseDoc
+from docarray.typing import TorchTensor
+import numpy as np
+
+
+class MyTensorsDoc(BaseDoc):
+ tensor: TorchTensor
+
+
+doc = MyTensorsDoc(tensor=np.zeros(512))
+doc.summary()
+```
+
+```bash
+📄 MyTensorsDoc : 0a10f88 ...
+╭─────────────────────┬────────────────────────────────────────────────────────╮
+│ Attribute │ Value │
+├─────────────────────┼────────────────────────────────────────────────────────┤
+│ tensor: TorchTensor │ TorchTensor of shape (512,), dtype: torch.float64 │
+╰─────────────────────┴────────────────────────────────────────────────────────╯
+```
+
+It also works in the other direction:
+
+```python
+from docarray import BaseDoc
+from docarray.typing import NdArray
+import torch
+
+
+class MyTensorsDoc(BaseDoc):
+ tensor: NdArray
+
+
+doc = MyTensorsDoc(tensor=torch.zeros(512))
+doc.summary()
+```
+
+```bash
+📄 MyTensorsDoc : 157e6f5 ...
+╭─────────────────┬────────────────────────────────────────────────────────────╮
+│ Attribute │ Value │
+├─────────────────┼────────────────────────────────────────────────────────────┤
+│ tensor: NdArray │ NdArray of shape (512,), dtype: float32 │
+╰─────────────────┴────────────────────────────────────────────────────────────╯
+```
+
+## `DocVec` with `AnyTensor`
+
+[`DocVec`][docarray.array.doc_vec.doc_vec.DocVec] can be used with a `BaseDoc` which has a field of [`AnyTensor`][docarray.typing.tensor.AnyTensor] or any other Union of tensor types.
+
+However, the `DocVec` needs to know the tensor type of the tensor field beforehand to create the correct column.
+
+You can specify these parameters with the `tensor_type` parameter of the [`DocVec`][docarray.vectorizer.doc_vec.DocVec] constructor:
+
+```python
+from docarray import BaseDoc, DocVec
+from docarray.typing import AnyTensor, NdArray
+
+import numpy as np
+
+
+class MyDoc(BaseDoc):
+ tensor: AnyTensor
+
+
+docs = DocVec[MyDoc](
+ [MyDoc(tensor=np.zeros(100)) for _ in range(10)], tensor_type=NdArray
+)
+
+assert isinstance(docs.tensor, NdArray)
+```
+
+!!! note
+ `NdArray` will be used by default if:
+
+ - you don't specify the `tensor_type` parameter
+ - your tensor field is a Union of tensor or [`AnyTensor`][docarray.typing.tensor.AnyTensor]
+
+## Compatibility of `TorchTensor` and `torch.compile()`
+
+PyTorch 2 [introduced compilation support](https://pytorch.org/blog/pytorch-2.0-release/) in the form of `torch.compile()`.
+
+Currently, **`torch.compile()` does not properly support subclasses of `torch.Tensor` such as [`TorchTensor`][docarray.typing.tensor.TorchTensor]**.
+The PyTorch team is currently working on a [fix for this issue](https://github.com/pytorch/pytorch/pull/105167#issuecomment-1678050808).
+
+For a workaround to this issue, see the [`TorchTensor` API reference][docarray.typing.tensor.TorchTensor].
diff --git a/docs/data_types/text/text.md b/docs/data_types/text/text.md
new file mode 100644
index 00000000000..5f3493c3415
--- /dev/null
+++ b/docs/data_types/text/text.md
@@ -0,0 +1,108 @@
+# 🔤 Text
+
+DocArray supports many different modalities including `Text`.
+This section will show you how to load and handle text data using DocArray.
+
+!!! tip
+ Check out our predefined [`TextDoc`](#getting-started-predefined-textdoc) to get started and play around with our text features.
+
+You can store text in DocArray like this:
+
+```python
+from docarray import BaseDoc
+
+
+class MyText(BaseDoc):
+ text: str = None
+
+
+doc = MyText(text='Hello world!')
+```
+
+Text can include any type of character, including emojis:
+
+```python
+doc.text = '👋 नमस्ते दुनिया! 你好世界！こんにちは世界！ Привет мир!'
+```
+
+## Load text file
+
+If your text data is too long to be written inline or if it is stored in a file, you can first define the URL as a [`TextUrl`][docarray.typing.url.text_url.TextUrl] and then load the text data.
+
+Let's first define a schema:
+
+```python
+from docarray import BaseDoc
+from docarray.typing import TextUrl
+
+
+class MyText(BaseDoc):
+ text: str = None
+ url: TextUrl = None
+```
+Next, instantiate a `MyText` object with a `url` attribute and load its content to the `text` field.
+
+```python
+doc = MyText(
+ url='https://www.w3.org/History/19921103-hypertext/hypertext/README.html',
+)
+doc.text = doc.url.load()
+
+assert doc.text.startswith('Output
+ ```text + 📄 Book : 321e9fd ... + ╭──────────────────────┬───────────────╮ + │ Attribute │ Value │ + ├──────────────────────┼───────────────┤ + │ title: str │ Title1 │ + │ author: str │ author1 │ + │ year: int │ 2020 │ + ╰──────────────────────┴───────────────╯ + 📄 Book : 16d2097 ... + ╭──────────────────────┬───────────────╮ + │ Attribute │ Value │ + ├──────────────────────┼───────────────┤ + │ title: str │ Title2 │ + │ author: str │ author2 │ + │ year: int │ 1234 │ + ╰──────────────────────┴───────────────╯ + ``` +
+
+
+## Getting started - Predefined `TextDoc`
+
+To get started and play around with your text data, DocArray provides a predefined [`TextDoc`][docarray.documents.text.TextDoc], which includes all of the previously mentioned functionalities:
+
+``` { .python }
+class TextDoc(BaseDoc):
+ text: Optional[str] = None
+ url: Optional[TextUrl] = None
+ embedding: Optional[AnyEmbedding] = None
+ bytes_: Optional[bytes] = None
+```
diff --git a/docs/data_types/video/key_frames.png b/docs/data_types/video/key_frames.png
new file mode 100644
index 00000000000..5c039e11c20
Binary files /dev/null and b/docs/data_types/video/key_frames.png differ
diff --git a/docs/data_types/video/mov_bbb.mp4 b/docs/data_types/video/mov_bbb.mp4
new file mode 100644
index 00000000000..530676502bd
Binary files /dev/null and b/docs/data_types/video/mov_bbb.mp4 differ
diff --git a/docs/data_types/video/mov_bbb_framerate_60.mp4 b/docs/data_types/video/mov_bbb_framerate_60.mp4
new file mode 100644
index 00000000000..46345c4f24f
Binary files /dev/null and b/docs/data_types/video/mov_bbb_framerate_60.mp4 differ
diff --git a/docs/data_types/video/video.md b/docs/data_types/video/video.md
new file mode 100644
index 00000000000..83b081c0cf0
--- /dev/null
+++ b/docs/data_types/video/video.md
@@ -0,0 +1,218 @@
+# 🎥 Video
+
+DocArray supports many modalities including `Video`.
+This section will show you how to load and handle video data using DocArray.
+
+Moreover, you will learn about DocArray's video-specific types, to represent your video data ranging from [`VideoUrl`][docarray.typing.url.VideoUrl] to [`VideoBytes`][docarray.typing.bytes.VideoBytes] and [`VideoNdArray`][docarray.typing.tensor.video.video_ndarray.VideoNdArray].
+
+!!! note
+ This requires an `av` dependency. You can install all necessary dependencies via:
+ ```cmd
+ pip install "docarray[video]"
+ ```
+
+## Load video data
+
+In DocArray video data is represented by a video tensor, an audio tensor, and the key frame indices.
+
+{: style='width: 600px; height: 330px'}
+
+!!! tip
+ Check out our predefined [`VideoDoc`](#getting-started-predefined-videodoc) to get started and play around with our video features.
+
+First, let's define a `MyVideo` class with all of those attributes and instantiate an object with a local or remote URL:
+
+```python
+from docarray import BaseDoc
+from docarray.typing import AudioNdArray, NdArray, VideoNdArray, VideoUrl
+
+
+class MyVideo(BaseDoc):
+ url: VideoUrl
+ video: VideoNdArray = None
+ audio: AudioNdArray = None
+ key_frame_indices: NdArray = None
+
+
+doc = MyVideo(
+ url='https://github.com/docarray/docarray/blob/main/tests/toydata/mov_bbb.mp4?raw=true'
+)
+```
+
+Now you can load the video file content by simply calling [`.load()`][docarray.typing.url.audio_url.AudioUrl.load] on your [`AudioUrl`][docarray.typing.url.audio_url.AudioUrl] instance.
+This will return a [NamedTuple](https://docs.python.org/3/library/typing.html#typing.NamedTuple) of a **video tensor**, an **audio tensor**, and the **key frame indices**:
+
+- The video tensor is a 4-dimensional array of shape `(n_frames, height, width, channels)`.
+ - The first dimension represents the frame id.
+ - The last three dimensions represent the image data of the corresponding frame.
+- If the video contains audio, it will be stored as an `AudioNdArray`.
+- Additionally, the key frame indices will be stored. A key frame is defined as the starting point of any smooth transition.
+
+```python
+doc.video, doc.audio, doc.key_frame_indices = doc.url.load()
+
+assert isinstance(doc.video, VideoNdArray)
+assert isinstance(doc.audio, AudioNdArray)
+assert isinstance(doc.key_frame_indices, NdArray)
+
+print(doc.video.shape)
+```
+
+``` { .text .no-copy }
+(250, 176, 320, 3)
+```
+
+For the given example you can infer from `doc.video`'s shape that the video contains 250 frames of size 176x320 in RGB mode.
+Based on the overall length of the video (10 seconds), you can infer the framerate is approximately 250/10 = 25 frames per second (fps).
+
+## VideoTensor
+
+DocArray offers several [`VideoTensor`s](../../../../api_references/typing/tensor/video) to store your data to:
+
+- [`VideoNdArray`][docarray.typing.tensor.video.video_ndarray.VideoNdArray]
+- [`VideoTorchTensor`][docarray.typing.tensor.video.VideoTorchTensor]
+- [`VideoTensorFlowTensor`][docarray.typing.tensor.video.VideoTensorFlowTensor]
+
+If you specify the type of your tensor as one of the above, it will be cast to that automatically:
+
+```python hl_lines="7 8 15 16"
+from docarray import BaseDoc
+from docarray.typing import VideoTensorFlowTensor, VideoTorchTensor, VideoUrl
+
+
+class MyVideo(BaseDoc):
+ url: VideoUrl
+ tf_tensor: VideoTensorFlowTensor = None
+ torch_tensor: VideoTorchTensor = None
+
+
+doc = MyVideo(
+ url='https://github.com/docarray/docarray/blob/main/tests/toydata/mov_bbb.mp4?raw=true'
+)
+
+doc.tf_tensor = doc.url.load().video
+doc.torch_tensor = doc.url.load().video
+
+assert isinstance(doc.tf_tensor, VideoTensorFlowTensor)
+assert isinstance(doc.torch_tensor, VideoTorchTensor)
+```
+
+## VideoBytes
+
+Alternatively, you can load your [`VideoUrl`][docarray.typing.url.VideoUrl] instance to [`VideoBytes`][docarray.typing.bytes.VideoBytes], and your [`VideoBytes`][docarray.typing.bytes.VideoBytes] instance to a [`VideoTensor`](../../../../api_references/typing/tensor/video) of your choice:
+
+```python hl_lines="15 16"
+from docarray import BaseDoc
+from docarray.typing import VideoTensor, VideoUrl, VideoBytes
+
+
+class MyVideo(BaseDoc):
+ url: VideoUrl
+ bytes_: VideoBytes = None
+ video: VideoTensor = None
+
+
+doc = MyVideo(
+ url='https://github.com/docarray/docarray/blob/main/tests/toydata/mov_bbb.mp4?raw=true'
+)
+
+doc.bytes_ = doc.url.load_bytes()
+doc.video = doc.url.load().video
+```
+
+Vice versa, you can also transform a [`VideoTensor`](../../../../api_references/typing/tensor/video) to [`VideoBytes`][docarray.typing.bytes.VideoBytes]:
+
+```python
+from docarray.typing import VideoBytes
+
+bytes_from_tensor = doc.video.to_bytes()
+
+assert isinstance(bytes_from_tensor, VideoBytes)
+```
+
+## Key frame extraction
+
+A key frame is defined as the starting point of any smooth transition.
+Given the key frame indices, you can access selected scenes:
+
+```python
+indices = doc.key_frame_indices
+first_scene = doc.video[indices[0] : indices[1]]
+
+assert (indices == [0, 95]).all()
+assert first_scene.shape == (95, 176, 320, 3)
+```
+
+Or you can access the first frame of all new scenes and display them in a notebook:
+
+``` { .python }
+from docarray.typing import ImageNdArray
+from pydantic import parse_obj_as
+
+
+key_frames = doc.video[doc.key_frame_indices]
+for frame in key_frames:
+ img = parse_obj_as(ImageNdArray, frame)
+ img.display()
+```
+
+Output
+ ``` { .text .no-copy } + 📄 Page : 13d909a ... + └── 💠 content: DocList[Sentence] + ├── 📄 Sentence : 6725382 ... + │ ╭────────────────┬─────────────────────╮ + │ │ Attribute │ Value │ + │ ├────────────────┼─────────────────────┤ + │ │ text: str │ First sentence │ + │ ╰────────────────┴─────────────────────╯ + ├── 📄 Sentence : 17a934c ... + │ ╭───────────────┬──────────────────────╮ + │ │ Attribute │ Value │ + │ ├───────────────┼──────────────────────┤ + │ │ text: str │ Second sentence │ + │ ╰───────────────┴──────────────────────╯ + └── ... 2 more Sentence documents + ``` +| hello.wav | -olleh.wav | -
|---|---|
| - | - |
+
diff --git a/docs/user_guide/storing/doc_store/store_s3.md b/docs/user_guide/storing/doc_store/store_s3.md
new file mode 100644
index 00000000000..cd26f1a358d
--- /dev/null
+++ b/docs/user_guide/storing/doc_store/store_s3.md
@@ -0,0 +1,142 @@
+# Store on S3
+When you want to use your [`DocList`][docarray.DocList] in another place, you can use the
+[`.push`][docarray.array.doc_list.pushpull.PushPullMixin.push] method to push the `DocList` to S3 and later use the
+[`.pull`][docarray.array.doc_list.pushpull.PushPullMixin.pull] function to pull its content back.
+
+!!! note
+ To store on S3, you need to install the extra dependency with the following line
+ ```cmd
+ pip install "docarray[aws]"
+ ```
+
+## Push & pull
+To use the store [`DocList`][docarray.DocList] on S3, you need to pass an S3 path to the function starting with `'s3://'`.
+
+In the following demo, we use `MinIO` as a local S3 service. You could use the following docker compose file to start the service in a Docker container.
+
+```yaml
+version: "3"
+services:
+ minio:
+ container_name: minio
+ image: "minio/minio:RELEASE.2023-03-13T19-46-17Z"
+ ports:
+ - "9005:9000"
+ command: server /data
+```
+Save the above file as `docker-compose.yml` and run the following line in the same folder as the file.
+```cmd
+docker compose up
+```
+
+```python
+from docarray import BaseDoc, DocList
+
+
+class SimpleDoc(BaseDoc):
+ text: str
+
+
+if __name__ == '__main__':
+ import boto3
+ from botocore.client import Config
+
+ BUCKET = 'tmp_bucket'
+ my_session = boto3.session.Session()
+ s3 = my_session.resource(
+ service_name='s3',
+ region_name="us-east-1",
+ use_ssl=False,
+ endpoint_url="http://localhost:9005",
+ aws_access_key_id="minioadmin",
+ aws_secret_access_key="minioadmin",
+ config=Config(signature_version="s3v4"),
+ )
+ # make a bucket
+ s3.create_bucket(Bucket=BUCKET)
+
+ store_docs = [SimpleDoc(text=f'doc {i}') for i in range(8)]
+ docs = DocList[SimpleDoc]()
+ docs.extend([SimpleDoc(text=f'doc {i}') for i in range(8)])
+
+ # .push() and .pull() use the default boto3 client
+ boto3.Session.client.__defaults__ = (
+ "us-east-1",
+ None,
+ False,
+ None,
+ "http://localhost:9005",
+ "minioadmin",
+ "minioadmin",
+ None,
+ Config(signature_version="s3v4"),
+ )
+ docs.push(f's3://{BUCKET}/simple_docs')
+ docs_pull = DocList[SimpleDoc].pull(f's3://{BUCKET}/simple_docs')
+```
+
+Under the bucket `tmp_bucket`, there is a file with the name of `simple_docs.docs` being created to store the `DocList`.
+
+!!! note
+ When using `.push()` and `.pull()`, `DocList` calls the default boto3 client. Be sure your default session is correctly set up.
+
+
+## Push & pull with streaming
+When you have a large amount of documents to push and pull, you could use the streaming function.
+[`.push_stream()`][docarray.array.doc_list.pushpull.PushPullMixin.push_stream] and
+[`.pull_stream()`][docarray.array.doc_list.pushpull.PushPullMixin.pull_stream] can help you to stream the
+[`DocList`][docarray.DocList] in order to save the memory usage. You set multiple [`DocList`][docarray.DocList] to pull from the same source as well. The usage is the same as using streaming with local files. Please refer to [Push & Pull with streaming with local files](store_file.md#push-pull-with-streaming).
+
+
+## Delete
+To delete the store, you need to use the static method [`.delete()`][docarray.store.s3.S3DocStore.delete] of [`S3DocStore`][docarray.store.s3.S3DocStore] class.
+
+```python hl_lines="44-47"
+from docarray import BaseDoc, DocList
+
+
+class SimpleDoc(BaseDoc):
+ text: str
+
+
+if __name__ == '__main__':
+ import boto3
+ from botocore.client import Config
+
+ BUCKET = 'tmp_bucket'
+ my_session = boto3.session.Session()
+ s3 = my_session.resource(
+ service_name='s3',
+ region_name="us-east-1",
+ use_ssl=False,
+ endpoint_url="http://localhost:9005",
+ aws_access_key_id="minioadmin",
+ aws_secret_access_key="minioadmin",
+ config=Config(signature_version="s3v4"),
+ )
+ # make a bucket
+ s3.create_bucket(Bucket=BUCKET)
+
+ store_docs = [SimpleDoc(text=f'doc {i}') for i in range(8)]
+ docs = DocList[SimpleDoc]()
+ docs.extend([SimpleDoc(text=f'doc {i}') for i in range(8)])
+
+ # .push() and .pull() use the default boto3 client
+ boto3.Session.client.__defaults__ = (
+ "us-east-1",
+ None,
+ False,
+ None,
+ "http://localhost:9005",
+ "minioadmin",
+ "minioadmin",
+ None,
+ Config(signature_version="s3v4"),
+ )
+ docs.push(f's3://{BUCKET}/simple_docs')
+
+ # delete bucket
+ from docarray.store import S3DocStore
+
+ success = S3DocStore.delete('{BUCKET}/simple_docs')
+```
diff --git a/docs/user_guide/storing/docindex.md b/docs/user_guide/storing/docindex.md
new file mode 100644
index 00000000000..7293c38597f
--- /dev/null
+++ b/docs/user_guide/storing/docindex.md
@@ -0,0 +1,190 @@
+# Introduction
+
+A Document Index lets you store your documents and search through them using vector similarity.
+
+This is useful if you want to store a bunch of data, and at a later point retrieve documents that are similar to
+some query that you provide.
+Relevant concrete examples are neural search applications, augmenting LLMs and chatbots with domain knowledge ([Retrieval-Augmented Generation](https://arxiv.org/abs/2005.11401)),
+or recommender systems.
+
+!!! question "How does vector similarity search work?"
+ Without going into too much detail, the idea behind vector similarity search is the following:
+
+ You represent every data point that you have (in our case, a document) as a _vector_, or _embedding_.
+ This vector should represent as much semantic information about your data as possible: Similar data points should
+ be represented by similar vectors.
+
+ These vectors (embeddings) are usually obtained by passing the data through a suitable neural network that has been
+ trained to produce such semantic representations - this is the _encoding_ step.
+
+ Once you have your vectors that represent your data, you can store them, for example in a vector database.
+
+ To perform similarity search, you take your input query and encode it in the same way as the data in your database.
+ Then, the database will search through the stored vectors and return those that are most similar to your query.
+ This similarity is measured by a _similarity metric_, which can be [cosine similarity](https://en.wikipedia.org/wiki/Cosine_similarity),
+ [Euclidean distance](https://en.wikipedia.org/wiki/Euclidean_distance), or any other metric that you can think of.
+
+ If you store a lot of data, performing this similarity computation for every data point in your database is expensive.
+ Therefore, vector databases usually perform _approximate nearest neighbor (ANN)_ search.
+ There are various algorithms for doing this, such as [HNSW](https://arxiv.org/abs/1603.09320), but in a nutshell,
+ they allow you to search through a large database of vectors very quickly, at the expense of a small loss in accuracy.
+
+DocArray's Document Index concept achieves this by providing a unified interface to a number of [vector databases](https://learn.microsoft.com/en-us/semantic-kernel/concepts-ai/vectordb).
+In fact, you can think of Document Index as an **[ORM](https://sqlmodel.tiangolo.com/db-to-code/) for vector databases**.
+
+Currently, DocArray supports the following vector databases:
+
+- [Weaviate](https://weaviate.io/) | [Docs](index_weaviate.md)
+- [Qdrant](https://qdrant.tech/) | [Docs](index_qdrant.md)
+- [Elasticsearch](https://www.elastic.co/elasticsearch/) v7 and v8 | [Docs](index_elastic.md)
+- [Epsilla](https://epsilla.com/) | [Docs](index_epsilla.md)
+- [Redis](https://redis.com/) | [Docs](index_redis.md)
+- [Milvus](https://milvus.io/) | [Docs](index_milvus.md)
+- [HNSWlib](https://github.com/nmslib/hnswlib) | [Docs](index_hnswlib.md)
+- InMemoryExactNNIndex | [Docs](index_in_memory.md)
+
+
+## Basic usage
+
+Let's learn the basic capabilities of Document Index with [InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex].
+This doesn't require a database server - rather, it saves your data locally.
+
+
+!!! note "Using a different vector database"
+ You can easily use Weaviate, Qdrant, Redis, Milvus or Elasticsearch instead -- their APIs are largely identical!
+ To do so, check their respective documentation sections.
+
+!!! note "InMemoryExactNNIndex in more detail"
+ The following section only covers the basics of InMemoryExactNNIndex.
+ For a deeper understanding, please look into its [documentation](index_in_memory.md).
+
+### Define document schema and create data
+The following code snippet defines a document schema using the `BaseDoc` class. Each document consists of a title (a string),
+a price (an integer), and an embedding (a 128-dimensional array). It also creates a list of ten documents with dummy titles,
+prices ranging from 0 to 9, and randomly generated embeddings.
+```python
+from docarray import BaseDoc, DocList
+from docarray.index import InMemoryExactNNIndex
+from docarray.typing import NdArray
+import numpy as np
+
+
+class MyDoc(BaseDoc):
+ title: str
+ price: int
+ embedding: NdArray[128]
+
+
+docs = DocList[MyDoc](
+ MyDoc(title=f"title #{i}", price=i, embedding=np.random.rand(128))
+ for i in range(10)
+)
+```
+
+### Initialize the Document Index and add data
+Here we initialize an `InMemoryExactNNIndex` instance with the document schema we defined previously, and add the created documents to this index.
+```python
+doc_index = InMemoryExactNNIndex[MyDoc]()
+doc_index.index(docs)
+```
+
+### Perform a vector similarity search
+Now, let's perform a similarity search on the document embeddings.
+As a result, we'll retrieve the ten most similar documents and their corresponding similarity scores.
+```python
+query = np.ones(128)
+retrieved_docs, scores = doc_index.find(query, search_field='embedding', limit=10)
+```
+
+### Filter documents
+In this snippet, we filter the indexed documents based on their price field, specifically retrieving documents with a price less than 5:
+```python
+query = {'price': {'$lt': 5}}
+filtered_docs = doc_index.filter(query, limit=10)
+```
+
+### Combine different search methods
+The final snippet combines the vector similarity search and filtering operations into a single query.
+We first perform a similarity search on the document embeddings and then apply a filter to return only those documents with a price greater than or equal to 2:
+```python
+query = (
+ doc_index.build_query() # get empty query object
+ .find(query=np.ones(128), search_field='embedding') # add vector similarity search
+ .filter(filter_query={'price': {'$gte': 2}}) # add filter search
+ .build() # build the query
+)
+retrieved_docs, scores = doc_index.execute_query(query)
+```
+
+### Using a predefined document as schema
+
+DocArray offers a number of predefined documents, like [ImageDoc][docarray.documents.ImageDoc] and [TextDoc][docarray.documents.TextDoc].
+If you try to use these directly as a schema for a Document Index, you will get unexpected behavior:
+Depending on the backend, an exception will be raised, or no vector index for ANN lookup will be built.
+
+The reason for this is that predefined documents don't hold information about the dimensionality of their `.embedding`
+field. But this is crucial information for any vector database to work properly!
+
+You can work around this problem by subclassing the predefined document and adding the dimensionality information:
+
+=== "Using type hint"
+ ```python
+ from docarray.documents import TextDoc
+ from docarray.typing import NdArray
+ from docarray.index import HnswDocumentIndex
+
+
+ class MyDoc(TextDoc):
+ embedding: NdArray[128]
+
+
+ db = HnswDocumentIndex[MyDoc]('test_db')
+ ```
+
+=== "Using Field()"
+ ```python
+ from docarray.documents import TextDoc
+ from docarray.typing import AnyTensor
+ from docarray.index import HnswDocumentIndex
+ from pydantic import Field
+
+
+ class MyDoc(TextDoc):
+ embedding: AnyTensor = Field(dim=128)
+
+
+ db = HnswDocumentIndex[MyDoc]('test_db3')
+ ```
+
+Once you have defined the schema of your Document Index in this way, the data that you index can be either the predefined Document type or your custom Document type.
+
+The [next section](#index) goes into more detail about data indexing, but note that if you have some `TextDoc`s, `ImageDoc`s etc. that you want to index, you _don't_ need to cast them to `MyDoc`:
+
+```python
+from docarray import DocList
+
+# data of type TextDoc
+data = DocList[TextDoc](
+ [
+ TextDoc(text='hello world', embedding=np.random.rand(128)),
+ TextDoc(text='hello world', embedding=np.random.rand(128)),
+ TextDoc(text='hello world', embedding=np.random.rand(128)),
+ ]
+)
+
+# you can index this into Document Index of type MyDoc
+db.index(data)
+```
+
+## Learn more
+The code snippets above just scratch the surface of what a Document Index can do.
+To learn more and get the most out of `DocArray`, take a look at the detailed guides for the vector database backends you're interested in:
+
+- [Weaviate](https://weaviate.io/) | [Docs](index_weaviate.md)
+- [Qdrant](https://qdrant.tech/) | [Docs](index_qdrant.md)
+- [Elasticsearch](https://www.elastic.co/elasticsearch/) v7 and v8 | [Docs](index_elastic.md)
+- [Epsilla](https://epsilla.com/) | [Docs](index_epsilla.md)
+- [Redis](https://redis.com/) | [Docs](index_redis.md)
+- [Milvus](https://milvus.io/) | [Docs](index_milvus.md)
+- [HNSWlib](https://github.com/nmslib/hnswlib) | [Docs](index_hnswlib.md)
+- InMemoryExactNNIndex | [Docs](index_in_memory.md)
diff --git a/docs/user_guide/storing/first_step.md b/docs/user_guide/storing/first_step.md
new file mode 100644
index 00000000000..9cd9b7e3e4d
--- /dev/null
+++ b/docs/user_guide/storing/first_step.md
@@ -0,0 +1,46 @@
+# Introduction
+
+In the previous sections we saw how to use [`BaseDoc`][docarray.base_doc.doc.BaseDoc], [`DocList`][docarray.array.doc_list.doc_list.DocList] and [`DocVec`][docarray.array.doc_vec.doc_vec.DocVec] to represent multimodal data and send it over the wire.
+In this section we will see how to store and persist this data.
+
+DocArray offers two ways of storing your data, each of which have their own documentation sections:
+
+1. **[Document Store](#document-store)** for simple long-term storage
+2. **[Document Index](#document-index)** for fast retrieval using vector similarity
+
+## Document Store
+
+[DocList][docarray.array.doc_list.doc_list.DocList] can be persisted using the
+[`.push()`][docarray.array.doc_list.pushpull.PushPullMixin.push] and
+[`.pull()`][docarray.array.doc_list.pushpull.PushPullMixin.pull] methods.
+Under the hood, [DocStore][docarray.store.abstract_doc_store.AbstractDocStore] is used to persist a `DocList`.
+You can either store your documents on-disk or upload them to [AWS S3](https://aws.amazon.com/s3/) or
+[minio](https://min.io).
+
+This section covers the following three topics:
+
+ - [Storing](doc_store/store_file.md) [`BaseDoc`][docarray.base_doc.doc.BaseDoc], [`DocList`][docarray.array.doc_list.doc_list.DocList] and [`DocVec`][docarray.array.doc_vec.doc_vec.DocVec] on-disk
+ - [Storing on S3](doc_store/store_s3.md)
+
+## Document Index
+
+A Document Index lets you store your documents and search through them using vector similarity.
+
+This is useful if you want to store a bunch of data, and at a later point retrieve documents that are similar to
+a query that you provide.
+Relevant concrete examples are neural search applications, augmenting LLMs and chatbots with domain knowledge ([Retrieval-Augmented Generation](https://arxiv.org/abs/2005.11401))]),
+or recommender systems.
+
+DocArray's Document Index concept achieves this by providing a unified interface to a number of [vector databases](https://learn.microsoft.com/en-us/semantic-kernel/concepts-ai/vectordb).
+In fact, you can think of Document Index as an **[ORM](https://sqlmodel.tiangolo.com/db-to-code/) for vector databases**.
+
+Currently, DocArray supports the following vector indexes. Some of them wrap vector databases (Weaviate, Qdrant, ElasticSearch) and act as a client for them, while others
+use a vector search library locally (HNSWLib, Exact NN search):
+
+- [Weaviate](https://weaviate.io/) | [Docs](index_weaviate.md)
+- [Qdrant](https://qdrant.tech/) | [Docs](index_qdrant.md)
+- [Elasticsearch](https://www.elastic.co/elasticsearch/) v7 and v8 | [Docs](index_elastic.md)
+- [Redis](https://redis.com/) | [Docs](index_redis.md)
+- [Milvus](https://milvus.io/) | [Docs](index_milvus.md)
+- [Hnswlib](https://github.com/nmslib/hnswlib) | [Docs](index_hnswlib.md)
+- InMemoryExactNNSearch | [Docs](index_in_memory.md)
diff --git a/docs/user_guide/storing/index_elastic.md b/docs/user_guide/storing/index_elastic.md
new file mode 100644
index 00000000000..89a104fefa6
--- /dev/null
+++ b/docs/user_guide/storing/index_elastic.md
@@ -0,0 +1,520 @@
+# Elasticsearch Document Index
+
+DocArray comes with two Document Indexes for [Elasticsearch](https://www.elastic.co/elasticsearch/):
+
+- [ElasticDocIndex][docarray.index.backends.elastic.ElasticDocIndex], based on [Elasticsearch 8](https://github.com/elastic/elasticsearch).
+- [ElasticV7DocIndex][docarray.index.backends.elasticv7.ElasticV7DocIndex], based on [Elasticsearch 7.10](https://www.elastic.co/downloads/past-releases/elasticsearch-7-10-0).
+
+!!! tip "Should you use ES v7 or v8?"
+ [Elasticsearch v8](https://www.elastic.co/blog/whats-new-elastic-8-0-0) is the current version of ES and offers
+ **native vector search (ANN) support**, alongside text and range search.
+
+ [Elasticsearch v7.10](https://www.elastic.co/downloads/past-releases/elasticsearch-7-10-0) can store vectors, but
+ **does _not_ support native ANN vector search**, but only exhaustive (i.e. slow) vector search, alongside text and range search.
+
+ Some users prefer to use ES v7.10 because it is available under a [different license](https://www.elastic.co/pricing/faq/licensing) to ES v8.0.0.
+
+!!! note "Installation"
+ To use [ElasticDocIndex][docarray.index.backends.elastic.ElasticDocIndex], you need to install the following dependencies:
+
+ ```console
+ pip install elasticsearch==8.6.2
+ pip install elastic-transport
+ ```
+
+ To use [ElasticV7DocIndex][docarray.index.backends.elasticv7.ElasticV7DocIndex], you need to install the following dependencies:
+
+ ```console
+ pip install elasticsearch==7.10.1
+ pip install elastic-transport
+ ```
+
+
+The following example is based on [ElasticDocIndex][docarray.index.backends.elastic.ElasticDocIndex],
+but will also work for [ElasticV7DocIndex][docarray.index.backends.elasticv7.ElasticV7DocIndex].
+
+
+## Basic usage
+This snippet demonstrates the basic usage of [ElasticDocIndex][docarray.index.backends.elastic.ElasticDocIndex]. It defines a document schema with a title and an embedding,
+creates ten dummy documents with random embeddings, initializes an instance of [ElasticDocIndex][docarray.index.backends.elastic.ElasticDocIndex] to index these documents,
+and performs a vector similarity search to retrieve the ten most similar documents to a given query vector.
+
+```python
+from docarray import BaseDoc, DocList
+from docarray.index import ElasticDocIndex # or ElasticV7DocIndex
+from docarray.typing import NdArray
+import numpy as np
+
+
+# Define the document schema.
+class MyDoc(BaseDoc):
+ title: str
+ embedding: NdArray[128]
+
+
+# Create dummy documents.
+docs = DocList[MyDoc](
+ MyDoc(title=f'title #{i}', embedding=np.random.rand(128)) for i in range(10)
+)
+
+# Initialize a new ElasticDocIndex instance and add the documents to the index.
+doc_index = ElasticDocIndex[MyDoc](index_name='my_index')
+doc_index.index(docs)
+
+# Perform a vector search.
+query = np.ones(128)
+retrieved_docs = doc_index.find(query, search_field='embedding', limit=10)
+```
+
+
+
+## Initialize
+
+
+You can use docker compose to create a local Elasticsearch service with the following `docker-compose.yml`.
+
+```yaml
+version: "3.3"
+services:
+ elastic:
+ image: docker.elastic.co/elasticsearch/elasticsearch:8.6.2
+ environment:
+ - xpack.security.enabled=false
+ - discovery.type=single-node
+ - ES_JAVA_OPTS=-Xmx1024m
+ ports:
+ - "9200:9200"
+ networks:
+ - elastic
+
+networks:
+ elastic:
+ name: elastic
+```
+
+Run the following command in the folder of the above `docker-compose.yml` to start the service:
+
+```bash
+docker compose up
+```
+
+### Schema definition
+
+To construct an index, you first need to define a schema in the form of a `Document`.
+
+There are a number of configurations you can pack into your schema:
+
+- Every field in your schema will become one column in the database
+- For vector fields, such as `NdArray`, `TorchTensor`, or `TensorflowTensor`, you need to specify a dimensionality to be able to perform vector search
+- You can override the default column type for every field by passing any [ES field data type](https://www.elastic.co/guide/en/elasticsearch/reference/8.6/mapping-types.html) to `field_name: Type = Field(col_type=...)`. You can see an example of this in the [section on keyword filters](#keyword-filter).
+
+Additionally, you can pass a `hosts` argument to the `__init__()` method to connect to an ES instance.
+By default, it is `http://localhost:9200`.
+
+```python
+import numpy as np
+from pydantic import Field
+
+from docarray import BaseDoc
+from docarray.index import ElasticDocIndex
+from docarray.typing import NdArray
+
+
+class SimpleDoc(BaseDoc):
+ # specify tensor field with dimensionality 128
+ tensor: NdArray[128]
+ # alternative and equivalent definition:
+ # tensor: NdArray = Field(dims=128)
+
+
+doc_index = ElasticDocIndex[SimpleDoc](hosts='http://localhost:9200')
+```
+
+## Index
+
+Now that you have a Document Index, you can add data to it, using the [`index()`][docarray.index.abstract.BaseDocIndex.index] method.
+The `.num_docs()` method returns the total number of documents in the index.
+
+```python
+from docarray import DocList
+
+# create some random data
+docs = DocList[SimpleDoc]([SimpleDoc(tensor=np.ones(128)) for _ in range(64)])
+
+doc_index.index(docs)
+
+print(f'number of docs in the index: {doc_index.num_docs()}')
+```
+
+As you can see, `DocList[SimpleDoc]` and `ElasticDocIndex[SimpleDoc]` both have `SimpleDoc` as a parameter.
+This means that they share the same schema, and in general, both the Document Index and the data that you want to store need to have compatible schemas.
+
+!!! question "When are two schemas compatible?"
+ The schemas of your Document Index and data need to be compatible with each other.
+
+ Let's say A is the schema of your Document Index and B is the schema of your data.
+ There are a few rules that determine if schema A is compatible with schema B.
+ If _any_ of the following are true, then A and B are compatible:
+
+ - A and B are the same class
+ - A and B have the same field names and field types
+ - A and B have the same field names, and, for every field, the type of B is a subclass of the type of A
+
+ In particular, this means that you can easily [index predefined documents](#using-a-predefined-document-as-schema) into a Document Index.
+
+
+
+## Vector search
+
+Now that you have indexed your data, you can perform vector similarity search using the [`find()`][docarray.index.abstract.BaseDocIndex.find] method.
+
+You can use the [`find()`][docarray.index.abstract.BaseDocIndex.find] function with a document of the type `MyDoc`
+to find similar documents within the Document Index:
+
+You can use the `limit` argument to configure how many documents to return.
+
+!!! note
+ [ElasticV7DocIndex][docarray.index.backends.elasticv7.ElasticV7DocIndex] uses Elasticsearch v7.10.1, which does not support approximate nearest neighbour algorithms such as HNSW.
+ This can lead to poor performance when the search involves many vectors.
+ [ElasticDocIndex][docarray.index.backends.elastic.ElasticDocIndex] does not have this limitation.
+
+=== "Search by Document"
+
+ ```python
+ # create a query document
+ query = SimpleDoc(tensor=np.ones(128))
+
+ # find similar documents
+ matches, scores = doc_index.find(query, search_field='tensor', limit=5)
+
+ print(f'{matches=}')
+ print(f'{matches.text=}')
+ print(f'{scores=}')
+ ```
+
+=== "Search by raw vector"
+
+ ```python
+ # create a query vector
+ query = np.random.rand(128)
+
+ # find similar documents
+ matches, scores = doc_index.find(query, search_field='tensor', limit=5)
+
+ print(f'{matches=}')
+ print(f'{matches.text=}')
+ print(f'{scores=}')
+ ```
+
+To peform a vector search, you need to specify a `search_field`. This is the field that serves as the
+basis of comparison between your query and the documents in the Document Index.
+
+In this example you only have one field (`tensor`) that is a vector, so you can trivially choose that one.
+In general, you could have multiple fields of type `NdArray` or `TorchTensor` or `TensorFlowTensor`, and you can choose
+which one to use for the search.
+
+The [`find()`][docarray.index.abstract.BaseDocIndex.find] method returns a named tuple containing the closest
+matching documents and their associated similarity scores.
+
+When searching on the subindex level, you can use the [`find_subindex()`][docarray.index.abstract.BaseDocIndex.find_subindex] method, which returns a named tuple containing the subindex documents, similarity scores and their associated root documents.
+
+How these scores are calculated depends on the backend, and can usually be [configured](#configuration).
+
+
+### Batched search
+
+You can also search for multiple documents at once, in a batch, using the [`find_batched()`][docarray.index.abstract.BaseDocIndex.find_batched] method.
+
+=== "Search by Documents"
+
+ ```python
+ # create some query Documents
+ queries = DocList[SimpleDoc](SimpleDoc(tensor=np.random.rand(128)) for i in range(3))
+
+ # find similar documents
+ matches, scores = doc_index.find_batched(queries, search_field='tensor', limit=5)
+
+ print(f'{matches=}')
+ print(f'{matches[0].text=}')
+ print(f'{scores=}')
+ ```
+
+=== "Search by raw vectors"
+
+ ```python
+ # create some query vectors
+ query = np.random.rand(3, 128)
+
+ # find similar documents
+ matches, scores = doc_index.find_batched(query, search_field='tensor', limit=5)
+
+ print(f'{matches=}')
+ print(f'{matches[0].text=}')
+ print(f'{scores=}')
+ ```
+
+The [`find_batched()`][docarray.index.abstract.BaseDocIndex.find_batched] method returns a named tuple containing
+a list of `DocList`s, one for each query, containing the closest matching documents and their similarity scores.
+
+
+
+## Filter
+
+You can filter your documents by using the `filter()` or `filter_batched()` method with a corresponding filter query.
+The query should follow [Elastic's query language](https://www.elastic.co/guide/en/elasticsearch/reference/current/query-filter-context.html).
+
+The `filter()` method accepts queries that follow the [Elasticsearch Query DSL](https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl.html) and consists of leaf and compound clauses.
+
+Using this, you can perform [keyword filters](#keyword-filter), [geolocation filters](#geolocation-filter) and [range filters](#range-filter).
+
+### Keyword filter
+
+To filter documents in your index by keyword, you can use `Field(col_type='keyword')` to enable keyword search for given fields:
+
+```python
+class NewsDoc(BaseDoc):
+ text: str
+ category: str = Field(col_type='keyword') # enable keyword filtering
+
+
+doc_index = ElasticDocIndex[NewsDoc]()
+index_docs = [
+ NewsDoc(id='0', text='this is a news for sport', category='sport'),
+ NewsDoc(id='1', text='this is a news for finance', category='finance'),
+ NewsDoc(id='2', text='this is another news for sport', category='sport'),
+]
+doc_index.index(index_docs)
+
+# search with filer
+query_filter = {'terms': {'category': ['sport']}}
+docs = doc_index.filter(query_filter)
+```
+
+### Geolocation filter
+
+To filter documents in your index by geolocation, you can use `Field(col_type='geo_point')` on a given field:
+
+```python
+class NewsDoc(BaseDoc):
+ text: str
+ location: dict = Field(col_type='geo_point') # enable geolocation filtering
+
+
+doc_index = ElasticDocIndex[NewsDoc]()
+index_docs = [
+ NewsDoc(text='this is from Berlin', location={'lon': 13.24, 'lat': 50.31}),
+ NewsDoc(text='this is from Beijing', location={'lon': 116.22, 'lat': 39.55}),
+ NewsDoc(text='this is from San Jose', location={'lon': -121.89, 'lat': 37.34}),
+]
+doc_index.index(index_docs)
+
+# filter the eastern hemisphere
+query = {
+ 'bool': {
+ 'filter': {
+ 'geo_bounding_box': {
+ 'location': {
+ 'top_left': {'lon': 0, 'lat': 90},
+ 'bottom_right': {'lon': 180, 'lat': 0},
+ }
+ }
+ }
+ }
+}
+
+docs = doc_index.filter(query)
+```
+
+### Range filter
+
+You can have [range field types](https://www.elastic.co/guide/en/elasticsearch/reference/8.6/range.html) in your document schema and set `Field(col_type='integer_range')`(or also `date_range`, etc.) to filter documents based on the range of the field.
+
+```python
+class NewsDoc(BaseDoc):
+ time_frame: dict = Field(
+ col_type='date_range', format='yyyy-MM-dd'
+ ) # enable range filtering
+
+
+doc_index = ElasticDocIndex[NewsDoc]()
+index_docs = [
+ NewsDoc(time_frame={'gte': '2023-01-01', 'lt': '2023-02-01'}),
+ NewsDoc(time_frame={'gte': '2023-02-01', 'lt': '2023-03-01'}),
+ NewsDoc(time_frame={'gte': '2023-03-01', 'lt': '2023-04-01'}),
+]
+doc_index.index(index_docs)
+
+query = {
+ 'bool': {
+ 'filter': {
+ 'range': {
+ 'time_frame': {
+ 'gte': '2023-02-05',
+ 'lt': '2023-02-10',
+ 'relation': 'contains',
+ }
+ }
+ }
+ }
+}
+
+docs = doc_index.filter(query)
+```
+
+
+## Text search
+
+In addition to vector similarity search, the Document Index interface offers methods for text search:
+[`text_search()`][docarray.index.abstract.BaseDocIndex.text_search],
+as well as the batched version [`text_search_batched()`][docarray.index.abstract.BaseDocIndex.text_search_batched].
+
+As in "pure" Elasticsearch, you can use text search directly on the field of type `str`:
+
+```python
+class NewsDoc(BaseDoc):
+ text: str
+
+
+doc_index = ElasticDocIndex[NewsDoc]()
+index_docs = [
+ NewsDoc(id='0', text='this is a news for sport'),
+ NewsDoc(id='1', text='this is a news for finance'),
+ NewsDoc(id='2', text='this is another news for sport'),
+]
+doc_index.index(index_docs)
+query = 'finance'
+
+# search with text
+docs, scores = doc_index.text_search(query, search_field='text')
+```
+
+
+## Hybrid search
+
+Document Index supports atomic operations for vector similarity search, text search and filter search.
+
+To combine these operations into a single, hybrid search query, you can use the query builder that is accessible
+through [`build_query()`][docarray.index.abstract.BaseDocIndex.build_query]:
+
+For example, you can build a hybrid serach query that performs range filtering, vector search and text search:
+
+```python
+class MyDoc(BaseDoc):
+ tens: NdArray[10] = Field(similarity='l2_norm')
+ num: int
+ text: str
+
+
+doc_index = ElasticDocIndex[MyDoc]()
+index_docs = [
+ MyDoc(id=f'{i}', tens=np.ones(10) * i, num=int(i / 2), text=f'text {int(i/2)}')
+ for i in range(10)
+]
+doc_index.index(index_docs)
+
+q = (
+ doc_index.build_query()
+ .filter({'range': {'num': {'lte': 3}}})
+ .find(index_docs[-1], search_field='tens')
+ .text_search('0', search_field='text')
+ .build()
+)
+docs, _ = doc_index.execute_query(q)
+```
+
+You can also manually build a valid ES query and directly pass it to the `execute_query()` method.
+
+
+## Access documents
+
+To access a document, you need to specify its `id`. You can also pass a list of `id`s to access multiple documents.
+
+```python
+# access a single Doc
+doc_index[index_docs[1].id]
+
+# access multiple Docs
+doc_index[index_docs[2].id, index_docs[3].id]
+```
+
+## Delete documents
+
+To delete documents, use the built-in function `del` with the `id` of the documents that you want to delete.
+You can also pass a list of `id`s to delete multiple documents.
+
+```python
+# delete a single Doc
+del doc_index[index_docs[1].id]
+
+# delete multiple Docs
+del doc_index[index_docs[2].id, index_docs[3].id]
+```
+
+
+## Configuration
+
+### DBConfig
+
+The following configs can be set in `DBConfig`:
+
+| Name | Description | Default |
+|-------------------|----------------------------------------------------------------------------------------------------------------------------------------|-------------------------|
+| `hosts` | Hostname of the Elasticsearch server | `http://localhost:9200` |
+| `es_config` | Other ES [configuration options](https://www.elastic.co/guide/en/elasticsearch/client/python-api/8.6/config.html) in a Dict and pass to `Elasticsearch` client constructor, e.g. `cloud_id`, `api_key` | `None` |
+| `index_name` | Elasticsearch index name, the name of Elasticsearch index object | `None`. Data will be stored in an index named after the Document type used as schema. |
+| `index_settings` | Other [index settings](https://www.elastic.co/guide/en/elasticsearch/reference/8.6/index-modules.html#index-modules-settings) in a Dict for creating the index | dict |
+| `index_mappings` | Other [index mappings](https://www.elastic.co/guide/en/elasticsearch/reference/8.6/mapping.html) in a Dict for creating the index | dict |
+| `default_column_config` | The default configurations for every column type. | dict |
+
+You can pass any of the above as keyword arguments to the `__init__()` method or pass an entire configuration object.
+See [here](docindex.md#configuration-options#customize-configurations) for more information.
+
+`default_column_config` is the default configurations for every column type. Since there are many column types in Elasticsearch, you can also consider changing the column config when defining the schema.
+
+```python
+class SimpleDoc(BaseDoc):
+ tensor: NdArray[128] = Field(similarity='l2_norm', m=32, num_candidates=5000)
+
+
+doc_index = ElasticDocIndex[SimpleDoc](index_name='my_index_1')
+```
+
+### RuntimeConfig
+
+The `RuntimeConfig` dataclass of `ElasticDocIndex` consists of `chunk_size`. You can change `chunk_size` for batch operations:
+
+```python
+doc_index = ElasticDocIndex[SimpleDoc](index_name='my_index_2')
+doc_index.configure(ElasticDocIndex.RuntimeConfig(chunk_size=1000))
+```
+
+You can pass the above as keyword arguments to the `configure()` method or pass an entire configuration object.
+See [here](docindex.md#configuration-options#customize-configurations) for more information.
+
+
+### Persistence
+
+You can hook into a database index that was persisted during a previous session by
+specifying the `index_name` and `hosts`:
+
+```python
+doc_index = ElasticDocIndex[MyDoc](
+ hosts='http://localhost:9200', index_name='previously_stored'
+)
+doc_index.index(index_docs)
+
+doc_index2 = ElasticDocIndex[MyDoc](
+ hosts='http://localhost:9200', index_name='previously_stored'
+)
+
+print(f'number of docs in the persisted index: {doc_index2.num_docs()}')
+```
+
+
+## Nested data and subindex search
+
+The examples provided primarily operate on a basic schema where each field corresponds to a straightforward type such as `str` or `NdArray`.
+However, it is also feasible to represent and store nested documents in a Document Index, including scenarios where a document
+contains a `DocList` of other documents.
+
+Go to the [Nested Data](nested_data.md) section to learn more.
\ No newline at end of file
diff --git a/docs/user_guide/storing/index_epsilla.md b/docs/user_guide/storing/index_epsilla.md
new file mode 100644
index 00000000000..425ebe48138
--- /dev/null
+++ b/docs/user_guide/storing/index_epsilla.md
@@ -0,0 +1,562 @@
+# Epsilla Document Index
+
+!!! note "Install dependencies"
+ To use [EpsillaDocumentIndex][docarray.index.backends.epsilla.EpsillaDocumentIndex], you need to install extra dependencies with the following command:
+
+ ```console
+ pip install "docarray[epsilla]"
+ pip install --upgrade pyepsilla
+ ```
+
+## Basic usage
+
+This snippet demonstrates the basic usage of
+[EpsillaDocumentIndex][docarray.index.backends.epsilla.EpsillaDocumentIndex]:
+
+1. Define a document schema with two fields: title and embedding
+2. Create ten dummy documents with random embeddings
+3. Set the db config and initialize the index
+4. Add dummy documents to the index
+5. Finally, perform a vector similarity search to retrieve the ten most similar documents to a given query vector
+
+```python
+from docarray import BaseDoc, DocList
+from docarray.index.backends.epsilla import EpsillaDocumentIndex
+from docarray.typing import NdArray
+from pydantic import Field
+import numpy as np
+
+
+# Define the document schema.
+class MyDoc(BaseDoc):
+ title: str
+ embedding: NdArray[128] = Field(is_embedding=True)
+
+
+# Create dummy documents.
+docs = DocList[MyDoc](
+ MyDoc(title=f'title #{i}', embedding=np.random.rand(128)) for i in range(10)
+)
+
+# db_config, see the initialize section below
+db_config = EpsillaDocumentIndex.DBConfig(
+ is_self_hosted=True,
+ protocol="http",
+ host="localhost",
+ port=8888,
+ db_path="/epsilla",
+ db_name="test",
+)
+
+# Initialize a new EpsillaDocumentIndex instance
+doc_index = EpsillaDocumentIndex[MyDoc](db_config=db_config)
+
+# Add the documents to the index.
+doc_index.index(docs)
+
+# Perform a vector search.
+query = MyDoc(title="test", embedding=np.ones(128))
+retrieved_docs = doc_index.find(query, limit=10, search_field="embedding")
+print(f'{retrieved_docs=}')
+retrieved_docs[0].summary()
+```
+
+The following sections will cover details of the individual steps.
+
+## Initialize
+
+### Start and connect to Epsilla
+
+To use [EpsillaDocumentIndex][docarray.index.backends.epsilla.EpsillaDocumentIndex], DocArray needs to hook into a
+running Epsilla service.
+There are multiple ways to start a Epsilla instance, depending on your use case.
+
+**Options - Overview**
+
+| Instance type | General use case | Configurability | Notes |
+| ------------------ | -------------------------- | --------------- | ------------------------------ |
+| **Epsilla Cloud ** | Development and production | Limited | **Recommended for most users** |
+| **Docker** | Self hosted | Full | |
+
+**Connect via Epsilla Cloud**
+
+Check out [Epsilla's documentation](https://epsilla-inc.gitbook.io/epsilladb/quick-start/epsilla-cloud) to create an
+instance, and for information on obtaining your credentials.
+
+**Connect via Docker (self-managed)**
+
+```bash
+docker pull epsilla/vectordb
+```
+
+Start the docker as the backend service
+
+```bash
+docker run --pull=always -d -p 8888:8888 epsilla/vectordb
+```
+
+### Connecting to Epsilla
+
+**Cloud instance**
+
+Check out [Epsilla's documentation](https://epsilla-inc.gitbook.io/epsilladb/quick-start/epsilla-cloud) for credentials.
+
+```python
+from docarray.index.backends.epsilla import EpsillaDocumentIndex
+
+db = EpsillaDocumentIndex.DBConfig(
+ is_self_hosted=False,
+ cloud_project_id="your-project-id",
+ cloud_db_id="your-database-id",
+ api_key="your-epsilla-api-key",
+)
+```
+
+**Self hosted**
+
+```python
+from docarray.index.backends.epsilla import EpsillaDocumentIndex
+
+db = EpsillaDocumentIndex.DBConfig(
+ is_self_hosted=True,
+ protocol=None,
+ host="localhost",
+ port=8888,
+ db_path=None,
+ db_name=None,
+)
+```
+
+### Create an instance
+
+Let's connect to a local Epsilla container and instantiate a `EpsillaDocumentIndex` instance for a given schema:
+
+```python
+from docarray import BaseDoc
+from docarray.index.backends.epsilla import EpsillaDocumentIndex
+from docarray.typing import NdArray
+from pydantic import Field
+
+
+# Define the document schema.
+class MyDoc(BaseDoc):
+ title: str
+ embedding: NdArray[128] = Field(is_embedding=True)
+
+
+# Set the database configuration.
+db_config = EpsillaDocumentIndex.DBConfig(
+ is_self_hosted=True,
+ protocol="http",
+ host="localhost",
+ port=8888,
+ db_path="/epsilla",
+ db_name="test",
+)
+
+# Initialize a new EpsillaDocumentIndex instance
+doc_index = EpsillaDocumentIndex[MyDoc](db_config=db_config)
+```
+
+### Schema definition
+
+In this code snippet, `EpsillaDocumentIndex` takes a schema of the form of `MyDoc`.
+The Document Index then _creates a column for each field in `MyDoc`_.
+
+The column types in the backend database are determined by the type hints of the document's fields.
+Optionally, you can [customize the database types for every field](#configuration).
+
+Most vector databases need to know the dimensionality of the vectors that will be stored.
+Here, that is automatically inferred from the type hint of the `embedding` field: `NdArray[128]` means that
+the database will store vectors with 128 dimensions.
+
+!!! note "PyTorch and TensorFlow support"
+ Instead of using `NdArray` you can use `TorchTensor` or `TensorFlowTensor` and the Document Index will handle that
+ for you. This is supported for all Document Index backends. No need to convert your tensors to NumPy arrays manually!
+
+### Using a predefined document as schema
+
+DocArray offers a number of predefined documents, like [ImageDoc][docarray.documents.ImageDoc]
+and [TextDoc][docarray.documents.TextDoc].
+If you try to use these directly as a schema for a Document Index, you will get unexpected behavior:
+Depending on the backend, an exception will be raised, or no vector index for ANN lookup will be built.
+
+The reason for this is that predefined documents don't hold information about the dimensionality of their `.embedding`
+field. But this is crucial information for any vector database to work properly!
+
+You can work around this problem by subclassing the predefined document and adding the dimensionality information:
+
+=== "Using type hint"
+
+```python
+from docarray.documents import TextDoc
+from docarray.typing import NdArray
+from docarray.index import EpsillaDocumentIndex
+from pydantic import Field
+
+
+class MyDoc(TextDoc):
+ embedding: NdArray[128] = Field(is_embedding=True)
+
+
+doc_index = EpsillaDocumentIndex[MyDoc]()
+```
+
+=== "Using Field()"
+
+```python
+from docarray.documents import TextDoc
+from docarray.typing import AnyTensor
+from docarray.index import EpsillaDocumentIndex
+from pydantic import Field
+
+
+class MyDoc(TextDoc):
+ embedding: AnyTensor = Field(dim=128, is_embedding=True)
+
+
+doc_index = EpsillaDocumentIndex[MyDoc]()
+```
+
+Once you have defined the schema of your Document Index in this way, the
+data that you index can be either the predefined Document type or your custom Document type.
+
+The [next section]( # index) goes into more detail about data indexing, but note that if you have some `TextDoc`
+, `ImageDoc` etc. that you want to index, you _don't_ need to cast them to `MyDoc`:
+
+```python
+from docarray import DocList
+
+data = DocList[MyDoc](
+ [
+ MyDoc(title='hello world', embedding=np.random.rand(128)),
+ MyDoc(title='hello world', embedding=np.random.rand(128)),
+ MyDoc(title='hello world', embedding=np.random.rand(128)),
+ ]
+)
+
+# you can index this into Document Index of type MyDoc
+doc_index.index(data)
+```
+
+## Index
+
+Now that you have a Document Index, you can add data to it, using
+the [`index()`][docarray.index.abstract.BaseDocIndex.index] method:
+
+```python
+from docarray import BaseDoc, DocList
+from docarray.index.backends.epsilla import EpsillaDocumentIndex
+from docarray.typing import NdArray
+from pydantic import Field
+import numpy as np
+
+
+class MyDoc(BaseDoc):
+ title: str
+ embedding: NdArray[128] = Field(is_embedding=True)
+
+
+# Create dummy documents.
+docs = DocList[MyDoc](
+ MyDoc(title=f'title #{i}', embedding=np.random.rand(128)) for i in range(10)
+)
+
+db_config = "..." # see the initialize section above
+
+doc_index = EpsillaDocumentIndex[MyDoc](db_config=db_config, index_name='mydoc_index')
+
+# add the data
+doc_index.index(docs)
+```
+
+That call to [`index()`][docarray.index.backends.epsilla.EpsillaDocumentIndex.index] stores all Documents in `docs` in
+the Document Index,
+ready to be retrieved in the next step.
+
+As you can see, `DocList[Document]` and `EpsillaDocumentIndex[Document]` both have `Document` as a parameter.
+This means that they share the same schema, and in general, both the Document Index and the data that you want to store
+need to have compatible schemas.
+
+!!! question "When are two schemas compatible?"
+ The schemas of your Document Index and data need to be compatible with each other.
+
+ Let's say A is the schema of your Document Index and B is the schema of your data.
+ There are a few rules that determine if schema A is compatible with schema B.
+ If _any_ of the following are true, then A and B are compatible:
+
+ - A and B are the same class
+ - A and B have the same field names and field types
+ - A and B have the same field names, and, for every field, the type of B is a subclass of the type of A
+
+ In particular, this means that you can easily [index predefined documents](#using-a-predefined-document-as-schema) into a Document Index.
+
+## Vector search
+
+Now that you have indexed your data, you can perform vector similarity search using
+the [`find()`][docarray.index.abstract.BaseDocIndex.find] method.
+
+You can perform a similarity search and find relevant documents by passing `MyDoc` or a raw vector to
+the [`find()`][docarray.index.abstract.BaseDocIndex.find] method:
+
+=== "Search by Document"
+
+ ```python
+ # create a query document
+ query = Document(
+ text="Hello world",
+ embedding=np.array([1, 2]),
+ file=np.random.rand(100),
+ )
+
+ # find similar documents
+ matches, scores = doc_index.find(query, limit=5)
+
+ print(f"{matches=}")
+ print(f"{matches.text=}")
+ print(f"{scores=}")
+ ```
+
+=== "Search by raw vector"
+
+ ```python
+ # create a query vector
+ query = np.random.rand(2)
+
+ # find similar documents
+ matches, scores = store.find(query, limit=5)
+
+ print(f'{matches=}')
+ print(f'{matches.text=}')
+ print(f'{scores=}')
+ ```
+
+The [`find()`][docarray.index.abstract.BaseDocIndex.find] method returns a named tuple containing the closest
+matching documents and their associated similarity scores.
+
+When searching on the subindex level, you can use
+the [`find_subindex()`][docarray.index.abstract.BaseDocIndex.find_subindex] method, which returns a named tuple
+containing the subindex documents, similarity scores and their associated root documents.
+
+How these scores are calculated depends on the backend, and can usually be [configured](#configuration).
+
+### Batched search
+
+You can also search for multiple documents at once, in a batch, using
+the [`find_batched()`][docarray.index.abstract.BaseDocIndex.find_batched] method.
+
+=== "Search by documents"
+
+ ```python
+ # create some query documents
+ queries = DocList[MyDoc](
+ Document(
+ text=f"Hello world {i}",
+ embedding=np.array([i, i + 1]),
+ file=np.random.rand(100),
+ )
+ for i in range(3)
+ )
+
+ # find similar documents
+ matches, scores = doc_index.find_batched(queries, limit=5)
+
+ print(f"{matches=}")
+ print(f"{matches[0].text=}")
+ print(f"{scores=}")
+ ```
+
+=== "Search by raw vectors"
+
+ ```python
+ # create some query vectors
+ query = np.random.rand(3, 2)
+
+ # find similar documents
+ matches, scores = doc_index.find_batched(query, limit=5)
+
+ print(f'{matches=}')
+ print(f'{matches[0].text=}')
+ print(f'{scores=}')
+ ```
+
+The [`find_batched()`][docarray.index.abstract.BaseDocIndex.find_batched] method returns a named tuple containing
+a list of `DocList`s, one for each query, containing the closest matching documents and their similarity scores.
+
+## Filter
+
+To perform filtering, follow the below syntax.
+
+This will perform a filtering on the field `title`:
+
+```python
+docs = doc_index.filter("title = 'test'", limit=5)
+```
+
+You can filter your documents by using the `filter()` or `filter_batched()` method with a corresponding filter query.
+The query should follow the [filters supported by Epsilla](https://epsilla-inc.gitbook.io/epsilladb/vector-database/search-the-top-k-semantically-similar-records#filter-expression).
+
+In the following example let's filter for all the books that are cheaper than 29 dollars:
+
+```python
+from docarray import BaseDoc, DocList
+from docarray.index.backends.epsilla import EpsillaDocumentIndex
+from docarray.typing import NdArray
+from pydantic import Field
+import numpy as np
+
+
+class Book(BaseDoc):
+ price: int
+ embedding: NdArray[10] = Field(is_embedding=True)
+
+
+books = DocList[Book](
+ [Book(price=i * 10, embedding=np.random.rand(10)) for i in range(10)]
+)
+db_config = "..." # see the initialize section above
+book_index = EpsillaDocumentIndex[Book](db_config=db_config, index_name='tmp_index')
+book_index.index(books)
+
+# filter for books that are cheaper than 29 dollars
+query = "price < 29"
+cheap_books = book_index.filter(filter_query=query)
+print(f"{cheap_books=}")
+cheap_books[0].summary()
+```
+
+## Text search
+
+!!! warning
+ The [EpsillaDocumentIndex][docarray.index.backends.epsilla.EpsillaDocumentIndex] implementation does not support text
+ search.
+
+## Hybrid search
+
+Document Index supports atomic operations for vector similarity search, text search and filter search.
+
+To combine these operations into a single, hybrid search query, you can use the query builder that is accessible
+through [`build_query()`][docarray.index.abstract.BaseDocIndex.build_query]:
+
+```python
+# Define the document schema.
+class SimpleSchema(BaseDoc):
+ year: int
+ price: int
+ embedding: NdArray[128]
+
+
+# Create dummy documents.
+docs = DocList[SimpleSchema](
+ SimpleSchema(year=2000 - i, price=i, embedding=np.random.rand(128))
+ for i in range(10)
+)
+
+doc_index = EpsillaDocumentIndex[SimpleSchema]()
+doc_index.index(docs)
+
+query = (
+ doc_index.build_query() # get empty query object
+ .filter(filter_query="year>1994") # pre-filtering
+ .find(
+ query=np.random.rand(128), search_field='embedding'
+ ) # add vector similarity search
+ .filter(filter_query="price<3") # post-filtering
+ .build()
+)
+# execute the combined query and return the results
+results = doc_index.execute_query(query)
+print(f'{results=}')
+```
+
+In the example above you can see how to form a hybrid query that combines vector similarity search and filtered search
+to obtain a combined set of results.
+
+The kinds of atomic queries that can be combined in this way depends on the backend.
+Some backends can combine text search and vector search, while others can perform filters and vectors search, etc.
+
+## Access documents
+
+To retrieve a document from a Document Index you don't necessarily need to perform a fancy search.
+
+You can also access data by the `id` that was assigned to each document:
+
+```python
+# prepare some data
+data = DocList[MyDoc](
+ MyDoc(embedding=np.random.rand(128), title=f'query {i}') for i in range(3)
+)
+
+# remember the Document ids and index the data
+ids = data.id
+doc_index.index(data)
+
+# access the documents by id
+doc = doc_index[ids[0]] # get by single id
+docs = doc_index[ids] # get by list of ids
+```
+
+## Delete documents
+
+In the same way you can access documents by `id`, you can also delete them:
+
+```python
+# prepare some data
+data = DocList[MyDoc](
+ MyDoc(embedding=np.random.rand(128), title=f'query {i}') for i in range(3)
+)
+
+# remember the Document ids and index the data
+ids = data.id
+doc_index.index(data)
+
+# access the documents by id
+del doc_index[ids[0]] # del by single id
+del doc_index[ids[1:]] # del by list of ids
+```
+
+## Count documents
+
+!!! warning
+ Unlike other index backends, Epsilla does not provide a count API. When using it with docarray, calling the `num_docs` method will raise errors.
+
+ ```python
+ # will raise errors
+ doc_index.num_docs()
+ ```
+
+If you need to count how many documents there are in the index, you can try to use the filter method.
+
+```python
+# use a larger limit as needed
+doc_index.filter(filter_query="", limit=100)
+```
+
+## Configuration
+
+### DBConfig
+
+The following configs can be set in `DBConfig`:
+
+| Name | Description | Default |
+| ------------------ | --------------------------------------------- | ------- |
+| `is_self_hosted` | If using Epsilla cloud or running self hosted | `false` |
+| `cloud_project_id` | If using Epsilla cloud; found in the console | `None` |
+| `cloud_db_id` | If using Epsilla cloud; found in the console | `None` |
+| `api_key` | If using Epsilla cloud; found in the console | `None` |
+| `host` | Address or 'localhost' | `None` |
+| `port` | The port number for the Epsilla server | 8888 |
+| `protocol` | Protocol to connect, e.g. 'http' | `None` |
+| `db_path` | Path to the database on disk | `None` |
+| `db_name` | Name of the database | `None` |
+
+You can pass any of the above as keyword arguments to the `__init__()` method or pass an entire configuration object.
+
+## Nested data and subindex search
+
+The examples provided primarily operate on a basic schema where each field corresponds to a straightforward type such
+as `str` or `NdArray`.
+However, it is also feasible to represent and store nested documents in a Document Index, including scenarios where a
+document contains a `DocList` of other documents.
+
+Go to the [Nested Data](nested_data.md) section to learn more.
diff --git a/docs/user_guide/storing/index_hnswlib.md b/docs/user_guide/storing/index_hnswlib.md
new file mode 100644
index 00000000000..e662cc220ae
--- /dev/null
+++ b/docs/user_guide/storing/index_hnswlib.md
@@ -0,0 +1,584 @@
+# Hnswlib Document Index
+
+!!! note
+ To use [HnswDocumentIndex][docarray.index.backends.hnswlib.HnswDocumentIndex], you need to install the extra dependency with the following command:
+
+ ```console
+ pip install "docarray[hnswlib]"
+ ```
+
+
+[HnswDocumentIndex][docarray.index.backends.hnswlib.HnswDocumentIndex] is a lightweight Document Index implementation
+that runs fully locally and is best suited for small- to medium-sized datasets.
+It stores vectors on disk in [hnswlib](https://github.com/nmslib/hnswlib), and stores all other data in [SQLite](https://www.sqlite.org/index.html).
+
+!!! note "Production readiness"
+ [HnswDocumentIndex][docarray.index.backends.hnswlib.HnswDocumentIndex] is a great starting point
+ for small- to medium-sized datasets, but it is not battle tested in production. If scalability, uptime, etc. are
+ important to you, we recommend you eventually transition to one of our database-backed Document Index implementations:
+
+ - [QdrantDocumentIndex][docarray.index.backends.qdrant.QdrantDocumentIndex]
+ - [WeaviateDocumentIndex][docarray.index.backends.weaviate.WeaviateDocumentIndex]
+ - [ElasticDocumentIndex][docarray.index.backends.elastic.ElasticDocIndex]
+ - [RedisDocumentIndex][docarray.index.backends.redis.RedisDocumentIndex]
+ - [MilvusDocumentIndex][docarray.index.backends.milvus.MilvusDocumentIndex]
+
+## Basic usage
+This snippet demonstrates the basic usage of [HnswDocumentIndex][docarray.index.backends.hnswlib.HnswDocumentIndex]. It defines a document schema with a title and an embedding,
+creates ten dummy documents with random embeddings, initializes an instance of [HnswDocumentIndex][docarray.index.backends.hnswlib.HnswDocumentIndex] to index these documents,
+and performs a vector similarity search to retrieve the ten most similar documents to a given query vector.
+
+```python
+from docarray import BaseDoc, DocList
+from docarray.index import HnswDocumentIndex
+from docarray.typing import NdArray
+import numpy as np
+
+# Define the document schema.
+class MyDoc(BaseDoc):
+ title: str
+ embedding: NdArray[128]
+
+# Create dummy documents.
+docs = DocList[MyDoc](MyDoc(title=f'title #{i}', embedding=np.random.rand(128)) for i in range(10))
+
+# Initialize a new HnswDocumentIndex instance and add the documents to the index.
+doc_index = HnswDocumentIndex[MyDoc](work_dir='./tmp_0')
+doc_index.index(docs)
+
+# Perform a vector search.
+query = np.ones(128)
+retrieved_docs = doc_index.find(query, search_field='embedding', limit=10)
+```
+
+## Initialize
+
+To create a Document Index, you first need a document class that defines the schema of your index:
+
+```python
+from docarray import BaseDoc
+from docarray.index import HnswDocumentIndex
+from docarray.typing import NdArray
+
+
+class MyDoc(BaseDoc):
+ embedding: NdArray[128]
+ text: str
+
+
+db = HnswDocumentIndex[MyDoc](work_dir='./tmp_1')
+```
+
+### Schema definition
+
+In this code snippet, `HnswDocumentIndex` takes a schema of the form of `MyDoc`.
+The Document Index then _creates a column for each field in `MyDoc`_.
+
+The column types in the backend database are determined by the type hints of the document's fields.
+Optionally, you can [customize the database types for every field](#configuration).
+
+Most vector databases need to know the dimensionality of the vectors that will be stored.
+Here, that is automatically inferred from the type hint of the `embedding` field: `NdArray[128]` means that
+the database will store vectors with 128 dimensions.
+
+!!! note "PyTorch and TensorFlow support"
+ Instead of using `NdArray` you can use `TorchTensor` or `TensorFlowTensor` and the Document Index will handle that
+ for you. This is supported for all Document Index backends. No need to convert your tensors to NumPy arrays manually!
+
+
+### Using a predefined document as schema
+
+DocArray offers a number of predefined documents, like [ImageDoc][docarray.documents.ImageDoc] and [TextDoc][docarray.documents.TextDoc].
+If you try to use these directly as a schema for a Document Index, you will get unexpected behavior:
+Depending on the backend, an exception will be raised, or no vector index for ANN lookup will be built.
+
+The reason for this is that predefined documents don't hold information about the dimensionality of their `.embedding`
+field. But this is crucial information for any vector database to work properly!
+
+You can work around this problem by subclassing the predefined document and adding the dimensionality information:
+
+=== "Using type hint"
+ ```python
+ from docarray.documents import TextDoc
+ from docarray.typing import NdArray
+ from docarray.index import HnswDocumentIndex
+
+
+ class MyDoc(TextDoc):
+ embedding: NdArray[128]
+
+
+ db = HnswDocumentIndex[MyDoc](work_dir='./tmp_2')
+ ```
+
+=== "Using Field()"
+ ```python
+ from docarray.documents import TextDoc
+ from docarray.typing import AnyTensor
+ from docarray.index import HnswDocumentIndex
+ from pydantic import Field
+
+
+ class MyDoc(TextDoc):
+ embedding: AnyTensor = Field(dim=128)
+
+
+ db = HnswDocumentIndex[MyDoc](work_dir='./tmp_3')
+ ```
+
+Once you have defined the schema of your Document Index in this way, the data that you index can be either the predefined Document type or your custom Document type.
+
+The [next section](#index) goes into more detail about data indexing, but note that if you have some `TextDoc`s, `ImageDoc`s etc. that you want to index, you _don't_ need to cast them to `MyDoc`:
+
+```python
+from docarray import DocList
+
+# data of type TextDoc
+data = DocList[TextDoc](
+ [
+ TextDoc(text='hello world', embedding=np.random.rand(128)),
+ TextDoc(text='hello world', embedding=np.random.rand(128)),
+ TextDoc(text='hello world', embedding=np.random.rand(128)),
+ ]
+)
+
+# you can index this into Document Index of type MyDoc
+db.index(data)
+```
+
+
+## Index
+
+Now that you have a Document Index, you can add data to it, using the [`index()`][docarray.index.abstract.BaseDocIndex.index] method:
+
+```python
+import numpy as np
+from docarray import DocList
+
+# create some random data
+docs = DocList[MyDoc](
+ [MyDoc(embedding=np.random.rand(128), text=f'text {i}') for i in range(100)]
+)
+
+# index the data
+db.index(docs)
+```
+
+That call to [`index()`][docarray.index.backends.hnswlib.HnswDocumentIndex.index] stores all Documents in `docs` in the Document Index,
+ready to be retrieved in the next step.
+
+As you can see, `DocList[MyDoc]` and `HnswDocumentIndex[MyDoc]` both have `MyDoc` as a parameter.
+This means that they share the same schema, and in general, both the Document Index and the data that you want to store need to have compatible schemas.
+
+!!! question "When are two schemas compatible?"
+ The schemas of your Document Index and data need to be compatible with each other.
+
+ Let's say A is the schema of your Document Index and B is the schema of your data.
+ There are a few rules that determine if schema A is compatible with schema B.
+ If _any_ of the following are true, then A and B are compatible:
+
+ - A and B are the same class
+ - A and B have the same field names and field types
+ - A and B have the same field names, and, for every field, the type of B is a subclass of the type of A
+
+ In particular, this means that you can easily [index predefined documents](#using-a-predefined-document-as-schema) into a Document Index.
+
+
+## Vector search
+
+Now that you have indexed your data, you can perform vector similarity search using the [`find()`][docarray.index.abstract.BaseDocIndex.find] method.
+
+You can use the [`find()`][docarray.index.abstract.BaseDocIndex.find] function with a document of the type `MyDoc`
+to find similar documents within the Document Index:
+
+=== "Search by Document"
+
+ ```python
+ # create a query document
+ query = MyDoc(embedding=np.random.rand(128), text='query')
+
+ # find similar documents
+ matches, scores = db.find(query, search_field='embedding', limit=5)
+
+ print(f'{matches=}')
+ print(f'{matches.text=}')
+ print(f'{scores=}')
+ ```
+
+=== "Search by raw vector"
+
+ ```python
+ # create a query vector
+ query = np.random.rand(128)
+
+ # find similar documents
+ matches, scores = db.find(query, search_field='embedding', limit=5)
+
+ print(f'{matches=}')
+ print(f'{matches.text=}')
+ print(f'{scores=}')
+ ```
+
+To peform a vector search, you need to specify a `search_field`. This is the field that serves as the
+basis of comparison between your query and the documents in the Document Index.
+
+In this example you only have one field (`embedding`) that is a vector, so you can trivially choose that one.
+In general, you could have multiple fields of type `NdArray` or `TorchTensor` or `TensorFlowTensor`, and you can choose
+which one to use for the search.
+
+The [`find()`][docarray.index.abstract.BaseDocIndex.find] method returns a named tuple containing the closest
+matching documents and their associated similarity scores.
+
+When searching on the subindex level, you can use the [`find_subindex()`][docarray.index.abstract.BaseDocIndex.find_subindex] method, which returns a named tuple containing the subindex documents, similarity scores and their associated root documents.
+
+How these scores are calculated depends on the backend, and can usually be [configured](#configuration).
+
+### Batched search
+
+You can also search for multiple documents at once, in a batch, using the [`find_batched()`][docarray.index.abstract.BaseDocIndex.find_batched] method.
+
+=== "Search by Documents"
+
+ ```python
+ # create some query Documents
+ queries = DocList[MyDoc](
+ MyDoc(embedding=np.random.rand(128), text=f'query {i}') for i in range(3)
+ )
+
+ # find similar documents
+ matches, scores = db.find_batched(queries, search_field='embedding', limit=5)
+
+ print(f'{matches=}')
+ print(f'{matches[0].text=}')
+ print(f'{scores=}')
+ ```
+
+=== "Search by raw vectors"
+
+ ```python
+ # create some query vectors
+ query = np.random.rand(3, 128)
+
+ # find similar documents
+ matches, scores = db.find_batched(query, search_field='embedding', limit=5)
+
+ print(f'{matches=}')
+ print(f'{matches[0].text=}')
+ print(f'{scores=}')
+ ```
+
+The [`find_batched()`][docarray.index.abstract.BaseDocIndex.find_batched] method returns a named tuple containing
+a list of `DocList`s, one for each query, containing the closest matching documents and their similarity scores.
+
+
+## Filter
+
+You can filter your documents by using the `filter()` or `filter_batched()` method with a corresponding filter query.
+The query should follow the query language of DocArray's [`filter_docs()`][docarray.utils.filter.filter_docs] function.
+
+In the following example let's filter for all the books that are cheaper than 29 dollars:
+
+```python
+from docarray import BaseDoc, DocList
+
+
+class Book(BaseDoc):
+ title: str
+ price: int
+
+
+books = DocList[Book]([Book(title=f'title {i}', price=i * 10) for i in range(10)])
+book_index = HnswDocumentIndex[Book](work_dir='./tmp_4')
+
+# filter for books that are cheaper than 29 dollars
+query = {'price': {'$lt': 29}}
+cheap_books = book_index.filter(query)
+
+assert len(cheap_books) == 3
+for doc in cheap_books:
+ doc.summary()
+```
+
+
+
+## Text search
+
+!!! note
+ The [HnswDocumentIndex][docarray.index.backends.hnswlib.HnswDocumentIndex] implementation does not support text search.
+
+ To see how to perform text search, you can check out other backends that offer support.
+
+In addition to vector similarity search, the Document Index interface offers methods for text search:
+[`text_search()`][docarray.index.abstract.BaseDocIndex.text_search],
+as well as the batched version [`text_search_batched()`][docarray.index.abstract.BaseDocIndex.text_search_batched].
+
+
+## Hybrid search
+
+Document Index supports atomic operations for vector similarity search, text search and filter search.
+
+To combine these operations into a single, hybrid search query, you can use the query builder that is accessible
+through [`build_query()`][docarray.index.abstract.BaseDocIndex.build_query]:
+
+```python
+# Define the document schema.
+class SimpleSchema(BaseDoc):
+ year: int
+ price: int
+ embedding: NdArray[128]
+
+# Create dummy documents.
+docs = DocList[SimpleSchema](SimpleSchema(year=2000-i, price=i, embedding=np.random.rand(128)) for i in range(10))
+
+doc_index = HnswDocumentIndex[SimpleSchema](work_dir='./tmp_5')
+doc_index.index(docs)
+
+query = (
+ doc_index.build_query() # get empty query object
+ .filter(filter_query={'year': {'$gt': 1994}}) # pre-filtering
+ .find(query=np.random.rand(128), search_field='embedding') # add vector similarity search
+ .filter(filter_query={'price': {'$lte': 3}}) # post-filtering
+ .build()
+)
+# execute the combined query and return the results
+results = doc_index.execute_query(query)
+print(f'{results=}')
+```
+
+In the example above you can see how to form a hybrid query that combines vector similarity search and filtered search
+to obtain a combined set of results.
+
+The kinds of atomic queries that can be combined in this way depends on the backend.
+Some backends can combine text search and vector search, while others can perform filters and vectors search, etc.
+
+
+## Access documents
+
+To retrieve a document from a Document Index you don't necessarily need to perform a fancy search.
+
+You can also access data by the `id` that was assigned to each document:
+
+```python
+# prepare some data
+data = DocList[MyDoc](
+ MyDoc(embedding=np.random.rand(128), text=f'query {i}') for i in range(3)
+)
+
+# remember the Document ids and index the data
+ids = data.id
+db.index(data)
+
+# access the Documents by id
+doc = db[ids[0]] # get by single id
+docs = db[ids] # get by list of ids
+```
+
+
+## Delete documents
+
+In the same way you can access Documents by `id`, you can also delete them:
+
+```python
+# prepare some data
+data = DocList[MyDoc](
+ MyDoc(embedding=np.random.rand(128), text=f'query {i}') for i in range(3)
+)
+
+# remember the Document ids and index the data
+ids = data.id
+db.index(data)
+
+# access the Documents by id
+del db[ids[0]] # del by single id
+del db[ids[1:]] # del by list of ids
+```
+
+## Update documents
+In order to update a Document inside the index, you only need to re-index it with the updated attributes.
+
+First, let's create a schema for our Document Index:
+```python
+import numpy as np
+from docarray import BaseDoc, DocList
+from docarray.typing import NdArray
+from docarray.index import HnswDocumentIndex
+class MyDoc(BaseDoc):
+ text: str
+ embedding: NdArray[128]
+```
+
+Now, we can instantiate our Index and add some data:
+```python
+docs = DocList[MyDoc](
+ [MyDoc(embedding=np.random.rand(128), text=f'I am the first version of Document {i}') for i in range(100)]
+)
+index = HnswDocumentIndex[MyDoc]()
+index.index(docs)
+assert index.num_docs() == 100
+```
+
+Let's retrieve our data and check its content:
+```python
+res = index.find(query=docs[0], search_field='embedding', limit=100)
+assert len(res.documents) == 100
+for doc in res.documents:
+ assert 'I am the first version' in doc.text
+```
+
+Then, let's update all of the text of these documents and re-index them:
+```python
+for i, doc in enumerate(docs):
+ doc.text = f'I am the second version of Document {i}'
+
+index.index(docs)
+assert index.num_docs() == 100
+```
+
+When we retrieve them again we can see that their text attribute has been updated accordingly:
+```python
+res = index.find(query=docs[0], search_field='embedding', limit=100)
+assert len(res.documents) == 100
+for doc in res.documents:
+ assert 'I am the second version' in doc.text
+```
+
+## Configuration
+
+This section lays out the configurations and options that are specific to [HnswDocumentIndex][docarray.index.backends.hnswlib.HnswDocumentIndex].
+
+### DBConfig
+
+The `DBConfig` of [HnswDocumentIndex][docarray.index.backends.hnswlib.HnswDocumentIndex] contains two argument:
+`work_dir` and `default_column_configs`
+
+`work_dir` is the location where all of the Index's data will be stored, namely the various HNSWLib indexes and the SQLite database.
+
+You can pass this directly to the constructor:
+
+```python
+from docarray import BaseDoc
+from docarray.index import HnswDocumentIndex
+from docarray.typing import NdArray
+
+
+class MyDoc(BaseDoc):
+ embedding: NdArray[128]
+ text: str
+
+
+db = HnswDocumentIndex[MyDoc](work_dir='./tmp_6')
+```
+
+To load existing data, you can specify a directory that stores data from a previous session.
+
+!!! note "Hnswlib file lock"
+ Hnswlib uses a file lock to prevent multiple processes from accessing the same index at the same time.
+ This means that if you try to open an index that is already open in another process, you will get an error.
+ To avoid this, you can specify a different `work_dir` for each process.
+
+`default_column_configs` contains the default mapping from Python types to column configurations.
+
+
+You can see in the [section below](#field-wise-configurations) how to override configurations for specific fields.
+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__`:
+
+```python
+import numpy as np
+
+
+db = HnswDocumentIndex[MyDoc](
+ work_dir='./tmp_7',
+ default_column_config={
+ np.ndarray: {
+ 'dim': -1,
+ 'index': True,
+ 'space': 'ip',
+ 'max_elements': 2048,
+ 'ef_construction': 100,
+ 'ef': 15,
+ 'M': 8,
+ 'allow_replace_deleted': True,
+ 'num_threads': 5,
+ },
+ None: {},
+ },
+)
+```
+
+This will set the default configuration for all vector fields to the one specified in the example above.
+
+!!! note
+ Even if your vectors come from PyTorch or TensorFlow, you can (and should) still use the `np.ndarray` configuration.
+ This is because all tensors are converted to `np.ndarray` under the hood.
+
+!!! note
+ max_elements is considered to have the initial maximum capacity of the index. However, the capacity of the index is doubled every time
+ that the number of Documents in the index exceeds this capacity. Expanding the capacity is an expensive operation, therefore it can be important to
+ choose an appropiate max_elements value at init time.
+
+For more information on these settings, see [below](#field-wise-configurations).
+
+Fields that are not vector fields (e.g. of type `str` or `int` etc.) do not offer any configuration, as they are simply
+stored as-is in a SQLite database.
+
+### Field-wise configuration
+
+There are various setting that you can tweak for every vector field that you index into Hnswlib.
+
+You pass all of those using the `field: Type = Field(...)` syntax:
+
+```python
+from pydantic import Field
+
+
+class Schema(BaseDoc):
+ tens: NdArray[100] = Field(max_elements=12, space='cosine')
+ tens_two: NdArray[10] = Field(M=4, space='ip')
+
+
+db = HnswDocumentIndex[Schema](work_dir='./tmp_8')
+```
+
+In the example above you can see how to configure two different vector fields, with two different sets of settings.
+
+In this way, you can pass [all options that Hnswlib supports](https://github.com/nmslib/hnswlib#api-description):
+
+| Keyword | Description | Default |
+|-------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------|
+| `max_elements` | Maximum number of vector that can be stored | 1024 |
+| `space` | Vector space (distance metric) the index operates in. Supports 'l2', 'ip', and 'cosine'. Output
+ ```text + get SimpleDoc(id='5a4b92af27aadbb852d636892506998b', text='doc 0'), get SimpleDoc(id='5a4b92af27aadbb852d636892506998b', text='doc 0') + get SimpleDoc(id='705e4f6acbab0a6ff10d11a07c03b24c', text='doc 1'), get SimpleDoc(id='705e4f6acbab0a6ff10d11a07c03b24c', text='doc 1') + get SimpleDoc(id='4fb5c01bd5f935bbe91cf73e271ad590', text='doc 2'), get SimpleDoc(id='4fb5c01bd5f935bbe91cf73e271ad590', text='doc 2') + get SimpleDoc(id='381498cef78f1d4f1d80415d67918940', text='doc 3'), get SimpleDoc(id='381498cef78f1d4f1d80415d67918940', text='doc 3') + get SimpleDoc(id='d968bc6fa235b1cfc69eded92926157e', text='doc 4'), get SimpleDoc(id='d968bc6fa235b1cfc69eded92926157e', text='doc 4') + get SimpleDoc(id='30bf347427a4bd50ce8ada1841320fe3', text='doc 5'), get SimpleDoc(id='30bf347427a4bd50ce8ada1841320fe3', text='doc 5') + get SimpleDoc(id='1389877ac97b3e6d0e8eb17568934708', text='doc 6'), get SimpleDoc(id='1389877ac97b3e6d0e8eb17568934708', text='doc 6') + get SimpleDoc(id='264b0eff2cd138d296f15c685e15bf23', text='doc 7'), get SimpleDoc(id='264b0eff2cd138d296f15c685e15bf23', text='doc 7') + ``` +**Note:** In contrast to the other backends, for HnswDocumentIndex `'cosine'` refers to **cosine distance**, not cosine similarity. To transform one to the other, you can use: `cos_sim = 1 - cos_dist`. For more details see [here](https://en.wikipedia.org/wiki/Cosine_similarity#Cosine_Distance). | 'l2' | +| `index` | Whether or not an index should be built for this field. | True | +| `ef_construction` | defines a construction time/accuracy trade-off | 200 | +| `ef` | parameter controlling query time/accuracy trade-off | 10 | +| `M` | parameter that defines the maximum number of outgoing connections in the graph | 16 | +| `allow_replace_deleted` | enables replacing of deleted elements with new added ones | True | +| `num_threads` | sets the number of cpu threads to use | 1 | + +!!! note + In HnswLibDocIndex `space='cosine'` refers to cosine distance, not to cosine similarity, as it does for the other backends. + +You can find more details on the parameters [here](https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md). + + +### Database location + +For `HnswDocumentIndex` you need to specify a `work_dir` where the data will be stored; for other backends you +usually specify a `host` and a `port` instead. + +In addition to a host and a port, most backends can also take an `index_name`, `table_name`, `collection_name` or similar. +This specifies the name of the index/table/collection that will be created in the database. +You don't have to specify this though: By default, this name will be taken from the name of the Document type that you use as schema. +For example, for `WeaviateDocumentIndex[MyDoc](...)` the data will be stored in a Weaviate Class of name `MyDoc`. + +In any case, if the location does not yet contain any data, we start from a blank slate. +If the location already contains data from a previous session, it will be accessible through the Document Index. + + + +## Nested data and subindex search + +The examples provided primarily operate on a basic schema where each field corresponds to a straightforward type such as `str` or `NdArray`. +However, it is also feasible to represent and store nested documents in a Document Index, including scenarios where a document +contains a `DocList` of other documents. + +Go to the [Nested Data](nested_data.md) section to learn more. diff --git a/docs/user_guide/storing/index_in_memory.md b/docs/user_guide/storing/index_in_memory.md new file mode 100644 index 00000000000..9b275b67063 --- /dev/null +++ b/docs/user_guide/storing/index_in_memory.md @@ -0,0 +1,515 @@ +# In-Memory Document Index + + +[InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex] stores all documents in memory using DocLists. +It is a great starting point for small datasets, where you may not want to launch a database server. + +For vector search and filtering [InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex] +utilizes DocArray's [`find()`][docarray.utils.find.find] and [`filter_docs()`][docarray.utils.filter.filter_docs] functions. + +!!! note "Production readiness" + [InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex] is a great starting point + for small- to medium-sized datasets, but it is not battle tested in production. If scalability, uptime, etc. are + important to you, we recommend you eventually transition to one of our database-backed Document Index implementations: + + - [QdrantDocumentIndex][docarray.index.backends.qdrant.QdrantDocumentIndex] + - [WeaviateDocumentIndex][docarray.index.backends.weaviate.WeaviateDocumentIndex] + - [ElasticDocumentIndex][docarray.index.backends.elastic.ElasticDocIndex] + - [RedisDocumentIndex][docarray.index.backends.redis.RedisDocumentIndex] + - [MilvusDocumentIndex][docarray.index.backends.milvus.MilvusDocumentIndex] + + + +## Basic usage +This snippet demonstrates the basic usage of [InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex]. It defines a document schema with a title and an embedding, +creates ten dummy documents with random embeddings, initializes an instance of [InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex] to index these documents, +and performs a vector similarity search to retrieve the ten most similar documents to a given query vector. + +```python +from docarray import BaseDoc, DocList +from docarray.index import InMemoryExactNNIndex +from docarray.typing import NdArray +import numpy as np + +# Define the document schema. +class MyDoc(BaseDoc): + title: str + embedding: NdArray[128] + +# Create dummy documents. +docs = DocList[MyDoc](MyDoc(title=f'title #{i}', embedding=np.random.rand(128)) for i in range(10)) + +# Initialize a new InMemoryExactNNIndex instance and add the documents to the index. +doc_index = InMemoryExactNNIndex[MyDoc]() +doc_index.index(docs) + +# Perform a vector search. +query = np.ones(128) +retrieved_docs, scores = doc_index.find(query, search_field='embedding', limit=10) +``` + +## Initialize + +To create a Document Index, you first need a document class that defines the schema of your index: + +```python +from docarray import BaseDoc +from docarray.index import InMemoryExactNNIndex +from docarray.typing import NdArray + + +class MyDoc(BaseDoc): + embedding: NdArray[128] + text: str + + +db = InMemoryExactNNIndex[MyDoc]() +``` + +### Schema definition + +In this code snippet, `InMemoryExactNNIndex` takes a schema of the form of `MyDoc`. +The Document Index then _creates a column for each field in `MyDoc`_. + +The column types in the backend database are determined by the type hints of the document's fields. +Optionally, you can [customize the database types for every field](#configuration). + +Most vector databases need to know the dimensionality of the vectors that will be stored. +Here, that is automatically inferred from the type hint of the `embedding` field: `NdArray[128]` means that +the database will store vectors with 128 dimensions. + +!!! note "PyTorch and TensorFlow support" + Instead of using `NdArray` you can use `TorchTensor` or `TensorFlowTensor` and the Document Index will handle that + for you. This is supported for all Document Index backends. No need to convert your tensors to NumPy arrays manually! + + +### Using a predefined document as schema + +DocArray offers a number of predefined documents, like [ImageDoc][docarray.documents.ImageDoc] and [TextDoc][docarray.documents.TextDoc]. +If you try to use these directly as a schema for a Document Index, you will get unexpected behavior: +Depending on the backend, an exception will be raised, or no vector index for ANN lookup will be built. + +The reason for this is that predefined documents don't hold information about the dimensionality of their `.embedding` +field. But this is crucial information for any vector database to work properly! + +You can work around this problem by subclassing the predefined document and adding the dimensionality information: + +=== "Using type hint" + ```python + from docarray.documents import TextDoc + from docarray.typing import NdArray + from docarray.index import InMemoryExactNNIndex + + + class MyDoc(TextDoc): + embedding: NdArray[128] + + + db = InMemoryExactNNIndex[MyDoc]() + ``` + +=== "Using Field()" + ```python + from docarray.documents import TextDoc + from docarray.typing import AnyTensor + from docarray.index import InMemoryExactNNIndex + from pydantic import Field + + + class MyDoc(TextDoc): + embedding: AnyTensor = Field(dim=128) + + + db = InMemoryExactNNIndex[MyDoc]() + ``` + +Once you have defined the schema of your Document Index in this way, the data that you index can be either the predefined Document type or your custom Document type. + +The [next section](#index) goes into more detail about data indexing, but note that if you have some `TextDoc`s, `ImageDoc`s etc. that you want to index, you _don't_ need to cast them to `MyDoc`: + +```python +from docarray import DocList + +# data of type TextDoc +data = DocList[MyDoc]( + [ + MyDoc(text='hello world', embedding=np.random.rand(128)), + MyDoc(text='hello world', embedding=np.random.rand(128)), + MyDoc(text='hello world', embedding=np.random.rand(128)), + ] +) + +# you can index this into Document Index of type MyDoc +db.index(data) +``` + +## Index + +Now that you have a Document Index, you can add data to it, using the [`index()`][docarray.index.abstract.BaseDocIndex.index] method: + +```python +import numpy as np +from docarray import DocList + +# create some random data +docs = DocList[MyDoc]( + [MyDoc(embedding=np.random.rand(128), text=f'text {i}') for i in range(100)] +) + +# index the data +db.index(docs) +``` + +That call to [`index()`][docarray.index.backends.in_memory.InMemoryExactNNIndex.index] stores all Documents in `docs` in the Document Index, +ready to be retrieved in the next step. + +As you can see, `DocList[MyDoc]` and `InMemoryExactNNIndex[MyDoc]` both have `MyDoc` as a parameter. +This means that they share the same schema, and in general, both the Document Index and the data that you want to store need to have compatible schemas. + +!!! question "When are two schemas compatible?" + The schemas of your Document Index and data need to be compatible with each other. + + Let's say A is the schema of your Document Index and B is the schema of your data. + There are a few rules that determine if schema A is compatible with schema B. + If _any_ of the following are true, then A and B are compatible: + + - A and B are the same class + - A and B have the same field names and field types + - A and B have the same field names, and, for every field, the type of B is a subclass of the type of A + + In particular, this means that you can easily [index predefined documents](#using-a-predefined-document-as-schema) into a Document Index. + + +## Vector search + +Now that you have indexed your data, you can perform vector similarity search using the [`find()`][docarray.index.abstract.BaseDocIndex.find] method. + +You can use the [`find()`][docarray.index.abstract.BaseDocIndex.find] function with a document of the type `MyDoc` +to find similar documents within the Document Index: + + +=== "Search by Document" + + ```python + # create a query document + query = MyDoc(embedding=np.random.rand(128), text='query') + + # find similar documents + matches, scores = db.find(query, search_field='embedding', limit=5) + + print(f'{matches=}') + print(f'{matches.text=}') + print(f'{scores=}') + ``` + +=== "Search by raw vector" + + ```python + # create a query vector + query = np.random.rand(128) + + # find similar documents + matches, scores = db.find(query, search_field='embedding', limit=5) + + print(f'{matches=}') + print(f'{matches.text=}') + print(f'{scores=}') + ``` + +To peform a vector search, you need to specify a `search_field`. This is the field that serves as the +basis of comparison between your query and the documents in the Document Index. + +In this example you only have one field (`embedding`) that is a vector, so you can trivially choose that one. +In general, you could have multiple fields of type `NdArray` or `TorchTensor` or `TensorFlowTensor`, and you can choose +which one to use for the search. + +The [`find()`][docarray.index.abstract.BaseDocIndex.find] method returns a named tuple containing the closest +matching documents and their associated similarity scores. + +When searching on the subindex level, you can use the [`find_subindex()`][docarray.index.abstract.BaseDocIndex.find_subindex] method, which returns a named tuple containing the subindex documents, similarity scores and their associated root documents. + +How these scores are calculated depends on the backend, and can usually be [configured](#configuration). + +### Batched search + +You can also search for multiple documents at once, in a batch, using the [`find_batched()`][docarray.index.abstract.BaseDocIndex.find_batched] method. + +=== "Search by documents" + + ```python + # create some query documents + queries = DocList[MyDoc]( + MyDoc(embedding=np.random.rand(128), text=f'query {i}') for i in range(3) + ) + + # find similar documents + matches, scores = db.find_batched(queries, search_field='embedding', limit=5) + + print(f'{matches=}') + print(f'{matches[0].text=}') + print(f'{scores=}') + ``` + +=== "Search by raw vectors" + + ```python + # create some query vectors + query = np.random.rand(3, 128) + + # find similar documents + matches, scores = db.find_batched(query, search_field='embedding', limit=5) + + print(f'{matches=}') + print(f'{matches[0].text=}') + print(f'{scores=}') + ``` + +The [`find_batched()`][docarray.index.abstract.BaseDocIndex.find_batched] method returns a named tuple containing +a list of `DocList`s, one for each query, containing the closest matching documents and their similarity scores. + + +## Filter + +To filter Documents, the `InMemoryExactNNIndex` uses DocArray's [`filter_docs()`][docarray.utils.filter.filter_docs] function. + +You can filter your documents by using the `filter()` or `filter_batched()` method with a corresponding filter query. +The query should follow the query language of DocArray's [`filter_docs()`][docarray.utils.filter.filter_docs] function. + +In the following example let's filter for all the books that are cheaper than 29 dollars: + +```python +from docarray import BaseDoc, DocList + + +class Book(BaseDoc): + title: str + price: int + + +books = DocList[Book]([Book(title=f'title {i}', price=i * 10) for i in range(10)]) +book_index = InMemoryExactNNIndex[Book](books) + +# filter for books that are cheaper than 29 dollars +query = {'price': {'$lte': 29}} +cheap_books = book_index.filter(query) + +assert len(cheap_books) == 3 +for doc in cheap_books: + doc.summary() +``` + +## Text search + +!!! note + The [InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex] implementation does not support text search. + + To see how to perform text search, you can check out other backends that offer support. + +In addition to vector similarity search, the Document Index interface offers methods for text search: +[`text_search()`][docarray.index.abstract.BaseDocIndex.text_search], +as well as the batched version [`text_search_batched()`][docarray.index.abstract.BaseDocIndex.text_search_batched]. + + + +## Hybrid search + +Document Index supports atomic operations for vector similarity search, text search and filter search. + +To combine these operations into a single, hybrid search query, you can use the query builder that is accessible +through [`build_query()`][docarray.index.abstract.BaseDocIndex.build_query]: + +```python +# Define the document schema. +class SimpleSchema(BaseDoc): + year: int + price: int + embedding: NdArray[128] + +# Create dummy documents. +docs = DocList[SimpleSchema](SimpleSchema(year=2000-i, price=i, embedding=np.random.rand(128)) for i in range(10)) + +doc_index = InMemoryExactNNIndex[SimpleSchema](docs) + +query = ( + doc_index.build_query() # get empty query object + .filter(filter_query={'year': {'$gt': 1994}}) # pre-filtering + .find(query=np.random.rand(128), search_field='embedding') # add vector similarity search + .filter(filter_query={'price': {'$lte': 3}}) # post-filtering + .build() +) +# execute the combined query and return the results +results = doc_index.execute_query(query) +print(f'{results=}') +``` + +In the example above you can see how to form a hybrid query that combines vector similarity search and filtered search +to obtain a combined set of results. + +The kinds of atomic queries that can be combined in this way depends on the backend. +Some backends can combine text search and vector search, while others can perform filters and vectors search, etc. + + +## Access documents + +To retrieve a document from a Document Index you don't necessarily need to perform a fancy search. + +You can also access data by the `id` that was assigned to each document: + +```python +# prepare some data +data = DocList[MyDoc]( + MyDoc(embedding=np.random.rand(128), text=f'query {i}') for i in range(3) +) + +# remember the Document ids and index the data +ids = data.id +db.index(data) + +# access the Documents by id +doc = db[ids[0]] # get by single id +docs = db[ids] # get by list of ids +``` + + +## Delete documents + +In the same way you can access Documents by `id`, you can also delete them: + +```python +# prepare some data +data = DocList[MyDoc]( + MyDoc(embedding=np.random.rand(128), text=f'query {i}') for i in range(3) +) + +# remember the Document ids and index the data +ids = data.id +db.index(data) + +# access the Documents by id +del db[ids[0]] # del by single id +del db[ids[1:]] # del by list of ids +``` + +## Update documents +In order to update a Document inside the index, you only need to re-index it with the updated attributes. + +First, let's create a schema for our Document Index: +```python +import numpy as np +from docarray import BaseDoc, DocList +from docarray.typing import NdArray +from docarray.index import InMemoryExactNNIndex +class MyDoc(BaseDoc): + text: str + embedding: NdArray[128] +``` + +Now, we can instantiate our Index and add some data: +```python +docs = DocList[MyDoc]( + [MyDoc(embedding=np.random.rand(128), text=f'I am the first version of Document {i}') for i in range(100)] +) +index = InMemoryExactNNIndex[MyDoc]() +index.index(docs) +assert index.num_docs() == 100 +``` + +Let's retrieve our data and check its content: +```python +res = index.find(query=docs[0], search_field='embedding', limit=100) +assert len(res.documents) == 100 +for doc in res.documents: + assert 'I am the first version' in doc.text +``` + +Then, let's update all of the text of these documents and re-index them: +```python +for i, doc in enumerate(docs): + doc.text = f'I am the second version of Document {i}' + +index.index(docs) +assert index.num_docs() == 100 +``` + +When we retrieve them again we can see that their text attribute has been updated accordingly +```python +res = index.find(query=docs[0], search_field='embedding', limit=100) +assert len(res.documents) == 100 +for doc in res.documents: + assert 'I am the second version' in doc.text +``` + + +## Configuration + +This section lays out the configurations and options that are specific to [InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex]. + +The `DBConfig` of [InMemoryExactNNIndex][docarray.index.backends.in_memory.InMemoryExactNNIndex] contains two entries: +`index_file_path` and `default_column_mapping`, the default mapping from Python types to column configurations. + +You can see in the [section below](#field-wise-configurations) how to override configurations for specific fields. +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__`:: + +```python +from collections import defaultdict +from docarray.typing.tensor.abstract_tensor import AbstractTensor +new_doc_index = InMemoryExactNNIndex[MyDoc]( + default_column_config=defaultdict( + dict, + { + AbstractTensor: {'space': 'cosine_sim'}, + }, + ) +) +``` + +This will set the default configuration for all vector fields to the one specified in the example above. + +For more information on these settings, see [below](#field-wise-configurations). + +Fields that are not vector fields (e.g. of type `str` or `int` etc.) do not offer any configuration. + + +### Field-wise configuration + +For a vector field you can adjust the `space` parameter. It can be one of: + +- `'cosine_sim'` (default) +- `'euclidean_dist'` +- `'sqeuclidean_dist'` + +You pass it using the `field: Type = Field(...)` syntax: + +```python +from docarray import BaseDoc +from pydantic import Field + + +class Schema(BaseDoc): + tensor_1: NdArray[100] = Field(space='euclidean_dist') + tensor_2: NdArray[100] = Field(space='sqeuclidean_dist') +``` + +In the example above you can see how to configure two different vector fields, with two different sets of settings. + + +### Persist and Load +You can pass an `index_file_path` argument to make sure that the index can be restored if persisted from that specific file. +```python +doc_index = InMemoryExactNNIndex[MyDoc](index_file_path='docs.bin') +doc_index.index(docs) + +doc_index.persist() + +# Initialize a new document index using the saved binary file +new_doc_index = InMemoryExactNNIndex[MyDoc](index_file_path='docs.bin') +``` + + +## Nested data and subindex search + +The examples provided primarily operate on a basic schema where each field corresponds to a straightforward type such as `str` or `NdArray`. +However, it is also feasible to represent and store nested documents in a Document Index, including scenarios where a document +contains a `DocList` of other documents. + +Go to the [Nested Data](nested_data.md) section to learn more. \ No newline at end of file diff --git a/docs/user_guide/storing/index_milvus.md b/docs/user_guide/storing/index_milvus.md new file mode 100644 index 00000000000..18431902cec --- /dev/null +++ b/docs/user_guide/storing/index_milvus.md @@ -0,0 +1,449 @@ +# Milvus Document Index + +!!! note "Install dependencies" + To use [MilvusDocumentIndex][docarray.index.backends.milvus.MilvusDocumentIndex], you need to install extra dependencies with the following command: + + ```console + pip install "docarray[milvus]" + ``` + +This is the user guide for the [MilvusDocumentIndex][docarray.index.backends.milvus.MilvusDocumentIndex], +focusing on special features and configurations of Milvus. + + +## Basic usage +This snippet demonstrates the basic usage of [MilvusDocumentIndex][docarray.index.backends.milvus.MilvusDocumentIndex]. It defines a document schema with a title and an embedding, +creates ten dummy documents with random embeddings, initializes an instance of [MilvusDocumentIndex][docarray.index.backends.milvus.MilvusDocumentIndex] to index these documents, +and performs a vector similarity search to retrieve the ten most similar documents to a given query vector. + +!!! note "Single Search Field Requirement" + In order to utilize vector search, it's necessary to define 'is_embedding' for one field only. + This is due to Milvus' configuration, which permits a single vector for each data object. + +```python +from docarray import BaseDoc, DocList +from docarray.index import MilvusDocumentIndex +from docarray.typing import NdArray +from pydantic import Field +import numpy as np + + +# Define the document schema. +class MyDoc(BaseDoc): + title: str + embedding: NdArray[128] = Field(is_embedding=True) + + +# Create dummy documents. +docs = DocList[MyDoc]( + MyDoc(title=f'title #{i}', embedding=np.random.rand(128)) for i in range(10) +) + +# Initialize a new MilvusDocumentIndex instance and add the documents to the index. +doc_index = MilvusDocumentIndex[MyDoc](index_name='tmp_index_1') +doc_index.index(docs) + +# Perform a vector search. +query = np.ones(128) +retrieved_docs = doc_index.find(query, limit=10) +``` + + +## Initialize + +First of all, you need to install and run Milvus. Download `docker-compose.yml` with the following command: + +```shell +wget https://github.com/milvus-io/milvus/releases/download/v2.2.11/milvus-standalone-docker-compose.yml -O docker-compose.yml +``` + +And start Milvus by running: +```shell +sudo docker compose up -d +``` + +Learn more on [Milvus documentation](https://milvus.io/docs/install_standalone-docker.md). + +Next, you can create a [MilvusDocumentIndex][docarray.index.backends.milvus.MilvusDocumentIndex] instance using: + +```python +from docarray import BaseDoc +from docarray.index import MilvusDocumentIndex +from docarray.typing import NdArray +from pydantic import Field + + +class MyDoc(BaseDoc): + embedding: NdArray[128] = Field(is_embedding=True) + text: str + + +doc_index = MilvusDocumentIndex[MyDoc](index_name='tmp_index_2') +``` + +### Schema definition +In this code snippet, `MilvusDocumentIndex` takes a schema of the form of `MyDoc`. +The Document Index then _creates a column for each field in `MyDoc`_. + +The column types in the backend database are determined by the type hints of the document's fields. +Optionally, you can [customize the database types for every field](#configuration). + +Most vector databases need to know the dimensionality of the vectors that will be stored. +Here, that is automatically inferred from the type hint of the `embedding` field: `NdArray[128]` means that +the database will store vectors with 128 dimensions. + +!!! note "PyTorch and TensorFlow support" + Instead of using `NdArray` you can use `TorchTensor` or `TensorFlowTensor` and the Document Index will handle that + for you. This is supported for all Document Index backends. No need to convert your tensors to NumPy arrays manually! + +### Using a predefined document as schema + +DocArray offers a number of predefined documents, like [ImageDoc][docarray.documents.ImageDoc] and [TextDoc][docarray.documents.TextDoc]. +If you try to use these directly as a schema for a Document Index, you will get unexpected behavior: +Depending on the backend, an exception will be raised, or no vector index for ANN lookup will be built. + +The reason for this is that predefined Documents don't hold information about the dimensionality of their `.embedding` +field. But this is crucial information for any vector database to work properly! + +You can work around this problem by subclassing the predefined document and adding the dimensionality information: + +=== "Using type hint" + ```python + from docarray.documents import TextDoc + from docarray.typing import NdArray + from docarray.index import MilvusDocumentIndex + from pydantic import Field + + + class MyDoc(TextDoc): + embedding: NdArray[128] = Field(is_embedding=True) + + + doc_index = MilvusDocumentIndex[MyDoc](index_name='tmp_index_3') + ``` + +=== "Using Field()" + ```python + from docarray.documents import TextDoc + from docarray.typing import AnyTensor + from docarray.index import MilvusDocumentIndex + from pydantic import Field + + + class MyDoc(TextDoc): + embedding: AnyTensor = Field(dim=128, is_embedding=True) + + + doc_index = MilvusDocumentIndex[MyDoc](index_name='tmp_index_4') + ``` + + +## Index + +Now that you have a Document Index, you can add data to it, using the [`index()`][docarray.index.abstract.BaseDocIndex.index] method: + +```python +import numpy as np +from docarray import DocList + + +class MyDoc(BaseDoc): + title: str + embedding: NdArray[128] = Field(is_embedding=True) + + +doc_index = MilvusDocumentIndex[MyDoc](index_name='tmp_index_5') + +# create some random data +docs = DocList[MyDoc]( + [MyDoc(embedding=np.random.rand(128), title=f'text {i}') for i in range(100)] +) + +# index the data +doc_index.index(docs) +``` + +That call to [`index()`][docarray.index.backends.milvus.MilvusDocumentIndex.index] stores all Documents in `docs` in the Document Index, +ready to be retrieved in the next step. + +As you can see, `DocList[MyDoc]` and `MilvusDocumentIndex[MyDoc]` both have `MyDoc` as a parameter. +This means that they share the same schema, and in general, both the Document Index and the data that you want to store need to have compatible schemas. + +!!! question "When are two schemas compatible?" + The schemas of your Document Index and data need to be compatible with each other. + + Let's say A is the schema of your Document Index and B is the schema of your data. + There are a few rules that determine if schema A is compatible with schema B. + If _any_ of the following are true, then A and B are compatible: + + - A and B are the same class + - A and B have the same field names and field types + - A and B have the same field names, and, for every field, the type of B is a subclass of the type of A + + In particular, this means that you can easily [index predefined Documents](#using-a-predefined-document-as-schema) into a Document Index. + + +## Vector search + +Now that you have indexed your data, you can perform vector similarity search using the [`find()`][docarray.index.abstract.BaseDocIndex.find] method. + +You can perform a similarity search and find relevant documents by passing `MyDoc` or a raw vector to +the [`find()`][docarray.index.abstract.BaseDocIndex.find] method: + +=== "Search by Document" + + ```python + # create a query document + query = MyDoc(embedding=np.random.rand(128), title='query') + + # find similar documents + matches, scores = doc_index.find(query, limit=5) + + print(f'{matches=}') + print(f'{matches.title=}') + print(f'{scores=}') + ``` + +=== "Search by raw vector" + + ```python + # create a query vector + query = np.random.rand(128) + + # find similar documents + matches, scores = doc_index.find(query, limit=5) + + print(f'{matches=}') + print(f'{matches.title=}') + print(f'{scores=}') + ``` + +The [`find()`][docarray.index.abstract.BaseDocIndex.find] method returns a named tuple containing the closest +matching documents and their associated similarity scores. + +When searching on the subindex level, you can use the [`find_subindex()`][docarray.index.abstract.BaseDocIndex.find_subindex] method, which returns a named tuple containing the subindex documents, similarity scores and their associated root documents. + +How these scores are calculated depends on the backend, and can usually be [configured](#configuration). + +### Batched search + +You can also search for multiple documents at once, in a batch, using the [`find_batched()`][docarray.index.abstract.BaseDocIndex.find_batched] method. + +=== "Search by documents" + + ```python + # create some query documents + queries = DocList[MyDoc]( + MyDoc(embedding=np.random.rand(128), text=f'query {i}') for i in range(3) + ) + + # find similar documents + matches, scores = doc_index.find_batched(queries, limit=5) + + print(f'{matches=}') + print(f'{matches[0].text=}') + print(f'{scores=}') + ``` + +=== "Search by raw vectors" + + ```python + # create some query vectors + query = np.random.rand(3, 128) + + # find similar documents + matches, scores = doc_index.find_batched(query, limit=5) + + print(f'{matches=}') + print(f'{matches[0].text=}') + print(f'{scores=}') + ``` + +The [`find_batched()`][docarray.index.abstract.BaseDocIndex.find_batched] method returns a named tuple containing +a list of `DocList`s, one for each query, containing the closest matching documents and their similarity scores. + + +## Filter + +You can filter your documents by using the `filter()` or `filter_batched()` method with a corresponding filter query. +The query should follow the [query language of the Milvus](https://milvus.io/docs/boolean.md). + +In the following example let's filter for all the books that are cheaper than 29 dollars: + +```python +from docarray import BaseDoc, DocList + + +class Book(BaseDoc): + price: int + embedding: NdArray[10] = Field(is_embedding=True) + + +books = DocList[Book]( + [Book(price=i * 10, embedding=np.random.rand(10)) for i in range(10)] +) +book_index = MilvusDocumentIndex[Book](index_name='tmp_index_6') +book_index.index(books) + +# filter for books that are cheaper than 29 dollars +query = 'price < 29' +cheap_books = book_index.filter(filter_query=query) + +assert len(cheap_books) == 3 +for doc in cheap_books: + doc.summary() +``` + +## Text search + +!!! note + The [MilvusDocumentIndex][docarray.index.backends.milvus.MilvusDocumentIndex] implementation does not support text search. + + To see how to perform text search, you can check out other backends that offer support. + +In addition to vector similarity search, the Document Index interface offers methods for text search: +[`text_search()`][docarray.index.abstract.BaseDocIndex.text_search], +as well as the batched version [`text_search_batched()`][docarray.index.abstract.BaseDocIndex.text_search_batched]. + + + +## Hybrid search + +Document Index supports atomic operations for vector similarity search, text search and filter search. + +To combine these operations into a single, hybrid search query, you can use the query builder that is accessible +through [`build_query()`][docarray.index.abstract.BaseDocIndex.build_query]: + +```python +# Define the document schema. +class SimpleSchema(BaseDoc): + price: int + embedding: NdArray[128] = Field(is_embedding=True) + + +# Create dummy documents. +docs = DocList[SimpleSchema]( + SimpleSchema(price=i, embedding=np.random.rand(128)) for i in range(10) +) + +doc_index = MilvusDocumentIndex[SimpleSchema](index_name='tmp_index_7') +doc_index.index(docs) + +query = ( + doc_index.build_query() # get empty query object + .find(query=np.random.rand(128)) # add vector similarity search + .filter(filter_query='price < 3') # add filter search + .build() +) +# execute the combined query and return the results +results = doc_index.execute_query(query) +print(f'{results=}') +``` + +In the example above you can see how to form a hybrid query that combines vector similarity search and filtered search +to obtain a combined set of results. + +The kinds of atomic queries that can be combined in this way depends on the backend. +Some backends can combine text search and vector search, while others can perform filters and vectors search, etc. + + +## Access documents + +To retrieve a document from a Document Index you don't necessarily need to perform a fancy search. + +You can also access data by the `id` that was assigned to each document: + +```python +# prepare some data +data = DocList[SimpleSchema]( + SimpleSchema(embedding=np.random.rand(128), price=i) for i in range(3) +) + +# remember the Document ids and index the data +ids = data.id +doc_index.index(data) + +# access the Documents by id +doc = doc_index[ids[0]] # get by single id +docs = doc_index[ids] # get by list of ids +``` + + +## Delete documents + +In the same way you can access Documents by `id`, you can also delete them: + +```python +# prepare some data +data = DocList[SimpleSchema]( + SimpleSchema(embedding=np.random.rand(128), price=i) for i in range(3) +) + +# remember the Document ids and index the data +ids = data.id +doc_index.index(data) + +# access the Documents by id +del doc_index[ids[0]] # del by single id +del doc_index[ids[1:]] # del by list of ids +``` + + +## Configuration + +This section lays out the configurations and options that are specific to [MilvusDocumentIndex][docarray.index.backends.milvus.MilvusDocumentIndex]. + +### DBConfig + +The following configs can be set in `DBConfig`: + +| Name | Description | Default | +|-------------------------|------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------| +| `host` | The host address for the Milvus server. | `localhost` | +| `port` | The port number for the Milvus server | 19530 | +| `index_name` | The name of the index in the Milvus database | `None`. Data will be stored in an index named after the Document type used as schema. | +| `user` | The username for the Milvus server | `None` | +| `password` | The password for the Milvus server | `None` | +| `token` | Token for secure connection | '' | +| `collection_description` | Description of the collection in the database | '' | +| `default_column_config` | The default configurations for every column type. | dict | + +You can pass any of the above as keyword arguments to the `__init__()` method or pass an entire configuration object. + + +### Field-wise configuration + + +`default_column_config` is the default configurations for every column type. Since there are many column types in Milvus, you can also consider changing the column config when defining the schema. + +```python +class SimpleDoc(BaseDoc): + tensor: NdArray[128] = Field( + is_embedding=True, index_type='IVF_FLAT', metric_type='L2' + ) + + +doc_index = MilvusDocumentIndex[SimpleDoc](index_name='tmp_index_10') +``` + + +### RuntimeConfig + +The `RuntimeConfig` dataclass of `MilvusDocumentIndex` consists of `batch_size` index/get/del operations. +You can change `batch_size` in the following way: + +```python +doc_index = MilvusDocumentIndex[SimpleDoc]() +doc_index.configure(MilvusDocumentIndex.RuntimeConfig(batch_size=128)) +``` + +You can pass the above as keyword arguments to the `configure()` method or pass an entire configuration object. + + +## Nested data and subindex search + +The examples provided primarily operate on a basic schema where each field corresponds to a straightforward type such as `str` or `NdArray`. +However, it is also feasible to represent and store nested documents in a Document Index, including scenarios where a document +contains a `DocList` of other documents. + +Go to the [Nested Data](nested_data.md) section to learn more. \ No newline at end of file diff --git a/docs/user_guide/storing/index_qdrant.md b/docs/user_guide/storing/index_qdrant.md new file mode 100644 index 00000000000..3d34b472a0c --- /dev/null +++ b/docs/user_guide/storing/index_qdrant.md @@ -0,0 +1,519 @@ +# Qdrant Document Index + +!!! note "Install dependencies" + To use [QdrantDocumentIndex][docarray.index.backends.qdrant.QdrantDocumentIndex], you need to install extra dependencies with the following command: + + ```console + pip install "docarray[qdrant]" + ``` + +The following is a starter script for using the [QdrantDocumentIndex][docarray.index.backends.qdrant.QdrantDocumentIndex], +based on the [Qdrant](https://qdrant.tech/) vector search engine. + + +## Basic usage +This snippet demonstrates the basic usage of [QdrantDocumentIndex][docarray.index.backends.qdrant.QdrantDocumentIndex]. It defines a document schema with a title and an embedding, +creates ten dummy documents with random embeddings, initializes an instance of [QdrantDocumentIndex][docarray.index.backends.qdrant.QdrantDocumentIndex] to index these documents, +and performs a vector similarity search to retrieve the ten most similar documents to a given query vector. + +```python +from docarray import BaseDoc, DocList +from docarray.index import QdrantDocumentIndex +from docarray.typing import NdArray +import numpy as np + + +# Define the document schema. +class MyDoc(BaseDoc): + title: str + embedding: NdArray[128] + + +# Create dummy documents. +docs = DocList[MyDoc]( + MyDoc(title=f'title #{i}', embedding=np.random.rand(128)) for i in range(10) +) + +# Initialize a new QdrantDocumentIndex instance and add the documents to the index. +doc_index = QdrantDocumentIndex[MyDoc](host='localhost') +doc_index.index(docs) + +# Perform a vector search. +query = np.ones(128) +retrieved_docs, scores = doc_index.find(query, search_field='embedding', limit=10) +``` + +## Initialize + +You can initialize [QdrantDocumentIndex][docarray.index.backends.qdrant.QdrantDocumentIndex] in three different ways: + + +**Connecting to a local Qdrant instance running as a Docker container** + +You can use docker compose to create a local Qdrant service with the following `docker-compose.yml`. + +```yaml +version: '3.8' + +services: + qdrant: + image: qdrant/qdrant:v1.1.2 + ports: + - "6333:6333" + - "6334:6334" + ulimits: # Only required for tests, as there are a lot of collections created + nofile: + soft: 65535 + hard: 65535 +``` + +Run the following command in the folder of the above `docker-compose.yml` to start the service: + +```bash +docker compose up +``` + +Next, you can create a [QdrantDocumentIndex][docarray.index.backends.qdrant.QdrantDocumentIndex] instance using: + +```python +qdrant_config = QdrantDocumentIndex.DBConfig('localhost') +doc_index = QdrantDocumentIndex[MyDoc](qdrant_config) + +# or just +doc_index = QdrantDocumentIndex[MyDoc](host='localhost') +``` + + +**Creating an in-memory Qdrant document index** +```python +qdrant_config = QdrantDocumentIndex.DBConfig(location=":memory:") +doc_index = QdrantDocumentIndex[MyDoc](qdrant_config) +``` + +**Connecting to Qdrant Cloud service** +```python +qdrant_config = QdrantDocumentIndex.DBConfig( + "https://YOUR-CLUSTER-URL.aws.cloud.qdrant.io", + api_key="
Hope the tests are going well :)
++ Hi! +
+ + diff --git a/tests/toydata/test.log b/tests/toydata/test.log new file mode 100644 index 00000000000..c0a0e904924 --- /dev/null +++ b/tests/toydata/test.log @@ -0,0 +1,6 @@ +2022-11-25 12:34:56 INFO: Program started +2022-11-25 12:35:01 ERROR: Unable to open file 'input.txt' +2022-11-25 12:35:10 INFO: 42 records processed successfully +2022-11-25 12:35:15 WARNING: Possible data corruption detected +2022-11-25 12:35:22 ERROR: Out of memory +2022-11-25 12:35:30 INFO: Program completed diff --git a/tests/toydata/test.md b/tests/toydata/test.md new file mode 100644 index 00000000000..320ec9ccc5f --- /dev/null +++ b/tests/toydata/test.md @@ -0,0 +1,4 @@ +# Hello +## This is a markdown file +### For testing stuff +Thank you! diff --git a/tests/unit/document/toydata/test.png b/tests/toydata/test.png similarity index 100% rename from tests/unit/document/toydata/test.png rename to tests/toydata/test.png diff --git a/tests/toydata/tetrahedron.mtl b/tests/toydata/tetrahedron.mtl new file mode 100644 index 00000000000..1bccd4474e4 --- /dev/null +++ b/tests/toydata/tetrahedron.mtl @@ -0,0 +1,22 @@ + +newmtl red +Ka 0.4449 0.0000 0.0000 +Kd 0.7714 0.0000 0.0000 +Ks 0.8857 0.0000 0.0000 +illum 2 +Ns 136.4300 + +newmtl lime +Ka 0.0000 0.5000 0.0000 +Kd 0.0000 1.0000 0.0000 +Ks 0.0000 0.5000 0.0000 +illum 2 +Ns 65.8900 + +newmtl gold +Ka 0.5265 0.2735 0.0122 +Kd 1.0000 0.5184 0.0286 +Ks 0.3000 0.3000 0.3000 +illum 2 +Ns 123.2600 + diff --git a/tests/toydata/tetrahedron.obj b/tests/toydata/tetrahedron.obj new file mode 100644 index 00000000000..40347bad7b7 --- /dev/null +++ b/tests/toydata/tetrahedron.obj @@ -0,0 +1,20 @@ +# tetrahedron.obj +# + +mtllib tetrahedron.mtl + +g tetrahedron + +v 1.00 1.00 1.00 +v 2.00 1.00 1.00 +v 1.00 2.00 1.00 +v 1.00 1.00 2.00 + +usemtl lime +f 1 3 2 +usemtl gold +f 1 4 3 +usemtl lime +f 1 2 4 +usemtl red +f 2 3 4 diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/unit/array/__init__.py b/tests/unit/array/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/unit/array/mixins/__init__.py b/tests/unit/array/mixins/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/unit/array/mixins/test-net.onnx b/tests/unit/array/mixins/test-net.onnx deleted file mode 100644 index f354c2d39f4..00000000000 Binary files a/tests/unit/array/mixins/test-net.onnx and /dev/null differ diff --git a/tests/unit/array/mixins/test_content.py b/tests/unit/array/mixins/test_content.py deleted file mode 100644 index df4bba9add8..00000000000 --- a/tests/unit/array/mixins/test_content.py +++ /dev/null @@ -1,68 +0,0 @@ -import numpy as np -import pytest - -from docarray import DocumentArray - - -@pytest.mark.parametrize('cls', [DocumentArray]) -@pytest.mark.parametrize( - 'content_attr', ['texts', 'embeddings', 'blobs', 'buffers', 'contents'] -) -def test_content_empty_getter_return_none(cls, content_attr): - da = cls() - assert getattr(da, content_attr) is None - - -@pytest.mark.parametrize('cls', [DocumentArray]) -@pytest.mark.parametrize( - 'content_attr', - [ - ('texts', ''), - ('embeddings', np.array([])), - ('blobs', np.array([])), - ('buffers', []), - ('contents', []), - ], -) -def test_content_empty_setter(cls, content_attr): - da = cls() - setattr(da, content_attr[0], content_attr[1]) - assert getattr(da, content_attr[0]) is None - - -@pytest.mark.parametrize('cls', [DocumentArray]) -@pytest.mark.parametrize( - 'content_attr', - [ - ('texts', ['s'] * 10), - ('blobs', np.random.random([10, 2])), - ('buffers', [b's'] * 10), - ], -) -def test_content_getter_setter(cls, content_attr): - da = cls.empty(10) - setattr(da, content_attr[0], content_attr[1]) - np.testing.assert_equal(da.contents, content_attr[1]) - da.contents = content_attr[1] - np.testing.assert_equal(da.contents, content_attr[1]) - np.testing.assert_equal(getattr(da, content_attr[0]), content_attr[1]) - da.contents = None - assert da.contents is None - - -@pytest.mark.parametrize('da_len', [0, 1, 2]) -def test_content_empty(da_len): - da = DocumentArray.empty(da_len) - assert not da.texts - assert not da.contents - assert not da.blobs - assert not da.buffers - - da.texts = ['hello'] * da_len - if da_len == 0: - assert not da.contents - else: - assert da.contents == ['hello'] * da_len - assert da.texts == ['hello'] * da_len - assert not da.blobs - assert not da.buffers diff --git a/tests/unit/array/mixins/test_embed.py b/tests/unit/array/mixins/test_embed.py deleted file mode 100644 index 1f25749fa50..00000000000 --- a/tests/unit/array/mixins/test_embed.py +++ /dev/null @@ -1,71 +0,0 @@ -import os - -import numpy as np -import onnxruntime -import paddle -import pytest -import tensorflow as tf -import torch - -from docarray import DocumentArray - -random_embed_models = { - 'keras': lambda: tf.keras.Sequential( - [tf.keras.layers.Dropout(0.5), tf.keras.layers.BatchNormalization()] - ), - 'pytorch': lambda: torch.nn.Sequential( - torch.nn.Dropout(0.5), torch.nn.BatchNorm1d(128) - ), - 'paddle': lambda: paddle.nn.Sequential( - paddle.nn.Dropout(0.5), paddle.nn.BatchNorm1D(128) - ), -} -cur_dir = os.path.dirname(os.path.abspath(__file__)) -torch.onnx.export( - random_embed_models['pytorch'](), - torch.rand(1, 128), - os.path.join(cur_dir, 'test-net.onnx'), - do_constant_folding=True, # whether to execute constant folding for optimization - input_names=['input'], # the model's input names - output_names=['output'], # the model's output names - dynamic_axes={ - 'input': {0: 'batch_size'}, # variable length axes - 'output': {0: 'batch_size'}, - }, -) - -random_embed_models['onnx'] = lambda: onnxruntime.InferenceSession( - os.path.join(cur_dir, 'test-net.onnx') -) - - -@pytest.mark.parametrize('framework', ['onnx', 'keras', 'pytorch', 'paddle']) -@pytest.mark.parametrize('da', [DocumentArray]) -@pytest.mark.parametrize('N', [2, 1000]) -@pytest.mark.parametrize('batch_size', [1, 256]) -@pytest.mark.parametrize('to_numpy', [True, False]) -def test_embedding_on_random_network(framework, da, N, batch_size, to_numpy): - docs = da.empty(N) - docs.blobs = np.random.random([N, 128]).astype(np.float32) - embed_model = random_embed_models[framework]() - docs.embed(embed_model, batch_size=batch_size, to_numpy=to_numpy) - - r = docs.embeddings - if hasattr(r, 'numpy'): - r = r.numpy() - embed1 = r.copy() - - # reset - docs.embeddings = np.random.random([N, 128]).astype(np.float32) - - # try it again, it should yield the same result - docs.embed(embed_model, batch_size=batch_size, to_numpy=to_numpy) - np.testing.assert_array_almost_equal(docs.embeddings, embed1) - - # reset - docs.embeddings = np.random.random([N, 128]).astype(np.float32) - - # now do this one by one - docs[: int(N / 2)].embed(embed_model, batch_size=batch_size, to_numpy=to_numpy) - docs[-int(N / 2) :].embed(embed_model, batch_size=batch_size, to_numpy=to_numpy) - np.testing.assert_array_almost_equal(docs.embeddings, embed1) diff --git a/tests/unit/array/mixins/test_empty.py b/tests/unit/array/mixins/test_empty.py deleted file mode 100644 index 480eec09f65..00000000000 --- a/tests/unit/array/mixins/test_empty.py +++ /dev/null @@ -1,8 +0,0 @@ -from docarray import DocumentArray - - -def test_empty_non_zero(): - da = DocumentArray.empty(10) - assert len(da) == 10 - da = DocumentArray.empty() - assert len(da) == 0 diff --git a/tests/unit/array/mixins/test_eval_class.py b/tests/unit/array/mixins/test_eval_class.py deleted file mode 100644 index 0c33d8309e4..00000000000 --- a/tests/unit/array/mixins/test_eval_class.py +++ /dev/null @@ -1,107 +0,0 @@ -import copy - -import numpy as np -import pytest - -from docarray import DocumentArray, Document - - -@pytest.mark.parametrize( - 'metric_fn, kwargs', - [ - ('r_precision', {}), - ('precision_at_k', {}), - ('hit_at_k', {}), - ('average_precision', {}), - ('reciprocal_rank', {}), - ('recall_at_k', {'max_rel': 9}), - ('f1_score_at_k', {'max_rel': 9}), - ('ndcg_at_k', {}), - ], -) -def test_eval_mixin_perfect_match(metric_fn, kwargs): - da = DocumentArray.empty(10) - da.embeddings = np.random.random([10, 256]) - da.match(da, exclude_self=True) - r = da.evaluate(da, metric=metric_fn, **kwargs) - assert isinstance(r, float) - assert r == 1.0 - for d in da: - assert d.evaluations[metric_fn].value == 1.0 - - -@pytest.mark.parametrize( - 'metric_fn, kwargs', - [ - ('r_precision', {}), - ('precision_at_k', {}), - ('hit_at_k', {}), - ('average_precision', {}), - ('reciprocal_rank', {}), - ('recall_at_k', {'max_rel': 9}), - ('f1_score_at_k', {'max_rel': 9}), - ('ndcg_at_k', {}), - ], -) -def test_eval_mixin_zero_match(metric_fn, kwargs): - da1 = DocumentArray.empty(10) - da1.embeddings = np.random.random([10, 256]) - da1.match(da1, exclude_self=True) - - da2 = copy.deepcopy(da1) - da2.embeddings = np.random.random([10, 256]) - da2.match(da2, exclude_self=True) - - r = da1.evaluate(da2, metric=metric_fn, **kwargs) - assert isinstance(r, float) - assert r == 1.0 - for d in da1: - d: Document - assert d.evaluations[metric_fn].value == 1.0 - - -def test_diff_len_should_raise(): - da1 = DocumentArray.empty(10) - da2 = DocumentArray.empty(5) - with pytest.raises(ValueError): - da1.evaluate(da2, metric='precision_at_k') - - -def test_diff_hash_fun_should_raise(): - da1 = DocumentArray.empty(10) - da2 = DocumentArray.empty(10) - with pytest.raises(ValueError): - da1.evaluate(da2, metric='precision_at_k') - - -def test_same_hash_same_len_fun_should_work(): - da1 = DocumentArray.empty(10) - da1.embeddings = np.random.random([10, 3]) - da1.match(da1) - da2 = DocumentArray.empty(10) - da2.embeddings = np.random.random([10, 3]) - da2.match(da2) - with pytest.raises(ValueError): - da1.evaluate(da2, metric='precision_at_k') - for d1, d2 in zip(da1, da2): - d1.id = d2.id - - da1.evaluate(da2, metric='precision_at_k') - - -def test_adding_noise(): - da = DocumentArray.empty(10) - - da.embeddings = np.random.random([10, 3]) - da.match(da, exclude_self=True) - - da2 = copy.deepcopy(da) - - for d in da2: - d.matches.extend(DocumentArray.empty(10)) - d.matches = d.matches.shuffle() - - assert da2.evaluate(da, metric='precision_at_k', k=10) < 1.0 - - for d in da2: - assert 0.0 < d.evaluations['precision_at_k'].value < 1.0 diff --git a/tests/unit/array/mixins/test_getset.py b/tests/unit/array/mixins/test_getset.py deleted file mode 100644 index 1a747f041c5..00000000000 --- a/tests/unit/array/mixins/test_getset.py +++ /dev/null @@ -1,156 +0,0 @@ -import numpy as np -import pytest -import scipy.sparse -import tensorflow as tf -import torch -from scipy.sparse import csr_matrix - -from docarray import DocumentArray, Document -from tests import random_docs - -rand_array = np.random.random([10, 3]) - - -def da_and_dam(): - rand_docs = random_docs(100) - da = DocumentArray() - da.extend(rand_docs) - return (da,) - - -@pytest.mark.parametrize( - 'array', - [ - rand_array, - torch.Tensor(rand_array), - tf.constant(rand_array), - csr_matrix(rand_array), - ], -) -def test_set_embeddings_multi_kind(array): - da = DocumentArray([Document() for _ in range(10)]) - da.embeddings = array - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_da_get_embeddings(da): - np.testing.assert_almost_equal(da._get_attributes('embedding'), da.embeddings) - np.testing.assert_almost_equal(da[:, 'embedding'], da.embeddings) - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_embeddings_setter_da(da): - emb = np.random.random((100, 128)) - da.embeddings = emb - np.testing.assert_almost_equal(da.embeddings, emb) - - for x, doc in zip(emb, da): - np.testing.assert_almost_equal(x, doc.embedding) - - da.embeddings = None - if hasattr(da, 'flush'): - da.flush() - assert not da.embeddings - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_embeddings_wrong_len(da): - embeddings = np.ones((2, 10)) - - with pytest.raises(ValueError): - da.embeddings = embeddings - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_blobs_getter_da(da): - blobs = np.random.random((100, 10, 10)) - da.blobs = blobs - assert len(da) == 100 - np.testing.assert_almost_equal(da.blobs, blobs) - - da.blobs = None - if hasattr(da, 'flush'): - da.flush() - assert not da.blobs - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_texts_getter_da(da): - assert len(da.texts) == 100 - assert da.texts == da[:, 'text'] - texts = ['text' for _ in range(100)] - da.texts = texts - assert da.texts == texts - - for x, doc in zip(texts, da): - assert x == doc.text - - da.texts = None - if hasattr(da, 'flush'): - da.flush() - - # unfortunately protobuf does not distinguish None and '' on string - # so non-set str field in Pb is '' - assert not da.texts - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_texts_wrong_len(da): - texts = ['hello'] - - with pytest.raises(ValueError): - da.texts = texts - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_blobs_wrong_len(da): - blobs = np.ones((2, 10, 10)) - - with pytest.raises(ValueError): - da.blobs = blobs - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_buffers_getter_setter(da): - with pytest.raises(ValueError): - da.buffers = [b'cc', b'bb', b'aa', b'dd'] - - da.buffers = [b'aa'] * len(da) - assert da.buffers == [b'aa'] * len(da) - - da.buffers = None - if hasattr(da, 'flush'): - da.flush() - - # unfortunately protobuf does not distinguish None and '' on string - # so non-set str field in Pb is '' - assert not da.buffers - - -def test_zero_embeddings(): - a = np.zeros([10, 6]) - da = DocumentArray.empty(10) - - # all zero, dense - da.embeddings = a - np.testing.assert_almost_equal(da.embeddings, a) - for d in da: - assert d.embedding.shape == (6,) - - # all zero, sparse - sp_a = scipy.sparse.coo_matrix(a) - da.embeddings = sp_a - np.testing.assert_almost_equal(da.embeddings.todense(), sp_a.todense()) - for d in da: - # scipy sparse row-vector can only be a (1, m) not squeezible - assert d.embedding.shape == (1, 6) - - # near zero, sparse - a = np.random.random([10, 6]) - a[a > 0.1] = 0 - sp_a = scipy.sparse.coo_matrix(a) - da.embeddings = sp_a - np.testing.assert_almost_equal(da.embeddings.todense(), sp_a.todense()) - for d in da: - # scipy sparse row-vector can only be a (1, m) not squeezible - assert d.embedding.shape == (1, 6) diff --git a/tests/unit/array/mixins/test_group.py b/tests/unit/array/mixins/test_group.py deleted file mode 100644 index 90316f3dfe4..00000000000 --- a/tests/unit/array/mixins/test_group.py +++ /dev/null @@ -1,90 +0,0 @@ -import pytest - -from docarray import DocumentArray, Document - - -def da_for_batching(): - da = DocumentArray.empty(100) - return (da,) - - -def docarray_for_split(): - da = DocumentArray() - da.append(Document(tags={'category': 'c'})) - da.append(Document(tags={'category': 'c'})) - da.append(Document(tags={'category': 'b'})) - da.append(Document(tags={'category': 'a'})) - da.append(Document(tags={'category': 'a'})) - return (da,) - - -def docarray_for_split_at_zero(): - da = DocumentArray() - da.append(Document(tags={'category': 0.0})) - da.append(Document(tags={'category': 0.0})) - da.append(Document(tags={'category': 1.0})) - da.append(Document(tags={'category': 2.0})) - da.append(Document(tags={'category': 2.0})) - return (da,) - - -def docarray_for_nest_split(): - da = DocumentArray() - da.append(Document(tags={'nest': {'category': 'c'}})) - da.append(Document(tags={'nest': {'category': 'c'}})) - da.append(Document(tags={'nest': {'category': 'b'}})) - da.append(Document(tags={'nest': {'category': 'a'}})) - da.append(Document(tags={'nest': {'category': 'a'}})) - return (da,) - - -@pytest.mark.parametrize('da', docarray_for_split()) -def test_split(da): - rv = da.split_by_tag('category') - assert isinstance(rv, dict) - assert sorted(list(rv.keys())) == ['a', 'b', 'c'] - # assure order is preserved c, b, a - assert list(rv.keys()) == ['c', 'b', 'a'] - # original input c, c, b, a, a - assert len(rv['c']) == 2 - assert len(rv['b']) == 1 - assert len(rv['a']) == 2 - rv = da.split_by_tag('random') - assert not rv # wrong tag returns empty dict - - -@pytest.mark.parametrize('da', docarray_for_split_at_zero()) -def test_split_at_zero(da): - rv = da.split_by_tag('category') - assert isinstance(rv, dict) - assert sorted(list(rv.keys())) == [0.0, 1.0, 2.0] - - -@pytest.mark.parametrize('da', docarray_for_nest_split()) -def test_dunder_split(da): - rv = da.split_by_tag('nest__category') - assert isinstance(rv, dict) - assert sorted(list(rv.keys())) == ['a', 'b', 'c'] - # assure order is preserved c, b, a - assert list(rv.keys()) == ['c', 'b', 'a'] - # original input c, c, b, a, a - assert len(rv['c']) == 2 - assert len(rv['b']) == 1 - assert len(rv['a']) == 2 - - assert len(da.split_by_tag('nest__random')) == 1 - - -@pytest.mark.parametrize('da', da_for_batching()) -@pytest.mark.parametrize('batch_size', [1, 5, 100, 200]) -@pytest.mark.parametrize('shuffle', [True, False]) -def test_batching(da, batch_size, shuffle): - all_ids = [] - for v in da.batch(batch_size=batch_size, shuffle=shuffle): - assert len(v) <= batch_size - all_ids.extend(v[:, 'id']) - - if shuffle: - assert all_ids != da[:, 'id'] - else: - assert all_ids == da[:, 'id'] diff --git a/tests/unit/array/mixins/test_io.py b/tests/unit/array/mixins/test_io.py deleted file mode 100644 index 21c5510daff..00000000000 --- a/tests/unit/array/mixins/test_io.py +++ /dev/null @@ -1,107 +0,0 @@ -import os -import uuid - -import numpy as np -import pytest - -from docarray import DocumentArray -from tests import random_docs - - -def da_and_dam(): - da = DocumentArray(random_docs(100)) - return (da,) - - -@pytest.mark.slow -@pytest.mark.parametrize('method', ['json', 'binary']) -@pytest.mark.parametrize('da', da_and_dam()) -def test_document_save_load(method, tmp_path, da): - tmp_file = os.path.join(tmp_path, 'test') - da.save(tmp_file, file_format=method) - da_r = type(da).load(tmp_file, file_format=method) - - assert type(da) is type(da_r) - assert len(da) == len(da_r) - for d, d_r in zip(da, da_r): - assert d.id == d_r.id - np.testing.assert_equal(d.embedding, d_r.embedding) - assert d.content == d_r.content - - -@pytest.mark.parametrize('flatten_tags', [True, False]) -@pytest.mark.parametrize('da', da_and_dam()) -def test_da_csv_write(flatten_tags, tmp_path, da): - tmpfile = os.path.join(tmp_path, 'test.csv') - da.save_csv(tmpfile, flatten_tags) - with open(tmpfile) as fp: - assert len([v for v in fp]) == len(da) + 1 - - -@pytest.mark.parametrize('da', [DocumentArray]) -def test_from_ndarray(da): - _da = da.from_ndarray(np.random.random([10, 256])) - assert len(_da) == 10 - - -@pytest.mark.parametrize('da', [DocumentArray]) -def test_from_files(da): - assert len(da.from_files(patterns='*.*', to_dataturi=True, size=1)) == 1 - - -cur_dir = os.path.dirname(os.path.abspath(__file__)) - - -@pytest.mark.parametrize('da', [DocumentArray]) -def test_from_ndjson(da): - with open(os.path.join(cur_dir, 'docs.jsonlines')) as fp: - _da = da.from_ndjson(fp) - assert len(_da) == 2 - - -@pytest.mark.parametrize('da_cls', [DocumentArray]) -def test_from_to_pd_dataframe(da_cls): - # simple - assert len(da_cls.from_dataframe(da_cls.empty(2).to_dataframe())) == 2 - - # more complicated - da = da_cls.empty(2) - da.embeddings = [[1, 2, 3], [4, 5, 6]] - da.blobs = [[1, 2], [2, 1]] - da[0].tags = {'hello': 'world'} - da2 = da_cls.from_dataframe(da.to_dataframe()) - assert da2[0].tags == {'hello': 'world'} - assert da2[1].tags == {} - - -@pytest.mark.parametrize('da_cls', [DocumentArray]) -def test_from_to_bytes(da_cls): - # simple - assert len(da_cls.load_binary(bytes(da_cls.empty(2)))) == 2 - - # more complicated - da = da_cls.empty(2) - da.embeddings = [[1, 2, 3], [4, 5, 6]] - da.blobs = [[1, 2], [2, 1]] - da[0].tags = {'hello': 'world'} - da2 = da_cls.load_binary(bytes(da)) - assert da2.blobs == [[1, 2], [2, 1]] - assert da2.embeddings == [[1, 2, 3], [4, 5, 6]] - assert da2[0].tags == {'hello': 'world'} - assert da2[1].tags == {} - - -@pytest.mark.parametrize('da_cls', [DocumentArray]) -@pytest.mark.parametrize('show_progress', [True, False]) -def test_push_pull_io(da_cls, show_progress): - da1 = da_cls.empty(10) - da1.embeddings = np.random.random([len(da1), 256]) - random_texts = [str(uuid.uuid1()) for _ in da1] - da1.texts = random_texts - - da1.push('myda', show_progress=show_progress) - - da2 = da_cls.pull('myda', show_progress=show_progress) - - assert len(da1) == len(da2) == 10 - assert da1.texts == da2.texts == random_texts diff --git a/tests/unit/array/mixins/test_magic.py b/tests/unit/array/mixins/test_magic.py deleted file mode 100644 index 76ade193d7c..00000000000 --- a/tests/unit/array/mixins/test_magic.py +++ /dev/null @@ -1,47 +0,0 @@ -import pytest - -from docarray import DocumentArray - -N = 100 - - -def da_and_dam(): - da = DocumentArray.empty(N) - return (da,) - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_iter_len_bool(da): - j = 0 - for _ in da: - j += 1 - assert j == N - assert j == len(da) - assert da - da.clear() - assert not da - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_repr(da): - assert f'length={N}' in repr(da) - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_iadd(da): - oid = id(da) - dap = DocumentArray.empty(10) - da += dap - assert len(da) == N + len(dap) - nid = id(da) - assert nid == oid - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_add(da): - oid = id(da) - dap = DocumentArray.empty(10) - da = da + dap - assert len(da) == N + len(dap) - nid = id(da) - assert nid != oid diff --git a/tests/unit/array/mixins/test_match.py b/tests/unit/array/mixins/test_match.py deleted file mode 100644 index 43592259593..00000000000 --- a/tests/unit/array/mixins/test_match.py +++ /dev/null @@ -1,517 +0,0 @@ -import copy - -import numpy as np -import paddle -import pytest -import scipy.sparse as sp -import tensorflow as tf -import torch -from scipy.sparse import csr_matrix, bsr_matrix, coo_matrix, csc_matrix -from scipy.spatial.distance import cdist as scipy_cdist - -from docarray import Document, DocumentArray - - -@pytest.fixture() -def doc_lists(): - d1 = Document(embedding=np.array([0, 0, 0])) - d2 = Document(embedding=np.array([3, 0, 0])) - d3 = Document(embedding=np.array([1, 0, 0])) - d4 = Document(embedding=np.array([2, 0, 0])) - - d1_m = Document(embedding=np.array([1, 0, 0])) - d2_m = Document(embedding=np.array([2, 0, 0])) - d3_m = Document(embedding=np.array([0, 0, 1])) - d4_m = Document(embedding=np.array([0, 0, 2])) - d5_m = Document(embedding=np.array([0, 0, 3])) - - return [d1, d2, d3, d4], [d1_m, d2_m, d3_m, d4_m, d5_m] - - -@pytest.fixture -def docarrays_for_embedding_distance_computation(doc_lists): - D1, D2 = doc_lists - da1 = DocumentArray(D1) - da2 = DocumentArray(D2) - return da1, da2 - - -@pytest.fixture -def docarrays_for_embedding_distance_computation_sparse(): - d1 = Document(embedding=sp.csr_matrix([0, 0, 0])) - d2 = Document(embedding=sp.csr_matrix([3, 0, 0])) - d3 = Document(embedding=sp.csr_matrix([1, 0, 0])) - d4 = Document(embedding=sp.csr_matrix([2, 0, 0])) - - d1_m = Document(embedding=sp.csr_matrix([1, 0, 0])) - d2_m = Document(embedding=sp.csr_matrix([2, 0, 0])) - d3_m = Document(embedding=sp.csr_matrix([0, 0, 1])) - d4_m = Document(embedding=sp.csr_matrix([0, 0, 2])) - d5_m = Document(embedding=sp.csr_matrix([0, 0, 3])) - - D1 = DocumentArray([d1, d2, d3, d4]) - D2 = DocumentArray([d1_m, d2_m, d3_m, d4_m, d5_m]) - return D1, D2 - - -@pytest.fixture -def embeddings(): - return np.array([[1, 0, 0], [2, 0, 0], [3, 0, 0]]) - - -def doc_lists_to_doc_arrays(doc_lists, *args, **kwargs): - doc_list1, doc_list2 = doc_lists - D1 = DocumentArray() - D1.extend(doc_list1) - D2 = DocumentArray() - D2.extend(doc_list2) - return D1, D2 - - -@pytest.mark.parametrize( - 'limit, batch_size', [(1, None), (2, None), (None, None), (1, 1), (1, 2), (2, 1)] -) -@pytest.mark.parametrize('only_id', [True, False]) -def test_matching_retrieves_correct_number( - doc_lists, - limit, - batch_size, - tmpdir, - only_id, -): - D1, D2 = doc_lists_to_doc_arrays( - doc_lists, - ) - D1.match( - D2, metric='sqeuclidean', limit=limit, batch_size=batch_size, only_id=only_id - ) - for m in D1[:, 'matches']: - if limit is None: - assert len(m) == len(D2) - else: - assert len(m) == limit - - -@pytest.mark.parametrize('metric', ['sqeuclidean', 'cosine']) -@pytest.mark.parametrize('only_id', [True, False]) -def test_matching_same_results_with_sparse( - docarrays_for_embedding_distance_computation, - docarrays_for_embedding_distance_computation_sparse, - metric, - only_id, -): - D1, D2 = docarrays_for_embedding_distance_computation - D1_sp, D2_sp = docarrays_for_embedding_distance_computation_sparse - - # use match with numpy arrays - D1.match(D2, metric=metric, only_id=only_id) - distances = [] - for m in D1[:, 'matches']: - for d in m: - distances.extend([d.scores[metric].value]) - - # use match with sparse arrays - D1_sp.match(D2_sp, metric=metric, is_sparse=True) - distances_sparse = [] - for m in D1[:, 'matches']: - for d in m: - distances_sparse.extend([d.scores[metric].value]) - - np.testing.assert_equal(distances, distances_sparse) - - -@pytest.mark.parametrize('metric', ['sqeuclidean', 'cosine']) -@pytest.mark.parametrize('only_id', [True, False]) -def test_matching_same_results_with_batch( - docarrays_for_embedding_distance_computation, metric, only_id -): - D1, D2 = docarrays_for_embedding_distance_computation - D1_batch = copy.deepcopy(D1) - D2_batch = copy.deepcopy(D2) - - # use match without batches - D1.match(D2, metric=metric, only_id=only_id) - distances = [] - for m in D1[:, 'matches']: - for d in m: - distances.extend([d.scores[metric].value]) - - # use match with batches - D1_batch.match(D2_batch, metric=metric, batch_size=10) - - distances_batch = [] - for m in D1[:, 'matches']: - for d in m: - distances_batch.extend([d.scores[metric].value]) - - np.testing.assert_equal(distances, distances_batch) - - -@pytest.mark.parametrize('metric', ['euclidean', 'cosine']) -@pytest.mark.parametrize('only_id', [True, False]) -def test_matching_scipy_cdist( - docarrays_for_embedding_distance_computation, metric, only_id -): - def scipy_cdist_metric(X, Y, *args): - return scipy_cdist(X, Y, metric=metric) - - D1, D2 = docarrays_for_embedding_distance_computation - D1_scipy = copy.deepcopy(D1) - - # match with our custom metric - D1.match(D2, metric=metric) - distances = [] - for m in D1[:, 'matches']: - for d in m: - distances.extend([d.scores[metric].value]) - - # match with callable cdist function from scipy - D1_scipy.match(D2, metric=scipy_cdist_metric, only_id=only_id) - distances_scipy = [] - for m in D1[:, 'matches']: - for d in m: - distances_scipy.extend([d.scores[metric].value]) - - np.testing.assert_equal(distances, distances_scipy) - - -@pytest.mark.parametrize( - 'normalization, metric', - [ - ((0, 1), 'sqeuclidean'), - (None, 'euclidean'), - ((0, 1), 'euclidean'), - (None, 'cosine'), - ((0, 1), 'cosine'), - ], -) -@pytest.mark.parametrize('use_scipy', [True, False]) -@pytest.mark.parametrize('only_id', [True, False]) -def test_matching_retrieves_closest_matches( - doc_lists, - normalization, - metric, - use_scipy, - only_id, -): - """ - Tests if match.values are returned 'low to high' if normalization is True or 'high to low' otherwise - """ - D1, D2 = doc_lists_to_doc_arrays( - doc_lists, - ) - D1.match( - D2, - metric=metric, - limit=3, - normalization=normalization, - use_scipy=use_scipy, - only_id=only_id, - ) - expected_sorted_values = [ - D1[0].matches[i].scores['sqeuclidean'].value for i in range(3) - ] - if normalization: - assert min(expected_sorted_values) >= 0 - assert max(expected_sorted_values) <= 1 - else: - assert expected_sorted_values == sorted(expected_sorted_values) - - -@pytest.mark.parametrize('buffer_pool_size', [1000, 3]) -@pytest.mark.parametrize('first_memmap', [True, False]) -@pytest.mark.parametrize('second_memmap', [True, False]) -@pytest.mark.parametrize('only_id', [True, False]) -def test_2arity_function( - first_memmap, second_memmap, doc_lists, tmpdir, buffer_pool_size, only_id -): - def dotp(x, y, *args): - return np.dot(x, np.transpose(y)) - - D1, D2 = doc_lists_to_doc_arrays( - doc_lists, - tmpdir, - first_memmap, - second_memmap, - buffer_pool_size=buffer_pool_size, - ) - D1.match(D2, metric=dotp, use_scipy=True, only_id=only_id) - - for d in D1: - for m in d.matches: - assert 'dotp' in m.scores - - -@pytest.mark.parametrize('only_id', [True, False]) -def test_match_inclusive(only_id): - """Call match function, while the other :class:`DocumentArray` is itself - or have same :class:`Document`. - """ - # The document array da1 match with itself. - da1 = DocumentArray( - [ - Document(embedding=np.array([1, 2, 3])), - Document(embedding=np.array([1, 0, 1])), - Document(embedding=np.array([1, 1, 2])), - ] - ) - - da1.match(da1, only_id=only_id) - assert len(da1) == 3 - traversed = da1.traverse_flat(traversal_paths='m,mm,mmm') - assert len(traversed) == 9 - # The document array da2 shares same documents with da1 - da2 = DocumentArray([Document(embedding=np.array([4, 1, 3])), da1[0], da1[1]]) - da1.match(da2, only_id=only_id) - assert len(da2) == 3 - traversed = da1.traverse_flat(traversal_paths='m,mm,mmm') - assert len(traversed) == 9 - - -@pytest.mark.parametrize('exclude_self, num_matches', [(True, 1), (False, 2)]) -@pytest.mark.parametrize('only_id', [True, False]) -def test_match_exclude_self(exclude_self, num_matches, only_id): - da1 = DocumentArray( - [ - Document(id='1', embedding=np.array([1, 2])), - Document(id='2', embedding=np.array([3, 4])), - ] - ) - da2 = DocumentArray( - [ - Document(id='1', embedding=np.array([1, 2])), - Document(id='2', embedding=np.array([3, 4])), - ] - ) - da1.match(da2, exclude_self=exclude_self, only_id=only_id) - for d in da1: - assert len(d.matches) == num_matches - - -@pytest.fixture() -def get_pair_document_array(): - da1 = DocumentArray( - [ - Document(id='1', embedding=np.array([1, 2])), - Document(id='2', embedding=np.array([3, 4])), - ] - ) - da2 = DocumentArray( - [ - Document(id='1', embedding=np.array([1, 2])), - Document(id='2', embedding=np.array([3, 4])), - Document(id='3', embedding=np.array([4, 5])), - ] - ) - yield da1, da2 - - -@pytest.mark.parametrize( - 'limit, expect_len, exclude_self', - [ - (2, 2, True), - (1, 1, True), - (3, 2, True), - (2, 2, False), - (1, 1, False), - (3, 3, False), - ], -) -def test_match_exclude_self_limit_2( - get_pair_document_array, exclude_self, limit, expect_len -): - da1, da2 = get_pair_document_array - da1.match(da2, exclude_self=exclude_self, limit=limit) - for d in da1: - assert len(d.matches) == expect_len - - -@pytest.mark.parametrize( - 'lhs, rhs', - [ - (DocumentArray(), DocumentArray()), - ( - DocumentArray( - [ - Document(embedding=np.array([3, 4])), - Document(embedding=np.array([4, 5])), - ] - ), - DocumentArray( - [ - Document(embedding=np.array([3, 4])), - Document(embedding=np.array([4, 5])), - ] - ), - ), - ( - DocumentArray(), - DocumentArray( - [ - Document(embedding=np.array([3, 4])), - Document(embedding=np.array([4, 5])), - ] - ), - ), - ( - ( - DocumentArray( - [ - Document(embedding=np.array([3, 4])), - Document(embedding=np.array([4, 5])), - ] - ) - ), - DocumentArray(), - ), - (None, DocumentArray()), - (DocumentArray(), None), - ], -) -def test_match_none(lhs, rhs): - if lhs is not None: - lhs.match(rhs) - if rhs is not None: - rhs.match(lhs) - - -@pytest.fixture() -def get_two_docarray(): - d1 = Document(embedding=np.array([0, 0, 0])) - d1c1 = Document(embedding=np.array([0, 1, 0])) - - d2 = Document(embedding=np.array([1, 0, 0])) - d2c1 = Document(embedding=np.array([1, 1, 0])) - d2c2 = Document(embedding=np.array([1, 0, 1])) - - d3 = Document(embedding=np.array([2, 1, 1])) - d3c1 = Document(embedding=np.array([2, 1, 0])) - d3c2 = Document(embedding=np.array([2, 0, 1])) - d3c3 = Document(embedding=np.array([2, 0, 0])) - - d4 = Document(embedding=np.array([3, 1, 1])) - d4c1 = Document(embedding=np.array([3, 1, 0])) - d4c2 = Document(embedding=np.array([3, 0, 1])) - d4c3 = Document(embedding=np.array([3, 0, 0])) - d4c4 = Document(embedding=np.array([3, 1, 1])) - - d1.chunks.extend([d1c1]) - d2.chunks.extend([d2c1, d2c2]) - d3.chunks.extend([d3c1, d3c2, d3c3]) - d4.chunks.extend([d4c1, d4c2, d4c3, d4c4]) - - da1 = DocumentArray([d1, d2]) - da2 = DocumentArray([d3, d4]) - yield da1, da2 - - -def test_match_with_traversal_path(get_two_docarray): - da1, da2 = get_two_docarray - da1.match(da2.traverse_flat('c')) - assert len(da1[0].matches) == len(da2[0].chunks) + len(da2[1].chunks) - - da2.match(da1.traverse_flat('c')) - assert len(da2[0].matches) == len(da1[0].chunks) + len(da1[1].chunks) - - -def test_match_on_two_sides_chunks(get_two_docarray): - da1, da2 = get_two_docarray - da2.traverse_flat('c').match(da1.traverse_flat('c')) - assert len(da2[0].matches) == 0 - assert len(da2[0].chunks[0].matches) == len(da1[0].chunks) + len(da1[1].chunks) - - da1.traverse_flat('c').match(da2.traverse_flat('c')) - assert len(da1[0].matches) == 0 - assert len(da1[0].chunks[0].matches) == len(da2[0].chunks) + len(da2[1].chunks) - - -@pytest.mark.parametrize('exclude_self', [True, False]) -@pytest.mark.parametrize('limit', [1, 2, 3]) -def test_exclude_self_should_keep_limit(limit, exclude_self): - da = DocumentArray( - [ - Document(embedding=np.array([3, 1, 0])), - Document(embedding=np.array([3, 0, 1])), - Document(embedding=np.array([3, 0, 0])), - Document(embedding=np.array([3, 1, 1])), - ] - ) - da.match(da, exclude_self=exclude_self, limit=limit) - for d in da: - assert len(d.matches) == limit - if exclude_self: - for m in d.matches: - assert d.id != m.id - - -@pytest.mark.parametrize('only_id', [True, False]) -def test_only_id(docarrays_for_embedding_distance_computation, only_id): - D1, D2 = docarrays_for_embedding_distance_computation - D1.match(D2, only_id=only_id) - for d in D1: - for m in d.matches: - assert (m.embedding is None) == only_id - assert m.id - - -@pytest.mark.parametrize( - 'match_kwargs', - [ - dict(limit=5, normalization=(1, 0), batch_size=10), - dict(normalization=(1, 0), batch_size=10), - dict(normalization=(1, 0)), - dict(), - ], -) -@pytest.mark.parametrize('nnz_ratio', [0.5, 1]) -def test_dense_vs_sparse_match(match_kwargs, nnz_ratio): - N = 100 - D = 256 - sp_embed = np.random.random([N, D]) - sp_embed[sp_embed > nnz_ratio] = 0 - - da1 = DocumentArray.empty(N) - da2 = DocumentArray.empty(N) - - # use sparse embedding - da1.embeddings = sp.coo_matrix(sp_embed) - da1.texts = [str(j) for j in range(N)] - size_sp = sum(d.nbytes for d in da1) - da1.match(da1, **match_kwargs) - - sparse_result = [m.text for m in da1[0].matches] - - # use dense embedding - da2.embeddings = sp_embed - da2.texts = [str(j) for j in range(N)] - size_dense = sum(d.nbytes for d in da2) - da2.match(da2, **match_kwargs) - dense_result = [m.text for m in da2[0].matches] - - assert sparse_result == dense_result - - print( - f'sparse DA: {size_sp} bytes is {size_sp / size_dense * 100:.0f}% of dense DA {size_dense} bytes' - ) - - -def get_ndarrays(): - a = np.random.random([10, 3]) - a[a > 0.5] = 0 - return [ - a, - torch.tensor(a), - tf.constant(a), - paddle.to_tensor(a), - csr_matrix(a), - bsr_matrix(a), - coo_matrix(a), - csc_matrix(a), - ] - - -@pytest.mark.parametrize('ndarray_val', get_ndarrays()) -def test_diff_framework_match(ndarray_val): - da = DocumentArray.empty(10) - da.embeddings = ndarray_val - da.match(da) diff --git a/tests/unit/array/mixins/test_parallel.py b/tests/unit/array/mixins/test_parallel.py deleted file mode 100644 index b05f1f215e6..00000000000 --- a/tests/unit/array/mixins/test_parallel.py +++ /dev/null @@ -1,112 +0,0 @@ -import os -import pytest - -from docarray import DocumentArray, Document - - -def foo(d: Document): - return ( - d.load_uri_to_image_blob() - .set_image_blob_normalization() - .set_image_blob_channel_axis(-1, 0) - .set_image_blob_shape((222, 222), 0) - ) - - -def foo_batch(da: DocumentArray): - for d in da: - foo(d) - return da - - -@pytest.mark.skipif( - 'GITHUB_WORKFLOW' in os.environ, - reason='this test somehow fail on Github CI, but it MUST run successfully on local', -) -@pytest.mark.parametrize( - 'da_cls', - [ - DocumentArray, - ], -) -@pytest.mark.parametrize('backend', ['process', 'thread']) -@pytest.mark.parametrize('num_worker', [1, 2, None]) -def test_parallel_map(pytestconfig, da_cls, backend, num_worker): - da = da_cls.from_files(f'{pytestconfig.rootdir}/**/*.jpeg')[:10] - - # use a generator - for d in da.map(foo, backend, num_worker=num_worker): - assert d.blob.shape == (3, 222, 222) - - da = da_cls.from_files(f'{pytestconfig.rootdir}/**/*.jpeg')[:10] - - # use as list, here the caveat is when using process backend you can not modify thing in-place - list(da.map(foo, backend, num_worker=num_worker)) - if backend == 'thread': - assert da.blobs.shape == (len(da), 3, 222, 222) - else: - assert da.blobs is None - - da = da_cls.from_files(f'{pytestconfig.rootdir}/**/*.jpeg')[:10] - da_new = da.apply(foo) - assert da_new.blobs.shape == (len(da_new), 3, 222, 222) - - -@pytest.mark.skipif( - 'GITHUB_WORKFLOW' in os.environ, - reason='this test somehow fail on Github CI, but it MUST run successfully on local', -) -@pytest.mark.parametrize( - 'da_cls', - [ - DocumentArray, - ], -) -@pytest.mark.parametrize('backend', ['thread']) -@pytest.mark.parametrize('num_worker', [1, 2, None]) -@pytest.mark.parametrize('b_size', [1, 2, 256]) -def test_parallel_map_batch(pytestconfig, da_cls, backend, num_worker, b_size): - da = da_cls.from_files(f'{pytestconfig.rootdir}/**/*.jpeg')[:10] - - # use a generator - for _da in da.map_batch( - foo_batch, batch_size=b_size, backend=backend, num_worker=num_worker - ): - for d in _da: - assert d.blob.shape == (3, 222, 222) - - da = da_cls.from_files(f'{pytestconfig.rootdir}/**/*.jpeg')[:10] - - # use as list, here the caveat is when using process backend you can not modify thing in-place - list( - da.map_batch( - foo_batch, batch_size=b_size, backend=backend, num_worker=num_worker - ) - ) - if backend == 'thread': - assert da.blobs.shape == (len(da), 3, 222, 222) - else: - assert da.blobs is None - - da_new = da.apply_batch(foo_batch, batch_size=b_size) - assert da_new.blobs.shape == (len(da_new), 3, 222, 222) - - -@pytest.mark.skipif( - 'GITHUB_WORKFLOW' in os.environ, - reason='this test somehow fail on Github CI, but it MUST run successfully on local', -) -@pytest.mark.parametrize( - 'da_cls', - [ - DocumentArray, - ], -) -def test_map_lambda(pytestconfig, da_cls): - da = da_cls.from_files(f'{pytestconfig.rootdir}/**/*.jpeg')[:10] - - for d in da: - assert d.blob is None - - for d in da.map(lambda x: x.load_uri_to_image_blob()): - assert d.blob is not None diff --git a/tests/unit/array/mixins/test_plot.py b/tests/unit/array/mixins/test_plot.py deleted file mode 100644 index 1ec701244b8..00000000000 --- a/tests/unit/array/mixins/test_plot.py +++ /dev/null @@ -1,72 +0,0 @@ -import json -import os -import random - -import numpy as np -import pytest - -from docarray import DocumentArray, Document - - -def test_sprite_image_generator(pytestconfig, tmpdir): - da = DocumentArray.from_files( - [ - f'{pytestconfig.rootdir}/**/*.png', - f'{pytestconfig.rootdir}/**/*.jpg', - f'{pytestconfig.rootdir}/**/*.jpeg', - ] - ) - da.plot_image_sprites(tmpdir / 'sprint_da.png') - assert os.path.exists(tmpdir / 'sprint_da.png') - - -def da_and_dam(): - embeddings = np.array([[1, 0, 0], [2, 0, 0], [3, 0, 0]]) - doc_array = DocumentArray( - [ - Document(embedding=x, tags={'label': random.randint(0, 5)}) - for x in embeddings - ] - ) - - return (doc_array,) - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_plot_embeddings(da): - p = da.plot_embeddings(start_server=False) - assert os.path.exists(p) - assert os.path.exists(os.path.join(p, 'config.json')) - with open(os.path.join(p, 'config.json')) as fp: - config = json.load(fp) - assert len(config['embeddings']) == 1 - assert config['embeddings'][0]['tensorShape'] == list(da.embeddings.shape) - - -def test_plot_embeddings_same_path(tmpdir): - da1 = DocumentArray.empty(100) - da1.embeddings = np.random.random([100, 5]) - p1 = da1.plot_embeddings(start_server=False, path=tmpdir) - da2 = DocumentArray.empty(768) - da2.embeddings = np.random.random([768, 5]) - p2 = da2.plot_embeddings(start_server=False, path=tmpdir) - assert p1 == p2 - assert os.path.exists(p1) - with open(os.path.join(p1, 'config.json')) as fp: - config = json.load(fp) - assert len(config['embeddings']) == 2 - - -def test_summary_homo_hetero(): - da = DocumentArray.empty(100) - da._get_attributes() - da.summary() - - da[0].pop('id') - da.summary() - - -def test_empty_get_attributes(): - da = DocumentArray.empty(10) - da[0].pop('id') - print(da[:, 'id']) diff --git a/tests/unit/array/mixins/test_reduce.py b/tests/unit/array/mixins/test_reduce.py deleted file mode 100644 index 0d0851b6c7c..00000000000 --- a/tests/unit/array/mixins/test_reduce.py +++ /dev/null @@ -1,209 +0,0 @@ -from copy import deepcopy - -import numpy as np - -from docarray import DocumentArray, Document - - -def test_reduce_doc(): - doc1 = Document( - text='doc1', - matches=[Document(id='m0'), Document(id='m2')], - chunks=[Document(id='c0'), Document(id='c1')], - ) - - doc2 = Document( - text='doc2', - embedding=np.zeros(3), - matches=[Document(id='m0'), Document(id='m1')], - chunks=[Document(id='c0'), Document(id='c2')], - ) - - DocumentArray._reduce_doc(doc1, doc2) - - assert doc1.text == 'doc1' - assert (doc1.embedding == np.zeros(3)).all() - for i in range(3): - assert f'c{i}' in doc1.chunks - assert f'm{i}' in doc1.matches - - -def test_reduce(): - da1, da2 = ( - DocumentArray( - [ - Document( - id='r0', - text='da1', - matches=[ - Document(id='r0m0'), - Document(id='r0m2'), - Document(id='r0m1'), - ], - ), - Document( - id='r2', - text='da1', - matches=[ - Document(id='r2m0'), - Document(id='r2m2'), - Document(id='r2m1'), - ], - ), - ] - ), - DocumentArray( - [ - Document( - id='r0', - text='da2', - matches=[ - Document(id='r0m0'), - Document(id='r0m1'), - Document(id='r0m3'), - ], - ), - Document( - id='r1', - text='da2', - matches=[ - Document(id='r1m0'), - Document(id='r1m1'), - Document(id='r1m2'), - Document(id='r1m3'), - ], - ), - Document( - id='r2', - text='da2', - matches=[ - Document(id='r2m3'), - ], - ), - ] - ), - ) - - da1.reduce(da2) - - for i in range(3): - assert f'r{i}' in da1 - for j in range(4): - assert f'r{i}m{j}' in da1[f'r{i}'].matches - - assert da1['r0'].text == 'da1' - assert da1['r1'].text == 'da2' - assert da1['r2'].text == 'da1' - - -def test_reduce_nested(): - da1, da2 = ( - DocumentArray( - [ - Document( - id='r1', - chunks=[ - Document( - id='c1', - chunks=[Document(id='c1c2')], - matches=[Document(id='c1m2')], - ), - ], - matches=[ - Document( - id='m1', chunks=[Document(id='m1c1'), Document(id='m1c2')] - ), - Document(id='m2'), - ], - ), - ] - ), - DocumentArray( - [ - Document( - id='r1', - chunks=[ - Document( - id='c1', - chunks=[Document(id='c1c1')], - matches=[Document(id='c1m1')], - ), - Document( - id='c2', - chunks=[Document(id='c2c1'), Document(id='c2c2')], - matches=[Document(id='c2m1'), Document(id='c2m2')], - ), - ], - matches=[ - Document( - id='m1', matches=[Document(id='m1m1'), Document(id='m1m2')] - ), - Document( - id='m2', - chunks=[Document(id='m2c1'), Document(id='m2c2')], - matches=[Document(id='m2m1'), Document(id='m2m2')], - ), - ], - ), - ] - ), - ) - - da1.reduce(da2) - for i in range(1, 3): - assert f'c{i}' in da1[0].chunks - assert f'm{i}' in da1[0].matches - for j in range(1, 3): - assert f'c{i}c{j}' in da1[0].chunks[f'c{i}'].chunks - assert f'c{i}m{j}' in da1[0].chunks[f'c{i}'].matches - - assert f'm{i}c{j}' in da1[0].matches[f'm{i}'].chunks - assert f'm{i}m{j}' in da1[0].matches[f'm{i}'].matches - - -def test_reduce_mat(): - docs = DocumentArray([Document(id=f'r{i}') for i in range(10)]) - doc_matrix = [deepcopy(docs) for _ in range(10)] - for i, da in enumerate(doc_matrix): - for doc in da: - doc.matches.append(Document(id=str(i))) - - reduced_da = doc_matrix[0].reduce_all(doc_matrix[1:]) - for doc in reduced_da: - for i in range(10): - assert str(i) in doc.matches - - -def test_reduce_data_props(): - da1, da2, da3 = ( - DocumentArray([Document(id='r', text='doc1', chunks=[Document(id='c1')])]), - DocumentArray( - [ - Document( - id='r', - text='doc2', - embedding=np.zeros(3), - chunks=[Document(id='c2')], - ) - ] - ), - DocumentArray( - [ - Document( - id='r', - text='doc3', - tags={'a': 'b'}, - matches=[Document(id='m3')], - ) - ] - ), - ) - da1.reduce_all([da2, da3]) - assert da1[0].text == 'doc1' - assert da1[0].id == 'r' - - # chunks and matches merged, not overridden - assert da1[0].chunks[0].id == 'c1' - assert da1[0].chunks[1].id == 'c2' - assert da1[0].matches[0].id == 'm3' - assert da1[0].tags == {'a': 'b'} diff --git a/tests/unit/array/mixins/test_sample.py b/tests/unit/array/mixins/test_sample.py deleted file mode 100644 index ddb15dc4fe1..00000000000 --- a/tests/unit/array/mixins/test_sample.py +++ /dev/null @@ -1,50 +0,0 @@ -import pytest - -from docarray import DocumentArray - - -def da_and_dam(N): - da = DocumentArray.empty(N) - return (da,) - - -@pytest.mark.parametrize('da', da_and_dam(100)) -def test_sample(da): - sampled = da.sample(1) - assert len(sampled) == 1 - sampled = da.sample(5) - assert len(sampled) == 5 - assert isinstance(sampled, DocumentArray) - with pytest.raises(ValueError): - da.sample(101) # can not sample with k greater than lenth of document array. - - -@pytest.mark.parametrize('da', da_and_dam(100)) -def test_sample_with_seed(da): - sampled_1 = da.sample(5, seed=1) - sampled_2 = da.sample(5, seed=1) - sampled_3 = da.sample(5, seed=2) - assert len(sampled_1) == len(sampled_2) == len(sampled_3) == 5 - assert sampled_1 == sampled_2 - assert sampled_1 != sampled_3 - - -@pytest.mark.parametrize('da', da_and_dam(100)) -def test_shuffle(da): - shuffled = da.shuffle() - assert len(shuffled) == len(da) - assert isinstance(shuffled, DocumentArray) - ids_before_shuffle = [d.id for d in da] - ids_after_shuffle = [d.id for d in shuffled] - assert ids_before_shuffle != ids_after_shuffle - assert sorted(ids_before_shuffle) == sorted(ids_after_shuffle) - - -@pytest.mark.parametrize('da', da_and_dam(100)) -def test_shuffle_with_seed(da): - shuffled_1 = da.shuffle(seed=1) - shuffled_2 = da.shuffle(seed=1) - shuffled_3 = da.shuffle(seed=2) - assert len(shuffled_1) == len(shuffled_2) == len(shuffled_3) == len(da) - assert shuffled_1 == shuffled_2 - assert shuffled_1 != shuffled_3 diff --git a/tests/unit/array/mixins/test_text.py b/tests/unit/array/mixins/test_text.py deleted file mode 100644 index 07d8d2d343a..00000000000 --- a/tests/unit/array/mixins/test_text.py +++ /dev/null @@ -1,98 +0,0 @@ -import numpy as np -import pytest - -from docarray import DocumentArray, Document - - -def da_and_dam(): - da = DocumentArray( - [ - Document(text='hello'), - Document(text='hello world'), - Document(text='goodbye world!'), - ] - ) - - return (da,) - - -@pytest.mark.parametrize('min_freq', [1, 2, 3]) -@pytest.mark.parametrize('da', da_and_dam()) -def test_da_vocabulary(da, min_freq): - vocab = da.get_vocabulary(min_freq) - if min_freq <= 1: - assert set(vocab.values()) == {2, 3, 4} # 0,1 are reserved - assert set(vocab.keys()) == {'hello', 'world', 'goodbye'} - elif min_freq == 2: - assert set(vocab.values()) == {2, 3} # 0,1 are reserved - assert set(vocab.keys()) == {'hello', 'world'} - elif min_freq == 3: - assert not vocab.values() - assert not vocab.keys() - - -@pytest.mark.parametrize('test_docs', da_and_dam()) -def test_da_text_to_blob_non_max_len(test_docs): - vocab = test_docs.get_vocabulary() - for d in test_docs: - d.convert_text_to_blob(vocab) - np.testing.assert_array_equal(test_docs[0].blob, [2]) - np.testing.assert_array_equal(test_docs[1].blob, [2, 3]) - np.testing.assert_array_equal(test_docs[2].blob, [4, 3]) - for d in test_docs: - d.convert_blob_to_text(vocab) - - assert test_docs[0].text == 'hello' - assert test_docs[1].text == 'hello world' - assert test_docs[2].text == 'goodbye world' - - -@pytest.mark.parametrize('test_docs', da_and_dam()) -def test_da_text_to_blob_max_len_3(test_docs): - vocab = test_docs.get_vocabulary() - for d in test_docs: - d.convert_text_to_blob(vocab, max_length=3) - np.testing.assert_array_equal(test_docs[0].blob, [0, 0, 2]) - np.testing.assert_array_equal(test_docs[1].blob, [0, 2, 3]) - np.testing.assert_array_equal(test_docs[2].blob, [0, 4, 3]) - for d in test_docs: - d.convert_blob_to_text(vocab) - - assert test_docs[0].text == 'hello' - assert test_docs[1].text == 'hello world' - assert test_docs[2].text == 'goodbye world' - - -@pytest.mark.parametrize('test_docs', da_and_dam()) -def test_da_text_to_blob_max_len_1(test_docs): - vocab = test_docs.get_vocabulary() - for d in test_docs: - d.convert_text_to_blob(vocab, max_length=1) - np.testing.assert_array_equal(test_docs[0].blob, [2]) - np.testing.assert_array_equal(test_docs[1].blob, [3]) - np.testing.assert_array_equal(test_docs[2].blob, [3]) - for d in test_docs: - d.convert_blob_to_text(vocab) - - assert test_docs[0].text == 'hello' - assert test_docs[1].text == 'world' - assert test_docs[2].text == 'world' - - -@pytest.mark.parametrize('da', da_and_dam()) -def test_convert_text_blob_random_text(da): - texts = ['a short phrase', 'word', 'this is a much longer sentence'] - da.clear() - da.extend(Document(text=t) for t in texts) - vocab = da.get_vocabulary() - - # encoding - for d in da: - d.convert_text_to_blob(vocab, max_length=10) - - # decoding - for d in da: - d.convert_blob_to_text(vocab) - - assert texts - assert da.texts == texts diff --git a/tests/unit/array/mixins/test_traverse.py b/tests/unit/array/mixins/test_traverse.py deleted file mode 100644 index 627b7de1db0..00000000000 --- a/tests/unit/array/mixins/test_traverse.py +++ /dev/null @@ -1,368 +0,0 @@ -import itertools -import types - -import numpy as np -import pytest - -from docarray import Document, DocumentArray -from tests import random_docs - -# some random prime number for sanity check -num_docs = 7 -num_chunks_per_doc = 11 -num_matches_per_doc = 3 -num_matches_per_chunk = 5 - - -@pytest.fixture -def doc_req(): - """Build a dummy request that has docs""" - ds = list(random_docs(num_docs, num_chunks_per_doc)) - # add some random matches - for d in ds: - for _ in range(num_matches_per_doc): - d.matches.append(Document(content='hello')) - for c in d.chunks: - for _ in range(num_matches_per_chunk): - c.matches.append(Document(content='world')) - yield DocumentArray(ds) - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_type(doc_req, filter_fn): - ds = doc_req.traverse('r', filter_fn=filter_fn) - assert isinstance(ds, types.GeneratorType) - assert isinstance(list(ds)[0], DocumentArray) - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_root(doc_req, filter_fn): - ds = list(doc_req.traverse('r', filter_fn=filter_fn)) - assert len(ds) == 1 - assert len(ds[0]) == num_docs - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse('c', filter_fn=filter_fn)) - assert len(ds) == num_docs - assert len(ds[0]) == num_chunks_per_doc - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_root_plus_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse('c,r', filter_fn=filter_fn)) - assert len(ds) == num_docs + 1 - assert len(ds[0]) == num_chunks_per_doc - assert len(ds[-1]) == num_docs - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_chunk_plus_root(doc_req, filter_fn): - ds = list(doc_req.traverse('r,c', filter_fn=filter_fn)) - assert len(ds) == 1 + num_docs - assert len(ds[-1]) == num_chunks_per_doc - assert len(ds[0]) == num_docs - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_match(doc_req, filter_fn): - ds = list(doc_req.traverse('m', filter_fn=filter_fn)) - assert len(ds) == num_docs - assert len(ds[0]) == num_matches_per_doc - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_match_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse('cm', filter_fn=filter_fn)) - assert len(ds) == num_docs * num_chunks_per_doc - assert len(ds[0]) == num_matches_per_chunk - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_root_match_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse('r,c,m,cm', filter_fn=filter_fn)) - assert len(ds) == 1 + num_docs + num_docs + num_docs * num_chunks_per_doc - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flatten_embedding(doc_req, filter_fn): - flattened_results = doc_req.traverse_flat('r,c', filter_fn=filter_fn) - ds = flattened_results.embeddings - assert ds.shape == (num_docs + num_chunks_per_doc * num_docs, 10) - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flatten_root(doc_req, filter_fn): - ds = list(doc_req.traverse_flat('r', filter_fn=filter_fn)) - assert len(ds) == num_docs - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flatten_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse_flat('c', filter_fn=filter_fn)) - assert len(ds) == num_docs * num_chunks_per_doc - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flatten_root_plus_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse_flat('c,r', filter_fn=filter_fn)) - assert len(ds) == num_docs + num_docs * num_chunks_per_doc - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flatten_match(doc_req, filter_fn): - ds = list(doc_req.traverse_flat('m', filter_fn=filter_fn)) - assert len(ds) == num_docs * num_matches_per_doc - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flatten_match_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse_flat('cm', filter_fn=filter_fn)) - assert len(ds) == num_docs * num_chunks_per_doc * num_matches_per_chunk - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flatten_root_match_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse_flat('r,c,m,cm', filter_fn=filter_fn)) - assert ( - len(ds) - == num_docs - + num_chunks_per_doc * num_docs - + num_matches_per_doc * num_docs - + num_docs * num_chunks_per_doc * num_matches_per_chunk - ) - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flattened_per_path_embedding(doc_req, filter_fn): - flattened_results = list(doc_req.traverse_flat_per_path('r,c', filter_fn=filter_fn)) - ds = flattened_results[0].embeddings - assert ds.shape == (num_docs, 10) - - ds = flattened_results[1].embeddings - assert ds.shape == (num_docs * num_chunks_per_doc, 10) - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flattened_per_path_root(doc_req, filter_fn): - ds = list(doc_req.traverse_flat_per_path('r', filter_fn=filter_fn)) - assert len(ds[0]) == num_docs - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flattened_per_path_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse_flat_per_path('c', filter_fn=filter_fn)) - assert len(ds[0]) == num_docs * num_chunks_per_doc - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flattened_per_path_root_plus_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse_flat_per_path('c,r', filter_fn=filter_fn)) - assert len(ds[0]) == num_docs * num_chunks_per_doc - assert len(ds[1]) == num_docs - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flattened_per_path_match(doc_req, filter_fn): - ds = list(doc_req.traverse_flat_per_path('m', filter_fn=filter_fn)) - assert len(ds[0]) == num_docs * num_matches_per_doc - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flattened_per_path_match_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse_flat_per_path('cm', filter_fn=filter_fn)) - assert len(ds[0]) == num_docs * num_chunks_per_doc * num_matches_per_chunk - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_flattened_per_path_root_match_chunk(doc_req, filter_fn): - ds = list(doc_req.traverse_flat_per_path('r,c,m,cm', filter_fn=filter_fn)) - assert len(ds[0]) == num_docs - assert len(ds[1]) == num_chunks_per_doc * num_docs - assert len(ds[2]) == num_matches_per_doc * num_docs - assert len(ds[3]) == num_docs * num_chunks_per_doc * num_matches_per_chunk - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_docuset_traverse_over_iterator_HACKY(filter_fn): - # HACKY USAGE DO NOT RECOMMEND: can also traverse over "runtime"-documentarray - ds = DocumentArray(random_docs(num_docs, num_chunks_per_doc)).traverse( - 'r', filter_fn=filter_fn - ) - assert len(list(list(ds)[0])) == num_docs - - ds = DocumentArray(random_docs(num_docs, num_chunks_per_doc)).traverse( - 'c', filter_fn=filter_fn - ) - ds = list(ds) - assert len(ds) == num_docs - assert len(ds[0]) == num_chunks_per_doc - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_docuset_traverse_over_iterator_CAVEAT(filter_fn): - # HACKY USAGE's CAVEAT: but it can not iterate over an iterator twice - ds = DocumentArray(random_docs(num_docs, num_chunks_per_doc)).traverse( - 'r,c', filter_fn=filter_fn - ) - # note that random_docs is a generator and can be only used once, - # therefore whoever comes first wil get iterated, and then it becomes empty - assert len(list(ds)) == 1 + num_docs - - ds = DocumentArray(random_docs(num_docs, num_chunks_per_doc)).traverse( - 'c,r', filter_fn=filter_fn - ) - assert len(list(ds)) == num_docs + 1 - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_doc_iter_method(filter_fn): - ds = list(random_docs(10)) - - for d in DocumentArray(ds): - assert d.text == 'hello world' - - for d in DocumentArray(ds).traverse_flat('c,r', filter_fn=filter_fn): - d.text = 'modified' - - for d in DocumentArray(ds): - assert d.text == 'modified' - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_matcharray(filter_fn): - doc = Document( - matches=[ - Document(id=f'm{i}', chunks=[Document(id=f'm{i}c{j}') for j in range(3)]) - for i in range(3) - ] - ) - flat_docs = doc.matches.traverse_flat('r,c', filter_fn=filter_fn) - assert isinstance(flat_docs, DocumentArray) - assert len(flat_docs) == 12 - - -@pytest.mark.parametrize('filter_fn', [(lambda d: True), None]) -def test_traverse_chunkarray(filter_fn): - doc = Document( - chunks=[ - Document(id=f'c{i}', matches=[Document(id=f'c{i}m{j}') for j in range(3)]) - for i in range(3) - ] - ) - flat_docs = doc.chunks.traverse_flat('r,m', filter_fn=filter_fn) - assert isinstance(flat_docs, DocumentArray) - assert len(flat_docs) == 12 - - -@pytest.mark.parametrize('use_dam', [True, False]) -@pytest.mark.parametrize( - ('filter_fn', 'docs_len'), - [ - (lambda d: False, 0), - (lambda d: d.text == 'hello', num_docs * num_matches_per_doc), - ( - lambda d: d.text == 'world', - num_docs * num_chunks_per_doc * num_matches_per_chunk, - ), - ( - lambda d: True, - num_docs - + num_docs * num_chunks_per_doc - + num_docs * num_matches_per_doc - + num_docs * num_chunks_per_doc * num_matches_per_chunk, - ), - ( - None, - num_docs - + num_docs * num_matches_per_doc - + num_docs * num_chunks_per_doc - + num_docs * num_chunks_per_doc * num_matches_per_chunk, - ), - ], -) -def test_filter_fn_traverse_flat(filter_fn, docs_len, doc_req, use_dam, tmp_path): - docs = doc_req - ds = list(docs.traverse_flat('r,c,m,cm', filter_fn=filter_fn)) - assert len(ds) == docs_len - assert all(isinstance(d, Document) for d in ds) - - -@pytest.mark.parametrize('use_dam', [True, False]) -@pytest.mark.parametrize( - ('filter_fn', 'docs_len'), - [ - (lambda d: False, [0, 0, 0, 0]), - (lambda d: d.text == 'hello', [0, 0, num_docs * num_matches_per_doc, 0]), - ( - lambda d: d.text == 'world', - [0, 0, 0, num_docs * num_chunks_per_doc * num_matches_per_chunk], - ), - ( - lambda d: True, - [ - num_docs, - num_docs * num_chunks_per_doc, - num_docs * num_matches_per_doc, - num_docs * num_chunks_per_doc * num_matches_per_chunk, - ], - ), - ( - None, - [ - num_docs, - num_docs * num_chunks_per_doc, - num_docs * num_matches_per_doc, - num_docs * num_chunks_per_doc * num_matches_per_chunk, - ], - ), - ], -) -def test_filter_fn_traverse_flat_per_path( - filter_fn, doc_req, docs_len, use_dam, tmp_path -): - docs = doc_req - ds = list(docs.traverse_flat_per_path('r,c,m,cm', filter_fn=filter_fn)) - assert len(ds) == 4 - for seq, length in zip(ds, docs_len): - assert isinstance(seq, DocumentArray) - assert len(list(seq)) == length - - -def test_traversal_path(): - da = DocumentArray([Document() for _ in range(6)]) - assert len(da) == 6 - - da.traverse_flat('r') - - -def test_traverse_flat_root_itself(): - da = DocumentArray([Document() for _ in range(100)]) - res = da.traverse_flat('r') - assert id(res) == id(da) - - -def da_and_dam(N): - da = DocumentArray(random_docs(N)) - return (da,) - - -@pytest.mark.parametrize('da', da_and_dam(100)) -def test_flatten(da): - daf = da.flatten() - assert len(daf) == 600 - assert isinstance(daf, DocumentArray) - assert len(set(d.id for d in daf)) == 600 - - # flattened DA can not be flattened again - daf = daf.flatten() - assert len(daf) == 600 - - -def test_flatten_no_copy(): - da = da_and_dam(100)[0] - daf = da.flatten() - new_text = 'hi i changed it!' - daf[53].text = new_text - assert da[daf[53].id].text == new_text diff --git a/tests/unit/array/test_advance_indexing.py b/tests/unit/array/test_advance_indexing.py deleted file mode 100644 index d87fdd80af9..00000000000 --- a/tests/unit/array/test_advance_indexing.py +++ /dev/null @@ -1,221 +0,0 @@ -import numpy as np -import pytest - -from docarray import DocumentArray, Document - - -@pytest.fixture -def docarray100(): - yield DocumentArray(Document(text=j) for j in range(100)) - - -def test_getter_int_str(docarray100): - # getter - assert docarray100[99].text == 99 - assert docarray100[np.int(99)].text == 99 - assert docarray100[-1].text == 99 - assert docarray100[0].text == 0 - # string index - assert docarray100[docarray100[0].id].text == 0 - assert docarray100[docarray100[99].id].text == 99 - assert docarray100[docarray100[-1].id].text == 99 - - with pytest.raises(IndexError): - docarray100[100] - - with pytest.raises(KeyError): - docarray100['adsad'] - - -def test_setter_int_str(docarray100): - # setter - docarray100[99] = Document(text='hello') - docarray100[0] = Document(text='world') - - assert docarray100[99].text == 'hello' - assert docarray100[-1].text == 'hello' - assert docarray100[0].text == 'world' - - docarray100[docarray100[2].id] = Document(text='doc2') - # string index - assert docarray100[docarray100[2].id].text == 'doc2' - - -def test_del_int_str(docarray100): - zero_id = docarray100[0].id - del docarray100[0] - assert len(docarray100) == 99 - assert zero_id not in docarray100 - - new_zero_id = docarray100[0].id - new_doc_zero = docarray100[0] - del docarray100[new_zero_id] - assert len(docarray100) == 98 - assert zero_id not in docarray100 - assert new_doc_zero not in docarray100 - - -def test_slice(docarray100): - # getter - assert len(docarray100[1:5]) == 4 - assert len(docarray100[1:100:5]) == 20 # 1 to 100, sep with 5 - - # setter - with pytest.raises(TypeError, match='can only assign an iterable'): - docarray100[1:5] = Document(text='repl') - - docarray100[1:5] = [Document(text=f'repl{j}') for j in range(4)] - for d in docarray100[1:5]: - assert d.text.startswith('repl') - assert len(docarray100) == 100 - - # del - zero_doc = docarray100[0] - twenty_doc = docarray100[20] - del docarray100[0:20] - assert len(docarray100) == 80 - assert zero_doc not in docarray100 - assert twenty_doc in docarray100 - - -def test_sequence_bool_index(docarray100): - # getter - mask = [True, False] * 50 - assert len(docarray100[mask]) == 50 - assert len(docarray100[[True, False]]) == 1 - - # setter - mask = [True, False] * 50 - docarray100[mask] = [Document(text=f'repl{j}') for j in range(50)] - - for idx, d in enumerate(docarray100): - if idx % 2 == 0: - # got replaced - assert d.text.startswith('repl') - else: - assert isinstance(d.text, int) - - # del - del docarray100[mask] - assert len(docarray100) == 50 - - del docarray100[mask] - assert len(docarray100) == 25 - - -@pytest.mark.parametrize('nparray', [lambda x: x, np.array, tuple]) -def test_sequence_int(docarray100, nparray): - # getter - idx = nparray([1, 3, 5, 7, -1, -2]) - assert len(docarray100[idx]) == len(idx) - - # setter - docarray100[idx] = [Document(text='repl') for _ in range(len(idx))] - for _id in idx: - assert docarray100[_id].text == 'repl' - - # del - idx = [-3, -4, -5, 9, 10, 11] - del docarray100[idx] - assert len(docarray100) == 100 - len(idx) - - -def test_sequence_str(docarray100): - # getter - idx = [d.id for d in docarray100[1, 3, 5, 7, -1, -2]] - - assert len(docarray100[idx]) == len(idx) - assert len(docarray100[tuple(idx)]) == len(idx) - - # setter - docarray100[idx] = [Document(text='repl') for _ in range(len(idx))] - idx = [d.id for d in docarray100[1, 3, 5, 7, -1, -2]] - for _id in idx: - assert docarray100[_id].text == 'repl' - - # del - idx = [d.id for d in docarray100[-3, -4, -5, 9, 10, 11]] - del docarray100[idx] - assert len(docarray100) == 100 - len(idx) - - -def test_docarray_list_tuple(docarray100): - assert isinstance(docarray100[99, 98], DocumentArray) - assert len(docarray100[99, 98]) == 2 - - -def test_path_syntax_indexing(): - da = DocumentArray().empty(3) - for d in da: - d.chunks = DocumentArray.empty(5) - d.matches = DocumentArray.empty(7) - for c in d.chunks: - c.chunks = DocumentArray.empty(3) - assert len(da['@c']) == 3 * 5 - assert len(da['@c:1']) == 3 - assert len(da['@c-1:']) == 3 - assert len(da['@c1']) == 3 - assert len(da['@c-2:']) == 3 * 2 - assert len(da['@c1:3']) == 3 * 2 - assert len(da['@c1:3c']) == (3 * 2) * 3 - assert len(da['@c1:3,c1:3c']) == (3 * 2) + (3 * 2) * 3 - assert len(da['@c 1:3 , c 1:3 c']) == (3 * 2) + (3 * 2) * 3 - assert len(da['@cc']) == 3 * 5 * 3 - assert len(da['@cc,m']) == 3 * 5 * 3 + 3 * 7 - assert len(da['@r:1cc,m']) == 1 * 5 * 3 + 3 * 7 - - -def test_attribute_indexing(): - da = DocumentArray.empty(10) - for v in da[:, 'id']: - assert v - da[:, 'mime_type'] = [f'type {j}' for j in range(10)] - for v in da[:, 'mime_type']: - assert v - del da[:, 'mime_type'] - for v in da[:, 'mime_type']: - assert not v - - da[:, ['text', 'mime_type']] = [ - [f'hello {j}' for j in range(10)], - [f'type {j}' for j in range(10)], - ] - da.summary() - - for v in da[:, ['mime_type', 'text']]: - for vv in v: - assert vv - - -def test_blob_attribute_selector(): - import scipy.sparse - - sp_embed = np.random.random([3, 10]) - sp_embed[sp_embed > 0.1] = 0 - sp_embed = scipy.sparse.coo_matrix(sp_embed) - - da = DocumentArray.empty(3) - - da[:, 'embedding'] = sp_embed - - assert da[:, 'embedding'].shape == (3, 10) - - for d in da: - assert d.embedding.shape == (1, 10) - - v1, v2 = da[:, ['embedding', 'id']] - assert isinstance(v1, scipy.sparse.coo_matrix) - assert isinstance(v2, list) - - v1, v2 = da[:, ['id', 'embedding']] - assert isinstance(v2, scipy.sparse.coo_matrix) - assert isinstance(v1, list) - - -def test_advance_selector_mixed(): - da = DocumentArray.empty(10) - da.embeddings = np.random.random([10, 3]) - da.match(da, exclude_self=True) - - assert len(da[:, ('id', 'embedding', 'matches')]) == 3 - assert len(da[:, ('id', 'embedding', 'matches')][0]) == 10 diff --git a/tests/unit/array/test_construct.py b/tests/unit/array/test_construct.py deleted file mode 100644 index 302a6359c2c..00000000000 --- a/tests/unit/array/test_construct.py +++ /dev/null @@ -1,62 +0,0 @@ -import pytest - -from docarray import Document, DocumentArray - - -@pytest.mark.parametrize('da_cls', [DocumentArray]) -def test_construct_docarray(da_cls): - da = da_cls() - assert len(da) == 0 - - da = da_cls(Document()) - assert len(da) == 1 - - da = da_cls([Document(), Document()]) - assert len(da) == 2 - - da = da_cls((Document(), Document())) - assert len(da) == 2 - - da = da_cls((Document() for _ in range(10))) - assert len(da) == 10 - - da1 = da_cls(da) - assert len(da1) == 10 - - -@pytest.mark.parametrize('da_cls', [DocumentArray]) -@pytest.mark.parametrize('is_copy', [True, False]) -def test_docarray_copy_singleton(da_cls, is_copy): - d = Document() - da = da_cls(d, copy=is_copy) - d.id = 'hello' - if is_copy: - assert da[0].id != 'hello' - else: - assert da[0].id == 'hello' - - -@pytest.mark.parametrize('da_cls', [DocumentArray]) -@pytest.mark.parametrize('is_copy', [True, False]) -def test_docarray_copy_da(da_cls, is_copy): - d1 = Document() - d2 = Document() - da = da_cls([d1, d2], copy=is_copy) - d1.id = 'hello' - if is_copy: - assert da[0].id != 'hello' - else: - assert da[0].id == 'hello' - - -@pytest.mark.parametrize('da_cls', [DocumentArray]) -@pytest.mark.parametrize('is_copy', [True, False]) -def test_docarray_copy_list(da_cls, is_copy): - d1 = Document() - d2 = Document() - da = da_cls([d1, d2], copy=is_copy) - d1.id = 'hello' - if is_copy: - assert da[0].id != 'hello' - else: - assert da[0].id == 'hello' diff --git a/tests/unit/array/test_from_to_bytes.py b/tests/unit/array/test_from_to_bytes.py deleted file mode 100644 index 2ac85ca9364..00000000000 --- a/tests/unit/array/test_from_to_bytes.py +++ /dev/null @@ -1,89 +0,0 @@ -import numpy as np -import pytest -import tensorflow as tf -import torch -from scipy.sparse import csr_matrix, coo_matrix, bsr_matrix, csc_matrix - -from docarray import DocumentArray -from docarray.math.ndarray import to_numpy_array -from tests import random_docs - - -def get_ndarrays_for_ravel(): - a = np.random.random([100, 3]) - a[a > 0.5] = 0 - return [ - (a, False), - (torch.tensor(a), False), - (tf.constant(a), False), - (torch.tensor(a).to_sparse(), True), - # (tf.sparse.from_dense(a), True), - (csr_matrix(a), True), - (bsr_matrix(a), True), - (coo_matrix(a), True), - (csc_matrix(a), True), - ] - - -@pytest.mark.parametrize('ndarray_val, is_sparse', get_ndarrays_for_ravel()) -@pytest.mark.parametrize('target_da', [DocumentArray.empty(100), random_docs(100)]) -@pytest.mark.parametrize( - 'protocol', ['protobuf', 'protobuf-array', 'pickle', 'pickle-array'] -) -@pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) -def test_to_from_bytes(target_da, protocol, compress, ndarray_val, is_sparse): - bstr = target_da.to_bytes(protocol=protocol, compress=compress) - print(protocol, compress, len(bstr)) - da2 = DocumentArray.from_bytes(bstr, protocol=protocol, compress=compress) - assert len(da2) == len(target_da) - - target_da.embeddings = ndarray_val - target_da.blobs = ndarray_val - bstr = target_da.to_bytes(protocol=protocol, compress=compress) - print(protocol, compress, len(bstr)) - da2 = DocumentArray.from_bytes(bstr, protocol=protocol, compress=compress) - assert len(da2) == len(target_da) - - np.testing.assert_almost_equal( - to_numpy_array(target_da.embeddings), to_numpy_array(da2.embeddings) - ) - np.testing.assert_almost_equal( - to_numpy_array(target_da.blobs), to_numpy_array(da2.blobs) - ) - - -@pytest.mark.parametrize('target_da', [DocumentArray.empty(100), random_docs(100)]) -@pytest.mark.parametrize( - 'protocol', ['protobuf', 'protobuf-array', 'pickle', 'pickle-array'] -) -@pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) -def test_save_bytes(target_da, protocol, compress, tmpfile): - target_da.save_binary(tmpfile, protocol=protocol, compress=compress) - target_da.save_binary(str(tmpfile), protocol=protocol, compress=compress) - - with open(tmpfile, 'wb') as fp: - target_da.save_binary(fp, protocol=protocol, compress=compress) - - DocumentArray.load_binary(tmpfile, protocol=protocol, compress=compress) - DocumentArray.load_binary(str(tmpfile), protocol=protocol, compress=compress) - with open(tmpfile, 'rb') as fp: - DocumentArray.load_binary(fp, protocol=protocol, compress=compress) - - -@pytest.mark.parametrize('target_da', [DocumentArray.empty(100), random_docs(100)]) -def test_from_to_protobuf(target_da): - DocumentArray.from_protobuf(target_da.to_protobuf()) - - -@pytest.mark.parametrize('target_da', [DocumentArray.empty(100), random_docs(100)]) -def test_from_to_safe_list(target_da): - DocumentArray.from_list(target_da.to_list()) - - -@pytest.mark.parametrize('protocol', ['protobuf', 'pickle']) -@pytest.mark.parametrize('show_progress', [True, False]) -def test_push_pull_show_progress(show_progress, protocol): - da = DocumentArray.empty(1000) - r = da.to_bytes(_show_progress=show_progress, protocol=protocol) - da_r = DocumentArray.from_bytes(r, _show_progress=show_progress, protocol=protocol) - assert da == da_r diff --git a/tests/unit/array/test_ravel_unravel.py b/tests/unit/array/test_ravel_unravel.py deleted file mode 100644 index 6ba872294f7..00000000000 --- a/tests/unit/array/test_ravel_unravel.py +++ /dev/null @@ -1,96 +0,0 @@ -import numpy as np -import paddle -import pytest -import tensorflow as tf - -from docarray.math.ndarray import to_numpy_array -import torch -from scipy.sparse import csr_matrix, coo_matrix, bsr_matrix, csc_matrix - -from docarray import DocumentArray, Document - - -def get_ndarrays_for_ravel(): - a = np.random.random([10, 3]) - a[a > 0.5] = 0 - return [ - (a, False), - (torch.tensor(a), False), - (tf.constant(a), False), - (paddle.to_tensor(a), False), - (torch.tensor(a).to_sparse(), True), - # (tf.sparse.from_dense(a), True), - (csr_matrix(a), True), - (bsr_matrix(a), True), - (coo_matrix(a), True), - (csc_matrix(a), True), - ] - - -@pytest.mark.parametrize('ndarray_val, is_sparse', get_ndarrays_for_ravel()) -@pytest.mark.parametrize('attr', ['embeddings', 'blobs']) -def test_ravel_embeddings_blobs(ndarray_val, attr, is_sparse): - da = DocumentArray.empty(10) - setattr(da, attr, ndarray_val) - ndav = getattr(da, attr) - - # test read/getter - assert type(ndav) is type(ndarray_val) - - if is_sparse: - if hasattr(ndav, 'todense'): - ndav = (ndav.todense(),) - ndarray_val = ndarray_val.todense() - if hasattr(ndav, 'to_dense'): - ndav = (ndav.to_dense(),) - ndarray_val = ndarray_val.to_dense() - if isinstance(ndav, tf.SparseTensor): - ndav = tf.sparse.to_dense(ndav) - ndarray_val = tf.sparse.to_dense(ndarray_val) - - if isinstance(ndav, tuple): - ndav = ndav[0] - if hasattr(ndav, 'numpy'): - ndav = ndav.numpy() - ndarray_val = ndarray_val.numpy() - - np.testing.assert_almost_equal(ndav, ndarray_val) - - -@pytest.mark.parametrize('sparse_cls', [csr_matrix, csc_matrix, bsr_matrix, coo_matrix]) -def test_bsr_coo_unravel(sparse_cls): - a = np.random.random([10, 72]) - a[a > 0.5] = 0 - - da = DocumentArray.empty(10) - for d, a_row in zip(da, a): - d.embedding = sparse_cls(a_row) - - np.testing.assert_almost_equal(a, da.embeddings.todense()) - - -def get_ndarrays(): - a = np.random.random([10, 3]) - a[a > 0.5] = 0 - return [ - (a, False), - (torch.tensor(a), False), - (tf.constant(a), False), - (paddle.to_tensor(a), False), - (torch.tensor(a).to_sparse(), True), - (tf.sparse.from_dense(a), True), - (csr_matrix(a), True), - (bsr_matrix(a), True), - (coo_matrix(a), True), - (csc_matrix(a), True), - ] - - -@pytest.mark.parametrize('ndarray_val, is_sparse', get_ndarrays()) -@pytest.mark.parametrize('attr', ['embedding', 'blob']) -def test_ndarray_force_numpy(ndarray_val, attr, is_sparse): - d = Document() - setattr(d, attr, ndarray_val) - ndav = to_numpy_array(getattr(d, attr)) - assert isinstance(ndav, np.ndarray) - assert ndav.shape == (10, 3) diff --git a/tests/unit/array/test_sequence.py b/tests/unit/array/test_sequence.py deleted file mode 100644 index 6ba6679f936..00000000000 --- a/tests/unit/array/test_sequence.py +++ /dev/null @@ -1,24 +0,0 @@ -import pytest - -from docarray import Document, DocumentArray - - -@pytest.mark.parametrize('da_cls', [DocumentArray]) -def test_insert(da_cls): - da = da_cls() - assert not len(da) - da.insert(0, Document(text='hello')) - da.insert(0, Document(text='world')) - assert len(da) == 2 - assert da[0].text == 'world' - assert da[1].text == 'hello' - - -@pytest.mark.parametrize('da_cls', [DocumentArray]) -def test_append_extend(da_cls): - da = da_cls() - da.append(Document()) - da.append(Document()) - assert len(da) == 2 - da.extend([Document(), Document()]) - assert len(da) == 4 diff --git a/tests/unit/document/__init__.py b/tests/unit/document/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/unit/document/test_converters.py b/tests/unit/document/test_converters.py deleted file mode 100644 index 37923f84e1e..00000000000 --- a/tests/unit/document/test_converters.py +++ /dev/null @@ -1,236 +0,0 @@ -import os - -import numpy as np -import pytest - -from docarray import Document -from docarray.document.generators import from_files -from docarray.helper import __windows__ - -cur_dir = os.path.dirname(os.path.abspath(__file__)) - - -def test_video_convert_pipe(pytestconfig, tmpdir): - num_d = 0 - fname = str(tmpdir / f'tmp{num_d}.mp4') - d = Document(uri=os.path.join(cur_dir, 'toydata/mov_bbb.mp4')) - d.load_uri_to_video_blob() - d.save_video_blob_to_file(fname) - assert os.path.exists(fname) - - -def test_audio_convert_pipe(pytestconfig, tmpdir): - num_d = 0 - for d in from_files(f'{cur_dir}/toydata/*.wav'): - fname = str(tmpdir / f'tmp{num_d}.wav') - d.load_uri_to_audio_blob() - d.blob = d.blob[::-1] - d.save_audio_blob_to_file(fname) - assert os.path.exists(fname) - num_d += 1 - assert num_d - - -def test_image_convert_pipe(pytestconfig): - for d in from_files(f'{pytestconfig.rootdir}/.github/**/*.png'): - ( - d.load_uri_to_image_blob() - .convert_uri_to_datauri() - .set_image_blob_shape((64, 64)) - .set_image_blob_normalization() - .set_image_blob_channel_axis(-1, 0) - ) - assert d.blob.shape == (3, 64, 64) - assert d.uri - - -def test_uri_to_blob(): - doc = Document(uri=os.path.join(cur_dir, 'toydata/test.png')) - doc.load_uri_to_image_blob() - assert isinstance(doc.blob, np.ndarray) - assert doc.blob.shape == (85, 152, 3) # h,w,c - assert doc.mime_type == 'image/png' - - -def test_datauri_to_blob(): - doc = Document(uri=os.path.join(cur_dir, 'toydata/test.png')) - doc.convert_uri_to_datauri() - assert not doc.blob - assert doc.mime_type == 'image/png' - - -def test_buffer_to_blob(): - doc = Document(uri=os.path.join(cur_dir, 'toydata/test.png')) - doc.load_uri_to_buffer() - doc.convert_buffer_to_image_blob() - assert isinstance(doc.blob, np.ndarray) - assert doc.mime_type == 'image/png' - assert doc.blob.shape == (85, 152, 3) # h,w,c - - -def test_convert_buffer_to_blob(): - rand_state = np.random.RandomState(0) - array = rand_state.random([10, 10]) - doc = Document(content=array.tobytes()) - assert doc.content_type == 'buffer' - intialiazed_buffer = doc.buffer - - doc.convert_buffer_to_blob() - assert doc.content_type == 'blob' - converted_buffer_in_one_of = doc.buffer - assert intialiazed_buffer != converted_buffer_in_one_of - np.testing.assert_almost_equal(doc.content.reshape([10, 10]), array) - - -@pytest.mark.parametrize('shape, channel_axis', [((3, 32, 32), 0), ((32, 32, 3), -1)]) -def test_image_normalize(shape, channel_axis): - doc = Document(content=np.random.randint(0, 255, shape, dtype=np.uint8)) - doc.set_image_blob_normalization(channel_axis=channel_axis) - assert doc.blob.ndim == 3 - assert doc.blob.shape == shape - assert doc.blob.dtype == np.float32 - - -@pytest.mark.parametrize( - 'arr_size, channel_axis, height, width', - [ - ([32, 28, 3], -1, 32, 28), # h, w, c (rgb) - ([3, 32, 28], 0, 32, 28), # c, h, w (rgb) - ([1, 32, 28], 0, 32, 28), # c, h, w, (greyscale) - ([32, 28, 1], -1, 32, 28), # h, w, c, (greyscale) - ], -) -def test_convert_image_blob_to_uri(arr_size, channel_axis, width, height): - doc = Document(content=np.random.randint(0, 255, arr_size)) - assert doc.blob.any() - assert not doc.uri - doc.set_image_blob_shape(channel_axis=channel_axis, shape=(width, height)) - - doc.convert_image_blob_to_uri(channel_axis=channel_axis) - assert doc.uri.startswith('data:image/png;base64,') - assert doc.mime_type == 'image/png' - assert doc.blob.any() # assure after conversion blob still exist. - - -@pytest.mark.xfail( - condition=__windows__, reason='x-python is not detected on windows CI' -) -@pytest.mark.parametrize( - 'uri, mimetype', - [ - (__file__, 'text/x-python'), - ('http://google.com/index.html', 'text/html'), - ('https://google.com/index.html', 'text/html'), - ], -) -def test_convert_uri_to_buffer(uri, mimetype): - d = Document(uri=uri) - assert not d.buffer - d.load_uri_to_buffer() - assert d.buffer - assert d.mime_type == mimetype - - -@pytest.mark.parametrize( - 'converter', ['convert_buffer_to_datauri', 'convert_content_to_datauri'] -) -def test_convert_buffer_to_uri(converter): - d = Document(content=open(__file__).read().encode(), mime_type='text/x-python') - assert d.buffer - getattr(d, converter)() - assert d.uri.startswith('data:text/x-python;') - - -@pytest.mark.parametrize( - 'converter', ['convert_text_to_datauri', 'convert_content_to_datauri'] -) -def test_convert_text_to_uri(converter): - d = Document(content=open(__file__).read()) - assert d.text - getattr(d, converter)() - assert d.uri.startswith('data:text/plain;') - - -@pytest.mark.xfail( - condition=__windows__, reason='x-python is not detected on windows CI' -) -@pytest.mark.parametrize( - 'uri, mimetype', - [ - pytest.param( - __file__, - 'text/x-python', - marks=pytest.mark.xfail( - condition=__windows__, reason='x-python is not detected on windows CI' - ), - ), - ('http://google.com/index.html', 'text/html'), - ('https://google.com/index.html', 'text/html'), - ], -) -def test_convert_uri_to_text(uri, mimetype): - doc = Document(uri=uri, mime_type=mimetype) - doc.load_uri_to_text() - if mimetype == 'text/html': - assert '' in doc.text - elif mimetype == 'text/x-python': - text_from_file = open(__file__).read() - assert doc.text == text_from_file - - -def test_convert_text_to_uri_and_back(): - text_from_file = open(__file__).read() - doc = Document(content=text_from_file, mime_type='text/x-python') - assert doc.text - assert doc.mime_type == 'text/x-python' - doc.convert_text_to_datauri() - doc.load_uri_to_text() - assert doc.mime_type == 'text/plain' - assert doc.text == text_from_file - - -def test_convert_text_diff_encoding(tmpfile): - otext = 'testä' - text = otext.encode('iso8859') - with open(tmpfile, 'wb') as fp: - fp.write(text) - with pytest.raises(UnicodeDecodeError): - d = Document(uri=str(tmpfile)).load_uri_to_text() - - d = Document(uri=str(tmpfile)).load_uri_to_text(charset='iso8859') - assert d.text == otext - - with open(tmpfile, 'w', encoding='iso8859') as fp: - fp.write(otext) - with pytest.raises(UnicodeDecodeError): - d = Document(uri=str(tmpfile)).load_uri_to_text() - - d = Document(uri=str(tmpfile)).load_uri_to_text(charset='iso8859') - assert d.text == otext - - -def test_convert_content_to_uri(): - d = Document(content=np.random.random([10, 10])) - with pytest.raises(NotImplementedError): - d.convert_content_to_datauri() - - -@pytest.mark.parametrize( - 'uri, mimetype', - [ - (__file__, 'text/x-python'), - ('http://google.com/index.html', 'text/html'), - ('https://google.com/index.html', 'text/html'), - ], -) -def test_convert_uri_to_data_uri(uri, mimetype): - doc = Document(uri=uri, mime_type=mimetype) - doc.convert_uri_to_datauri() - assert doc.uri.startswith(f'data:{mimetype}') - assert doc.mime_type == mimetype - - -def test_glb_converters(): - doc = Document(uri=os.path.join(cur_dir, 'toydata/test.glb')) - doc.load_uri_to_point_cloud_blob(2000) - assert doc.blob.shape == (2000, 3) diff --git a/tests/unit/document/test_docdata.py b/tests/unit/document/test_docdata.py deleted file mode 100644 index ff39b24bb68..00000000000 --- a/tests/unit/document/test_docdata.py +++ /dev/null @@ -1,263 +0,0 @@ -import numpy as np -import pytest - -from docarray import Document, DocumentArray -from docarray.array.chunk import ChunkArray -from docarray.array.match import MatchArray -from docarray.score import NamedScore - - -@pytest.mark.parametrize('init_args', [None, dict(id=123), Document()]) -@pytest.mark.parametrize('copy', [True, False]) -def test_construct_doc(init_args, copy): - Document(init_args, copy) - - -def test_doc_hash_identical(): - d1 = Document(text='hello') - d2 = Document(text='hello') - assert hash(d1) != hash(d2) - assert d1 != d2 - d1.id = d2.id - assert hash(d1) == hash(d2) - assert d1 == d2 - - -def test_doc_hash_complicate_content(): - d1 = Document(text='hello', embedding=np.array([1, 2, 3]), id=1) - d2 = Document(text='hello', embedding=np.array([1, 2, 3]), id=1) - assert d1 == d2 - assert hash(d1) == hash(d2) - - -def test_pop_field(): - d1 = Document(text='hello', embedding=np.array([1, 2, 3]), id=1) - assert d1.non_empty_fields == ('id', 'mime_type', 'text', 'embedding') - d1.pop('text') - assert d1.non_empty_fields == ('id', 'mime_type', 'embedding') - d1.pop('id', 'embedding', 'mime_type') - assert d1.non_empty_fields == tuple() - - d1.pop('foobar') - with pytest.raises(AttributeError): - assert d1.foobar - - -def test_clear_fields(): - d1 = Document(text='hello', embedding=np.array([1, 2, 3]), id=1) - d1.clear() - assert d1.non_empty_fields == tuple() - - -def test_exclusive_content(): - d = Document(text='hello') - assert d.content_type == 'text' - d.buffer = b'123' - assert d.buffer - assert not d.text - assert not d.blob - assert d.content_type == 'buffer' - d.blob = [1, 2, 3] - assert d.blob - assert not d.buffer - assert not d.text - assert d.content_type == 'blob' - d.text = 'hello' - assert d.text - assert not d.buffer - assert not d.blob - assert d.content_type == 'text' - - -def test_content_setter(): - d = Document() - assert not d.content_type - d.content = 'hello' - assert d.content_type == 'text' - d.content = None - assert not d.content_type - - -def test_chunks_matches_setter(): - d = Document(chunks=[Document()], matches=[Document(), Document()]) - assert len(d.chunks) == 1 - assert len(d.matches) == 2 - assert isinstance(d.chunks, DocumentArray) - assert isinstance(d.chunks, ChunkArray) - assert isinstance(d.matches, DocumentArray) - assert isinstance(d.matches, MatchArray) - - -def test_empty_doc_chunks_matches(): - assert isinstance(Document().chunks, DocumentArray) - assert isinstance(Document().matches, DocumentArray) - assert isinstance(Document().matches, MatchArray) - assert isinstance(Document().chunks, ChunkArray) - - d = Document() - d.chunks.append(Document()) - assert isinstance(d.chunks, ChunkArray) - - d.chunks = [Document(), Document()] - assert isinstance(d.chunks, ChunkArray) - - -def test_chunk_match_increase_granularity(): - d = Document() - d.chunks.append(Document()) - assert d.chunks[0].granularity == 1 - assert id(d.chunks.reference_doc) == id(d) - d.matches.append(Document()) - assert d.matches[0].adjacency == 1 - assert id(d.matches.reference_doc) == id(d) - - d = d.chunks[0] - d.chunks.append(Document()) - assert d.chunks[0].granularity == 2 - assert id(d.chunks.reference_doc) == id(d) - - d.matches.append(Document()) - assert d.matches[0].adjacency == 1 - assert id(d.matches.reference_doc) == id(d) - - -def test_offset(): - d1 = Document(offset=1.0) - d2 = Document() - d2.offset = 1.0 - assert d1.offset == d2.offset == 1.0 - - -def test_exclusive_content_2(): - d = Document(text='hello', buffer=b'sda') - assert len(d.non_empty_fields) == 3 - d.content = b'sda' - assert d.content == b'sda' - assert 'buffer' in d.non_empty_fields - d = Document(content='hello') - assert d.content_type == 'text' - d = Document(content=b'hello') - assert d.content_type == 'buffer' - d = Document(content=[1, 2, 3]) - assert d.content_type == 'blob' - - -def test_get_attr_values(): - d = Document( - **{ - 'id': '123', - 'text': 'document', - 'feature1': 121, - 'name': 'name', - 'tags': {'id': 'identity', 'a': 'b', 'c': 'd', 'e': [0, 1, {'f': 'g'}]}, - } - ) - d.scores['metric'] = NamedScore(value=42) - - required_keys = [ - 'id', - 'text', - 'tags__name', - 'tags__feature1', - 'scores__metric__value', - 'tags__c', - 'tags__id', - 'tags__e__2__f', - ] - res = d._get_attributes(*required_keys) - assert len(res) == len(required_keys) - assert res[required_keys.index('id')] == '123' - assert res[required_keys.index('tags__feature1')] == 121 - assert res[required_keys.index('tags__name')] == 'name' - assert res[required_keys.index('text')] == 'document' - assert res[required_keys.index('tags__c')] == 'd' - assert res[required_keys.index('tags__id')] == 'identity' - assert res[required_keys.index('scores__metric__value')] == 42 - assert res[required_keys.index('tags__e__2__f')] == 'g' - - required_keys_2 = ['tags', 'text'] - res2 = d._get_attributes(*required_keys_2) - assert len(res2) == 2 - assert res2[required_keys_2.index('text')] == 'document' - assert res2[required_keys_2.index('tags')] == d.tags - - d = Document({'id': '123', 'tags': {'outterkey': {'innerkey': 'real_value'}}}) - required_keys_3 = ['tags__outterkey__innerkey'] - res3 = d._get_attributes(*required_keys_3) - assert res3 == 'real_value' - - d = Document(content=np.array([1, 2, 3])) - res4 = np.stack(d._get_attributes(*['blob'])) - np.testing.assert_equal(res4, np.array([1, 2, 3])) - - -def test_set_get_mime(): - a = Document() - a.mime_type = 'jpg' - assert a.mime_type == 'image/jpeg' - b = Document() - b.mime_type = 'jpeg' - assert b.mime_type == 'image/jpeg' - c = Document() - c.mime_type = '.jpg' - assert c.mime_type == 'image/jpeg' - - -def test_doc_content(): - d = Document() - assert d.content is None - d.text = 'abc' - assert d.content == 'abc' - c = np.random.random([10, 10]) - d.blob = c - np.testing.assert_equal(d.content, c) - d.buffer = b'123' - assert d.buffer == b'123' - - -def test_dict_constructor(): - - d1 = Document( - uri='https://jina.ai', mime_type='text/plain', granularity=1, adjacency=3 - ) - - d2 = Document( - dict(uri='https://jina.ai', mime_type='text/plain', granularity=1, adjacency=3) - ) - - d3 = Document( - { - 'uri': 'https://jina.ai', - 'mime_type': 'text/plain', - 'granularity': 1, - 'adjacency': 3, - } - ) - - assert d1 != d2 - d1.id = None - d2.id = None - d3.id = None - assert d1 == d2 == d3 - - -def test_unknown_fields_behavior(): - d = Document(hello='world') - assert d.tags == {'hello': 'world'} - - d = Document(hello='world', unknown_fields_handler='drop') - assert d.tags == {} - - with pytest.raises(AttributeError): - d = Document(hello='world', unknown_fields_handler='raise') - - -def test_content_setter_as_proxy(): - d = Document(content='hello') - assert d.content == 'hello' - - assert 'content' not in d.non_empty_fields - assert 'text' in d.non_empty_fields - d.content = [1, 2, 3] - assert 'blob' in d.non_empty_fields - assert 'text' not in d.non_empty_fields diff --git a/tests/unit/document/test_feature_hashing.py b/tests/unit/document/test_feature_hashing.py deleted file mode 100644 index edd81932709..00000000000 --- a/tests/unit/document/test_feature_hashing.py +++ /dev/null @@ -1,20 +0,0 @@ -import pytest - -from docarray import DocumentArray -from docarray.math.ndarray import to_numpy_array - - -@pytest.mark.parametrize('n_dim', [2, 4, 100]) -@pytest.mark.parametrize('sparse', [True, False]) -@pytest.mark.parametrize('metric', ['jaccard', 'cosine']) -def test_feature_hashing(n_dim, sparse, metric): - da = DocumentArray.empty(3) - da.texts = ['hello world', 'world, bye', 'hello bye'] - da.apply(lambda d: d.embed_feature_hashing(n_dim=n_dim, sparse=sparse)) - assert da.embeddings.shape == (3, n_dim) - da.embeddings = to_numpy_array(da.embeddings) - da.match(da, metric=metric, use_scipy=True) - result = da['@m', ('id', f'scores__{metric}__value')] - assert len(result) == 2 - assert result[1][0] == 0.0 - assert result[1][1] > 0.0 diff --git a/tests/unit/document/test_io.py b/tests/unit/document/test_io.py deleted file mode 100644 index 4f226d6b954..00000000000 --- a/tests/unit/document/test_io.py +++ /dev/null @@ -1,306 +0,0 @@ -import os -import shutil - -import numpy as np -import pytest - -from docarray import Document, DocumentArray -from docarray.document.generators import ( - from_files, - from_ndarray, - from_lines, - from_csv, - from_huggingface_datasets, -) - -cur_dir = os.path.dirname(os.path.abspath(__file__)) - - -@pytest.fixture(scope='function') -def filepath(tmpdir): - input_filepath = os.path.join(tmpdir, 'input_file.csv') - with open(input_filepath, 'w') as input_file: - input_file.writelines(["1\n", "2\n", "3\n"]) - return input_filepath - - -@pytest.fixture(scope='function') -def dataset_configs(): - config = { - 'adversarial': { - 'dataset_path': 'adversarial_qa', - 'name': 'adversarialQA', - 'split': 'test', - }, - 'tweet_eval': { - 'dataset_path': 'tweet_eval', - 'name': 'emoji', - 'split': 'train+test', - }, - } - return config - - -def test_input_lines_with_filepath(filepath): - result = list(from_lines(filepath=filepath, size=2)) - assert len(result) == 2 - assert isinstance(result[0], Document) - - -def test_input_csv_from_file(): - with open(os.path.join(cur_dir, 'toydata/docs.csv')) as fp: - result = list(from_csv(fp)) - assert len(result) == 2 - assert isinstance(result[0], Document) - assert result[0].tags['source'] == 'testsrc' - - -def test_input_csv_from_lines(): - with open(os.path.join(cur_dir, 'toydata/docs.csv')) as fp: - result = list(from_lines(fp, line_format='csv')) - assert len(result) == 2 - assert isinstance(result[0], Document) - assert result[0].tags['source'] == 'testsrc' - - -def test_input_csv_from_lines_field_resolver(): - with open(os.path.join(cur_dir, 'toydata/docs.csv')) as fp: - result = list( - from_lines(fp, line_format='csv', field_resolver={'question': 'text'}) - ) - assert len(result) == 2 - assert isinstance(result[0], Document) - assert result[0].tags['source'] == 'testsrc' - assert not result[0].uri - assert result[0].text - - -@pytest.mark.parametrize( - 'da', - [ - DocumentArray, - ], -) -def test_input_csv_from_strings(da): - result = da.from_csv(os.path.join(cur_dir, 'toydata/docs.csv')) - assert len(result) == 2 - assert isinstance(result[0], Document) - assert result[0].tags['source'] == 'testsrc' - - -def test_input_lines_with_empty_filepath_and_lines(): - with pytest.raises(ValueError): - lines = from_lines(lines=None, filepath=None) - for _ in lines: - pass - - -def test_input_lines_with_jsonlines_docs(): - result = list(from_lines(filepath=os.path.join(cur_dir, 'toydata/docs.jsonlines'))) - assert len(result) == 2 - assert result[0].text == "a" - assert result[1].text == "b" - - -@pytest.mark.parametrize( - 'size, sampling_rate', - [ - (None, None), - (1, None), - (None, 0.5), - ], -) -@pytest.mark.parametrize( - 'da', - [ - DocumentArray, - ], -) -def test_input_lines_with_jsonlines_file(size, sampling_rate, da): - result = da.from_lines( - filepath=os.path.join(cur_dir, 'toydata/docs.jsonlines'), - size=size, - sampling_rate=sampling_rate, - ) - - assert len(result) == size if size is not None else 2 - if sampling_rate is None: - assert result[0].text == "a" - if size is None: - assert result[1].text == "b" - - -@pytest.mark.parametrize( - 'size, sampling_rate', - [ - (None, None), - (1, None), - (None, 0.5), - ], -) -def test_input_lines_with_jsonslines(size, sampling_rate): - with open(os.path.join(cur_dir, 'toydata/docs.jsonlines')) as fp: - lines = fp.readlines() - result = list( - from_lines( - lines=lines, line_format='json', size=size, sampling_rate=sampling_rate - ) - ) - assert len(result) == size if size is not None else 2 - if sampling_rate is None: - assert result[0].text == "a" - if size is None: - assert result[1].text == "b" - - -def test_input_lines_with_jsonlines_docs_groundtruth(): - result = list( - from_lines(filepath=os.path.join(cur_dir, 'toydata/docs_groundtruth.jsonlines')) - ) - assert len(result) == 2 - assert result[0][0].text == "a" - assert result[0][1].text == "b" - assert result[1][0].text == "c" - assert result[1][1].text == "d" - - -@pytest.mark.parametrize( - 'size, sampling_rate', - [ - (None, None), - (1, None), - (None, 0.5), - ], -) -def test_input_huggingface_datasets_from_path(dataset_configs, size, sampling_rate): - result = list( - from_huggingface_datasets( - dataset_configs['adversarial']['dataset_path'], - size=size, - name=dataset_configs['adversarial']['name'], - sampling_rate=sampling_rate, - split=dataset_configs['adversarial']['split'], - ) - ) - - if size is not None: - assert len(result) == size - - assert isinstance(result[0], Document) - - -def test_input_huggingface_datasets_with_tweet_dataset(dataset_configs): - result = list( - from_huggingface_datasets( - dataset_configs['tweet_eval']['dataset_path'], - name=dataset_configs['tweet_eval']['name'], - split=dataset_configs['tweet_eval']['split'], - ) - ) - assert isinstance(result[0], Document) - assert result[0].text - - -@pytest.mark.parametrize( - 'da', - [ - DocumentArray, - ], -) -def test_input_huggingface_datasets_from_csv_file(dataset_configs, da): - field_resolver = {'question': 'text'} - result = da.from_huggingface_datasets( - 'csv', - field_resolver=field_resolver, - data_files=os.path.join(cur_dir, 'toydata/docs.csv'), - split='train', - ) - - assert len(result) == 2 - assert isinstance(result[0], Document) - assert result[0].text == 'What are the symptoms?' - assert result[0].tags['source'] == 'testsrc' - - -@pytest.mark.parametrize( - 'da', - [ - DocumentArray, - ], -) -def test_input_huggingface_datasets_with_field_resolver(dataset_configs, da): - field_resolver = {'question': 'text'} - result = da.from_huggingface_datasets( - dataset_configs['adversarial']['dataset_path'], - field_resolver=field_resolver, - name=dataset_configs['adversarial']['name'], - split=dataset_configs['adversarial']['split'], - ) - - assert isinstance(result[0], Document) - assert result[0].text - assert 'title' in result[0].tags - - -def test_input_huggingface_datasets_with_filter_fields(dataset_configs): - field_resolver = {'question': 'text'} - result = list( - from_huggingface_datasets( - dataset_configs['adversarial']['dataset_path'], - field_resolver=field_resolver, - filter_fields=True, - name=dataset_configs['adversarial']['name'], - split=dataset_configs['adversarial']['split'], - ) - ) - assert isinstance(result[0], Document) - assert result[0].text - assert not 'title' in result[0].tags - - -def test_input_huggingface_datasets_with_no_split(dataset_configs): - with pytest.raises(ValueError): - result = from_huggingface_datasets( - dataset_configs['adversarial']['dataset_path'], - name=dataset_configs['adversarial']['name'], - ) - for _ in result: - pass - - -def test_input_huggingface_datasets_with_filter_fields_and_no_resolver(dataset_configs): - with pytest.raises(ValueError): - result = from_huggingface_datasets( - dataset_configs['adversarial']['dataset_path'], - name=dataset_configs['adversarial']['name'], - filter_fields=True, - ) - for _ in result: - pass - - -@pytest.mark.parametrize( - 'patterns, recursive, size, sampling_rate, read_mode', - [ - ('*.*', True, None, None, None), - ('*.*', False, None, None, None), - ('*.*', True, 2, None, None), - ('*.*', True, 2, None, 'rb'), - ('*.*', True, None, 0.5, None), - ], -) -def test_input_files(patterns, recursive, size, sampling_rate, read_mode): - DocumentArray( - from_files( - patterns=patterns, - recursive=recursive, - size=size, - sampling_rate=sampling_rate, - read_mode=read_mode, - ) - ) - - -def test_from_files_with_uri(): - for d in from_files(patterns='*.*', to_dataturi=True, size=10): - assert d.uri.startswith('data:') diff --git a/tests/unit/document/test_namedscore.py b/tests/unit/document/test_namedscore.py deleted file mode 100644 index c128a2d014a..00000000000 --- a/tests/unit/document/test_namedscore.py +++ /dev/null @@ -1,11 +0,0 @@ -import pytest - -from docarray.score import NamedScore - - -@pytest.mark.parametrize( - 'init_args', [None, dict(value=123, description='hello'), NamedScore()] -) -@pytest.mark.parametrize('copy', [True, False]) -def test_construct_ns(init_args, copy): - NamedScore(init_args, copy) diff --git a/tests/unit/document/test_ndarray.py b/tests/unit/document/test_ndarray.py deleted file mode 100644 index 91bb76d1410..00000000000 --- a/tests/unit/document/test_ndarray.py +++ /dev/null @@ -1,35 +0,0 @@ -import numpy as np -import paddle -import pytest -import tensorflow as tf -import torch -from scipy.sparse import csr_matrix, coo_matrix, bsr_matrix, csc_matrix - -from docarray import Document - - -def get_ndarrays(): - a = np.random.random([10, 3]) - a[a > 0.5] = 0 - return [ - (a, False), - (torch.tensor(a), False), - (tf.constant(a), False), - (paddle.to_tensor(a), False), - (torch.tensor(a).to_sparse(), True), - (tf.sparse.from_dense(a), True), - (csr_matrix(a), True), - (bsr_matrix(a), True), - (coo_matrix(a), True), - (csc_matrix(a), True), - ] - - -@pytest.mark.parametrize('ndarray_val, is_sparse', get_ndarrays()) -@pytest.mark.parametrize('attr', ['embedding', 'blob']) -def test_ndarray_force_numpy(ndarray_val, attr, is_sparse): - d = Document() - setattr(d, attr, ndarray_val) - assert type(getattr(Document.from_protobuf(d.to_protobuf()), attr)) is type( - ndarray_val - ) diff --git a/tests/unit/document/test_pickle.py b/tests/unit/document/test_pickle.py deleted file mode 100644 index 10df0ad29e8..00000000000 --- a/tests/unit/document/test_pickle.py +++ /dev/null @@ -1,35 +0,0 @@ -import pickle - -import pytest - -from docarray import Document -from docarray.document.data import DocumentData -from docarray.base import BaseDCType -from tests import random_docs - - -@pytest.mark.parametrize('cls', BaseDCType.__subclasses__()) -def test_pickle_dump_load(cls): - r = pickle.loads(pickle.dumps(cls())) - isinstance(r, cls) - - -def test_pickle_dump_load_real_doc(): - for d in random_docs(10): - dr = pickle.loads(pickle.dumps(d)) - assert dr == d - assert dr.embedding is not None - assert len(dr.chunks) == len(d.chunks) - - -def test_pickle_rely_on_data_class_and_document_class(): - # TODO (Han): This is not really a designed behavior, but atm I see no harm - # of having it, and no real usecases that against it. - - d = Document() - d.id = 'hello' - setattr(d, 'foo', 'bar') - assert getattr(d, 'foo') == 'bar' - r_d = Document.from_bytes(d.to_bytes(protocol='pickle')) - assert r_d.id == d.id - assert getattr(r_d, 'foo') == 'bar' diff --git a/tests/unit/document/test_porting.py b/tests/unit/document/test_porting.py deleted file mode 100644 index e90b927fc67..00000000000 --- a/tests/unit/document/test_porting.py +++ /dev/null @@ -1,22 +0,0 @@ -import pytest - -from docarray import Document, DocumentArray -from tests import random_docs - - -@pytest.mark.parametrize('protocol', ['protobuf', 'pickle']) -@pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) -def test_to_from_bytes(protocol, compress): - d = Document(embedding=[1, 2, 3, 4, 5], text='hello') - bstr = d.to_bytes(protocol=protocol, compress=compress) - print(protocol, compress, len(bstr)) - d2 = Document.from_bytes(bstr, protocol=protocol, compress=compress) - assert d2.non_empty_fields == d.non_empty_fields - - -@pytest.mark.parametrize('target', [DocumentArray.empty(10), random_docs(10)]) -def test_dict_json(target): - for d in target: - d1 = Document.from_dict(d.to_dict()) - d2 = Document.from_json(d.to_json()) - assert d1 == d2 diff --git a/tests/unit/document/test_protobuf.py b/tests/unit/document/test_protobuf.py deleted file mode 100644 index d844e12aeac..00000000000 --- a/tests/unit/document/test_protobuf.py +++ /dev/null @@ -1,34 +0,0 @@ -import numpy as np -import pytest - -from docarray import Document -from docarray.score import NamedScore - - -@pytest.mark.parametrize( - 'doc', - [ - Document(tags={'hello': 'world', 'sad': {'nest': 123}, 'hello12': 1.2}), - Document(scores={'hello': NamedScore(value=1.0, description='hello')}), - Document(location=[1.0, 2.0, 3.0]), - Document(chunks=[Document()], matches=[Document(), Document()]), - ], -) -def test_to_from_protobuf(doc): - docr = Document.from_protobuf(doc.to_protobuf()) - assert docr == doc - - -def test_to_protobuf(): - with pytest.raises(TypeError): - Document(text='hello', embedding=np.array([1, 2, 3]), id=1).to_protobuf() - - with pytest.raises(AttributeError): - Document(tags=1).to_protobuf() - - assert ( - Document(text='hello', embedding=np.array([1, 2, 3])).to_protobuf().text - == 'hello' - ) - assert Document(tags={'hello': 'world'}).to_protobuf().tags - assert len(Document(chunks=[Document(), Document()]).to_protobuf().chunks) == 2 diff --git a/tests/unit/document/test_repr_str.py b/tests/unit/document/test_repr_str.py deleted file mode 100644 index d793440aec4..00000000000 --- a/tests/unit/document/test_repr_str.py +++ /dev/null @@ -1,37 +0,0 @@ -import pytest - -from docarray import Document -from docarray.array.chunk import ChunkArray -from docarray.array.match import MatchArray -from docarray.score import NamedScore - - -@pytest.mark.parametrize( - 'obj', - [ - Document(), - NamedScore(), - MatchArray([Document()], Document()), - ChunkArray([Document()], Document()), - ], -) -def test_builtin_str_repr_no_content(obj): - print(obj) - print(f'{obj!r}') - - -@pytest.mark.parametrize( - 'obj', - [ - Document(content='123', chunks=[Document(content='abc')]), - NamedScore( - op_name='operation', - value=10.0, - ref_id='10' * 16, - description='score description', - ), - ], -) -def test_builtin_str_repr_has_content(obj): - print(obj) - print(f'{obj!r}') diff --git a/tests/unit/document/test_summary.py b/tests/unit/document/test_summary.py deleted file mode 100644 index 6516dbc50c2..00000000000 --- a/tests/unit/document/test_summary.py +++ /dev/null @@ -1,29 +0,0 @@ -import os - -from docarray import Document - -cur_dir = os.path.dirname(os.path.abspath(__file__)) - - -def test_single_doc_summary(): - # empty doc - Document().summary() - # nested doc - Document( - chunks=[ - Document(), - Document(chunks=[Document()]), - Document(), - ], - matches=[Document(), Document()], - ).summary() - - -def test_plot_image(): - d = Document(uri=os.path.join(cur_dir, 'toydata/test.png')) - d.plot() - - d.load_uri_to_image_blob() - d.uri = None - - d.plot() diff --git a/tests/unit/document/toydata/docs.jsonlines b/tests/unit/document/toydata/docs.jsonlines deleted file mode 100644 index 58b3fbfe534..00000000000 --- a/tests/unit/document/toydata/docs.jsonlines +++ /dev/null @@ -1,2 +0,0 @@ -{"text": "a"} -{"text": "b"} diff --git a/tests/unit/document/toydata/mov_bbb.mp4 b/tests/unit/document/toydata/mov_bbb.mp4 deleted file mode 100644 index 0a4dd5b4017..00000000000 Binary files a/tests/unit/document/toydata/mov_bbb.mp4 and /dev/null differ diff --git a/tests/unit/document/toydata/olleh.wav b/tests/unit/document/toydata/olleh.wav deleted file mode 100644 index 28523a7750f..00000000000 Binary files a/tests/unit/document/toydata/olleh.wav and /dev/null differ diff --git a/tests/unit/test_placeholder.py b/tests/unit/test_placeholder.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/units/__init__.py b/tests/units/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/array/__init__.py b/tests/units/array/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/array/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/array/stack/__init__.py b/tests/units/array/stack/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/array/stack/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/array/stack/storage/__init__.py b/tests/units/array/stack/storage/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/array/stack/storage/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/array/stack/storage/test_array_stack_with_optional.py b/tests/units/array/stack/storage/test_array_stack_with_optional.py new file mode 100644 index 00000000000..182b0178593 --- /dev/null +++ b/tests/units/array/stack/storage/test_array_stack_with_optional.py @@ -0,0 +1,88 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Optional + +import numpy as np +import pytest + +from docarray import BaseDoc, DocList, DocVec +from docarray.typing import NdArray + + +class Nested(BaseDoc): + tensor: NdArray + + +class Image(BaseDoc): + features: Optional[Nested] = None + + +def test_optional_field(): + docs = DocVec[Image]([Image() for _ in range(10)]) + + assert docs.features is None + + docs.features = DocList[Nested]([Nested(tensor=np.zeros(10)) for _ in range(10)]) + + assert docs.features.tensor.shape == (10, 10) + + for doc in docs: + assert doc.features.tensor.shape == (10,) + + +def test_set_none(): + docs = DocVec[Image]( + [Image(features=Nested(tensor=np.zeros(10))) for _ in range(10)] + ) + assert docs.features.tensor.shape == (10, 10) + + docs.features = None + + assert docs.features is None + + for doc in docs: + assert doc.features is None + + +def test_set_doc(): + docs = DocVec[Image]( + [Image(features=Nested(tensor=np.zeros(10))) for _ in range(10)] + ) + assert docs.features.tensor.shape == (10, 10) + + for doc in docs: + doc.features = Nested(tensor=np.ones(10)) + + with pytest.raises(ValueError): + doc.features = None + + +def test_set_doc_none(): + docs = DocVec[Image]([Image() for _ in range(10)]) + + assert docs.features is None + + for doc in docs: + with pytest.raises(ValueError): + doc.features = Nested(tensor=np.ones(10)) + + +def test_no_uniform_none(): + with pytest.raises(ValueError): + DocVec[Image]([Image(), Image(features=Nested(tensor=np.zeros(10)))]) + + with pytest.raises(ValueError): + DocVec[Image]([Image(features=Nested(tensor=np.zeros(10))), Image()]) diff --git a/tests/units/array/stack/storage/test_storage.py b/tests/units/array/stack/storage/test_storage.py new file mode 100644 index 00000000000..b91585d3737 --- /dev/null +++ b/tests/units/array/stack/storage/test_storage.py @@ -0,0 +1,104 @@ +import numpy as np + +from docarray import BaseDoc +from docarray.array import DocVec +from docarray.array.doc_vec.column_storage import ColumnStorageView +from docarray.typing import AnyTensor + + +def test_column_storage_init(): + class InnerDoc(BaseDoc): + price: int + + class MyDoc(BaseDoc): + tensor: AnyTensor + name: str + doc: InnerDoc + + docs = [ + MyDoc(tensor=np.zeros(10), name='hello', doc=InnerDoc(price=i)) + for i in range(4) + ] + + storage = DocVec[MyDoc](docs)._storage + + assert (storage.tensor_columns['tensor'] == np.zeros((4, 10))).all() + for name in storage.any_columns['name']: + assert name == 'hello' + inner_docs = storage.doc_columns['doc'] + assert isinstance(inner_docs, DocVec) + for i, doc in enumerate(inner_docs): + assert isinstance(doc, InnerDoc) + assert doc.price == i + + +def test_column_storage_view(): + class MyDoc(BaseDoc): + tensor: AnyTensor + name: str + + docs = [MyDoc(tensor=np.zeros((10, 10)), name='hello', id=str(i)) for i in range(4)] + + storage = DocVec[MyDoc](docs)._storage + + view = ColumnStorageView(0, storage) + + assert view['id'] == '0' + assert (view['tensor'] == np.zeros(10)).all() + assert view['name'] == 'hello' + + view['id'] = '1' + view['tensor'] = np.ones(10) + view['name'] = 'byebye' + + assert storage.any_columns['id'][0] == '1' + assert (storage.tensor_columns['tensor'][0] == np.ones(10)).all() + assert storage.any_columns['name'][0] == 'byebye' + + +def test_column_storage_to_dict(): + class MyDoc(BaseDoc): + tensor: AnyTensor + name: str + + docs = [MyDoc(tensor=np.zeros((10, 10)), name='hello', id=str(i)) for i in range(4)] + + storage = DocVec[MyDoc](docs)._storage + + view = ColumnStorageView(0, storage) + + dict_view = view.to_dict() + + assert dict_view['id'] == '0' + assert (dict_view['tensor'] == np.zeros(10)).all() + assert np.may_share_memory(dict_view['tensor'], view['tensor']) + assert dict_view['name'] == 'hello' + + +def test_storage_view_dict_like(): + class MyDoc(BaseDoc): + tensor: AnyTensor + name: str + + docs = [MyDoc(tensor=np.zeros((10, 10)), name='hello', id=str(i)) for i in range(4)] + + storage = DocVec[MyDoc](docs)._storage + + view = ColumnStorageView(0, storage) + + assert list(view.keys()) == ['id', 'name', 'tensor'] + + # since boolean value of np array is ambiguous, we iterate manually + for val_view, val_reference in zip(view.values(), ['0', 'hello', np.zeros(10)]): + if isinstance(val_view, np.ndarray): + assert (val_view == val_reference).all() + else: + assert val_view == val_reference + for item_view, item_reference in zip( + view.items(), [('id', '0'), ('name', 'hello'), ('tensor', np.zeros(10))] + ): + if isinstance(item_view[1], np.ndarray): + assert item_view[0] == item_reference[0] + assert (item_view[1] == item_reference[1]).all() + else: + assert item_view == item_reference diff --git a/tests/units/array/stack/test_array_stacked.py b/tests/units/array/stack/test_array_stacked.py new file mode 100644 index 00000000000..b1b385840dd --- /dev/null +++ b/tests/units/array/stack/test_array_stacked.py @@ -0,0 +1,680 @@ +from typing import Dict, Optional, Union + +import numpy as np +import pytest +import torch +from pydantic import parse_obj_as + +from docarray import BaseDoc, DocList +from docarray.array import DocVec +from docarray.documents import ImageDoc +from docarray.exceptions.exceptions import UnusableObjectError +from docarray.typing import AnyEmbedding, AnyTensor, NdArray, TorchTensor + + +@pytest.fixture() +def batch(): + class ImageDoc(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + batch = DocVec[ImageDoc]( + [ImageDoc(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + return batch + + +@pytest.fixture() +def nested_batch(): + class ImageDoc(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + class MMdoc(BaseDoc): + img: DocList[ImageDoc] + + batch = DocList[MMdoc]( + [ + MMdoc( + img=DocList[ImageDoc]( + [ImageDoc(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + ) + for _ in range(10) + ] + ) + + return batch.to_doc_vec() + + +def test_create_from_list_docs(): + list_ = [ImageDoc(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + da_stacked = DocVec[ImageDoc](docs=list_, tensor_type=TorchTensor) + assert len(da_stacked) == 10 + assert da_stacked.tensor.shape == tuple([10, 3, 224, 224]) + + +def test_len(batch): + assert len(batch) == 10 + + +def test_create_from_None(): + with pytest.raises(ValueError): + DocVec[ImageDoc]([]) + + +def test_getitem(batch): + for i in range(len(batch)): + assert (batch[i].tensor == torch.zeros(3, 224, 224)).all() + + +def test_iterator(batch): + for doc in batch: + assert (doc.tensor == torch.zeros(3, 224, 224)).all() + + +def test_stack_setter(): + class ImageDoc(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + batch = DocList[ImageDoc]( + [ImageDoc(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + batch = batch.to_doc_vec() + batch.tensor = torch.ones(10, 3, 224, 224) + + assert (batch.tensor == torch.ones(10, 3, 224, 224)).all() + + for i, doc in enumerate(batch): + assert (doc.tensor == batch.tensor[i]).all() + + +def test_stack_setter_np(): + class ImageDoc(BaseDoc): + tensor: NdArray[3, 224, 224] + + batch = DocList[ImageDoc]( + [ImageDoc(tensor=np.zeros((3, 224, 224))) for _ in range(10)] + ) + + batch = batch.to_doc_vec() + batch.tensor = np.ones((10, 3, 224, 224)) + + assert (batch.tensor == np.ones((10, 3, 224, 224))).all() + + for i, doc in enumerate(batch): + assert (doc.tensor == batch.tensor[i]).all() + + +def test_stack_optional(batch): + assert ( + batch._storage.tensor_columns['tensor'] == torch.zeros(10, 3, 224, 224) + ).all() + assert (batch.tensor == torch.zeros(10, 3, 224, 224)).all() + + +def test_stack_numpy(): + class ImageDoc(BaseDoc): + tensor: NdArray[3, 224, 224] + + batch = DocList[ImageDoc]( + [ImageDoc(tensor=np.zeros((3, 224, 224))) for _ in range(10)] + ) + + batch = batch.to_doc_vec() + + assert ( + batch._storage.tensor_columns['tensor'] == np.zeros((10, 3, 224, 224)) + ).all() + assert (batch.tensor == np.zeros((10, 3, 224, 224))).all() + assert ( + batch.tensor.ctypes.data == batch._storage.tensor_columns['tensor'].ctypes.data + ) + + +def test_stack(batch): + assert ( + batch._storage.tensor_columns['tensor'] == torch.zeros(10, 3, 224, 224) + ).all() + assert (batch.tensor == torch.zeros(10, 3, 224, 224)).all() + assert batch._storage.tensor_columns['tensor'].data_ptr() == batch.tensor.data_ptr() + + for doc, tensor in zip(batch, batch.tensor): + assert doc.tensor.data_ptr() == tensor.data_ptr() + + for i in range(len(batch)): + assert batch[i].tensor.data_ptr() == batch.tensor[i].data_ptr() + + +def test_stack_mod_nested_document(): + class ImageDoc(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + class MMdoc(BaseDoc): + img: ImageDoc + + batch = DocList[MMdoc]( + [MMdoc(img=ImageDoc(tensor=torch.zeros(3, 224, 224))) for _ in range(10)] + ) + + batch = batch.to_doc_vec() + + assert ( + batch._storage.doc_columns['img']._storage.tensor_columns['tensor'] + == torch.zeros(10, 3, 224, 224) + ).all() + + assert (batch.img.tensor == torch.zeros(10, 3, 224, 224)).all() + + assert ( + batch._storage.doc_columns['img']._storage.tensor_columns['tensor'].data_ptr() + == batch.img.tensor.data_ptr() + ) + + +def test_stack_nested_DocArray(nested_batch): + for i in range(len(nested_batch)): + assert ( + nested_batch[i].img._storage.tensor_columns['tensor'] + == torch.zeros(10, 3, 224, 224) + ).all() + assert (nested_batch[i].img.tensor == torch.zeros(10, 3, 224, 224)).all() + assert ( + nested_batch[i].img._storage.tensor_columns['tensor'].data_ptr() + == nested_batch[i].img.tensor.data_ptr() + ) + + +def test_convert_to_da(batch): + class ImageDoc(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + batch = DocList[ImageDoc]( + [ImageDoc(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + batch = batch.to_doc_vec() + da = batch.to_doc_list() + + for doc in da: + assert (doc.tensor == torch.zeros(3, 224, 224)).all() + + +def test_unstack_nested_document(): + class ImageDoc(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + class MMdoc(BaseDoc): + img: ImageDoc + + batch = DocList[MMdoc]( + [MMdoc(img=ImageDoc(tensor=torch.zeros(3, 224, 224))) for _ in range(10)] + ) + + batch = batch.to_doc_vec() + + da = batch.to_doc_list() + + for doc in da: + assert (doc.img.tensor == torch.zeros(3, 224, 224)).all() + + +def test_unstack_nested_DocArray(nested_batch): + batch = nested_batch.to_doc_list() + for i in range(len(batch)): + assert isinstance(batch[i].img, DocList) + for doc in batch[i].img: + assert (doc.tensor == torch.zeros(3, 224, 224)).all() + + +def test_stack_call(): + class ImageDoc(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + da = DocList[ImageDoc]( + [ImageDoc(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + da = da.to_doc_vec() + + assert len(da) == 10 + + assert da.tensor.shape == (10, 3, 224, 224) + + +def test_stack_union(): + class ImageDoc(BaseDoc): + tensor: Union[NdArray[3, 224, 224], TorchTensor[3, 224, 224]] + + batch = DocList[ImageDoc]( + [ImageDoc(tensor=np.zeros((3, 224, 224))) for _ in range(10)] + ) + batch[3].tensor = np.zeros((3, 224, 224)) + + # union fields aren't actually doc_vec + # just checking that there is no error + batch.to_doc_vec() + + +@pytest.mark.parametrize( + 'tensor_type,tensor', + [(TorchTensor, torch.zeros(3, 224, 224)), (NdArray, np.zeros((3, 224, 224)))], +) +def test_any_tensor_with_torch(tensor_type, tensor): + class ImageDoc(BaseDoc): + tensor: AnyTensor + + da = DocVec[ImageDoc]( + [ImageDoc(tensor=tensor) for _ in range(10)], + tensor_type=tensor_type, + ) + + for i in range(len(da)): + assert (da[i].tensor == tensor).all() + + assert 'tensor' in da._storage.tensor_columns.keys() + assert isinstance(da._storage.tensor_columns['tensor'], tensor_type) + + +def test_any_tensor_with_optional(): + tensor = torch.zeros(3, 224, 224) + + class ImageDoc(BaseDoc): + tensor: Optional[AnyTensor] = None + + class TopDoc(BaseDoc): + img: ImageDoc + + da = DocVec[TopDoc]( + [TopDoc(img=ImageDoc(tensor=tensor)) for _ in range(10)], + tensor_type=TorchTensor, + ) + + for i in range(len(da)): + assert (da.img[i].tensor == tensor).all() + + assert 'tensor' in da.img._storage.tensor_columns.keys() + assert isinstance(da.img._storage.tensor_columns['tensor'], TorchTensor) + + +def test_dict_stack(): + class MyDoc(BaseDoc): + my_dict: Dict[str, int] + + da = DocVec[MyDoc]([MyDoc(my_dict={'a': 1, 'b': 2}) for _ in range(10)]) + + da.my_dict + + +def test_get_from_slice_stacked(): + class Doc(BaseDoc): + text: str + tensor: NdArray + + N = 10 + + da = DocVec[Doc]( + [Doc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N)] + ) + + da_sliced = da[0:10:2] + assert isinstance(da_sliced, DocVec) + + tensors = da_sliced.tensor + assert tensors.shape == (5, 3, 224, 224) + + texts = da_sliced.text + assert len(texts) == 5 + for i, text in enumerate(texts): + assert text == f'hello{i * 2}' + + +def test_stack_embedding(): + class MyDoc(BaseDoc): + embedding: AnyEmbedding + + da = DocVec[MyDoc]([MyDoc(embedding=np.zeros(10)) for _ in range(10)]) + + assert 'embedding' in da._storage.tensor_columns.keys() + assert (da.embedding == np.zeros((10, 10))).all() + + +@pytest.mark.parametrize('tensor_backend', [TorchTensor, NdArray]) +def test_stack_none(tensor_backend): + class MyDoc(BaseDoc): + tensor: Optional[AnyTensor] = None + + da = DocVec[MyDoc]( + [MyDoc(tensor=None) for _ in range(10)], tensor_type=tensor_backend + ) + + assert 'tensor' in da._storage.tensor_columns.keys() + + +def test_to_device(): + da = DocVec[ImageDoc]([ImageDoc(tensor=torch.zeros(3, 5))], tensor_type=TorchTensor) + assert da.tensor.device == torch.device('cpu') + da.to('meta') + assert da.tensor.device == torch.device('meta') + + +def test_to_device_with_nested_da(): + class Video(BaseDoc): + images: DocVec[ImageDoc] + + da_image = DocVec[ImageDoc]( + [ImageDoc(tensor=torch.zeros(3, 5))], tensor_type=TorchTensor + ) + + da = DocVec[Video]([Video(images=da_image)]) + assert da.images[0].tensor.device == torch.device('cpu') + da.to('meta') + assert da.images[0].tensor.device == torch.device('meta') + + +def test_to_device_nested(): + class MyDoc(BaseDoc): + tensor: TorchTensor + docs: ImageDoc + + da = DocVec[MyDoc]( + [MyDoc(tensor=torch.zeros(3, 5), docs=ImageDoc(tensor=torch.zeros(3, 5)))], + tensor_type=TorchTensor, + ) + assert da.tensor.device == torch.device('cpu') + assert da.docs.tensor.device == torch.device('cpu') + da.to('meta') + assert da.tensor.device == torch.device('meta') + assert da.docs.tensor.device == torch.device('meta') + + +def test_to_device_numpy(): + da = DocVec[ImageDoc]([ImageDoc(tensor=np.zeros((3, 5)))], tensor_type=NdArray) + with pytest.raises(NotImplementedError): + da.to('meta') + + +def test_keep_dtype_torch(): + class MyDoc(BaseDoc): + tensor: TorchTensor + + da = DocList[MyDoc]( + [MyDoc(tensor=torch.zeros([2, 4], dtype=torch.int32)) for _ in range(3)] + ) + assert da[0].tensor.dtype == torch.int32 + + da = da.to_doc_vec() + assert da[0].tensor.dtype == torch.int32 + assert da.tensor.dtype == torch.int32 + + +def test_keep_dtype_np(): + class MyDoc(BaseDoc): + tensor: NdArray + + da = DocList[MyDoc]( + [MyDoc(tensor=np.zeros([2, 4], dtype=np.int32)) for _ in range(3)] + ) + assert da[0].tensor.dtype == np.int32 + + da = da.to_doc_vec() + assert da[0].tensor.dtype == np.int32 + assert da.tensor.dtype == np.int32 + + +def test_del_item(batch): + assert len(batch) == 10 + assert batch.tensor.shape[0] == 10 + with pytest.raises(NotImplementedError): + del batch[2] + + +def test_np_scalar(): + class MyDoc(BaseDoc): + scalar: NdArray + + da = DocList[MyDoc]([MyDoc(scalar=np.array(2.0)) for _ in range(3)]) + assert all(doc.scalar.ndim == 0 for doc in da) + assert all(doc.scalar == 2.0 for doc in da) + + stacked_da = da.to_doc_vec() + assert type(stacked_da.scalar) == NdArray + + assert all(type(doc.scalar) == NdArray for doc in stacked_da) + assert all(doc.scalar.ndim == 1 for doc in stacked_da) + assert all(doc.scalar == 2.0 for doc in stacked_da) + + # Make sure they share memory + stacked_da.scalar[0] = 3.0 + assert stacked_da[0].scalar == 3.0 + + +def test_torch_scalar(): + class MyDoc(BaseDoc): + scalar: TorchTensor + + da = DocList[MyDoc]( + [MyDoc(scalar=torch.tensor(2.0)) for _ in range(3)], + ) + assert all(doc.scalar.ndim == 0 for doc in da) + assert all(doc.scalar == 2.0 for doc in da) + stacked_da = da.to_doc_vec(tensor_type=TorchTensor) + assert type(stacked_da.scalar) == TorchTensor + + assert all(type(doc.scalar) == TorchTensor for doc in stacked_da) + assert all(doc.scalar.ndim == 1 for doc in stacked_da) # TODO failing here + assert all(doc.scalar == 2.0 for doc in stacked_da) + + stacked_da.scalar[0] = 3.0 + assert stacked_da[0].scalar == 3.0 + + +def test_np_nan(): + class MyDoc(BaseDoc): + scalar: Optional[NdArray] = None + + da = DocList[MyDoc]([MyDoc() for _ in range(3)]) + assert all(doc.scalar is None for doc in da) + assert all(doc.scalar == doc.scalar for doc in da) + stacked_da = da.to_doc_vec() + assert stacked_da.scalar is None + + assert all(doc.scalar is None for doc in stacked_da) + # Stacking them turns them into np.nan + + +def test_from_storage(): + class ImageDoc(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + batch = DocVec[ImageDoc]( + [ImageDoc(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + DocVec[ImageDoc].from_columns_storage(batch._storage) + + +def test_validate_from_da(): + class ImageDoc(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + batch = DocList[ImageDoc]( + [ImageDoc(tensor=torch.zeros(3, 224, 224)) for _ in range(10)] + ) + + da = parse_obj_as(DocVec[ImageDoc], batch) + + assert isinstance(da, DocVec) + for d in da: + assert isinstance(d, ImageDoc) + + +def test_validation_column_tensor(batch): + batch.tensor = torch.zeros(10, 3, 224, 244) + assert isinstance(batch.tensor, TorchTensor) + + +def test_validation_column_tensor_fail(batch): + with pytest.raises(ValueError): + batch.tensor = ['hello'] * 10 + + with pytest.raises(ValueError): + batch.tensor = torch.zeros(11, 3, 224, 244) + + +@pytest.fixture() +def batch_nested_doc(): + class Inner(BaseDoc): + hello: str + + class Doc(BaseDoc): + inner: Inner + + batch = DocVec[Doc]([Doc(inner=Inner(hello='hello')) for _ in range(10)]) + return batch, Doc, Inner + + +def test_validation_column_doc(batch_nested_doc): + batch, Doc, Inner = batch_nested_doc + + batch.inner = DocList[Inner]([Inner(hello='hello') for _ in range(10)]) + assert isinstance(batch.inner, DocVec) + for d in batch.inner: + assert isinstance(d, Inner) + + +def test_validation_list_doc(batch_nested_doc): + batch, Doc, Inner = batch_nested_doc + + batch.inner = [Inner(hello='hello') for _ in range(10)] + assert isinstance(batch.inner, DocVec) + for d in batch.inner: + assert isinstance(d, Inner) + + +def test_validation_col_doc_fail(batch_nested_doc): + batch, Doc, Inner = batch_nested_doc + + with pytest.raises(ValueError): + batch.inner = ['hello'] * 10 + + with pytest.raises(ValueError): + batch.inner = DocList[Inner]([Inner(hello='hello') for _ in range(11)]) + + +def test_doc_view_update(batch): + batch[0].tensor = 12 * torch.ones(3, 224, 224) + assert (batch.tensor[0] == 12 * torch.ones(3, 224, 224)).all() + + +def test_doc_view_nested(batch_nested_doc): + batch, Doc, Inner = batch_nested_doc + batch[0].inner = Inner(hello='world') + assert batch.inner[0].hello == 'world' + + +def test_type_error_no_doc_type(): + + with pytest.raises(TypeError): + DocVec([BaseDoc() for _ in range(10)]) + + +def test_doc_view_dict(batch: DocVec[ImageDoc]): + doc_view = batch[0] + assert doc_view.is_view() + d = doc_view.dict() + assert d['tensor'].shape == (3, 224, 224) + assert d['id'] == doc_view.id + + doc_view_two = batch[1] + assert doc_view_two.is_view() + d = doc_view_two.dict() + assert d['tensor'].shape == (3, 224, 224) + assert d['id'] == doc_view_two.id + + +def test_doc_vec_equality(): + class Text(BaseDoc): + text: str + + da = DocVec[Text]([Text(text='hello') for _ in range(10)]) + da2 = DocList[Text]([Text(text='hello') for _ in range(10)]) + + assert da != da2 + assert da == da2.to_doc_vec() + + +@pytest.mark.parametrize('tensor_type', [TorchTensor, NdArray]) +def test_doc_vec_equality_tensor(tensor_type): + class Text(BaseDoc): + tens: tensor_type + + da = DocVec[Text]( + [Text(tens=[1, 2, 3, 4]) for _ in range(10)], tensor_type=tensor_type + ) + da2 = DocVec[Text]( + [Text(tens=[1, 2, 3, 4]) for _ in range(10)], tensor_type=tensor_type + ) + assert da == da2 + + da2 = DocVec[Text]( + [Text(tens=[1, 2, 3, 4, 5]) for _ in range(10)], tensor_type=tensor_type + ) + assert da != da2 + + +@pytest.mark.tensorflow +def test_doc_vec_equality_tf(): + from docarray.typing import TensorFlowTensor + + class Text(BaseDoc): + tens: TensorFlowTensor + + da = DocVec[Text]( + [Text(tens=[1, 2, 3, 4]) for _ in range(10)], tensor_type=TensorFlowTensor + ) + da2 = DocVec[Text]( + [Text(tens=[1, 2, 3, 4]) for _ in range(10)], tensor_type=TensorFlowTensor + ) + assert da == da2 + + da2 = DocVec[Text]( + [Text(tens=[1, 2, 3, 4, 5]) for _ in range(10)], tensor_type=TensorFlowTensor + ) + assert da != da2 + + +def test_doc_vec_nested(batch_nested_doc): + batch, Doc, Inner = batch_nested_doc + batch2 = DocVec[Doc]([Doc(inner=Inner(hello='hello')) for _ in range(10)]) + + assert batch == batch2 + + +def test_doc_vec_tensor_type(): + class ImageDoc(BaseDoc): + tensor: AnyTensor + + da = DocVec[ImageDoc]([ImageDoc(tensor=np.zeros((3, 224, 224))) for _ in range(10)]) + + da2 = DocVec[ImageDoc]( + [ImageDoc(tensor=torch.zeros(3, 224, 224)) for _ in range(10)], + tensor_type=TorchTensor, + ) + + assert da != da2 + + +def teste_unusable_state_raises_exception(): + from docarray import DocVec + from docarray.documents import ImageDoc + + docs = DocVec[ImageDoc]([ImageDoc(url='http://url.com/foo.png') for _ in range(10)]) + + docs.to_doc_list() + + with pytest.raises(UnusableObjectError): + docs.url + + with pytest.raises(UnusableObjectError): + docs.url = 'hi' diff --git a/tests/units/array/stack/test_array_stacked_jax.py b/tests/units/array/stack/test_array_stacked_jax.py new file mode 100644 index 00000000000..86f1399a40d --- /dev/null +++ b/tests/units/array/stack/test_array_stacked_jax.py @@ -0,0 +1,301 @@ +from typing import Optional, Union + +import pytest + +from docarray import BaseDoc, DocList +from docarray.array import DocVec +from docarray.typing import ( + AnyEmbedding, + AnyTensor, + AudioTensor, + ImageTensor, + NdArray, + VideoTensor, +) +from docarray.utils._internal.misc import is_jax_available + +jax_available = is_jax_available() +if jax_available: + import jax.numpy as jnp + + from docarray.typing import JaxArray + + +@pytest.fixture() +@pytest.mark.jax +def batch(): + + import jax.numpy as jnp + + class Image(BaseDoc): + tensor: JaxArray[3, 224, 224] + + batch = DocList[Image]([Image(tensor=jnp.zeros((3, 224, 224))) for _ in range(10)]) + + return batch.to_doc_vec() + + +@pytest.fixture() +@pytest.mark.jax +def nested_batch(): + class Image(BaseDoc): + tensor: JaxArray[3, 224, 224] + + class MMdoc(BaseDoc): + img: DocList[Image] + + batch = DocVec[MMdoc]( + [ + MMdoc( + img=DocList[Image]( + [Image(tensor=jnp.zeros((3, 224, 224))) for _ in range(10)] + ) + ) + for _ in range(10) + ] + ) + + return batch + + +@pytest.mark.jax +def test_len(batch): + assert len(batch) == 10 + + +@pytest.mark.jax +def test_getitem(batch): + for i in range(len(batch)): + item = batch[i] + assert isinstance(item.tensor, JaxArray) + assert jnp.allclose(item.tensor.tensor, jnp.zeros((3, 224, 224))) + + +@pytest.mark.jax +def test_get_slice(batch): + sliced = batch[0:2] + assert isinstance(sliced, DocVec) + assert len(sliced) == 2 + + +@pytest.mark.jax +def test_iterator(batch): + for doc in batch: + assert jnp.allclose(doc.tensor.tensor, jnp.zeros((3, 224, 224))) + + +@pytest.mark.jax +def test_set_after_stacking(): + class Image(BaseDoc): + tensor: JaxArray[3, 224, 224] + + batch = DocVec[Image]([Image(tensor=jnp.zeros((3, 224, 224))) for _ in range(10)]) + + batch.tensor = jnp.ones((10, 3, 224, 224)) + assert jnp.allclose(batch.tensor.tensor, jnp.ones((10, 3, 224, 224))) + for i, doc in enumerate(batch): + assert jnp.allclose(doc.tensor.tensor, batch.tensor.tensor[i]) + + +@pytest.mark.jax +def test_stack_optional(batch): + assert jnp.allclose( + batch._storage.tensor_columns['tensor'].tensor, jnp.zeros((10, 3, 224, 224)) + ) + assert jnp.allclose(batch.tensor.tensor, jnp.zeros((10, 3, 224, 224))) + + +@pytest.mark.jax +def test_stack_mod_nested_document(): + class Image(BaseDoc): + tensor: JaxArray[3, 224, 224] + + class MMdoc(BaseDoc): + img: Image + + batch = DocList[MMdoc]( + [MMdoc(img=Image(tensor=jnp.zeros((3, 224, 224)))) for _ in range(10)] + ).to_doc_vec() + + assert jnp.allclose( + batch._storage.doc_columns['img']._storage.tensor_columns['tensor'].tensor, + jnp.zeros((10, 3, 224, 224)), + ) + + assert jnp.allclose(batch.img.tensor.tensor, jnp.zeros((10, 3, 224, 224))) + + +@pytest.mark.jax +def test_stack_nested_DocArray(nested_batch): + for i in range(len(nested_batch)): + assert jnp.allclose( + nested_batch[i].img._storage.tensor_columns['tensor'].tensor, + jnp.zeros((10, 3, 224, 224)), + ) + + assert jnp.allclose( + nested_batch[i].img.tensor.tensor, jnp.zeros((10, 3, 224, 224)) + ) + + +@pytest.mark.jax +def test_convert_to_da(batch): + da = batch.to_doc_list() + + for doc in da: + assert jnp.allclose(doc.tensor.tensor, jnp.zeros((3, 224, 224))) + + +@pytest.mark.jax +def test_unstack_nested_document(): + class Image(BaseDoc): + tensor: JaxArray[3, 224, 224] + + class MMdoc(BaseDoc): + img: Image + + batch = DocVec[MMdoc]( + [MMdoc(img=Image(tensor=jnp.zeros((3, 224, 224)))) for _ in range(10)] + ) + assert isinstance(batch.img._storage.tensor_columns['tensor'], JaxArray) + da = batch.to_doc_list() + + for doc in da: + assert jnp.allclose(doc.img.tensor.tensor, jnp.zeros((3, 224, 224))) + + +@pytest.mark.jax +def test_unstack_nested_DocArray(nested_batch): + batch = nested_batch.to_doc_list() + for i in range(len(batch)): + assert isinstance(batch[i].img, DocList) + for doc in batch[i].img: + assert jnp.allclose(doc.tensor.tensor, jnp.zeros((3, 224, 224))) + + +@pytest.mark.jax +def test_stack_call(): + class Image(BaseDoc): + tensor: JaxArray[3, 224, 224] + + da = DocList[Image]([Image(tensor=jnp.zeros((3, 224, 224))) for _ in range(10)]) + + da = da.to_doc_vec() + + assert len(da) == 10 + + assert da.tensor.tensor.shape == (10, 3, 224, 224) + + +@pytest.mark.jax +def test_stack_union(): + class Image(BaseDoc): + tensor: Union[JaxArray[3, 224, 224], NdArray[3, 224, 224]] + + DocVec[Image]( + [Image(tensor=jnp.zeros((3, 224, 224))) for _ in range(10)], + tensor_type=JaxArray, + ) + + # union fields aren't actually doc_vec + # just checking that there is no error + + +@pytest.mark.jax +def test_setitem_tensor(batch): + batch[3].tensor.tensor = jnp.zeros((3, 224, 224)) + + +@pytest.mark.jax +@pytest.mark.skip('not working yet') +def test_setitem_tensor_direct(batch): + batch[3].tensor = jnp.zeros((3, 224, 224)) + + +@pytest.mark.jax +@pytest.mark.parametrize( + 'cls_tensor', [ImageTensor, AudioTensor, VideoTensor, AnyEmbedding, AnyTensor] +) +def test_generic_tensors_with_jnp(cls_tensor): + tensor = jnp.zeros((3, 224, 224)) + + class Image(BaseDoc): + tensor: cls_tensor + + da = DocVec[Image]( + [Image(tensor=tensor) for _ in range(10)], + tensor_type=JaxArray, + ) + + for i in range(len(da)): + assert jnp.allclose(da[i].tensor.tensor, tensor) + + assert 'tensor' in da._storage.tensor_columns.keys() + assert isinstance(da._storage.tensor_columns['tensor'], JaxArray) + + +@pytest.mark.jax +@pytest.mark.parametrize( + 'cls_tensor', [ImageTensor, AudioTensor, VideoTensor, AnyEmbedding, AnyTensor] +) +def test_generic_tensors_with_optional(cls_tensor): + tensor = jnp.zeros((3, 224, 224)) + + class Image(BaseDoc): + tensor: Optional[cls_tensor] = None + + class TopDoc(BaseDoc): + img: Image + + da = DocVec[TopDoc]( + [TopDoc(img=Image(tensor=tensor)) for _ in range(10)], + tensor_type=JaxArray, + ) + + for i in range(len(da)): + assert jnp.allclose(da.img[i].tensor.tensor, tensor) + + assert 'tensor' in da.img._storage.tensor_columns.keys() + assert isinstance(da.img._storage.tensor_columns['tensor'], JaxArray) + assert isinstance(da.img._storage.tensor_columns['tensor'].tensor, jnp.ndarray) + + +@pytest.mark.jax +def test_get_from_slice_stacked(): + class Doc(BaseDoc): + text: str + tensor: JaxArray + + da = DocVec[Doc]( + [Doc(text=f'hello{i}', tensor=jnp.zeros((3, 224, 224))) for i in range(10)] + ) + + da_sliced = da[0:10:2] + assert isinstance(da_sliced, DocVec) + + tensors = da_sliced.tensor.tensor + assert tensors.shape == (5, 3, 224, 224) + + +@pytest.mark.jax +def test_stack_none(): + class MyDoc(BaseDoc): + tensor: Optional[AnyTensor] = None + + da = DocVec[MyDoc]([MyDoc(tensor=None) for _ in range(10)], tensor_type=JaxArray) + assert 'tensor' in da._storage.tensor_columns.keys() + + +@pytest.mark.jax +def test_keep_dtype_jnp(): + class MyDoc(BaseDoc): + tensor: JaxArray + + da = DocList[MyDoc]( + [MyDoc(tensor=jnp.zeros([2, 4], dtype=jnp.int32)) for _ in range(3)] + ) + assert da[0].tensor.tensor.dtype == jnp.int32 + + da = da.to_doc_vec() + assert da[0].tensor.tensor.dtype == jnp.int32 + assert da.tensor.tensor.dtype == jnp.int32 diff --git a/tests/units/array/stack/test_array_stacked_tf.py b/tests/units/array/stack/test_array_stacked_tf.py new file mode 100644 index 00000000000..da055fcd8ee --- /dev/null +++ b/tests/units/array/stack/test_array_stacked_tf.py @@ -0,0 +1,303 @@ +from typing import Optional, Union + +import pytest + +from docarray import BaseDoc, DocList +from docarray.array import DocVec +from docarray.typing import ( + AnyEmbedding, + AnyTensor, + AudioTensor, + ImageTensor, + NdArray, + VideoTensor, +) +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + import tensorflow._api.v2.experimental.numpy as tnp + + from docarray.typing import TensorFlowTensor + + +@pytest.fixture() +def batch(): + class Image(BaseDoc): + tensor: TensorFlowTensor[3, 224, 224] + + import tensorflow as tf + + batch = DocList[Image]([Image(tensor=tf.zeros((3, 224, 224))) for _ in range(10)]) + + return batch.to_doc_vec() + + +@pytest.fixture() +def nested_batch(): + class Image(BaseDoc): + tensor: TensorFlowTensor[3, 224, 224] + + class MMdoc(BaseDoc): + img: DocList[Image] + + import tensorflow as tf + + batch = DocVec[MMdoc]( + [ + MMdoc( + img=DocList[Image]( + [Image(tensor=tf.zeros((3, 224, 224))) for _ in range(10)] + ) + ) + for _ in range(10) + ] + ) + + return batch + + +@pytest.mark.tensorflow +def test_len(batch): + assert len(batch) == 10 + + +@pytest.mark.tensorflow +def test_getitem(batch): + for i in range(len(batch)): + item = batch[i] + assert isinstance(item.tensor, TensorFlowTensor) + assert tnp.allclose(item.tensor.tensor, tf.zeros((3, 224, 224))) + + +@pytest.mark.tensorflow +def test_get_slice(batch): + sliced = batch[0:2] + assert isinstance(sliced, DocVec) + assert len(sliced) == 2 + + +@pytest.mark.tensorflow +def test_iterator(batch): + for doc in batch: + assert tnp.allclose(doc.tensor.tensor, tf.zeros((3, 224, 224))) + + +@pytest.mark.tensorflow +def test_set_after_stacking(): + class Image(BaseDoc): + tensor: TensorFlowTensor[3, 224, 224] + + batch = DocVec[Image]([Image(tensor=tf.zeros((3, 224, 224))) for _ in range(10)]) + + batch.tensor = tf.ones((10, 3, 224, 224)) + assert tnp.allclose(batch.tensor.tensor, tf.ones((10, 3, 224, 224))) + for i, doc in enumerate(batch): + assert tnp.allclose(doc.tensor.tensor, batch.tensor.tensor[i]) + + +@pytest.mark.tensorflow +def test_stack_optional(batch): + assert tnp.allclose( + batch._storage.tensor_columns['tensor'].tensor, tf.zeros((10, 3, 224, 224)) + ) + assert tnp.allclose(batch.tensor.tensor, tf.zeros((10, 3, 224, 224))) + + +@pytest.mark.tensorflow +def test_stack_mod_nested_document(): + class Image(BaseDoc): + tensor: TensorFlowTensor[3, 224, 224] + + class MMdoc(BaseDoc): + img: Image + + batch = DocList[MMdoc]( + [MMdoc(img=Image(tensor=tf.zeros((3, 224, 224)))) for _ in range(10)] + ).to_doc_vec() + + assert tnp.allclose( + batch._storage.doc_columns['img']._storage.tensor_columns['tensor'].tensor, + tf.zeros((10, 3, 224, 224)), + ) + + assert tnp.allclose(batch.img.tensor.tensor, tf.zeros((10, 3, 224, 224))) + + +@pytest.mark.tensorflow +def test_stack_nested_DocArray(nested_batch): + for i in range(len(nested_batch)): + assert tnp.allclose( + nested_batch[i].img._storage.tensor_columns['tensor'].tensor, + tf.zeros((10, 3, 224, 224)), + ) + + assert tnp.allclose( + nested_batch[i].img.tensor.tensor, tf.zeros((10, 3, 224, 224)) + ) + + +@pytest.mark.tensorflow +def test_convert_to_da(batch): + da = batch.to_doc_list() + + for doc in da: + assert tnp.allclose(doc.tensor.tensor, tf.zeros((3, 224, 224))) + + +@pytest.mark.tensorflow +def test_unstack_nested_document(): + class Image(BaseDoc): + tensor: TensorFlowTensor[3, 224, 224] + + class MMdoc(BaseDoc): + img: Image + + batch = DocVec[MMdoc]( + [MMdoc(img=Image(tensor=tf.zeros((3, 224, 224)))) for _ in range(10)] + ) + assert isinstance(batch.img._storage.tensor_columns['tensor'], TensorFlowTensor) + da = batch.to_doc_list() + + for doc in da: + assert tnp.allclose(doc.img.tensor.tensor, tf.zeros((3, 224, 224))) + + +@pytest.mark.tensorflow +def test_unstack_nested_DocArray(nested_batch): + batch = nested_batch.to_doc_list() + for i in range(len(batch)): + assert isinstance(batch[i].img, DocList) + for doc in batch[i].img: + assert tnp.allclose(doc.tensor.tensor, tf.zeros((3, 224, 224))) + + +@pytest.mark.tensorflow +def test_stack_call(): + class Image(BaseDoc): + tensor: TensorFlowTensor[3, 224, 224] + + da = DocList[Image]([Image(tensor=tf.zeros((3, 224, 224))) for _ in range(10)]) + + da = da.to_doc_vec() + + assert len(da) == 10 + + assert da.tensor.tensor.shape == (10, 3, 224, 224) + + +@pytest.mark.tensorflow +def test_stack_union(): + class Image(BaseDoc): + tensor: Union[TensorFlowTensor[3, 224, 224], NdArray[3, 224, 224]] + + DocVec[Image]( + [Image(tensor=tf.zeros((3, 224, 224))) for _ in range(10)], + tensor_type=TensorFlowTensor, + ) + + # union fields aren't actually doc_vec + # just checking that there is no error + + +@pytest.mark.tensorflow +def test_setitem_tensor(batch): + batch[3].tensor.tensor = tf.zeros((3, 224, 224)) + + +@pytest.mark.skip('not working yet') +@pytest.mark.tensorflow +def test_setitem_tensor_direct(batch): + batch[3].tensor = tf.zeros((3, 224, 224)) + + +@pytest.mark.parametrize( + 'cls_tensor', [ImageTensor, AudioTensor, VideoTensor, AnyEmbedding, AnyTensor] +) +@pytest.mark.tensorflow +def test_generic_tensors_with_tf(cls_tensor): + tensor = tf.zeros((3, 224, 224)) + + class Image(BaseDoc): + tensor: cls_tensor + + da = DocVec[Image]( + [Image(tensor=tensor) for _ in range(10)], + tensor_type=TensorFlowTensor, + ) + + for i in range(len(da)): + assert tnp.allclose(da[i].tensor.tensor, tensor) + + assert 'tensor' in da._storage.tensor_columns.keys() + assert isinstance(da._storage.tensor_columns['tensor'], TensorFlowTensor) + + +@pytest.mark.parametrize( + 'cls_tensor', [ImageTensor, AudioTensor, VideoTensor, AnyEmbedding, AnyTensor] +) +@pytest.mark.tensorflow +def test_generic_tensors_with_optional(cls_tensor): + tensor = tf.zeros((3, 224, 224)) + + class Image(BaseDoc): + tensor: Optional[cls_tensor] + + class TopDoc(BaseDoc): + img: Image + + da = DocVec[TopDoc]( + [TopDoc(img=Image(tensor=tensor)) for _ in range(10)], + tensor_type=TensorFlowTensor, + ) + + for i in range(len(da)): + assert tnp.allclose(da.img[i].tensor.tensor, tensor) + + assert 'tensor' in da.img._storage.tensor_columns.keys() + assert isinstance(da.img._storage.tensor_columns['tensor'], TensorFlowTensor) + assert isinstance(da.img._storage.tensor_columns['tensor'].tensor, tf.Tensor) + + +@pytest.mark.tensorflow +def test_get_from_slice_stacked(): + class Doc(BaseDoc): + text: str + tensor: TensorFlowTensor + + da = DocVec[Doc]( + [Doc(text=f'hello{i}', tensor=tf.zeros((3, 224, 224))) for i in range(10)] + ) + + da_sliced = da[0:10:2] + assert isinstance(da_sliced, DocVec) + + tensors = da_sliced.tensor.tensor + assert tensors.shape == (5, 3, 224, 224) + + +@pytest.mark.tensorflow +def test_stack_none(): + class MyDoc(BaseDoc): + tensor: Optional[AnyTensor] = None + + da = DocVec[MyDoc]( + [MyDoc(tensor=None) for _ in range(10)], tensor_type=TensorFlowTensor + ) + assert 'tensor' in da._storage.tensor_columns.keys() + + +@pytest.mark.tensorflow +def test_keep_dtype_tf(): + class MyDoc(BaseDoc): + tensor: TensorFlowTensor + + da = DocList[MyDoc]( + [MyDoc(tensor=tf.zeros([2, 4], dtype=tf.int32)) for _ in range(3)] + ) + assert da[0].tensor.tensor.dtype == tf.int32 + + da = da.to_doc_vec() + assert da[0].tensor.tensor.dtype == tf.int32 + assert da.tensor.tensor.dtype == tf.int32 diff --git a/tests/units/array/stack/test_init.py b/tests/units/array/stack/test_init.py new file mode 100644 index 00000000000..232c9276002 --- /dev/null +++ b/tests/units/array/stack/test_init.py @@ -0,0 +1,48 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np + +from docarray import BaseDoc +from docarray.array.doc_vec.doc_vec import DocVec +from docarray.typing import AnyTensor, NdArray + + +def test_da_init(): + class MyDoc(BaseDoc): + tensor: AnyTensor + name: str + + docs = [MyDoc(tensor=np.zeros(10), name='hello') for _ in range(4)] + + da = DocVec[MyDoc](docs, tensor_type=NdArray) + + assert (da._storage.tensor_columns['tensor'] == np.zeros((4, 10))).all() + assert da._storage.any_columns['name'] == ['hello' for _ in range(4)] + + +def test_da_iter(): + class MyDoc(BaseDoc): + tensor: AnyTensor + name: str + + docs = [MyDoc(tensor=i * np.zeros((10, 10)), name=f'hello{i}') for i in range(4)] + + da = DocVec[MyDoc](docs, tensor_type=NdArray) + + for i, doc in enumerate(da): + assert isinstance(doc, MyDoc) + assert (doc.tensor == i * np.zeros((10, 10))).all() + assert doc.name == f'hello{i}' diff --git a/tests/units/array/stack/test_proto.py b/tests/units/array/stack/test_proto.py new file mode 100644 index 00000000000..d46766cde30 --- /dev/null +++ b/tests/units/array/stack/test_proto.py @@ -0,0 +1,348 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +from typing import Dict, Optional, Union + +import numpy as np +import pytest +import torch + +from docarray import BaseDoc, DocList +from docarray.array import DocVec +from docarray.typing import NdArray, TorchTensor + + +@pytest.fixture() +def batch(): + class Image(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + batch = DocList[Image]([Image(tensor=torch.zeros(3, 224, 224)) for _ in range(10)]) + + return batch.to_doc_vec() + + +@pytest.mark.proto +def test_proto_stacked_mode_torch(batch): + batch.from_protobuf(batch.to_protobuf()) + + +@pytest.mark.proto +def test_proto_stacked_mode_numpy(): + class MyDoc(BaseDoc): + tensor: NdArray[3, 224, 224] + + da = DocList[MyDoc]([MyDoc(tensor=np.zeros((3, 224, 224))) for _ in range(10)]) + + da = da.to_doc_vec() + + da.from_protobuf(da.to_protobuf()) + + +@pytest.mark.proto +def test_stacked_proto(): + class CustomDocument(BaseDoc): + image: NdArray + + da = DocList[CustomDocument]( + [CustomDocument(image=np.zeros((3, 224, 224))) for _ in range(10)] + ).to_doc_vec() + + da2 = DocVec[CustomDocument].from_protobuf(da.to_protobuf()) + + assert isinstance(da2, DocVec) + assert da.doc_type == da2.doc_type + assert (da2.image == da.image).all() + + +@pytest.mark.proto +def test_proto_none_tensor_column(): + class MyOtherDoc(BaseDoc): + embedding: Union[NdArray, None] = None + other_embedding: NdArray + third_embedding: Union[NdArray, None] = None + + da = DocVec[MyOtherDoc]( + [ + MyOtherDoc( + other_embedding=np.random.random(512), + ), + MyOtherDoc(other_embedding=np.random.random(512)), + ] + ) + assert da._storage.tensor_columns['embedding'] is None + assert da._storage.tensor_columns['other_embedding'] is not None + assert da._storage.tensor_columns['third_embedding'] is None + + proto = da.to_protobuf() + da_after = DocVec[MyOtherDoc].from_protobuf(proto) + + assert da_after._storage.tensor_columns['embedding'] is None + assert da_after._storage.tensor_columns['other_embedding'] is not None + assert ( + da_after._storage.tensor_columns['other_embedding'] + == da._storage.tensor_columns['other_embedding'] + ).all() + assert da_after._storage.tensor_columns['third_embedding'] is None + + +@pytest.mark.proto +def test_proto_none_doc_column(): + class InnerDoc(BaseDoc): + embedding: NdArray + + class MyDoc(BaseDoc): + inner: Union[InnerDoc, None] = None + other_inner: Union[InnerDoc, None] = None + + da = DocVec[MyDoc]( + [ + MyDoc(other_inner=InnerDoc(embedding=np.random.random(512))), + MyDoc(other_inner=InnerDoc(embedding=np.random.random(512))), + ] + ) + assert da._storage.doc_columns['inner'] is None + assert len(da._storage.doc_columns['other_inner']) == 2 + + proto = da.to_protobuf() + da_after = DocVec[MyDoc].from_protobuf(proto) + + assert da_after._storage.doc_columns['inner'] is None + assert len(da._storage.doc_columns['other_inner']) == 2 + assert (da.other_inner.embedding == da_after.other_inner.embedding).all() + + +@pytest.mark.proto +def test_proto_none_docvec_column(): + class InnerDoc(BaseDoc): + embedding: NdArray + + class MyDoc(BaseDoc): + inner_l: Union[DocList[InnerDoc], None] = None + inner_v: Union[DocVec[InnerDoc], None] = None + inner_exists_v: Union[DocVec[InnerDoc], None] = None + inner_exists_l: Union[DocList[InnerDoc], None] = None + + def _make_inner_list(): + return DocList[InnerDoc]( + [ + InnerDoc(embedding=np.random.random(512)), + InnerDoc(embedding=np.random.random(512)), + ] + ) + + da = DocVec[MyDoc]( + [ + MyDoc( + inner_exists_l=_make_inner_list(), + inner_exists_v=_make_inner_list().to_doc_vec(), + ), + MyDoc( + inner_exists_l=_make_inner_list(), + inner_exists_v=_make_inner_list().to_doc_vec(), + ), + ] + ) + assert da._storage.docs_vec_columns['inner_l'] is None + assert da._storage.docs_vec_columns['inner_v'] is None + assert len(da._storage.docs_vec_columns['inner_exists_l']) == 2 + assert len(da._storage.docs_vec_columns['inner_exists_v']) == 2 + assert da.inner_exists_l[0].embedding.shape == (2, 512) + assert da.inner_exists_l[1].embedding.shape == (2, 512) + assert da.inner_exists_v[0].embedding.shape == (2, 512) + assert da.inner_exists_v[1].embedding.shape == (2, 512) + + proto = da.to_protobuf() + da_after = DocVec[MyDoc].from_protobuf(proto) + + assert da_after._storage.docs_vec_columns['inner_l'] is None + assert da_after._storage.docs_vec_columns['inner_v'] is None + assert len(da._storage.docs_vec_columns['inner_exists_l']) == 2 + assert len(da._storage.docs_vec_columns['inner_exists_v']) == 2 + assert ( + da.inner_exists_l[0].embedding == da_after.inner_exists_l[0].embedding + ).all() + assert ( + da.inner_exists_l[1].embedding == da_after.inner_exists_l[1].embedding + ).all() + assert ( + da.inner_exists_v[0].embedding == da_after.inner_exists_v[0].embedding + ).all() + assert ( + da.inner_exists_v[1].embedding == da_after.inner_exists_v[1].embedding + ).all() + + +@pytest.mark.proto +def test_proto_any_column(): + class MyDoc(BaseDoc): + embedding: NdArray + text: str + d: Dict + + da = DocVec[MyDoc]( + [ + MyDoc( + embedding=np.random.random(512), + text='hi', + d={'a': 1}, + ), + MyDoc(embedding=np.random.random(512), text='there', d={'b': 2}), + ] + ) + assert da._storage.tensor_columns['embedding'].shape == (2, 512) + assert da._storage.any_columns['text'] == ['hi', 'there'] + assert da._storage.any_columns['d'] == [{'a': 1}, {'b': 2}] + + proto = da.to_protobuf() + da_after = DocVec[MyDoc].from_protobuf(proto) + + assert da_after.doc_type == da.doc_type + assert da._storage.tensor_columns['embedding'].shape == (2, 512) + assert ( + da_after._storage.tensor_columns['embedding'] + == da._storage.tensor_columns['embedding'] + ).all() + assert da._storage.any_columns['text'] == ['hi', 'there'] + assert da._storage.any_columns['d'] == [{'a': 1}, {'b': 2}] + + assert (da_after.embedding == da.embedding).all() + assert da_after.text == da.text + assert da_after.d == da.d + + +@pytest.mark.proto +def test_proto_none_any_column(): + class MyDoc(BaseDoc): + text: Optional[str] = None + d: Optional[Dict] = None + + da = DocVec[MyDoc]( + [ + MyDoc(), + MyDoc(), + ] + ) + assert da._storage.any_columns['text'] == [None, None] + assert da._storage.any_columns['d'] == [None, None] + + proto = da.to_protobuf() + da_after = DocVec[MyDoc].from_protobuf(proto) + + assert da_after._storage.any_columns['text'] == [None, None] + assert da_after._storage.any_columns['d'] == [None, None] + + +@pytest.mark.skipif('GITHUB_WORKFLOW' in os.environ, reason='Flaky in Github') +@pytest.mark.proto +@pytest.mark.parametrize('tensor_type', [NdArray, TorchTensor]) +def test_proto_tensor_type(tensor_type): + class InnerDoc(BaseDoc): + embedding: tensor_type + + class MyDoc(BaseDoc): + tensor: tensor_type + inner: InnerDoc + inner_v: DocVec[InnerDoc] + + def _get_rand_tens(): + arr = np.random.random(512) + return tensor_type.from_ndarray(arr) if tensor_type == TorchTensor else arr + + da = DocVec[MyDoc]( + [ + MyDoc( + tensor=_get_rand_tens(), + inner=InnerDoc(embedding=_get_rand_tens()), + inner_v=DocVec[InnerDoc]([InnerDoc(embedding=_get_rand_tens())]), + ), + MyDoc( + tensor=_get_rand_tens(), + inner=InnerDoc(embedding=_get_rand_tens()), + inner_v=DocVec[InnerDoc]([InnerDoc(embedding=_get_rand_tens())]), + ), + ] + ) + assert isinstance(da.tensor, tensor_type) + assert da.tensor.shape == (2, 512) + assert isinstance(da.inner.embedding, tensor_type) + assert da.inner.embedding.shape == (2, 512) + assert isinstance(da.inner_v[0].embedding, tensor_type) + assert da.inner_v[0].embedding.shape == (1, 512) + + proto = da.to_protobuf() + da_after = DocVec[MyDoc].from_protobuf(proto, tensor_type=tensor_type) + + assert isinstance(da_after.tensor, tensor_type) + assert (da.tensor == da_after.tensor).all() + assert isinstance(da_after.inner.embedding, tensor_type) + assert (da.inner.embedding == da_after.inner.embedding).all() + assert isinstance(da_after.inner_v[0].embedding, tensor_type) + assert (da.inner_v[0].embedding == da_after.inner_v[0].embedding).all() + + +@pytest.mark.tensorflow +def test_proto_tensor_type_tf(): + import tensorflow as tf + + from docarray.typing import TensorFlowTensor + + class InnerDoc(BaseDoc): + embedding: TensorFlowTensor + + class MyDoc(BaseDoc): + tensor: TensorFlowTensor + inner: InnerDoc + inner_v: DocVec[InnerDoc] + + def _get_rand_tens(): + arr = np.random.random(512) + return TensorFlowTensor.from_ndarray(arr) + + da = DocVec[MyDoc]( + [ + MyDoc( + tensor=_get_rand_tens(), + inner=InnerDoc(embedding=_get_rand_tens()), + inner_v=DocVec[InnerDoc]([InnerDoc(embedding=_get_rand_tens())]), + ), + MyDoc( + tensor=_get_rand_tens(), + inner=InnerDoc(embedding=_get_rand_tens()), + inner_v=DocVec[InnerDoc]([InnerDoc(embedding=_get_rand_tens())]), + ), + ] + ) + assert isinstance(da.tensor, TensorFlowTensor) + assert len(da.tensor) == 2 + assert isinstance(da.inner.embedding, TensorFlowTensor) + assert len(da.inner.embedding) == 2 + assert isinstance(da.inner_v[0].embedding, TensorFlowTensor) + assert len(da.inner_v[0].embedding) == 1 + + proto = da.to_protobuf() + da_after = DocVec[MyDoc].from_protobuf(proto, tensor_type=TensorFlowTensor) + + assert isinstance(da_after.tensor, TensorFlowTensor) + assert tf.math.reduce_all(tf.equal(da.tensor.tensor, da_after.tensor.tensor)) + assert isinstance(da_after.inner.embedding, TensorFlowTensor) + assert tf.math.reduce_all( + tf.equal(da.inner.embedding.tensor, da_after.inner.embedding.tensor) + ) + assert isinstance(da_after.inner_v[0].embedding, TensorFlowTensor) + assert tf.math.reduce_all( + tf.equal(da.inner_v[0].embedding.tensor, da_after.inner_v[0].embedding.tensor) + ) diff --git a/tests/units/array/test_array.py b/tests/units/array/test_array.py new file mode 100644 index 00000000000..8e51cc1c37e --- /dev/null +++ b/tests/units/array/test_array.py @@ -0,0 +1,524 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Optional, TypeVar, Union + +import numpy as np +import pytest +import torch +from pydantic import parse_obj_as + +from docarray import BaseDoc, DocList +from docarray.typing import ImageUrl, NdArray, TorchTensor +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + + from docarray.typing import TensorFlowTensor + + +@pytest.fixture() +def da(): + class Text(BaseDoc): + text: str + + return DocList[Text]([Text(text=f'hello {i}') for i in range(10)]) + + +def test_iterate(da): + for doc, doc2 in zip(da, da): + assert doc.id == doc2.id + + +def test_append(): + class Text(BaseDoc): + text: str + + da = DocList[Text]([]) + + da.append(Text(text='hello', id='1')) + + assert len(da) == 1 + assert da[0].id == '1' + + +def test_extend(): + class Text(BaseDoc): + text: str + + da = DocList[Text]([Text(text='hello', id=str(i)) for i in range(10)]) + + da.extend([Text(text='hello', id=str(10 + i)) for i in range(10)]) + + assert len(da) == 20 + for da, i in zip(da, range(20)): + assert da.id == str(i) + + +def test_extend_itself(): + class Text(BaseDoc): + text: str + + da = DocList[Text]([Text(text='hello', id=str(i)) for i in range(10)]) + + da.extend(da) + + assert len(da) == 20 + + +def test_slice(da): + da2 = da[0:5] + assert type(da2) == da.__class__ + assert len(da2) == 5 + + +def test_document_array(): + class Text(BaseDoc): + text: str + + da = DocList([Text(text='hello') for _ in range(10)]) + + assert len(da) == 10 + + +def test_empty_array(): + da = DocList() + len(da) == 0 + + +def test_document_array_fixed_type(): + class Text(BaseDoc): + text: str + + da = DocList[Text]([Text(text='hello') for _ in range(10)]) + + assert len(da) == 10 + + +def test_ndarray_equality(): + class Text(BaseDoc): + tensor: NdArray + + arr1 = Text(tensor=np.zeros(5)) + arr2 = Text(tensor=np.zeros(5)) + arr3 = Text(tensor=np.ones(5)) + arr4 = Text(tensor=np.zeros(4)) + + assert arr1 == arr2 + assert arr1 != arr3 + assert arr1 != arr4 + + +def test_tensor_equality(): + class Text(BaseDoc): + tensor: TorchTensor + + torch1 = Text(tensor=torch.zeros(128)) + torch2 = Text(tensor=torch.zeros(128)) + torch3 = Text(tensor=torch.zeros(126)) + torch4 = Text(tensor=torch.ones(128)) + + assert torch1 == torch2 + assert torch1 != torch3 + assert torch1 != torch4 + + +def test_documentarray(): + class Text(BaseDoc): + text: str + + da1 = DocList([Text(text='hello')]) + da2 = DocList([Text(text='hello')]) + + assert da1 == da2 + assert da1 == [Text(text='hello') for _ in range(len(da1))] + assert da2 == [Text(text='hello') for _ in range(len(da2))] + + +@pytest.mark.tensorflow +def test_tensorflowtensor_equality(): + class Text(BaseDoc): + tensor: TensorFlowTensor + + tensor1 = Text(tensor=tf.constant([1, 2, 3, 4, 5, 6])) + tensor2 = Text(tensor=tf.constant([1, 2, 3, 4, 5, 6])) + tensor3 = Text(tensor=tf.constant([[1.0, 2.0], [3.0, 5.0]])) + tensor4 = Text(tensor=tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])) + + assert tensor1 == tensor2 + assert tensor1 != tensor3 + assert tensor1 != tensor4 + + +def test_text_tensor(): + class Text1(BaseDoc): + tensor: NdArray + + class Text2(BaseDoc): + tensor: TorchTensor + + arr_tensor1 = Text1(tensor=np.zeros(2)) + arr_tensor2 = Text2(tensor=torch.zeros(2)) + + assert arr_tensor1 == arr_tensor2 + + +def test_get_bulk_attributes_function(): + class Mmdoc(BaseDoc): + text: str + tensor: NdArray + + N = 10 + + da = DocList[Mmdoc]( + (Mmdoc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N)) + ) + + tensors = da._get_data_column('tensor') + + assert len(tensors) == N + for tensor in tensors: + assert tensor.shape == (3, 224, 224) + + texts = da._get_data_column('text') + + assert len(texts) == N + for i, text in enumerate(texts): + assert text == f'hello{i}' + + +def test_set_attributes(): + class InnerDoc(BaseDoc): + text: str + + class Mmdoc(BaseDoc): + inner: InnerDoc + + N = 10 + + da = DocList[Mmdoc]((Mmdoc(inner=InnerDoc(text=f'hello{i}')) for i in range(N))) + + list_docs = [InnerDoc(text=f'hello{i}') for i in range(N)] + da._set_data_column('inner', list_docs) + + for doc, list_doc in zip(da, list_docs): + assert doc.inner == list_doc + + +def test_get_bulk_attributes(): + class Mmdoc(BaseDoc): + text: str + tensor: NdArray + + N = 10 + + da = DocList[Mmdoc]( + (Mmdoc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N)) + ) + + tensors = da.tensor + + assert len(tensors) == N + for tensor in tensors: + assert tensor.shape == (3, 224, 224) + + texts = da.text + + assert len(texts) == N + for i, text in enumerate(texts): + assert text == f'hello{i}' + + +def test_get_bulk_attributes_document(): + class InnerDoc(BaseDoc): + text: str + + class Mmdoc(BaseDoc): + inner: InnerDoc + + N = 10 + + da = DocList[Mmdoc]((Mmdoc(inner=InnerDoc(text=f'hello{i}')) for i in range(N))) + + assert isinstance(da.inner, DocList) + + +def test_get_bulk_attributes_optional_type(): + class Mmdoc(BaseDoc): + text: str + tensor: Optional[NdArray] + + N = 10 + + da = DocList[Mmdoc]( + (Mmdoc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N)) + ) + + tensors = da.tensor + + assert len(tensors) == N + for tensor in tensors: + assert tensor.shape == (3, 224, 224) + + texts = da.text + + assert len(texts) == N + for i, text in enumerate(texts): + assert text == f'hello{i}' + + +def test_get_bulk_attributes_union_type(): + class Mmdoc(BaseDoc): + text: str + tensor: Union[NdArray, TorchTensor] + + N = 10 + + da = DocList[Mmdoc]( + (Mmdoc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N)) + ) + + tensors = da.tensor + + assert len(tensors) == N + assert isinstance(tensors, list) + for tensor in tensors: + assert tensor.shape == (3, 224, 224) + + texts = da.text + + assert len(texts) == N + for i, text in enumerate(texts): + assert text == f'hello{i}' + + +@pytest.mark.tensorflow +def test_get_bulk_attributes_union_type_nested(): + class MyDoc(BaseDoc): + embedding: Union[Optional[TorchTensor], Optional[NdArray]] + embedding2: Optional[Union[TorchTensor, NdArray, TensorFlowTensor]] + embedding3: Optional[Optional[TorchTensor]] + embedding4: Union[ + Optional[Union[TorchTensor, NdArray, TensorFlowTensor]], TorchTensor + ] + + da = DocList[MyDoc]( + [ + MyDoc( + embedding=torch.rand(10), + embedding2=torch.rand(10), + embedding3=torch.rand(10), + embedding4=torch.rand(10), + ) + for _ in range(10) + ] + ) + + for attr in ['embedding', 'embedding2', 'embedding3', 'embedding4']: + tensors = getattr(da, attr) + assert len(tensors) == 10 + assert isinstance(tensors, list) + for tensor in tensors: + assert tensor.shape == (10,) + + +def test_get_from_slice(): + class Doc(BaseDoc): + text: str + tensor: NdArray + + N = 10 + + da = DocList[Doc]( + (Doc(text=f'hello{i}', tensor=np.zeros((3, 224, 224))) for i in range(N)) + ) + + da_sliced = da[0:10:2] + assert isinstance(da_sliced, DocList) + + tensors = da_sliced.tensor + assert len(tensors) == 5 + for tensor in tensors: + assert tensor.shape == (3, 224, 224) + + texts = da_sliced.text + assert len(texts) == 5 + for i, text in enumerate(texts): + assert text == f'hello{i * 2}' + + +def test_del_item(da): + assert len(da) == 10 + del da[2] + assert len(da) == 9 + assert da.text == [ + 'hello 0', + 'hello 1', + 'hello 3', + 'hello 4', + 'hello 5', + 'hello 6', + 'hello 7', + 'hello 8', + 'hello 9', + ] + del da[0:2] + assert len(da) == 7 + assert da.text == [ + 'hello 3', + 'hello 4', + 'hello 5', + 'hello 6', + 'hello 7', + 'hello 8', + 'hello 9', + ] + + +def test_generic_type_var(): + T = TypeVar('T', bound=BaseDoc) + + def f(a: DocList[T]) -> DocList[T]: + return a + + def g(a: DocList['BaseDoc']) -> DocList['BaseDoc']: + return a + + a = DocList() + f(a) + g(a) + + +def test_construct(): + class Text(BaseDoc): + text: str + + docs = [Text(text=f'hello {i}') for i in range(10)] + [BaseDoc()] + + da = DocList[Text].construct(docs) + + assert type(da[-1]) == BaseDoc + + +def test_reverse(): + class Text(BaseDoc): + text: str + + docs = [Text(text=f'hello {i}') for i in range(10)] + + da = DocList[Text](docs) + da.reverse() + assert da[-1].text == 'hello 0' + assert da[0].text == 'hello 9' + + +class Image(BaseDoc): + tensor: Optional[NdArray] = None + url: ImageUrl + + +def test_remove(): + images = [Image(url=f'http://url.com/foo_{i}.png') for i in range(3)] + da = DocList[Image](images) + da.remove(images[1]) + assert len(da) == 2 + assert da[0] == images[0] + assert da[1] == images[2] + + +def test_pop(): + images = [Image(url=f'http://url.com/foo_{i}.png') for i in range(3)] + da = DocList[Image](images) + popped = da.pop(1) + assert len(da) == 2 + assert popped == images[1] + assert da[0] == images[0] + assert da[1] == images[2] + + +def test_sort(): + images = [ + Image(url=f'http://url.com/foo_{i}.png', tensor=NdArray(i)) for i in [2, 0, 1] + ] + da = DocList[Image](images) + da.sort(key=lambda img: len(img.tensor)) + assert len(da) == 3 + assert da[0].url == 'http://url.com/foo_0.png' + assert da[1].url == 'http://url.com/foo_1.png' + + +def test_optional_field(): + from typing import Optional + + from docarray import BaseDoc, DocList + from docarray.typing import ImageUrl, NdArray + + class Nested(BaseDoc): + tensor: NdArray + + class Image(BaseDoc): + url: ImageUrl + features: Optional[Nested] = None + + docs = DocList[Image]([Image(url='http://url.com/foo.png') for _ in range(10)]) + + assert docs.features == [None for _ in range(10)] + assert isinstance(docs.features, list) + assert not isinstance(docs.features, DocList) + + +def test_validate_list_dict(): + images = [ + dict(url=f'http://url.com/foo_{i}.png', tensor=NdArray(i)) for i in [2, 0, 1] + ] + + # docs = DocList[Image]([Image(url=image['url'], tensor=image['tensor']) for image in images]) + + docs = parse_obj_as(DocList[Image], images) + + assert docs.url == [ + 'http://url.com/foo_2.png', + 'http://url.com/foo_0.png', + 'http://url.com/foo_1.png', + ] + + +def test_legacy_doc(): + from docarray.documents.legacy import LegacyDocument + + newDoc = LegacyDocument() + da = DocList[LegacyDocument]([newDoc]) + da.summary() + + +def test_parameterize_list(): + from docarray import DocList, BaseDoc + + with pytest.raises(TypeError) as excinfo: + da = DocList[BaseDoc()] + assert da is None + + assert str(excinfo.value) == 'Expecting a type, got object instead' + + +def test_not_double_subcriptable(): + from docarray import DocList + from docarray.documents import TextDoc + + with pytest.raises(TypeError) as excinfo: + da = DocList[TextDoc][TextDoc] + assert da is None diff --git a/tests/units/array/test_array_from_to_bytes.py b/tests/units/array/test_array_from_to_bytes.py new file mode 100644 index 00000000000..0ab952ce4a7 --- /dev/null +++ b/tests/units/array/test_array_from_to_bytes.py @@ -0,0 +1,148 @@ +import pytest + +from docarray import BaseDoc, DocList, DocVec +from docarray.documents import ImageDoc +from docarray.typing import NdArray, TorchTensor + + +class MyDoc(BaseDoc): + embedding: NdArray + text: str + image: ImageDoc + + +@pytest.mark.parametrize( + 'protocol', ['pickle-array', 'protobuf-array', 'protobuf', 'pickle'] +) +@pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) +@pytest.mark.parametrize('show_progress', [False, True]) +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_to_bytes(protocol, compress, show_progress, array_cls): + da = array_cls[MyDoc]( + [ + MyDoc( + embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png') + ), + MyDoc(embedding=[5, 4, 3, 2, 1], text='hello world', image=ImageDoc()), + ] + ) + bytes_da = da.to_bytes( + protocol=protocol, compress=compress, show_progress=show_progress + ) + da2 = array_cls[MyDoc].from_bytes( + bytes_da, protocol=protocol, compress=compress, show_progress=show_progress + ) + assert len(da2) == 2 + assert len(da) == len(da2) + for d1, d2 in zip(da, da2): + assert d1.embedding.tolist() == d2.embedding.tolist() + assert d1.text == d2.text + assert d1.image.url == d2.image.url + assert da[1].image.url is None + assert da2[1].image.url is None + + +@pytest.mark.parametrize( + 'protocol', ['pickle-array', 'protobuf-array', 'protobuf', 'pickle'] +) +@pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) +@pytest.mark.parametrize('show_progress', [False, True]) # [False, True]) +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_to_base64(protocol, compress, show_progress, array_cls): + da = array_cls[MyDoc]( + [ + MyDoc( + embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png') + ), + MyDoc(embedding=[5, 4, 3, 2, 1], text='hello world', image=ImageDoc()), + ] + ) + bytes_da = da.to_base64( + protocol=protocol, compress=compress, show_progress=show_progress + ) + da2 = array_cls[MyDoc].from_base64( + bytes_da, protocol=protocol, compress=compress, show_progress=show_progress + ) + assert len(da2) == 2 + assert len(da) == len(da2) + for d1, d2 in zip(da, da2): + assert d1.embedding.tolist() == d2.embedding.tolist() + assert d1.text == d2.text + assert d1.image.url == d2.image.url + + assert da[1].image.url is None + assert da2[1].image.url is None + + +# test_from_to_base64('protobuf', 'lz4', False, DocVec) +class MyTensorTypeDocNdArray(BaseDoc): + embedding: NdArray + text: str + image: ImageDoc + + +class MyTensorTypeDocTorchTensor(BaseDoc): + embedding: TorchTensor + text: str + image: ImageDoc + + +@pytest.mark.parametrize( + 'doc_type, tensor_type', + [(MyTensorTypeDocNdArray, NdArray), (MyTensorTypeDocTorchTensor, TorchTensor)], +) +@pytest.mark.parametrize('protocol', ['protobuf-array', 'pickle-array']) +def test_from_to_base64_tensor_type(doc_type, tensor_type, protocol): + da = DocVec[doc_type]( + [ + doc_type( + embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png') + ), + doc_type(embedding=[5, 4, 3, 2, 1], text='hello world', image=ImageDoc()), + ], + tensor_type=tensor_type, + ) + bytes_da = da.to_base64(protocol=protocol) + da2 = DocVec[doc_type].from_base64( + bytes_da, tensor_type=tensor_type, protocol=protocol + ) + assert da2.tensor_type == tensor_type + assert isinstance(da2.embedding, tensor_type) + + +@pytest.mark.parametrize('tensor_type', [NdArray, TorchTensor]) +def test_from_to_bytes_tensor_type(tensor_type): + da = DocVec[MyDoc]( + [ + MyDoc( + embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png') + ), + MyDoc(embedding=[5, 4, 3, 2, 1], text='hello world', image=ImageDoc()), + ], + tensor_type=tensor_type, + ) + bytes_da = da.to_bytes() + da2 = DocVec[MyDoc].from_bytes(bytes_da, tensor_type=tensor_type) + assert da2.tensor_type == tensor_type + assert isinstance(da2.embedding, tensor_type) + + +def test_union_type_error(tmp_path): + from typing import Union + + from docarray.documents import TextDoc + + class CustomDoc(BaseDoc): + ud: Union[TextDoc, ImageDoc] = TextDoc(text='union type') + + docs = DocList[CustomDoc]([CustomDoc(ud=TextDoc(text='union type'))]) + + with pytest.raises(ValueError): + docs.from_bytes(docs.to_bytes()) + + class BasisUnion(BaseDoc): + ud: Union[int, str] + + docs_basic = DocList[BasisUnion]([BasisUnion(ud="hello")]) + docs_copy = DocList[BasisUnion].from_bytes(docs_basic.to_bytes()) + assert docs_copy == docs_basic diff --git a/tests/units/array/test_array_from_to_csv.py b/tests/units/array/test_array_from_to_csv.py new file mode 100644 index 00000000000..07d353ffc0f --- /dev/null +++ b/tests/units/array/test_array_from_to_csv.py @@ -0,0 +1,166 @@ +import os +from typing import Optional + +import pytest + +from docarray import BaseDoc, DocList, DocVec +from docarray.documents import ImageDoc +from tests import TOYDATA_DIR + + +@pytest.fixture() +def nested_doc_cls(): + class MyDoc(BaseDoc): + count: Optional[int] = None + text: str + + class MyDocNested(MyDoc): + image: ImageDoc + image2: ImageDoc + + return MyDocNested + + +def test_to_from_csv(tmpdir, nested_doc_cls): + da = DocList[nested_doc_cls]( + [ + nested_doc_cls( + count=0, + text='hello', + image=ImageDoc(url='aux.png'), + image2=ImageDoc(url='aux.png'), + ), + nested_doc_cls(text='hello world', image=ImageDoc(), image2=ImageDoc()), + ] + ) + tmp_file = str(tmpdir / 'tmp.csv') + da.to_csv(tmp_file) + assert os.path.isfile(tmp_file) + + da_from = DocList[nested_doc_cls].from_csv(tmp_file) + assert isinstance(da_from, DocList) + for doc1, doc2 in zip(da, da_from): + assert doc1 == doc2 + + +def test_from_csv_nested(nested_doc_cls): + da = DocList[nested_doc_cls].from_csv( + file_path=str(TOYDATA_DIR / 'docs_nested.csv') + ) + assert isinstance(da, DocList) + assert len(da) == 3 + + for i, doc in enumerate(da): + assert doc.count.__class__ == int + assert doc.count == int(f'{i}{i}{i}') + + assert doc.text.__class__ == str + assert doc.text == f'hello {i}' + + assert doc.image.__class__ == ImageDoc + assert doc.image.tensor is None + assert doc.image.embedding is None + assert doc.image.bytes_ is None + + assert doc.image2.__class__ == ImageDoc + assert doc.image2.tensor is None + assert doc.image2.embedding is None + assert doc.image2.bytes_ is None + + assert da[0].image2.url == 'image_10.png' + assert da[1].image2.url is None + assert da[2].image2.url is None + + +@pytest.fixture() +def nested_doc(): + class Inner(BaseDoc): + img: Optional[ImageDoc] = None + + class Middle(BaseDoc): + img: Optional[ImageDoc] = None + inner: Optional[Inner] = None + + class Outer(BaseDoc): + img: Optional[ImageDoc] = None + middle: Optional[Middle] = None + + doc = Outer( + img=ImageDoc(), middle=Middle(img=ImageDoc(), inner=Inner(img=ImageDoc())) + ) + return doc + + +def test_from_csv_without_schema_raise_exception(): + with pytest.raises(TypeError, match='no document schema defined'): + DocList.from_csv(file_path=str(TOYDATA_DIR / 'docs_nested.csv')) + + +def test_from_csv_with_wrong_schema_raise_exception(nested_doc): + with pytest.raises(ValueError, match='Column names do not match the schema'): + DocList[nested_doc.__class__].from_csv(file_path=str(TOYDATA_DIR / 'docs.csv')) + + +def test_from_remote_csv_file(): + remote_url = 'https://github.com/docarray/docarray/blob/main/tests/toydata/books.csv?raw=true' + + class Book(BaseDoc): + title: str + author: str + year: int + + books = DocList[Book].from_csv(file_path=remote_url) + assert isinstance(books, DocList) + + assert len(books) == 3 + + +def test_doc_list_error(tmpdir): + class Book(BaseDoc): + title: str + + # not testing DocVec bc it already fails here (as it should!) + docs = DocList([Book(title='hello'), Book(title='world')]) + tmp_file = str(tmpdir / 'tmp.csv') + with pytest.raises(TypeError): + docs.to_csv(tmp_file) + + +def test_union_type_error(tmp_path): + from typing import Union + + from docarray.documents import TextDoc + + class CustomDoc(BaseDoc): + ud: Union[TextDoc, ImageDoc] = TextDoc(text='union type') + + docs = DocList[CustomDoc]([CustomDoc(ud=TextDoc(text='union type'))]) + + with pytest.raises(ValueError): + docs.to_csv(str(tmp_path) + ".csv") + DocList[CustomDoc].from_csv(str(tmp_path) + ".csv") + + class BasisUnion(BaseDoc): + ud: Union[int, str] + + docs_basic = DocList[BasisUnion]([BasisUnion(ud="hello")]) + docs_basic.to_csv(str(tmp_path) + ".csv") + docs_copy = DocList[BasisUnion].from_csv(str(tmp_path) + ".csv") + assert docs_copy == docs_basic + + +def test_to_from_csv_docvec_raises(): + class Book(BaseDoc): + title: str + author: str + year: int + + books = DocVec[Book]( + [Book(title='It\'s me, hi', author='I\'m the problem it\'s me', year=2022)] + ) + + with pytest.raises(NotImplementedError): + books.to_csv('dummy/file/path') + + with pytest.raises(NotImplementedError): + DocVec[Book].from_csv('dummy/file/path') diff --git a/tests/units/array/test_array_from_to_json.py b/tests/units/array/test_array_from_to_json.py new file mode 100644 index 00000000000..2324652c6d0 --- /dev/null +++ b/tests/units/array/test_array_from_to_json.py @@ -0,0 +1,180 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Optional, Dict, List + +import numpy as np +import pytest +import torch + +from docarray import BaseDoc, DocList, DocVec +from docarray.documents import ImageDoc +from docarray.typing import NdArray, TorchTensor + + +class MyDoc(BaseDoc): + embedding: NdArray + text: str + image: ImageDoc + + +def test_from_to_json_doclist(): + da = DocList[MyDoc]( + [ + MyDoc( + embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png') + ), + MyDoc(embedding=[5, 4, 3, 2, 1], text='hello world', image=ImageDoc()), + ] + ) + json_da = da.to_json() + da2 = DocList[MyDoc].from_json(json_da) + assert len(da2) == 2 + assert len(da) == len(da2) + for d1, d2 in zip(da, da2): + assert d1.embedding.tolist() == d2.embedding.tolist() + assert d1.text == d2.text + assert d1.image.url == d2.image.url + assert da[1].image.url is None + assert da2[1].image.url is None + + +@pytest.mark.parametrize('tensor_type', [TorchTensor, NdArray]) +def test_from_to_json_docvec(tensor_type): + def generate_docs(tensor_type): + class InnerDoc(BaseDoc): + tens: tensor_type + + class MyDoc(BaseDoc): + text: str + num: Optional[int] = None + tens: tensor_type + tens_none: Optional[tensor_type] = None + inner: InnerDoc + inner_none: Optional[InnerDoc] = None + inner_vec: DocVec[InnerDoc] + inner_vec_none: Optional[DocVec[InnerDoc]] = None + + def _rand_vec_gen(tensor_type): + arr = np.random.rand(5) + if tensor_type == TorchTensor: + arr = torch.from_numpy(arr).to(torch.float32) + return arr + + inner = InnerDoc(tens=_rand_vec_gen(tensor_type)) + inner_vec = DocVec[InnerDoc]([inner, inner], tensor_type=tensor_type) + vec = DocVec[MyDoc]( + [ + MyDoc( + text=str(i), + num=None, + tens=_rand_vec_gen(tensor_type), + inner=inner, + inner_none=None, + inner_vec=inner_vec, + inner_vec_none=None, + ) + for i in range(5) + ], + tensor_type=tensor_type, + ) + return vec + + v = generate_docs(tensor_type) + json_str = v.to_json() + + v_after = DocVec[v.doc_type].from_json(json_str, tensor_type=tensor_type) + + assert v_after.tensor_type == v.tensor_type + assert set(v_after._storage.columns.keys()) == set(v._storage.columns.keys()) + assert v_after._storage == v._storage + + +@pytest.mark.tensorflow +def test_from_to_json_docvec_tf(): + from docarray.typing import TensorFlowTensor + + def generate_docs(): + class InnerDoc(BaseDoc): + tens: TensorFlowTensor + + class MyDoc(BaseDoc): + text: str + num: Optional[int] = None + tens: TensorFlowTensor + tens_none: Optional[TensorFlowTensor] = None + inner: InnerDoc + inner_none: Optional[InnerDoc] = None + inner_vec: DocVec[InnerDoc] + inner_vec_none: Optional[DocVec[InnerDoc]] = None + + inner = InnerDoc(tens=np.random.rand(5)) + inner_vec = DocVec[InnerDoc]([inner, inner], tensor_type=TensorFlowTensor) + vec = DocVec[MyDoc]( + [ + MyDoc( + text=str(i), + num=None, + tens=np.random.rand(5), + inner=inner, + inner_none=None, + inner_vec=inner_vec, + inner_vec_none=None, + ) + for i in range(5) + ], + tensor_type=TensorFlowTensor, + ) + return vec + + v = generate_docs() + json_str = v.to_json() + + v_after = DocVec[v.doc_type].from_json(json_str, tensor_type=TensorFlowTensor) + + assert v_after.tensor_type == v.tensor_type + assert set(v_after._storage.columns.keys()) == set(v._storage.columns.keys()) + assert v_after._storage == v._storage + + +def test_union_type(): + from typing import Union + + from docarray.documents import TextDoc + + class CustomDoc(BaseDoc): + ud: Union[TextDoc, ImageDoc] = TextDoc(text='union type') + + docs = DocList[CustomDoc]([CustomDoc(ud=TextDoc(text='union type'))]) + + docs_copy = docs.from_json(docs.to_json()) + assert docs == docs_copy + + +@pytest.mark.parametrize('tensor_type', [NdArray, TorchTensor]) +def test_from_to_json_tensor_type(tensor_type): + da = DocVec[MyDoc]( + [ + MyDoc( + embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png') + ), + MyDoc(embedding=[5, 4, 3, 2, 1], text='hello world', image=ImageDoc()), + ], + tensor_type=tensor_type, + ) + json_da = da.to_json() + da2 = DocVec[MyDoc].from_json(json_da, tensor_type=tensor_type) + assert da2.tensor_type == tensor_type + assert isinstance(da2.embedding, tensor_type) diff --git a/tests/units/array/test_array_from_to_pandas.py b/tests/units/array/test_array_from_to_pandas.py new file mode 100644 index 00000000000..440398562ff --- /dev/null +++ b/tests/units/array/test_array_from_to_pandas.py @@ -0,0 +1,162 @@ +from typing import List, Optional + +import pandas as pd +import pytest + +from docarray import BaseDoc, DocList, DocVec +from docarray.documents import ImageDoc +from docarray.typing import NdArray, TorchTensor + + +@pytest.fixture() +def nested_doc_cls(): + class MyDoc(BaseDoc): + count: Optional[int] = None + text: str + + class MyDocNested(MyDoc): + image: ImageDoc + lst: List[str] + + return MyDocNested + + +@pytest.mark.parametrize('doc_vec', [False, True]) +def test_to_from_pandas_df(nested_doc_cls, doc_vec): + da = DocList[nested_doc_cls]( + [ + nested_doc_cls( + count=0, + text='hello', + image=ImageDoc(url='aux.png'), + lst=["hello", "world"], + ), + nested_doc_cls( + text='hello world', image=ImageDoc(), lst=["hello", "world"] + ), + ] + ) + if doc_vec: + da = da.to_doc_vec() + df = da.to_dataframe() + assert isinstance(df, pd.DataFrame) + assert len(df) == 2 + assert ( + df.columns + == [ + 'id', + 'count', + 'text', + 'image__id', + 'image__url', + 'image__tensor', + 'image__embedding', + 'image__bytes_', + 'lst', + ] + ).all() + + if doc_vec: + da_from_df = DocVec[nested_doc_cls].from_dataframe(df) + assert isinstance(da_from_df, DocVec) + else: + da_from_df = DocList[nested_doc_cls].from_dataframe(df) + assert isinstance(da_from_df, DocList) + for doc1, doc2 in zip(da, da_from_df): + assert doc1 == doc2 + + +@pytest.fixture() +def nested_doc(): + class Inner(BaseDoc): + img: Optional[ImageDoc] = None + + class Middle(BaseDoc): + img: Optional[ImageDoc] = None + inner: Optional[Inner] = None + + class Outer(BaseDoc): + img: Optional[ImageDoc] = None + middle: Optional[Middle] = None + + doc = Outer( + img=ImageDoc(), middle=Middle(img=ImageDoc(), inner=Inner(img=ImageDoc())) + ) + return doc + + +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_pandas_without_schema_raise_exception(array_cls): + with pytest.raises(TypeError, match='no document schema defined'): + df = pd.DataFrame( + columns=['title', 'count'], data=[['title 0', 0], ['title 1', 1]] + ) + array_cls.from_dataframe(df=df) + + +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_pandas_with_wrong_schema_raise_exception(nested_doc, array_cls): + with pytest.raises(ValueError, match='Column names do not match the schema'): + df = pd.DataFrame( + columns=['title', 'count'], data=[['title 0', 0], ['title 1', 1]] + ) + array_cls[nested_doc.__class__].from_dataframe(df=df) + + +def test_doc_list_error(): + class Book(BaseDoc): + title: str + + # not testing DocVec bc it already fails here (as it should!) + docs = DocList([Book(title='hello'), Book(title='world')]) + with pytest.raises(TypeError): + docs.to_dataframe() + + +@pytest.mark.proto +def test_union_type_error(): + from typing import Union + + from docarray.documents import TextDoc + + class CustomDoc(BaseDoc): + ud: Union[TextDoc, ImageDoc] = TextDoc(text='union type') + + docs = DocList[CustomDoc]([CustomDoc(ud=TextDoc(text='union type'))]) + + with pytest.raises(ValueError): + DocList[CustomDoc].from_dataframe(docs.to_dataframe()) + + class BasisUnion(BaseDoc): + ud: Union[int, str] + + docs_basic = DocList[BasisUnion]([BasisUnion(ud="hello")]) + docs_copy = DocList[BasisUnion].from_dataframe(docs_basic.to_dataframe()) + assert docs_copy == docs_basic + + +@pytest.mark.parametrize('tensor_type', [NdArray, TorchTensor]) +@pytest.mark.parametrize('tensor_len', [0, 5]) +def test_from_to_pandas_tensor_type(tensor_type, tensor_len): + class MyDoc(BaseDoc): + embedding: tensor_type + text: str + image: ImageDoc + + da = DocVec[MyDoc]( + [ + MyDoc( + embedding=list(range(tensor_len)), + text='hello', + image=ImageDoc(url='aux.png'), + ), + MyDoc( + embedding=list(range(tensor_len)), text='hello world', image=ImageDoc() + ), + ], + tensor_type=tensor_type, + ) + df_da = da.to_dataframe() + da2 = DocVec[MyDoc].from_dataframe(df_da, tensor_type=tensor_type) + assert da2.tensor_type == tensor_type + assert isinstance(da2.embedding, tensor_type) diff --git a/tests/units/array/test_array_proto.py b/tests/units/array/test_array_proto.py new file mode 100644 index 00000000000..8b6cc172725 --- /dev/null +++ b/tests/units/array/test_array_proto.py @@ -0,0 +1,167 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytest +from typing import Dict, List + +from docarray import BaseDoc, DocList +from docarray.base_doc import AnyDoc +from docarray.documents import ImageDoc, TextDoc +from docarray.typing import NdArray + + +@pytest.mark.proto +def test_simple_proto(): + class CustomDoc(BaseDoc): + text: str + tensor: NdArray + + da = DocList( + [CustomDoc(text='hello', tensor=np.zeros((3, 224, 224))) for _ in range(10)] + ) + + new_da = DocList[CustomDoc].from_protobuf(da.to_protobuf()) + + for doc1, doc2 in zip(da, new_da): + assert doc1.text == doc2.text + assert (doc1.tensor == doc2.tensor).all() + + +@pytest.mark.proto +def test_nested_proto(): + class CustomDocument(BaseDoc): + text: TextDoc + image: ImageDoc + + da = DocList[CustomDocument]( + [ + CustomDocument( + text=TextDoc(text='hello'), + image=ImageDoc(tensor=np.zeros((3, 224, 224))), + ) + for _ in range(10) + ] + ) + + DocList[CustomDocument].from_protobuf(da.to_protobuf()) + + +@pytest.mark.proto +def test_nested_proto_any_doc(): + class CustomDocument(BaseDoc): + text: TextDoc + image: ImageDoc + + da = DocList[CustomDocument]( + [ + CustomDocument( + text=TextDoc(text='hello'), + image=ImageDoc(tensor=np.zeros((3, 224, 224))), + ) + for _ in range(10) + ] + ) + + DocList.from_protobuf(da.to_protobuf()) + + +@pytest.mark.proto +def test_any_doc_list_proto(): + doc = AnyDoc(hello='world') + pt = DocList([doc]).to_protobuf() + docs = DocList.from_protobuf(pt) + assert docs[0].hello == 'world' + + +@pytest.mark.proto +def test_any_nested_doc_list_proto(): + from docarray import BaseDoc, DocList + + class TextDocWithId(BaseDoc): + id: str + text: str + + class ResultTestDoc(BaseDoc): + matches: DocList[TextDocWithId] + + index_da = DocList[TextDocWithId]( + [TextDocWithId(id=f'{i}', text=f'ID {i}') for i in range(10)] + ) + + out_da = DocList[ResultTestDoc]([ResultTestDoc(matches=index_da[0:2])]) + pb = out_da.to_protobuf() + docs = DocList.from_protobuf(pb) + assert docs[0].matches[0].id == '0' + assert len(docs[0].matches) == 2 + assert len(docs) == 1 + + +@pytest.mark.proto +def test_union_type_error(): + from typing import Union + + class CustomDoc(BaseDoc): + ud: Union[TextDoc, ImageDoc] = TextDoc(text='union type') + + docs = DocList[CustomDoc]([CustomDoc(ud=TextDoc(text='union type'))]) + + with pytest.raises(ValueError): + DocList[CustomDoc].from_protobuf(docs.to_protobuf()) + + class BasisUnion(BaseDoc): + ud: Union[int, str] + + docs_basic = DocList[BasisUnion]([BasisUnion(ud="hello")]) + docs_copy = DocList[BasisUnion].from_protobuf(docs_basic.to_protobuf()) + assert docs_copy == docs_basic + + +class MySimpleDoc(BaseDoc): + title: str + + +class MyComplexDoc(BaseDoc): + content_dict_doclist: Dict[str, DocList[MySimpleDoc]] + content_dict_list: Dict[str, List[MySimpleDoc]] + aux_dict: Dict[str, int] + + +def test_to_from_proto_complex(): + da = DocList[MyComplexDoc]( + [ + MyComplexDoc( + content_dict_doclist={ + 'test1': DocList[MySimpleDoc]( + [MySimpleDoc(title='123'), MySimpleDoc(title='456')] + ) + }, + content_dict_list={ + 'test1': [MySimpleDoc(title='123'), MySimpleDoc(title='456')] + }, + aux_dict={'a': 0}, + ) + ] + ) + da2 = DocList[MyComplexDoc].from_protobuf(da.to_protobuf()) + assert len(da2) == 1 + d2 = da2[0] + assert d2.aux_dict == {'a': 0} + assert len(d2.content_dict_doclist['test1']) == 2 + assert d2.content_dict_doclist['test1'][0].title == '123' + assert d2.content_dict_doclist['test1'][1].title == '456' + assert len(d2.content_dict_list['test1']) == 2 + assert d2.content_dict_list['test1'][0].title == '123' + assert d2.content_dict_list['test1'][1].title == '456' diff --git a/tests/units/array/test_array_save_load.py b/tests/units/array/test_array_save_load.py new file mode 100644 index 00000000000..b5ee6b616e4 --- /dev/null +++ b/tests/units/array/test_array_save_load.py @@ -0,0 +1,137 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import numpy as np +import pytest + +from docarray import BaseDoc, DocList, DocVec +from docarray.documents import ImageDoc +from docarray.typing import NdArray, TorchTensor + + +class MyDoc(BaseDoc): + embedding: NdArray + text: str + image: ImageDoc + + +@pytest.mark.slow +@pytest.mark.parametrize( + 'protocol', ['pickle-array', 'protobuf-array', 'protobuf', 'pickle', 'json-array'] +) +@pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) +@pytest.mark.parametrize('show_progress', [False, True]) +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_array_save_load_binary(protocol, compress, tmp_path, show_progress, array_cls): + tmp_file = os.path.join(tmp_path, 'test') + + da = array_cls[MyDoc]( + [ + MyDoc( + embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png') + ), + MyDoc(embedding=[5, 4, 3, 2, 1], text='hello world', image=ImageDoc()), + ] + ) + + da.save_binary( + tmp_file, protocol=protocol, compress=compress, show_progress=show_progress + ) + + da2 = array_cls[MyDoc].load_binary( + tmp_file, protocol=protocol, compress=compress, show_progress=show_progress + ) + + assert len(da2) == 2 + assert len(da) == len(da2) + for d1, d2 in zip(da, da2): + assert d1.embedding.tolist() == d2.embedding.tolist() + assert d1.text == d2.text + assert d1.image.url == d2.image.url + assert da[1].image.url is None + assert da2[1].image.url is None + + +@pytest.mark.slow +@pytest.mark.parametrize( + 'protocol', ['pickle-array', 'protobuf-array', 'protobuf', 'pickle', 'json-array'] +) +@pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) +@pytest.mark.parametrize('show_progress', [False, True]) +@pytest.mark.parametrize('to_doc_vec', [True, False]) +def test_array_save_load_binary_streaming( + protocol, compress, tmp_path, show_progress, to_doc_vec +): + tmp_file = os.path.join(tmp_path, 'test') + array_cls = DocVec if to_doc_vec else DocList + + da = DocList[MyDoc]() + + def _extend_da(num_docs=100): + for _ in range(num_docs): + da.extend( + [ + MyDoc( + embedding=np.random.rand(3, 2), + text='hello', + image=ImageDoc(url='aux.png'), + ), + ] + ) + + _extend_da() + if to_doc_vec: + da = da.to_doc_vec() + + da.save_binary( + tmp_file, protocol=protocol, compress=compress, show_progress=show_progress + ) + + da_after = array_cls[MyDoc].load_binary( + tmp_file, protocol=protocol, compress=compress, show_progress=show_progress + ) + + for i, doc in enumerate(da_after): + assert doc.id == da[i].id + assert doc.text == da[i].text + assert doc.image.url == da[i].image.url + + assert i == 99 + + +@pytest.mark.parametrize('tensor_type', [NdArray, TorchTensor]) +def test_save_load_tensor_type(tensor_type, tmp_path): + tmp_file = os.path.join(tmp_path, 'test123') + + class MyDoc(BaseDoc): + embedding: tensor_type + text: str + image: ImageDoc + + da = DocVec[MyDoc]( + [ + MyDoc( + embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png') + ), + MyDoc(embedding=[5, 4, 3, 2, 1], text='hello world', image=ImageDoc()), + ], + tensor_type=tensor_type, + ) + da.save_binary(tmp_file) + da2 = DocVec[MyDoc].load_binary(tmp_file, tensor_type=tensor_type) + assert da2.tensor_type == tensor_type + assert isinstance(da2.embedding, tensor_type) diff --git a/tests/units/array/test_batching.py b/tests/units/array/test_batching.py new file mode 100644 index 00000000000..0387b7a2b91 --- /dev/null +++ b/tests/units/array/test_batching.py @@ -0,0 +1,61 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytest + +from docarray import BaseDoc, DocList +from docarray.typing import NdArray + + +@pytest.mark.parametrize('shuffle', [False, True]) +@pytest.mark.parametrize('stack', [False, True]) +@pytest.mark.parametrize('batch_size,n_batches', [(16, 7), (10, 10)]) +def test_batch(shuffle, stack, batch_size, n_batches): + class MyDoc(BaseDoc): + id: int + tensor: NdArray + + t_shape = (32, 32) + da = DocList[MyDoc]( + [ + MyDoc( + id=str(i), + tensor=np.zeros(t_shape), + ) + for i in range(100) + ] + ) + if stack: + da = da.to_doc_vec() + + batches = list(da._batch(batch_size=batch_size, shuffle=shuffle)) + assert len(batches) == n_batches + + for i, batch in enumerate(batches): + if i < n_batches - 1: + assert len(batch) == batch_size + if stack: + assert batch.tensor.shape == (batch_size, *t_shape) + else: + assert len(batch) <= batch_size + + non_shuffled_ids = [ + i for i in range(i * batch_size, min((i + 1) * batch_size, len(da))) + ] + if not shuffle: + assert batch.id == non_shuffled_ids + else: + assert not (batch.id == non_shuffled_ids) diff --git a/tests/units/array/test_doclist_schema.py b/tests/units/array/test_doclist_schema.py new file mode 100644 index 00000000000..02a5f562807 --- /dev/null +++ b/tests/units/array/test_doclist_schema.py @@ -0,0 +1,22 @@ +import pytest +from docarray import BaseDoc, DocList +from docarray.utils._internal.pydantic import is_pydantic_v2 + + +@pytest.mark.skipif(not is_pydantic_v2, reason='Feature only available for Pydantic V2') +def test_schema_nested(): + # check issue https://github.com/docarray/docarray/issues/1521 + + class Doc1Test(BaseDoc): + aux: str + + class DocDocTest(BaseDoc): + docs: DocList[Doc1Test] + + assert 'Doc1Test' in DocDocTest.schema()['$defs'] + d = DocDocTest(docs=DocList[Doc1Test]([Doc1Test(aux='aux')])) + + assert isinstance(d.docs, DocList) + for dd in d.docs: + assert isinstance(dd, Doc1Test) + assert d.docs.aux == ['aux'] diff --git a/tests/units/array/test_generic_array.py b/tests/units/array/test_generic_array.py new file mode 100644 index 00000000000..92d77d2a405 --- /dev/null +++ b/tests/units/array/test_generic_array.py @@ -0,0 +1,34 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from docarray import BaseDoc, DocList +from docarray.base_doc import AnyDoc + + +def test_generic_init(): + class Text(BaseDoc): + text: str + + da = DocList[Text]([]) + da.doc_type == Text + + assert isinstance(da, DocList) + + +def test_normal_access_init(): + da = DocList([]) + da.doc_type == AnyDoc + + assert isinstance(da, DocList) diff --git a/tests/units/array/test_indexing.py b/tests/units/array/test_indexing.py new file mode 100644 index 00000000000..f733b1e6630 --- /dev/null +++ b/tests/units/array/test_indexing.py @@ -0,0 +1,265 @@ +import numpy as np +import pytest +import torch + +from docarray import DocList, DocVec +from docarray.documents import TextDoc +from docarray.typing import TorchTensor + + +@pytest.fixture() +def da(): + texts = [f'hello {i}' for i in range(10)] + tensors = [torch.ones((4,)) * i for i in range(10)] + return DocList[TextDoc]( + [TextDoc(text=text, embedding=tens) for text, tens in zip(texts, tensors)], + ) + + +@pytest.fixture() +def da_to_set(): + texts = [f'hello {2*i}' for i in range(5)] + tensors = [torch.ones((4,)) * i * 2 for i in range(5)] + return DocList[TextDoc]( + [TextDoc(text=text, embedding=tens) for text, tens in zip(texts, tensors)], + ) + + +########### +# getitem +########### + + +@pytest.mark.parametrize('stack', [True, False]) +def test_simple_getitem(stack, da): + if stack: + da = da.to_doc_vec(tensor_type=TorchTensor) + + assert torch.all(da[0].embedding == torch.zeros((4,))) + assert da[0].text == 'hello 0' + + +@pytest.mark.parametrize('stack', [True, False]) +def test_get_none(stack, da): + if stack: + da = da.to_doc_vec(tensor_type=TorchTensor) + + assert da[None] is da + + +@pytest.mark.parametrize('stack', [True, False]) +@pytest.mark.parametrize('index', [(1, 2, 3, 4, 6), [1, 2, 3, 4, 6]]) +def test_iterable_getitem(stack, da, index): + if stack: + da = da.to_doc_vec(tensor_type=TorchTensor) + + indexed_da = da[index] + + for pos, d in zip(index, indexed_da): + assert d.text == f'hello {pos}' + assert torch.all(d.embedding == torch.ones((4,)) * pos) + + +@pytest.mark.parametrize('stack', [True, False]) +@pytest.mark.parametrize('index_dtype', [torch.int64]) +def test_torchtensor_getitem(stack, da, index_dtype): + if stack: + da = da.to_doc_vec(tensor_type=TorchTensor) + + index = torch.tensor([1, 2, 3, 4, 6], dtype=index_dtype) + + indexed_da = da[index] + + for pos, d in zip(index, indexed_da): + assert d.text == f'hello {pos}' + assert torch.all(d.embedding == torch.ones((4,)) * pos) + + +@pytest.mark.parametrize('stack', [True, False]) +@pytest.mark.parametrize('index_dtype', [int, np.int_, np.int32, np.int64]) +def test_nparray_getitem(stack, da, index_dtype): + if stack: + da = da.to_doc_vec(tensor_type=TorchTensor) + + index = np.array([1, 2, 3, 4, 6], dtype=index_dtype) + + indexed_da = da[index] + for pos, d in zip(index, indexed_da): + assert d.text == f'hello {pos}' + assert torch.all(d.embedding == torch.ones((4,)) * pos) + + +@pytest.mark.parametrize('stack', [True, False]) +@pytest.mark.parametrize( + 'index', + [ + [False, True, True, True, True, False, True, False, False, False], + (False, True, True, True, True, False, True, False, False, False), + torch.tensor([0, 1, 1, 1, 1, 0, 1, 0, 0, 0], dtype=torch.bool), + np.array([0, 1, 1, 1, 1, 0, 1, 0, 0, 0], dtype=bool), + ], +) +def test_boolmask_getitem(stack, da, index): + if stack: + da = da.to_doc_vec(tensor_type=TorchTensor) + + indexed_da = da[index] + + mask_true_idx = [1, 2, 3, 4, 6] + + for pos, d in zip(mask_true_idx, indexed_da): + assert d.text == f'hello {pos}' + assert torch.all(d.embedding == torch.ones((4,)) * pos) + + +########### +# setitem +########### + + +@pytest.mark.parametrize('stack_left', [True, False]) +def test_simple_setitem(stack_left, da, da_to_set): + if stack_left: + da = da.to_doc_vec(tensor_type=TorchTensor) + + da[0] = da_to_set[0] + + assert torch.all(da[0].embedding == da_to_set[0].embedding) + assert da[0].text == da_to_set[0].text + + +@pytest.mark.parametrize('stack_left', [True, False]) +@pytest.mark.parametrize('stack_right', [True, False]) +@pytest.mark.parametrize('index', [(1, 2, 3, 4, 6), [1, 2, 3, 4, 6]]) +def test_iterable_setitem(stack_left, stack_right, da, da_to_set, index): + if stack_left: + da = da.to_doc_vec(tensor_type=TorchTensor) + if stack_right: + da_to_set = da_to_set.to_doc_vec(tensor_type=TorchTensor) + + da[index] = da_to_set + + i_da_to_set = 0 + for i, d in enumerate(da): + if i in index: + d_reference = da_to_set[i_da_to_set] + assert d.text == d_reference.text + assert torch.all(d.embedding == d_reference.embedding) + i_da_to_set += 1 + else: + assert d.text == f'hello {i}' + assert torch.all(d.embedding == torch.ones((4,)) * i) + + +@pytest.mark.parametrize('stack_left', [True, False]) +@pytest.mark.parametrize('stack_right', [True, False]) +@pytest.mark.parametrize('index_dtype', [torch.int64]) +def test_torchtensor_setitem(stack_left, stack_right, da, da_to_set, index_dtype): + if stack_left: + da = da.to_doc_vec(tensor_type=TorchTensor) + if stack_right: + da_to_set = da_to_set.to_doc_vec(tensor_type=TorchTensor) + + index = torch.tensor([1, 2, 3, 4, 6], dtype=index_dtype) + + da[index] = da_to_set + + i_da_to_set = 0 + for i, d in enumerate(da): + if i in index: + d_reference = da_to_set[i_da_to_set] + assert d.text == d_reference.text + assert torch.all(d.embedding == d_reference.embedding) + i_da_to_set += 1 + else: + assert d.text == f'hello {i}' + assert torch.all(d.embedding == torch.ones((4,)) * i) + + +@pytest.mark.parametrize('stack_left', [True, False]) +@pytest.mark.parametrize('stack_right', [True, False]) +@pytest.mark.parametrize('index_dtype', [int, np.int_, np.int32, np.int64]) +def test_nparray_setitem(stack_left, stack_right, da, da_to_set, index_dtype): + if stack_left: + da = da.to_doc_vec(tensor_type=TorchTensor) + if stack_right: + da_to_set = da_to_set.to_doc_vec(tensor_type=TorchTensor) + + index = np.array([1, 2, 3, 4, 6], dtype=index_dtype) + + da[index] = da_to_set + + i_da_to_set = 0 + for i, d in enumerate(da): + if i in index: + d_reference = da_to_set[i_da_to_set] + assert d.text == d_reference.text + assert torch.all(d.embedding == d_reference.embedding) + i_da_to_set += 1 + else: + assert d.text == f'hello {i}' + assert torch.all(d.embedding == torch.ones((4,)) * i) + + +@pytest.mark.parametrize('stack_left', [True, False]) +@pytest.mark.parametrize('stack_right', [True, False]) +@pytest.mark.parametrize( + 'index', + [ + [False, True, True, True, True, False, True, False, False, False], + (False, True, True, True, True, False, True, False, False, False), + torch.tensor([0, 1, 1, 1, 1, 0, 1, 0, 0, 0], dtype=torch.bool), + np.array([0, 1, 1, 1, 1, 0, 1, 0, 0, 0], dtype=bool), + ], +) +def test_boolmask_setitem(stack_left, stack_right, da, da_to_set, index): + if stack_left: + da = da.to_doc_vec(tensor_type=TorchTensor) + if stack_right: + da_to_set = da_to_set.to_doc_vec(tensor_type=TorchTensor) + + da[index] = da_to_set + + mask_true_idx = [1, 2, 3, 4, 6] + i_da_to_set = 0 + for i, d in enumerate(da): + if i in mask_true_idx: + d_reference = da_to_set[i_da_to_set] + assert d.text == d_reference.text + assert torch.all(d.embedding == d_reference.embedding) + i_da_to_set += 1 + else: + assert d.text == f'hello {i}' + assert torch.all(d.embedding == torch.ones((4,)) * i) + + +def test_setitem_update_column(): + texts = [f'hello {i}' for i in range(10)] + tensors = [torch.ones((4,)) * (i + 1) for i in range(10)] + da = DocVec[TextDoc]( + [TextDoc(text=text, embedding=tens) for text, tens in zip(texts, tensors)], + tensor_type=TorchTensor, + ) + + da[0] = TextDoc(text='hello', embedding=torch.zeros((4,))) + + assert da[0].text == 'hello' + assert (da[0].embedding == torch.zeros((4,))).all() + assert (da.embedding[0] == torch.zeros((4,))).all() + + assert da._storage.any_columns['text'][0] == 'hello' + assert (da._storage.tensor_columns['embedding'][0] == torch.zeros((4,))).all() + assert (da._storage.tensor_columns['embedding'][0] == torch.zeros((4,))).all() + + +@pytest.mark.parametrize( + 'index', + [ + [False, True, True, True, True, False, True, False, False, False], + (False, True, True, True, True, False, True, False, False, False), + torch.tensor([0, 1, 1, 1, 1, 0, 1, 0, 0, 0], dtype=torch.bool), + np.array([0, 1, 1, 1, 1, 0, 1, 0, 0, 0], dtype=bool), + ], +) +def test_del_getitem(da, index): + del da[index] diff --git a/tests/units/array/test_traverse.py b/tests/units/array/test_traverse.py new file mode 100644 index 00000000000..4c513148bd4 --- /dev/null +++ b/tests/units/array/test_traverse.py @@ -0,0 +1,118 @@ +from typing import Optional + +import pytest +import torch + +from docarray import BaseDoc, DocList +from docarray.array.any_array import AnyDocArray +from docarray.documents import TextDoc +from docarray.typing import TorchTensor + +num_docs = 5 +num_sub_docs = 2 +num_sub_sub_docs = 3 + + +@pytest.fixture +def multi_model_docs(): + class SubSubDoc(BaseDoc): + sub_sub_text: TextDoc + sub_sub_tensor: TorchTensor[2] + + class SubDoc(BaseDoc): + sub_text: TextDoc + sub_da: DocList[SubSubDoc] + + class MultiModalDoc(BaseDoc): + mm_text: TextDoc + mm_tensor: Optional[TorchTensor[3, 2, 2]] = None + mm_da: DocList[SubDoc] + + docs = DocList[MultiModalDoc]( + [ + MultiModalDoc( + mm_text=TextDoc(text=f'hello{i}'), + mm_da=[ + SubDoc( + sub_text=TextDoc(text=f'sub_{i}_1'), + sub_da=DocList[SubSubDoc]( + [ + SubSubDoc( + sub_sub_text=TextDoc(text='subsub'), + sub_sub_tensor=torch.zeros(2), + ) + for _ in range(num_sub_sub_docs) + ] + ), + ) + for _ in range(num_sub_docs) + ], + ) + for i in range(num_docs) + ] + ) + + return docs + + +@pytest.mark.parametrize( + 'access_path,len_result', + [ + ('mm_text', num_docs), # List of 5 Text objs + ('mm_text__text', num_docs), # List of 5 strings + ('mm_da', num_docs * num_sub_docs), # List of 5 * 2 SubDoc objs + ('mm_da__sub_text', num_docs * num_sub_docs), # List of 5 * 2 Text objs + ( + 'mm_da__sub_da', + num_docs * num_sub_docs * num_sub_sub_docs, + ), # List of 5 * 2 * 3 SubSubDoc objs + ( + 'mm_da__sub_da__sub_sub_text', + num_docs * num_sub_docs * num_sub_sub_docs, + ), # List of 5 * 2 * 3 Text objs + ], +) +def test_traverse_flat(multi_model_docs, access_path, len_result): + traversed = multi_model_docs.traverse_flat(access_path) + assert len(traversed) == len_result + + +def test_traverse_stacked_da(): + class Image(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + batch = DocList[Image]( + [ + Image( + tensor=torch.zeros(3, 224, 224), + ) + for _ in range(2) + ] + ) + + batch_stacked = batch.to_doc_vec() + tensors = batch_stacked.traverse_flat(access_path='tensor') + + assert tensors.shape == (2, 3, 224, 224) + assert isinstance(tensors, torch.Tensor) + + +@pytest.mark.parametrize( + 'input_list,output_list', + [ + ([1, 2, 3], [1, 2, 3]), + ([[1], [2], [3]], [1, 2, 3]), + ([[[1]], [[2]], [[3]]], [[1], [2], [3]]), + ], +) +def test_flatten_one_level(input_list, output_list): + flattened = AnyDocArray._flatten_one_level(sequence=input_list) + assert flattened == output_list + + +def test_flatten_one_level_list_of_da(): + doc = BaseDoc() + input_list = [DocList([doc, doc, doc])] + + flattened = AnyDocArray._flatten_one_level(sequence=input_list) + assert flattened == [doc, doc, doc] diff --git a/tests/units/computation_backends/__init__.py b/tests/units/computation_backends/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/computation_backends/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/computation_backends/backend_comparisons/__init__.py b/tests/units/computation_backends/backend_comparisons/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/computation_backends/backend_comparisons/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/computation_backends/backend_comparisons/test_metrics.py b/tests/units/computation_backends/backend_comparisons/test_metrics.py new file mode 100644 index 00000000000..f899bc44d39 --- /dev/null +++ b/tests/units/computation_backends/backend_comparisons/test_metrics.py @@ -0,0 +1,74 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch + +from docarray.computation.numpy_backend import NumpyCompBackend +from docarray.computation.torch_backend import TorchCompBackend + +np_metrics = NumpyCompBackend.Metrics +torch_metrics = TorchCompBackend.Metrics + + +def test_cosine_sim_compare(): + a = torch.rand(128) + b = torch.rand(128) + torch.testing.assert_close( + torch_metrics.cosine_sim(a, b), + torch.from_numpy(np_metrics.cosine_sim(a.numpy(), b.numpy())), + ) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + torch.testing.assert_close( + torch_metrics.cosine_sim(a, b), + torch.from_numpy(np_metrics.cosine_sim(a.numpy(), b.numpy())), + ) + + +def test_euclidean_dist_compare(): + a = torch.rand(128) + b = torch.rand(128) + torch.testing.assert_close( + torch_metrics.euclidean_dist(a, b), + torch.from_numpy(np_metrics.euclidean_dist(a.numpy(), b.numpy())).to( + torch.float32 + ), + ) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + torch.testing.assert_close( + torch_metrics.euclidean_dist(a, b), + torch.from_numpy(np_metrics.euclidean_dist(a.numpy(), b.numpy())), + ) + + +def test_sqeuclidean_dist_compare(): + a = torch.rand(128) + b = torch.rand(128) + torch.testing.assert_close( + torch_metrics.sqeuclidean_dist(a, b), + torch.from_numpy(np_metrics.sqeuclidean_dist(a.numpy(), b.numpy())).to( + torch.float32 + ), + ) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + torch.testing.assert_close( + torch_metrics.sqeuclidean_dist(a, b), + torch.from_numpy(np_metrics.sqeuclidean_dist(a.numpy(), b.numpy())), + ) diff --git a/tests/units/computation_backends/jax_backend/__init__.py b/tests/units/computation_backends/jax_backend/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/computation_backends/jax_backend/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/computation_backends/jax_backend/test_basics.py b/tests/units/computation_backends/jax_backend/test_basics.py new file mode 100644 index 00000000000..db064430c9b --- /dev/null +++ b/tests/units/computation_backends/jax_backend/test_basics.py @@ -0,0 +1,163 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import pytest + +from docarray.utils._internal.misc import is_jax_available + +jax_available = is_jax_available() +if jax_available: + print("is jax available", jax_available) + import jax + import jax.numpy as jnp + + from docarray.computation.jax_backend import JaxCompBackend + from docarray.typing import JaxArray + + jax.config.update("jax_enable_x64", True) + + +@pytest.mark.jax +@pytest.mark.parametrize( + 'shape,result', + [ + ((5), 1), + ((1, 5), 2), + ((5, 5), 2), + ((), 0), + ], +) +def test_n_dim(shape, result): + + array = JaxArray(jnp.zeros(shape)) + assert JaxCompBackend.n_dim(array) == result + + +@pytest.mark.jax +@pytest.mark.parametrize( + 'shape,result', + [ + ((10,), (10,)), + ((5, 5), (5, 5)), + ((), ()), + ], +) +def test_shape(shape, result): + array = JaxArray(jnp.zeros(shape)) + shape = JaxCompBackend.shape(array) + assert shape == result + assert type(shape) == tuple + + +@pytest.mark.jax +def test_to_device(): + array = JaxArray(jnp.zeros((3))) + array = JaxCompBackend.to_device(array, 'cpu') + assert array.tensor.device().platform.endswith('cpu') + + +@pytest.mark.jax +@pytest.mark.parametrize( + 'dtype,result_type', + [ + ('int64', 'int64'), + ('float64', 'float64'), + ('int8', 'int8'), + ('double', 'float64'), + ], +) +def test_dtype(dtype, result_type): + array = JaxArray(jnp.array([1, 2, 3], dtype=dtype)) + assert JaxCompBackend.dtype(array) == result_type + + +@pytest.mark.jax +def test_empty(): + array = JaxCompBackend.empty((10, 3)) + assert array.tensor.shape == (10, 3) + + +@pytest.mark.jax +def test_empty_dtype(): + tf_tensor = JaxCompBackend.empty((10, 3), dtype=jnp.int32) + assert tf_tensor.tensor.shape == (10, 3) + assert tf_tensor.tensor.dtype == jnp.int32 + + +@pytest.mark.jax +def test_empty_device(): + tensor = JaxCompBackend.empty((10, 3), device='cpu') + assert tensor.tensor.shape == (10, 3) + assert tensor.tensor.device().platform.endswith('cpu') + + +@pytest.mark.jax +def test_squeeze(): + tensor = JaxArray(jnp.zeros(shape=(1, 1, 3, 1))) + squeezed = JaxCompBackend.squeeze(tensor) + assert squeezed.tensor.shape == (3,) + + +@pytest.mark.jax +@pytest.mark.parametrize( + 'data_input,t_range,x_range,data_result', + [ + ( + [0, 1, 2, 3, 4, 5], + (0, 10), + None, + [0, 2, 4, 6, 8, 10], + ), + ( + [0, 1, 2, 3, 4, 5], + (0, 10), + (0, 10), + [0, 1, 2, 3, 4, 5], + ), + ( + [[0.0, 1.0], [0.0, 1.0]], + (0, 10), + None, + [[0.0, 10.0], [0.0, 10.0]], + ), + ], +) +def test_minmax_normalize(data_input, t_range, x_range, data_result): + array = JaxArray(jnp.array(data_input)) + output = JaxCompBackend.minmax_normalize( + tensor=array, t_range=t_range, x_range=x_range + ) + assert jnp.allclose(output.tensor, jnp.array(data_result)) + + +@pytest.mark.jax +def test_reshape(): + tensor = JaxArray(jnp.zeros((3, 224, 224))) + reshaped = JaxCompBackend.reshape(tensor, (224, 224, 3)) + assert reshaped.tensor.shape == (224, 224, 3) + + +@pytest.mark.jax +def test_stack(): + t0 = JaxArray(jnp.zeros((3, 224, 224))) + t1 = JaxArray(jnp.ones((3, 224, 224))) + + stacked1 = JaxCompBackend.stack([t0, t1], dim=0) + assert isinstance(stacked1, JaxArray) + assert stacked1.tensor.shape == (2, 3, 224, 224) + + stacked2 = JaxCompBackend.stack([t0, t1], dim=-1) + assert isinstance(stacked2, JaxArray) + assert stacked2.tensor.shape == (3, 224, 224, 2) diff --git a/tests/units/computation_backends/jax_backend/test_metrics.py b/tests/units/computation_backends/jax_backend/test_metrics.py new file mode 100644 index 00000000000..50dc6339d63 --- /dev/null +++ b/tests/units/computation_backends/jax_backend/test_metrics.py @@ -0,0 +1,81 @@ +import pytest + +from docarray.utils._internal.misc import is_jax_available + +jax_available = is_jax_available() +if jax_available: + import jax + import jax.numpy as jnp + + from docarray.computation.jax_backend import JaxCompBackend + from docarray.typing import JaxArray + + metrics = JaxCompBackend.Metrics +else: + metrics = None + + +@pytest.mark.jax +def test_cosine_sim_jax(): + a = JaxArray(jax.random.uniform(jax.random.PRNGKey(0), shape=(128,))) + b = JaxArray(jax.random.uniform(jax.random.PRNGKey(1), shape=(128,))) + assert metrics.cosine_sim(a, b).tensor.shape == (1,) + assert metrics.cosine_sim(a, b).tensor == metrics.cosine_sim(b, a).tensor + + assert jnp.allclose(metrics.cosine_sim(a, a).tensor, jnp.ones((1,))) + + a = JaxArray(jax.random.uniform(jax.random.PRNGKey(2), shape=(10, 3))) + b = JaxArray(jax.random.uniform(jax.random.PRNGKey(3), shape=(5, 3))) + assert metrics.cosine_sim(a, b).tensor.shape == (10, 5) + assert metrics.cosine_sim(b, a).tensor.shape == (5, 10) + diag_dists = jnp.diagonal(metrics.cosine_sim(b, b).tensor) # self-comparisons + assert jnp.allclose(diag_dists, jnp.ones((5,))) + + +@pytest.mark.jax +@pytest.mark.skip +def test_euclidean_dist_jax(): + a = JaxArray(jax.random.normal(jax.random.PRNGKey(0), shape=(128,))) + b = JaxArray(jax.random.normal(jax.random.PRNGKey(1), shape=(128,))) + assert metrics.euclidean_dist(a, b).tensor.shape == (1,) + assert jnp.allclose( + metrics.euclidean_dist(a, b).tensor, metrics.euclidean_dist(b, a).tensor + ) + + assert jnp.allclose(metrics.euclidean_dist(a, a).tensor, jnp.zeros((1,))) + + a = JaxArray(jnp.zeros((1, 1))) + b = JaxArray(jnp.ones((4, 1))) + assert metrics.euclidean_dist(a, b).tensor.shape == (4,) + assert jnp.allclose( + metrics.euclidean_dist(a, b).tensor, metrics.euclidean_dist(b, a).tensor + ) + assert jnp.allclose(metrics.euclidean_dist(a, a).tensor, jnp.zeros((1,))) + + a = JaxArray(jnp.array([0.0, 2.0, 0.0])) + b = JaxArray(jnp.array([0.0, 0.0, 2.0])) + desired_output_singleton = jnp.sqrt(jnp.array([2.0**2.0 + 2.0**2.0])) + assert jnp.allclose(metrics.euclidean_dist(a, b).tensor, desired_output_singleton) + + a = JaxArray(jnp.array([[0.0, 2.0, 0.0], [0.0, 0.0, 2.0]])) + b = JaxArray(jnp.array([[0.0, 0.0, 2.0], [0.0, 2.0, 0.0]])) + desired_output_singleton = jnp.array([[2.828427, 0.0], [0.0, 2.828427]]) + + assert jnp.allclose(metrics.euclidean_dist(a, b).tensor, desired_output_singleton) + + +@pytest.mark.jax +def test_sqeuclidea_dist_jnp(): + a = JaxArray(jax.random.uniform(jax.random.PRNGKey(0), shape=(128,))) + b = JaxArray(jax.random.uniform(jax.random.PRNGKey(1), shape=(128,))) + assert metrics.sqeuclidean_dist(a, b).tensor.shape == (1,) + assert jnp.allclose( + metrics.sqeuclidean_dist(a, b).tensor, metrics.euclidean_dist(a, b).tensor ** 2 + ) + + a = JaxArray(jax.random.uniform(jax.random.PRNGKey(2), shape=(10, 3))) + b = JaxArray(jax.random.uniform(jax.random.PRNGKey(3), shape=(5, 3))) + assert metrics.sqeuclidean_dist(a, b).tensor.shape == (10, 5) + assert jnp.allclose( + metrics.sqeuclidean_dist(a, b).tensor, metrics.euclidean_dist(a, b).tensor ** 2 + ) diff --git a/tests/units/computation_backends/jax_backend/test_retrieval.py b/tests/units/computation_backends/jax_backend/test_retrieval.py new file mode 100644 index 00000000000..7d827e2d383 --- /dev/null +++ b/tests/units/computation_backends/jax_backend/test_retrieval.py @@ -0,0 +1,81 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import pytest + +from docarray.utils._internal.misc import is_jax_available + +jax_available = is_jax_available() +if jax_available: + import jax.numpy as jnp + + from docarray.computation.jax_backend import JaxCompBackend + from docarray.typing import JaxArray + + metrics = JaxCompBackend.Metrics +else: + metrics = None + + +@pytest.mark.jax +def test_top_k_descending_false(): + top_k = JaxCompBackend.Retrieval.top_k + + a = JaxArray(jnp.array([1, 4, 2, 7, 4, 9, 2])) + vals, indices = top_k(a, 3, descending=False) + + assert vals.tensor.shape == (1, 3) + assert indices.tensor.shape == (1, 3) + assert jnp.allclose(jnp.squeeze(vals.tensor), jnp.array([1, 2, 2])) + assert jnp.allclose(jnp.squeeze(indices.tensor), jnp.array([0, 2, 6])) or ( + jnp.allclose(jnp.squeeze.indices.tensor), + jnp.array([0, 6, 2]), + ) + + a = JaxArray(jnp.array([[1, 4, 2, 7, 4, 9, 2], [11, 6, 2, 7, 3, 10, 4]])) + vals, indices = top_k(a, 3, descending=False) + assert vals.tensor.shape == (2, 3) + assert indices.tensor.shape == (2, 3) + assert jnp.allclose(vals.tensor[0], jnp.array([1, 2, 2])) + assert jnp.allclose(indices.tensor[0], jnp.array([0, 2, 6])) or jnp.allclose( + indices.tensor[0], jnp.array([0, 6, 2]) + ) + assert jnp.allclose(vals.tensor[1], jnp.array([2, 3, 4])) + assert jnp.allclose(indices.tensor[1], jnp.array([2, 4, 6])) + + +@pytest.mark.jax +def test_top_k_descending_true(): + top_k = JaxCompBackend.Retrieval.top_k + + a = JaxArray(jnp.array([1, 4, 2, 7, 4, 9, 2])) + vals, indices = top_k(a, 3, descending=True) + + assert vals.tensor.shape == (1, 3) + assert indices.tensor.shape == (1, 3) + assert jnp.allclose(jnp.squeeze(vals.tensor), jnp.array([9, 7, 4])) + assert jnp.allclose(jnp.squeeze(indices.tensor), jnp.array([5, 3, 1])) + + a = JaxArray(jnp.array([[1, 4, 2, 7, 4, 9, 2], [11, 6, 2, 7, 3, 10, 4]])) + vals, indices = top_k(a, 3, descending=True) + + assert vals.tensor.shape == (2, 3) + assert indices.tensor.shape == (2, 3) + + assert jnp.allclose(vals.tensor[0], jnp.array([9, 7, 4])) + assert jnp.allclose(indices.tensor[0], jnp.array([5, 3, 1])) + + assert jnp.allclose(vals.tensor[1], jnp.array([11, 10, 7])) + assert jnp.allclose(indices.tensor[1], jnp.array([0, 5, 3])) diff --git a/tests/units/computation_backends/numpy_backend/__init__.py b/tests/units/computation_backends/numpy_backend/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/computation_backends/numpy_backend/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/computation_backends/numpy_backend/test_basics.py b/tests/units/computation_backends/numpy_backend/test_basics.py new file mode 100644 index 00000000000..7ab511db9ad --- /dev/null +++ b/tests/units/computation_backends/numpy_backend/test_basics.py @@ -0,0 +1,119 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytest +from pydantic import parse_obj_as + +from docarray.computation.numpy_backend import NumpyCompBackend +from docarray.typing import NdArray + + +def test_to_device(): + with pytest.raises(NotImplementedError): + NumpyCompBackend.to_device(np.random.rand(10, 3), 'meta') + + +@pytest.mark.parametrize( + 'array,result', + [ + (np.zeros((5)), 1), + (np.zeros((1, 5)), 2), + (np.zeros((5, 5)), 2), + (np.zeros(()), 0), + ], +) +def test_n_dim(array, result): + assert NumpyCompBackend.n_dim(array) == result + + +@pytest.mark.parametrize( + 'array,result', + [ + (np.zeros((10,)), (10,)), + (np.zeros((5, 5)), (5, 5)), + (np.zeros(()), ()), + ], +) +def test_shape(array, result): + shape = NumpyCompBackend.shape(array) + assert shape == result + assert type(shape) == tuple + + +def test_device(): + array = np.array([1, 2, 3]) + assert NumpyCompBackend.device(array) is None + + +@pytest.mark.parametrize('dtype', [np.int64, np.float64, int, float]) +def test_dtype(dtype): + array = np.array([1, 2, 3], dtype=dtype) + assert NumpyCompBackend.dtype(array) == dtype + + +def test_empty(): + array = NumpyCompBackend.empty((10, 3)) + assert array.shape == (10, 3) + + +def test_empty_dtype(): + tensor = NumpyCompBackend.empty((10, 3), dtype=np.int32) + assert tensor.shape == (10, 3) + assert tensor.dtype == np.int32 + + +def test_empty_device(): + with pytest.raises(NotImplementedError): + NumpyCompBackend.empty((10, 3), device='meta') + + +def test_squeeze(): + tensor = np.zeros(shape=(1, 1, 3, 1)) + squeezed = NumpyCompBackend.squeeze(tensor) + assert squeezed.shape == (3,) + + +@pytest.mark.parametrize( + 'array,t_range,x_range,result', + [ + (np.array([0, 1, 2, 3, 4, 5]), (0, 10), None, np.array([0, 2, 4, 6, 8, 10])), + (np.array([0, 1, 2, 3, 4, 5]), (0, 10), (0, 10), np.array([0, 1, 2, 3, 4, 5])), + ( + np.array([[0.0, 1.0], [0.0, 1.0]]), + (0, 10), + None, + np.array([[0.0, 10.0], [0.0, 10.0]]), + ), + ], +) +def test_minmax_normalize(array, t_range, x_range, result): + output = NumpyCompBackend.minmax_normalize( + tensor=array, t_range=t_range, x_range=x_range + ) + assert np.allclose(output, result) + + +def test_stack(): + t0 = parse_obj_as(NdArray, np.zeros((3, 224, 224))) + t1 = parse_obj_as(NdArray, np.ones((3, 224, 224))) + + stacked1 = NumpyCompBackend.stack([t0, t1], dim=0) + assert isinstance(stacked1, np.ndarray) + assert stacked1.shape == (2, 3, 224, 224) + + stacked2 = NumpyCompBackend.stack([t0, t1], dim=-1) + assert isinstance(stacked2, np.ndarray) + assert stacked2.shape == (3, 224, 224, 2) diff --git a/tests/units/computation_backends/numpy_backend/test_metrics.py b/tests/units/computation_backends/numpy_backend/test_metrics.py new file mode 100644 index 00000000000..2cd0369e876 --- /dev/null +++ b/tests/units/computation_backends/numpy_backend/test_metrics.py @@ -0,0 +1,67 @@ +import numpy as np + +from docarray.computation.numpy_backend import NumpyCompBackend + +metrics = NumpyCompBackend.Metrics + + +def test_cosine_sim_np(): + a = np.random.rand(128) + b = np.random.rand(128) + assert metrics.cosine_sim(a, b).shape == (1,) + assert metrics.cosine_sim(a, b) == metrics.cosine_sim(b, a) + np.testing.assert_array_almost_equal(metrics.cosine_sim(a, a), np.ones((1,))) + + a = np.random.rand(10, 3) + b = np.random.rand(5, 3) + assert metrics.cosine_sim(a, b).shape == (10, 5) + assert metrics.cosine_sim(b, a).shape == (5, 10) + diag_dists = np.diagonal(metrics.cosine_sim(b, b)) # self-comparisons + np.testing.assert_array_almost_equal(diag_dists, np.ones((5,))) + + +def test_euclidean_dist_np(): + a = np.random.rand(128) + b = np.random.rand(128) + assert metrics.euclidean_dist(a, b).shape == (1,) + assert metrics.euclidean_dist(a, b) == metrics.euclidean_dist(b, a) + np.testing.assert_array_almost_equal(metrics.euclidean_dist(a, a), np.zeros((1,))) + + a = np.random.rand(10, 3) + b = np.random.rand(5, 3) + assert metrics.euclidean_dist(a, b).shape == (10, 5) + assert metrics.euclidean_dist(b, a).shape == (5, 10) + diag_dists = np.diagonal(metrics.euclidean_dist(b, b)) # self-comparisons + np.testing.assert_array_almost_equal(diag_dists, np.zeros((5,))) + + a = np.array([0.0, 2.0, 0.0]) + b = np.array([0.0, 0.0, 2.0]) + desired_output_singleton = np.sqrt(np.array([2.0**2.0 + 2.0**2.0])) + np.testing.assert_array_almost_equal( + metrics.euclidean_dist(a, b), desired_output_singleton + ) + + a = np.array([[0.0, 2.0, 0.0], [0.0, 0.0, 2.0]]) + b = np.array([[0.0, 0.0, 2.0], [0.0, 2.0, 0.0]]) + desired_output_singleton = np.array( + [[desired_output_singleton.item(), 0.0], [0.0, desired_output_singleton.item()]] + ) + np.testing.assert_array_almost_equal( + metrics.euclidean_dist(a, b), desired_output_singleton + ) + + +def test_sqeuclidea_dist_np(): + a = np.random.rand(128) + b = np.random.rand(128) + assert metrics.sqeuclidean_dist(a, b).shape == (1,) + np.testing.assert_array_almost_equal( + metrics.sqeuclidean_dist(a, b), metrics.euclidean_dist(a, b) ** 2 + ) + + a = np.random.rand(10, 3) + b = np.random.rand(5, 3) + assert metrics.sqeuclidean_dist(a, b).shape == (10, 5) + np.testing.assert_array_almost_equal( + metrics.sqeuclidean_dist(a, b), metrics.euclidean_dist(a, b) ** 2 + ) diff --git a/tests/units/computation_backends/numpy_backend/test_retrieval.py b/tests/units/computation_backends/numpy_backend/test_retrieval.py new file mode 100644 index 00000000000..5fa693dde61 --- /dev/null +++ b/tests/units/computation_backends/numpy_backend/test_retrieval.py @@ -0,0 +1,42 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np + +from docarray.computation.numpy_backend import NumpyCompBackend + + +def test_topk_numpy(): + top_k = NumpyCompBackend.Retrieval.top_k + + a = np.array([1, 4, 2, 7, 4, 9, 2]) + vals, indices = top_k(a, 3) + assert vals.shape == (1, 3) + assert indices.shape == (1, 3) + assert (vals.squeeze() == np.array([1, 2, 2])).all() + assert (indices.squeeze() == np.array([0, 2, 6])).all() or ( + indices.squeeze() == np.array([0, 6, 2]) + ).all() + + a = np.array([[1, 4, 2, 7, 4, 9, 2], [11, 6, 2, 7, 3, 10, 4]]) + vals, indices = top_k(a, 3) + assert vals.shape == (2, 3) + assert indices.shape == (2, 3) + assert (vals[0] == np.array([1, 2, 2])).all() + assert (indices[0] == np.array([0, 2, 6])).all() or ( + indices[0] == np.array([0, 6, 2]) + ).all() + assert (vals[1] == np.array([2, 3, 4])).all() + assert (indices[1] == np.array([2, 4, 6])).all() diff --git a/tests/units/computation_backends/tensorflow_backend/__init__.py b/tests/units/computation_backends/tensorflow_backend/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/computation_backends/tensorflow_backend/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/computation_backends/tensorflow_backend/test_basics.py b/tests/units/computation_backends/tensorflow_backend/test_basics.py new file mode 100644 index 00000000000..6747eecb87e --- /dev/null +++ b/tests/units/computation_backends/tensorflow_backend/test_basics.py @@ -0,0 +1,159 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytest + +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + + from docarray.computation.tensorflow_backend import TensorFlowCompBackend + from docarray.typing import TensorFlowTensor + + +@pytest.mark.tensorflow +@pytest.mark.parametrize( + 'shape,result', + [ + ((5), 1), + ((1, 5), 2), + ((5, 5), 2), + ((), 0), + ], +) +def test_n_dim(shape, result): + array = TensorFlowTensor(tf.zeros(shape)) + assert TensorFlowCompBackend.n_dim(array) == result + + +@pytest.mark.tensorflow +@pytest.mark.parametrize( + 'shape,result', + [ + ((10,), (10,)), + ((5, 5), (5, 5)), + ((), ()), + ], +) +def test_shape(shape, result): + array = TensorFlowTensor(tf.zeros(shape)) + shape = TensorFlowCompBackend.shape(array) + assert shape == result + assert type(shape) == tuple + + +@pytest.mark.tensorflow +def test_to_device(): + array = TensorFlowTensor(tf.constant([1, 2, 3])) + array = TensorFlowCompBackend.to_device(array, 'CPU:0') + assert array.tensor.device.endswith('CPU:0') + + +@pytest.mark.tensorflow +@pytest.mark.parametrize( + 'dtype,result_type', + [ + ('int64', 'int64'), + ('float64', 'float64'), + ('int8', 'int8'), + ('double', 'float64'), + ], +) +def test_dtype(dtype, result_type): + array = TensorFlowTensor(tf.constant([1, 2, 3], dtype=getattr(tf, dtype))) + assert TensorFlowCompBackend.dtype(array) == result_type + + +@pytest.mark.tensorflow +def test_empty(): + array = TensorFlowCompBackend.empty((10, 3)) + assert array.tensor.shape == (10, 3) + + +@pytest.mark.tensorflow +def test_empty_dtype(): + tf_tensor = TensorFlowCompBackend.empty((10, 3), dtype=tf.int32) + assert tf_tensor.tensor.shape == (10, 3) + assert tf_tensor.tensor.dtype == tf.int32 + + +@pytest.mark.tensorflow +def test_empty_device(): + tensor = TensorFlowCompBackend.empty((10, 3), device='CPU:0') + assert tensor.tensor.shape == (10, 3) + assert tensor.tensor.device.endswith('CPU:0') + + +@pytest.mark.tensorflow +def test_squeeze(): + tensor = TensorFlowTensor(tf.zeros(shape=(1, 1, 3, 1))) + squeezed = TensorFlowCompBackend.squeeze(tensor) + assert squeezed.tensor.shape == (3,) + + +@pytest.mark.tensorflow +@pytest.mark.parametrize( + 'data_input,t_range,x_range,data_result', + [ + ( + [0, 1, 2, 3, 4, 5], + (0, 10), + None, + [0, 2, 4, 6, 8, 10], + ), + ( + [0, 1, 2, 3, 4, 5], + (0, 10), + (0, 10), + [0, 1, 2, 3, 4, 5], + ), + ( + [[0.0, 1.0], [0.0, 1.0]], + (0, 10), + None, + [[0.0, 10.0], [0.0, 10.0]], + ), + ], +) +def test_minmax_normalize(data_input, t_range, x_range, data_result): + array = TensorFlowTensor(tf.constant(data_input)) + output = TensorFlowCompBackend.minmax_normalize( + tensor=array, t_range=t_range, x_range=x_range + ) + assert np.allclose(output.tensor, tf.constant(data_result)) + + +@pytest.mark.tensorflow +def test_reshape(): + tensor = TensorFlowTensor(tf.zeros((3, 224, 224))) + reshaped = TensorFlowCompBackend.reshape(tensor, (224, 224, 3)) + assert reshaped.tensor.shape == (224, 224, 3) + + +@pytest.mark.tensorflow +def test_stack(): + t0 = TensorFlowTensor(tf.zeros((3, 224, 224))) + t1 = TensorFlowTensor(tf.ones((3, 224, 224))) + + stacked1 = TensorFlowCompBackend.stack([t0, t1], dim=0) + assert isinstance(stacked1, TensorFlowTensor) + assert stacked1.tensor.shape == (2, 3, 224, 224) + + stacked2 = TensorFlowCompBackend.stack([t0, t1], dim=-1) + assert isinstance(stacked2, TensorFlowTensor) + assert stacked2.tensor.shape == (3, 224, 224, 2) diff --git a/tests/units/computation_backends/tensorflow_backend/test_metrics.py b/tests/units/computation_backends/tensorflow_backend/test_metrics.py new file mode 100644 index 00000000000..196297adf0b --- /dev/null +++ b/tests/units/computation_backends/tensorflow_backend/test_metrics.py @@ -0,0 +1,82 @@ +import pytest + +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + + from docarray.computation.tensorflow_backend import TensorFlowCompBackend + from docarray.typing import TensorFlowTensor + + metrics = TensorFlowCompBackend.Metrics +else: + metrics = None + + +@pytest.mark.tensorflow +def test_cosine_sim_tf(): + a = TensorFlowTensor(tf.random.normal((128,))) + b = TensorFlowTensor(tf.random.normal((128,))) + assert metrics.cosine_sim(a, b).tensor.shape == (1,) + assert metrics.cosine_sim(a, b).tensor == metrics.cosine_sim(b, a).tensor + tf.experimental.numpy.allclose(metrics.cosine_sim(a, a).tensor, tf.ones(1)) + + a = TensorFlowTensor(tf.random.normal((10, 3))) + b = TensorFlowTensor(tf.random.normal((5, 3))) + assert metrics.cosine_sim(a, b).tensor.shape == (10, 5) + assert metrics.cosine_sim(b, a).tensor.shape == (5, 10) + diag_dists = tf.linalg.diag(metrics.cosine_sim(b, b).tensor) # self-comparisons + tf.experimental.numpy.allclose(diag_dists, tf.ones(5)) + + +@pytest.mark.tensorflow +def test_euclidean_dist_tf(): + a = TensorFlowTensor(tf.random.normal((128,))) + b = TensorFlowTensor(tf.random.normal((128,))) + assert metrics.euclidean_dist(a, b).tensor.shape == (1,) + assert metrics.euclidean_dist(a, b).tensor == metrics.euclidean_dist(b, a).tensor + tf.experimental.numpy.allclose(metrics.euclidean_dist(a, a).tensor, tf.zeros(1)) + + a = TensorFlowTensor(tf.zeros((1, 1))) + b = TensorFlowTensor(tf.ones((4, 1))) + assert metrics.euclidean_dist(a, b).tensor.shape == (4,) + tf.experimental.numpy.allclose( + metrics.euclidean_dist(a, b).tensor, metrics.euclidean_dist(b, a).tensor + ) + tf.experimental.numpy.allclose(metrics.euclidean_dist(a, a).tensor, tf.zeros(1)) + + a = TensorFlowTensor(tf.constant([0.0, 2.0, 0.0])) + b = TensorFlowTensor(tf.constant([0.0, 0.0, 2.0])) + desired_output_singleton: tf.Tensor = tf.math.sqrt( + tf.constant([2.0**2.0 + 2.0**2.0]) + ) + tf.experimental.numpy.allclose( + metrics.euclidean_dist(a, b).tensor, desired_output_singleton + ) + + a = TensorFlowTensor(tf.constant([[0.0, 2.0, 0.0], [0.0, 0.0, 2.0]])) + b = TensorFlowTensor(tf.constant([[0.0, 0.0, 2.0], [0.0, 2.0, 0.0]])) + desired_output_singleton = tf.constant([[2.828427, 0.0], [0.0, 2.828427]]) + tf.experimental.numpy.allclose( + metrics.euclidean_dist(a, b).tensor, desired_output_singleton + ) + + +@pytest.mark.tensorflow +def test_sqeuclidean_dist_torch(): + a = TensorFlowTensor(tf.random.normal((128,))) + b = TensorFlowTensor(tf.random.normal((128,))) + assert metrics.sqeuclidean_dist(a, b).tensor.shape == (1,) + tf.experimental.numpy.allclose( + metrics.sqeuclidean_dist(a, b).tensor, + metrics.euclidean_dist(a, b).tensor ** 2, + ) + + a = TensorFlowTensor(tf.random.normal((1, 1))) + b = TensorFlowTensor(tf.random.normal((4, 1))) + assert metrics.sqeuclidean_dist(b, a).tensor.shape == (4,) + tf.experimental.numpy.allclose( + metrics.sqeuclidean_dist(a, b).tensor, + metrics.euclidean_dist(a, b).tensor ** 2, + ) diff --git a/tests/units/computation_backends/tensorflow_backend/test_retrieval.py b/tests/units/computation_backends/tensorflow_backend/test_retrieval.py new file mode 100644 index 00000000000..f4d40e7a317 --- /dev/null +++ b/tests/units/computation_backends/tensorflow_backend/test_retrieval.py @@ -0,0 +1,78 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import pytest + +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + import tensorflow._api.v2.experimental.numpy as tnp + + from docarray.computation.tensorflow_backend import TensorFlowCompBackend + from docarray.typing import TensorFlowTensor + + +@pytest.mark.tensorflow +def test_top_k_descending_false(): + top_k = TensorFlowCompBackend.Retrieval.top_k + + a = TensorFlowTensor(tf.constant([1, 4, 2, 7, 4, 9, 2])) + vals, indices = top_k(a, 3, descending=False) + + assert vals.tensor.shape == (1, 3) + assert indices.tensor.shape == (1, 3) + assert tnp.allclose(tnp.squeeze(vals.tensor), tf.constant([1, 2, 2])) + assert tnp.allclose(tnp.squeeze(indices.tensor), tf.constant([0, 2, 6])) or ( + tnp.allclose(tnp.squeeze.indices.tensor), + tf.constant([0, 6, 2]), + ) + + a = TensorFlowTensor(tf.constant([[1, 4, 2, 7, 4, 9, 2], [11, 6, 2, 7, 3, 10, 4]])) + vals, indices = top_k(a, 3, descending=False) + assert vals.tensor.shape == (2, 3) + assert indices.tensor.shape == (2, 3) + assert tnp.allclose(vals.tensor[0], tf.constant([1, 2, 2])) + assert tnp.allclose(indices.tensor[0], tf.constant([0, 2, 6])) or tnp.allclose( + indices.tensor[0], tf.constant([0, 6, 2]) + ) + assert tnp.allclose(vals.tensor[1], tf.constant([2, 3, 4])) + assert tnp.allclose(indices.tensor[1], tf.constant([2, 4, 6])) + + +@pytest.mark.tensorflow +def test_top_k_descending_true(): + top_k = TensorFlowCompBackend.Retrieval.top_k + + a = TensorFlowTensor(tf.constant([1, 4, 2, 7, 4, 9, 2])) + vals, indices = top_k(a, 3, descending=True) + + assert vals.tensor.shape == (1, 3) + assert indices.tensor.shape == (1, 3) + assert tnp.allclose(tnp.squeeze(vals.tensor), tf.constant([9, 7, 4])) + assert tnp.allclose(tnp.squeeze(indices.tensor), tf.constant([5, 3, 1])) + + a = TensorFlowTensor(tf.constant([[1, 4, 2, 7, 4, 9, 2], [11, 6, 2, 7, 3, 10, 4]])) + vals, indices = top_k(a, 3, descending=True) + + assert vals.tensor.shape == (2, 3) + assert indices.tensor.shape == (2, 3) + + assert tnp.allclose(vals.tensor[0], tf.constant([9, 7, 4])) + assert tnp.allclose(indices.tensor[0], tf.constant([5, 3, 1])) + + assert tnp.allclose(vals.tensor[1], tf.constant([11, 10, 7])) + assert tnp.allclose(indices.tensor[1], tf.constant([0, 5, 3])) diff --git a/tests/units/computation_backends/torch_backend/__init__.py b/tests/units/computation_backends/torch_backend/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/computation_backends/torch_backend/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/computation_backends/torch_backend/test_basics.py b/tests/units/computation_backends/torch_backend/test_basics.py new file mode 100644 index 00000000000..b0b98980b7e --- /dev/null +++ b/tests/units/computation_backends/torch_backend/test_basics.py @@ -0,0 +1,161 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytest +import torch + +from docarray.computation.torch_backend import TorchCompBackend + + +def test_to_device(): + t = torch.rand(10, 3) + assert t.device == torch.device('cpu') + t = TorchCompBackend.to_device(t, 'meta') + assert t.device == torch.device('meta') + + +@pytest.mark.parametrize( + 'array,result', + [ + (torch.zeros((5)), 1), + (torch.zeros((1, 5)), 2), + (torch.zeros((5, 5)), 2), + (torch.zeros(()), 0), + ], +) +def test_n_dim(array, result): + assert TorchCompBackend.n_dim(array) == result + + +@pytest.mark.parametrize( + 'array,result', + [ + (torch.zeros((10,)), (10,)), + (torch.zeros((5, 5)), (5, 5)), + (torch.zeros(()), ()), + ], +) +def test_shape(array, result): + shape = TorchCompBackend.shape(array) + assert shape == result + assert type(shape) == tuple + + +@pytest.mark.parametrize('dtype', [torch.int64, torch.float64, torch.int, torch.float]) +def test_dtype(dtype): + tensor = torch.tensor([1, 2, 3], dtype=dtype) + assert TorchCompBackend.dtype(tensor) == dtype + + +def test_device(): + tensor = torch.tensor([1, 2, 3]) + assert TorchCompBackend.device(tensor) == 'cpu' + + +def test_empty(): + tensor = TorchCompBackend.empty((10, 3)) + assert tensor.shape == (10, 3) + + +def test_empty_dtype(): + tensor = TorchCompBackend.empty((10, 3), dtype=torch.int32) + assert tensor.shape == (10, 3) + assert tensor.dtype == torch.int32 + + +def test_empty_device(): + tensor = TorchCompBackend.empty((10, 3), device='meta') + assert tensor.shape == (10, 3) + assert tensor.device == torch.device('meta') + + +def test_squeeze(): + tensor = torch.zeros(size=(1, 1, 3, 1)) + squeezed = TorchCompBackend.squeeze(tensor) + assert squeezed.shape == (3,) + + +@pytest.mark.parametrize( + 'array,t_range,x_range,result', + [ + ( + torch.tensor([0, 1, 2, 3, 4, 5]), + (0, 10), + None, + torch.tensor([0, 2, 4, 6, 8, 10]), + ), + ( + torch.tensor([0, 1, 2, 3, 4, 5]), + (0, 10), + (0, 10), + torch.tensor([0, 1, 2, 3, 4, 5]), + ), + ( + torch.tensor([[0.0, 1.0], [0.0, 1.0]]), + (0, 10), + None, + torch.tensor([[0.0, 10.0], [0.0, 10.0]]), + ), + ], +) +def test_minmax_normalize(array, t_range, x_range, result): + output = TorchCompBackend.minmax_normalize( + tensor=array, t_range=t_range, x_range=x_range + ) + assert torch.allclose(output, result) + + +def test_reshape(): + a = torch.tensor([[[1, 2, 3], [4, 5, 6]]]) + b = TorchCompBackend.reshape(a, (2, 3)) + assert torch.equal(b, torch.tensor([[1, 2, 3], [4, 5, 6]])) + + +def test_copy(): + a = torch.tensor([1, 2, 3]) + b = TorchCompBackend.copy(a) + assert torch.equal(a, b) + + +def test_stack(): + a = torch.tensor([1, 2, 3]) + b = torch.tensor([4, 5, 6]) + stacked = TorchCompBackend.stack([a, b], dim=0) + assert torch.equal(stacked, torch.tensor([[1, 2, 3], [4, 5, 6]])) + + +def test_empty_all(): + shape = (2, 3) + dtype = torch.float32 + device = 'cpu' + a = TorchCompBackend.empty(shape, dtype, device) + assert a.shape == shape and a.dtype == dtype and a.device.type == device + + +def test_to_numpy(): + a = torch.tensor([1, 2, 3]) + b = TorchCompBackend.to_numpy(a) + assert np.array_equal(b, np.array(a)) + + +def test_none_value(): + assert torch.isnan(TorchCompBackend.none_value()) + + +def test_detach(): + a = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) + b = TorchCompBackend.detach(a) + assert not b.requires_grad diff --git a/tests/units/computation_backends/torch_backend/test_metrics.py b/tests/units/computation_backends/torch_backend/test_metrics.py new file mode 100644 index 00000000000..c353e4f77bd --- /dev/null +++ b/tests/units/computation_backends/torch_backend/test_metrics.py @@ -0,0 +1,65 @@ +import torch + +from docarray.computation.torch_backend import TorchCompBackend + +metrics = TorchCompBackend.Metrics + + +def test_cosine_sim_torch(): + a = torch.rand(128) + b = torch.rand(128) + assert metrics.cosine_sim(a, b).shape == (1,) + assert metrics.cosine_sim(a, b) == metrics.cosine_sim(b, a) + torch.testing.assert_close(metrics.cosine_sim(a, a), torch.ones(1)) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + assert metrics.cosine_sim(a, b).shape == (10, 5) + assert metrics.cosine_sim(b, a).shape == (5, 10) + diag_dists = torch.diagonal(metrics.cosine_sim(b, b)) # self-comparisons + torch.testing.assert_allclose(diag_dists, torch.ones(5)) + + +def test_euclidean_dist_torch(): + a = torch.rand(128) + b = torch.rand(128) + assert metrics.euclidean_dist(a, b).shape == (1,) + assert metrics.euclidean_dist(a, b) == metrics.euclidean_dist(b, a) + torch.testing.assert_close(metrics.euclidean_dist(a, a), torch.zeros(1)) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + assert metrics.euclidean_dist(a, b).shape == (10, 5) + assert metrics.euclidean_dist(b, a).shape == (5, 10) + diag_dists = torch.diagonal(metrics.euclidean_dist(b, b)) # self-comparisons + torch.testing.assert_allclose(diag_dists, torch.zeros(5)) + + a = torch.tensor([0.0, 2.0, 0.0]) + b = torch.tensor([0.0, 0.0, 2.0]) + desired_output_singleton = torch.sqrt(torch.tensor([2.0**2.0 + 2.0**2.0])) + torch.testing.assert_close(metrics.euclidean_dist(a, b), desired_output_singleton) + + a = torch.tensor([[0.0, 2.0, 0.0], [0.0, 0.0, 2.0]]) + b = torch.tensor([[0.0, 0.0, 2.0], [0.0, 2.0, 0.0]]) + desired_output_singleton = torch.tensor( + [[desired_output_singleton.item(), 0.0], [0.0, desired_output_singleton.item()]] + ) + torch.testing.assert_close(metrics.euclidean_dist(a, b), desired_output_singleton) + + +def test_sqeuclidean_dist_torch(): + a = torch.rand(128) + b = torch.rand(128) + assert metrics.sqeuclidean_dist(a, b).shape == (1,) + torch.testing.assert_close( + metrics.sqeuclidean_dist(a, b), + metrics.euclidean_dist(a, b) ** 2, + ) + + a = torch.rand(10, 3) + b = torch.rand(5, 3) + assert metrics.sqeuclidean_dist(a, b).shape == (10, 5) + torch.testing.assert_close( + metrics.sqeuclidean_dist(a, b), + metrics.euclidean_dist(a, b) ** 2, + ) diff --git a/tests/units/computation_backends/torch_backend/test_retrieval.py b/tests/units/computation_backends/torch_backend/test_retrieval.py new file mode 100644 index 00000000000..56fc63afc18 --- /dev/null +++ b/tests/units/computation_backends/torch_backend/test_retrieval.py @@ -0,0 +1,42 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch + +from docarray.computation.torch_backend import TorchCompBackend + + +def test_topk(): + top_k = TorchCompBackend.Retrieval.top_k + + a = torch.tensor([1, 4, 2, 7, 4, 9, 2]) + vals, indices = top_k(a, 3) + assert vals.shape == (1, 3) + assert indices.shape == (1, 3) + assert (vals.squeeze() == torch.tensor([1, 2, 2])).all() + assert (indices.squeeze() == torch.tensor([0, 2, 6])).all() or ( + indices.squeeze() == torch.tensor([0, 6, 2]) + ).all() + + a = torch.tensor([[1, 4, 2, 7, 4, 9, 2], [11, 6, 2, 7, 3, 10, 4]]) + vals, indices = top_k(a, 3) + assert vals.shape == (2, 3) + assert indices.shape == (2, 3) + assert (vals[0] == torch.tensor([1, 2, 2])).all() + assert (indices[0] == torch.tensor([0, 2, 6])).all() or ( + indices[0] == torch.tensor([0, 6, 2]) + ).all() + assert (vals[1] == torch.tensor([2, 3, 4])).all() + assert (indices[1] == torch.tensor([2, 4, 6])).all() diff --git a/tests/units/document/__init__.py b/tests/units/document/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/document/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/document/proto/__init__.py b/tests/units/document/proto/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/document/proto/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/document/proto/test_document_proto.py b/tests/units/document/proto/test_document_proto.py new file mode 100644 index 00000000000..0fc16482c6d --- /dev/null +++ b/tests/units/document/proto/test_document_proto.py @@ -0,0 +1,387 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Dict, List, Optional, Set, Tuple + +import numpy as np +import pytest +import torch + +from docarray import DocList +from docarray.base_doc import AnyDoc, BaseDoc +from docarray.documents.image import ImageDoc +from docarray.typing import NdArray, TorchTensor +from docarray.utils._internal.misc import is_tf_available + +if is_tf_available(): + import tensorflow as tf + + +@pytest.mark.proto +def test_proto_simple(): + class CustomDoc(BaseDoc): + text: str + + doc = CustomDoc(text='hello') + + CustomDoc.from_protobuf(doc.to_protobuf()) + + +@pytest.mark.proto +def test_proto_ndarray(): + class CustomDoc(BaseDoc): + tensor: NdArray + + tensor = np.zeros((3, 224, 224)) + doc = CustomDoc(tensor=tensor) + + new_doc = CustomDoc.from_protobuf(doc.to_protobuf()) + + assert (new_doc.tensor == tensor).all() + + +@pytest.mark.proto +def test_proto_with_nested_doc(): + class CustomInnerDoc(BaseDoc): + tensor: NdArray + + class CustomDoc(BaseDoc): + text: str + inner: CustomInnerDoc + + doc = CustomDoc(text='hello', inner=CustomInnerDoc(tensor=np.zeros((3, 224, 224)))) + + CustomDoc.from_protobuf(doc.to_protobuf()) + + +@pytest.mark.proto +def test_proto_with_chunks_doc(): + class CustomInnerDoc(BaseDoc): + tensor: NdArray + + class CustomDoc(BaseDoc): + text: str + chunks: DocList[CustomInnerDoc] + + doc = CustomDoc( + text='hello', + chunks=DocList[CustomInnerDoc]( + [CustomInnerDoc(tensor=np.zeros((3, 224, 224))) for _ in range(5)], + ), + ) + + new_doc = CustomDoc.from_protobuf(doc.to_protobuf()) + + for chunk1, chunk2 in zip(doc.chunks, new_doc.chunks): + assert (chunk1.tensor == chunk2.tensor).all() + + +@pytest.mark.proto +def test_proto_with_nested_doc_pytorch(): + class CustomInnerDoc(BaseDoc): + tensor: TorchTensor + + class CustomDoc(BaseDoc): + text: str + inner: CustomInnerDoc + + doc = CustomDoc( + text='hello', inner=CustomInnerDoc(tensor=torch.zeros((3, 224, 224))) + ) + + CustomDoc.from_protobuf(doc.to_protobuf()) + + +@pytest.mark.proto +def test_proto_with_chunks_doc_pytorch(): + class CustomInnerDoc(BaseDoc): + tensor: TorchTensor + + class CustomDoc(BaseDoc): + text: str + chunks: DocList[CustomInnerDoc] + + doc = CustomDoc( + text='hello', + chunks=DocList[CustomInnerDoc]( + [CustomInnerDoc(tensor=torch.zeros((3, 224, 224))) for _ in range(5)], + ), + ) + + new_doc = CustomDoc.from_protobuf(doc.to_protobuf()) + + for chunk1, chunk2 in zip(doc.chunks, new_doc.chunks): + assert (chunk1.tensor == chunk2.tensor).all() + + +@pytest.mark.proto +def test_optional_field_in_doc(): + class CustomDoc(BaseDoc): + text: Optional[str] = None + + CustomDoc.from_protobuf(CustomDoc().to_protobuf()) + + +@pytest.mark.proto +def test_optional_field_nested_in_doc(): + class InnerDoc(BaseDoc): + title: str + + class CustomDoc(BaseDoc): + text: Optional[InnerDoc] = None + + CustomDoc.from_protobuf(CustomDoc().to_protobuf()) + + +@pytest.mark.proto +def test_integer_field(): + class Meow(BaseDoc): + age: int + wealth: float + registered: bool + + d = Meow(age=30, wealth=100.5, registered=True) + rebuilt_doc = Meow.from_protobuf(d.to_protobuf()) + assert rebuilt_doc.age == 30 + assert rebuilt_doc.wealth == 100.5 + assert rebuilt_doc.registered + + +@pytest.mark.proto +def test_list_set_dict_tuple_field(): + class MyDoc(BaseDoc): + list_: List + dict_: Dict + tuple_: Tuple + set_: Set + + d = MyDoc( + list_=[0, 1, 2], dict_={'a': 0, 'b': 1}, tuple_=tuple([0, 1]), set_={0, 1} + ) + rebuilt_doc = MyDoc.from_protobuf(d.to_protobuf()) + assert rebuilt_doc.list_ == [0, 1, 2] + assert rebuilt_doc.dict_ == {'a': 0, 'b': 1} + assert rebuilt_doc.tuple_ == (0, 1) + assert rebuilt_doc.set_ == {0, 1} + + +@pytest.mark.proto +@pytest.mark.parametrize( + 'dtype', + [ + np.uint, + np.uint8, + np.uint64, + int, + np.int8, + np.int64, + float, + np.float16, + np.longfloat, + np.double, + ], +) +def test_ndarray_dtype(dtype): + class MyDoc(BaseDoc): + tensor: NdArray + + doc = MyDoc(tensor=np.ndarray([1, 2, 3], dtype=dtype)) + assert doc.tensor.dtype == dtype + assert MyDoc.from_protobuf(doc.to_protobuf()).tensor.dtype == dtype + assert MyDoc.parse_obj(doc.dict()).tensor.dtype == dtype + + +@pytest.mark.proto +@pytest.mark.parametrize( + 'dtype', + [ + torch.uint8, + torch.int, + torch.int8, + torch.int64, + torch.float, + torch.float64, + torch.double, + ], +) +def test_torch_dtype(dtype): + class MyDoc(BaseDoc): + tensor: TorchTensor + + doc = MyDoc(tensor=torch.zeros([5, 5], dtype=dtype)) + assert doc.tensor.dtype == dtype + assert MyDoc.from_protobuf(doc.to_protobuf()).tensor.dtype == dtype + assert MyDoc.parse_obj(doc.dict()).tensor.dtype == dtype + + +@pytest.mark.proto +def test_nested_dict(): + class MyDoc(BaseDoc): + data: Dict + + doc = MyDoc(data={'data': (1, 2)}) + + MyDoc.from_protobuf(doc.to_protobuf()) + + +@pytest.mark.proto +def test_nested_dict_error(): + class MyDoc(BaseDoc): + data: Dict + + doc = MyDoc(data={0: (1, 2)}) + + with pytest.raises(ValueError, match="Protobuf only support string as key"): + doc.to_protobuf() + + +@pytest.mark.proto +def test_tuple_complex(): + class MyDoc(BaseDoc): + data: Tuple + + doc = MyDoc(data=(1, 2)) + + doc2 = MyDoc.from_protobuf(doc.to_protobuf()) + + assert doc2.data == (1, 2) + + +@pytest.mark.proto +def test_list_complex(): + class MyDoc(BaseDoc): + data: List + + doc = MyDoc(data=[(1, 2)]) + + doc2 = MyDoc.from_protobuf(doc.to_protobuf()) + + assert doc2.data == [(1, 2)] + + +@pytest.mark.proto +def test_nested_tensor_list(): + class MyDoc(BaseDoc): + data: List + + doc = MyDoc(data=[np.zeros(10)]) + + doc2 = MyDoc.from_protobuf(doc.to_protobuf()) + + assert isinstance(doc2.data[0], np.ndarray) + assert isinstance(doc2.data[0], NdArray) + + assert (doc2.data[0] == np.zeros(10)).all() + + +@pytest.mark.proto +def test_nested_tensor_dict(): + class MyDoc(BaseDoc): + data: Dict + + doc = MyDoc(data={'hello': np.zeros(10)}) + + doc2 = MyDoc.from_protobuf(doc.to_protobuf()) + + assert isinstance(doc2.data['hello'], np.ndarray) + assert isinstance(doc2.data['hello'], NdArray) + + assert (doc2.data['hello'] == np.zeros(10)).all() + + +@pytest.mark.proto +def test_super_complex_nested(): + class MyDoc(BaseDoc): + data: Dict + + data = {'hello': (torch.zeros(55), 1, 'hi', [torch.ones(55), np.zeros(10), (1, 2)])} + doc = MyDoc(data=data) + + doc2 = MyDoc.from_protobuf(doc.to_protobuf()) + + (doc2.data['hello'][3][0] == torch.ones(55)).all() + + +@pytest.mark.tensorflow +def test_super_complex_nested_tensorflow(): + class MyDoc(BaseDoc): + data: Dict + + data = {'hello': (torch.zeros(55), 1, 'hi', [tf.ones(55), np.zeros(10), (1, 2)])} + doc = MyDoc(data=data) + + MyDoc.from_protobuf(doc.to_protobuf()) + + +@pytest.mark.proto +def test_any_doc_proto(): + doc = AnyDoc(hello='world') + pt = doc.to_protobuf() + doc2 = AnyDoc.from_protobuf(pt) + assert doc2.hello == 'world' + + +@pytest.mark.proto +def test_nested_list(): + from typing import List + + from docarray import BaseDoc, DocList + from docarray.documents import TextDoc + + class TextDocWithId(TextDoc): + id: str + + class ResultTestDoc(BaseDoc): + matches: List[TextDocWithId] + + da = DocList[ResultTestDoc]( + [ + ResultTestDoc(matches=[TextDocWithId(id=f'{i}') for _ in range(10)]) + for i in range(10) + ] + ) + + DocList[ResultTestDoc].from_protobuf(da.to_protobuf()) + + +@pytest.mark.proto +def test_nested_dict_typed(): + from docarray import BaseDoc, DocList + from docarray.documents import TextDoc + + class TextDocWithId(TextDoc): + id: str + + class ResultTestDoc(BaseDoc): + matches: Dict[str, TextDocWithId] + + da = DocList[ResultTestDoc]( + [ + ResultTestDoc(matches={f'{i}': TextDocWithId(id=f'{i}') for _ in range(10)}) + for i in range(10) + ] + ) + + DocList[ResultTestDoc].from_protobuf(da.to_protobuf()) + + +def test_image_doc_proto(): + + doc = ImageDoc(url="aux.png") + pt = doc.to_protobuf() + assert "aux.png" in str(pt) + d2 = ImageDoc.from_protobuf(pt) + + assert doc.url == d2.url diff --git a/tests/units/document/proto/test_proto_based_object.py b/tests/units/document/proto/test_proto_based_object.py new file mode 100644 index 00000000000..69849dc99f6 --- /dev/null +++ b/tests/units/document/proto/test_proto_based_object.py @@ -0,0 +1,48 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytest + +from docarray.proto import DocProto, NodeProto +from docarray.typing import NdArray + + +@pytest.mark.proto +def test_ndarray(): + original_ndarray = np.zeros((3, 224, 224)) + + custom_ndarray = NdArray._docarray_from_native(original_ndarray) + + tensor = NdArray.from_protobuf(custom_ndarray.to_protobuf()) + + assert (tensor == original_ndarray).all() + + +@pytest.mark.proto +def test_document_proto_set(): + data = {} + + nested_item1 = NodeProto(text='hello') + + ndarray = NdArray._docarray_from_native(np.zeros((3, 224, 224))) + nd_proto = ndarray.to_protobuf() + + nested_item2 = NodeProto(ndarray=nd_proto) + + data['a'] = nested_item1 + data['b'] = nested_item2 + + DocProto(data=data) diff --git a/tests/units/document/test_any_document.py b/tests/units/document/test_any_document.py new file mode 100644 index 00000000000..7a235b45fed --- /dev/null +++ b/tests/units/document/test_any_document.py @@ -0,0 +1,92 @@ +from typing import Dict, List + +import numpy as np +import pytest + +from docarray import DocList +from docarray.base_doc import AnyDoc, BaseDoc +from docarray.typing import NdArray + + +def test_any_doc(): + class InnerDocument(BaseDoc): + text: str + tensor: NdArray + + class CustomDoc(BaseDoc): + inner: InnerDocument + text: str + + doc = CustomDoc( + text='bye', inner=InnerDocument(text='hello', tensor=np.zeros((3, 224, 224))) + ) + + any_doc = AnyDoc(**doc.__dict__) + + assert any_doc.text == doc.text + assert any_doc.inner.text == doc.inner.text + assert (any_doc.inner.tensor == doc.inner.tensor).all() + + +@pytest.mark.parametrize('protocol', ['proto', 'json']) +def test_any_document_from_to(protocol): + class InnerDoc(BaseDoc): + text: str + t: Dict[str, str] + + class DocTest(BaseDoc): + text: str + tags: Dict[str, int] + l_: List[int] + d: InnerDoc + ld: DocList[InnerDoc] + + inner_doc = InnerDoc(text='I am inner', t={'a': 'b'}) + da = DocList[DocTest]( + [ + DocTest( + text='type1', + tags={'type': 1}, + l_=[1, 2], + d=inner_doc, + ld=DocList[InnerDoc]([inner_doc]), + ), + DocTest( + text='type2', + tags={'type': 2}, + l_=[1, 2], + d=inner_doc, + ld=DocList[InnerDoc]([inner_doc]), + ), + ] + ) + + from docarray.base_doc import AnyDoc + + if protocol == 'proto': + aux = DocList[AnyDoc].from_protobuf(da.to_protobuf()) + else: + aux = DocList[AnyDoc].from_json(da.to_json()) + assert len(aux) == 2 + assert len(aux.id) == 2 + for i, d in enumerate(aux): + assert d.tags['type'] == i + 1 + assert d.text == f'type{i + 1}' + assert d.l_ == [1, 2] + if protocol == 'proto': + assert isinstance(d.d, AnyDoc) + assert d.d.text == 'I am inner' # inner Document is a Dict + assert d.d.t == {'a': 'b'} + else: + assert isinstance(d.d, dict) + assert d.d['text'] == 'I am inner' # inner Document is a Dict + assert d.d['t'] == {'a': 'b'} + assert len(d.ld) == 1 + if protocol == 'proto': + assert isinstance(d.ld[0], AnyDoc) + assert d.ld[0].text == 'I am inner' + assert d.ld[0].t == {'a': 'b'} + else: + assert isinstance(d.ld[0], dict) + assert d.ld[0]['text'] == 'I am inner' + assert d.ld[0]['t'] == {'a': 'b'} diff --git a/tests/units/document/test_base_document.py b/tests/units/document/test_base_document.py new file mode 100644 index 00000000000..2bd80af3763 --- /dev/null +++ b/tests/units/document/test_base_document.py @@ -0,0 +1,189 @@ +from typing import Any, List, Optional, Tuple + +import numpy as np +import orjson +import pytest + +from docarray import DocList, DocVec +from docarray.base_doc.doc import BaseDoc +from docarray.base_doc.io.json import orjson_dumps_and_decode +from docarray.typing import NdArray +from docarray.typing.tensor.abstract_tensor import AbstractTensor +from docarray.utils._internal.pydantic import is_pydantic_v2 + + +def test_base_document_init(): + doc = BaseDoc() + + assert doc.id is not None + + +def test_update(): + class MyDocument(BaseDoc): + content: str + title: Optional[str] = None + tags_: List + + doc1 = MyDocument( + content='Core content of the document', title='Title', tags_=['python', 'AI'] + ) + doc2 = MyDocument(content='Core content updated', tags_=['docarray']) + + doc1.update(doc2) + assert doc1.content == 'Core content updated' + assert doc1.title == 'Title' + assert doc1.tags_ == ['python', 'AI', 'docarray'] + + +def test_equal_nested_docs(): + import numpy as np + + from docarray import BaseDoc, DocList + from docarray.typing import NdArray + + class SimpleDoc(BaseDoc): + simple_tens: NdArray[10] + + class NestedDoc(BaseDoc): + docs: DocList[SimpleDoc] + + nested_docs = NestedDoc( + docs=DocList[SimpleDoc]([SimpleDoc(simple_tens=np.ones(10)) for j in range(2)]), + ) + + assert nested_docs == nested_docs + + +@pytest.fixture +def nested_docs(): + class SimpleDoc(BaseDoc): + simple_tens: NdArray[10] + + class NestedDoc(BaseDoc): + docs: DocList[SimpleDoc] + hello: str = 'world' + + nested_docs = NestedDoc( + docs=DocList[SimpleDoc]([SimpleDoc(simple_tens=np.ones(10)) for j in range(2)]), + ) + + return nested_docs + + +@pytest.fixture +def nested_docs_docvec(): + class SimpleDoc(BaseDoc): + simple_tens: NdArray[10] + + class NestedDoc(BaseDoc): + docs: DocVec[SimpleDoc] + hello: str = 'world' + + nested_docs = NestedDoc( + docs=DocList[SimpleDoc]([SimpleDoc(simple_tens=np.ones(10)) for j in range(2)]), + ) + + return nested_docs + + +def test_nested_to_dict(nested_docs): + d = nested_docs.dict() + assert (d['docs'][0]['simple_tens'] == np.ones(10)).all() + assert isinstance(d['docs'], list) + assert not isinstance(d['docs'], DocList) + + +def test_nested_docvec_to_dict(nested_docs_docvec): + d = nested_docs_docvec.dict() + assert (d['docs'][0]['simple_tens'] == np.ones(10)).all() + + +def test_nested_to_dict_exclude(nested_docs): + d = nested_docs.dict(exclude={'docs'}) + assert 'docs' not in d.keys() + + +def test_nested_to_dict_exclude_set(nested_docs): + d = nested_docs.dict(exclude={'hello'}) + assert 'hello' not in d.keys() + + +def test_nested_to_dict_exclude_dict(nested_docs): + d = nested_docs.dict(exclude={'hello': True}) + assert 'hello' not in d.keys() + + +def test_nested_to_json(nested_docs): + d = nested_docs.json() + nested_docs.__class__.parse_raw(d) + + +@pytest.fixture +def nested_none_docs(): + class SimpleDoc(BaseDoc): + simple_tens: NdArray[10] + + class NestedDoc(BaseDoc): + docs: Optional[DocList[SimpleDoc]] = None + hello: str = 'world' + + nested_docs = NestedDoc() + + return nested_docs + + +def test_nested_none_to_dict(nested_none_docs): + d = nested_none_docs.dict() + assert d == {'docs': None, 'hello': 'world', 'id': nested_none_docs.id} + + +def test_nested_none_to_json(nested_none_docs): + d = nested_none_docs.json() + d = nested_none_docs.__class__.parse_raw(d) + assert d.dict() == {'docs': None, 'hello': 'world', 'id': nested_none_docs.id} + + +def test_get_get_field_inner_type(): + class MyDoc(BaseDoc): + tuple_: Tuple + + field_type = MyDoc._get_field_inner_type("tuple_") + + assert field_type == Any + + +@pytest.mark.skipif( + is_pydantic_v2, reason="syntax only working with pydantic v1 for now" +) +def test_subclass_config(): + class MyDoc(BaseDoc): + x: str + + class Config(BaseDoc.Config): + arbitrary_types_allowed = True # just an example setting + + assert MyDoc.Config.json_loads == orjson.loads + assert MyDoc.Config.json_dumps == orjson_dumps_and_decode + assert ( + MyDoc.Config.json_encoders[AbstractTensor](3) == 3 + ) # dirty check that it is identity + assert MyDoc.Config.validate_assignment + assert not MyDoc.Config._load_extra_fields_from_protobuf + assert MyDoc.Config.arbitrary_types_allowed + + +@pytest.mark.skipif(not (is_pydantic_v2), reason="syntax only working with pydantic v2") +def test_subclass_config_v2(): + class MyDoc(BaseDoc): + x: str + + model_config = BaseDoc.ConfigDocArray( + arbitrary_types_allowed=True + ) # just an example setting + + assert ( + MyDoc.model_config['json_encoders'][AbstractTensor](3) == 3 + ) # dirty check that it is identity + assert MyDoc.model_config['validate_assignment'] + assert not MyDoc.model_config['_load_extra_fields_from_protobuf'] + assert MyDoc.model_config['arbitrary_types_allowed'] diff --git a/tests/units/document/test_doc_wo_id.py b/tests/units/document/test_doc_wo_id.py new file mode 100644 index 00000000000..4e2a8bba118 --- /dev/null +++ b/tests/units/document/test_doc_wo_id.py @@ -0,0 +1,31 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from docarray import DocList +from docarray.base_doc.doc import BaseDocWithoutId + + +def test_doc_list(): + class A(BaseDocWithoutId): + text: str + + cls_doc_list = DocList[A] + + da = cls_doc_list([A(text='hey here')]) + + assert isinstance(da, DocList) + for d in da: + assert isinstance(d, A) + assert not hasattr(d, 'id') diff --git a/tests/units/document/test_docs_operators.py b/tests/units/document/test_docs_operators.py new file mode 100644 index 00000000000..36cfc258811 --- /dev/null +++ b/tests/units/document/test_docs_operators.py @@ -0,0 +1,40 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from docarray.documents.text import TextDoc + + +def test_text_document_operators(): + doc = TextDoc(text='text', url='http://url.com') + + assert doc == 'text' + assert doc != 'http://url.com' + + doc2 = TextDoc(id=doc.id, text='text', url='http://url.com') + assert doc == doc2 + + doc3 = TextDoc(id='other-id', text='text', url='http://url.com') + assert doc == doc3 + + assert 't' in doc + assert 'a' not in doc + + t = TextDoc(text='this is my text document') + assert 'text' in t + assert 'docarray' not in t + + text = TextDoc() + assert text is not None + assert text.text is None diff --git a/tests/units/document/test_from_to_bytes.py b/tests/units/document/test_from_to_bytes.py new file mode 100644 index 00000000000..9ee971eb5c5 --- /dev/null +++ b/tests/units/document/test_from_to_bytes.py @@ -0,0 +1,117 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import pytest +from typing import Dict, List + +from docarray import BaseDoc, DocList +from docarray.documents import ImageDoc +from docarray.typing import NdArray + + +class MyDoc(BaseDoc): + embedding: NdArray + text: str + image: ImageDoc + + +class MySimpleDoc(BaseDoc): + title: str + + +class MyComplexDoc(BaseDoc): + content_dict_doclist: Dict[str, DocList[MySimpleDoc]] + content_dict_list: Dict[str, List[MySimpleDoc]] + aux_dict: Dict[str, int] + + +@pytest.mark.parametrize('protocol', ['protobuf', 'pickle']) +@pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) +def test_to_from_bytes(protocol, compress): + d = MyDoc(embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png')) + + assert d.text == 'hello' + assert d.embedding.tolist() == [1, 2, 3, 4, 5] + assert d.image.url == 'aux.png' + bstr = d.to_bytes(protocol=protocol, compress=compress) + d2 = MyDoc.from_bytes(bstr, protocol=protocol, compress=compress) + assert d2.text == 'hello' + assert d2.embedding.tolist() == [1, 2, 3, 4, 5] + assert d2.image.url == 'aux.png' + + +@pytest.mark.parametrize('protocol', ['protobuf', 'pickle']) +@pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) +def test_to_from_base64(protocol, compress): + d = MyDoc(embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png')) + + assert d.text == 'hello' + assert d.embedding.tolist() == [1, 2, 3, 4, 5] + assert d.image.url == 'aux.png' + bstr = d.to_base64(protocol=protocol, compress=compress) + d2 = MyDoc.from_base64(bstr, protocol=protocol, compress=compress) + assert d2.text == 'hello' + assert d2.embedding.tolist() == [1, 2, 3, 4, 5] + assert d2.image.url == 'aux.png' + + +@pytest.mark.parametrize('protocol', ['protobuf', 'pickle']) +@pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) +def test_to_from_bytes_complex(protocol, compress): + d = MyComplexDoc( + content_dict_doclist={ + 'test1': DocList[MySimpleDoc]( + [MySimpleDoc(title='123'), MySimpleDoc(title='456')] + ) + }, + content_dict_list={ + 'test1': [MySimpleDoc(title='123'), MySimpleDoc(title='456')] + }, + aux_dict={'a': 0}, + ) + bstr = d.to_bytes(protocol=protocol, compress=compress) + d2 = MyComplexDoc.from_bytes(bstr, protocol=protocol, compress=compress) + assert d2.aux_dict == {'a': 0} + assert len(d2.content_dict_doclist['test1']) == 2 + assert d2.content_dict_doclist['test1'][0].title == '123' + assert d2.content_dict_doclist['test1'][1].title == '456' + assert len(d2.content_dict_list['test1']) == 2 + assert d2.content_dict_list['test1'][0].title == '123' + assert d2.content_dict_list['test1'][1].title == '456' + + +@pytest.mark.parametrize('protocol', ['protobuf', 'pickle']) +@pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) +def test_to_from_base64_complex(protocol, compress): + d = MyComplexDoc( + content_dict_doclist={ + 'test1': DocList[MySimpleDoc]( + [MySimpleDoc(title='123'), MySimpleDoc(title='456')] + ) + }, + content_dict_list={ + 'test1': [MySimpleDoc(title='123'), MySimpleDoc(title='456')] + }, + aux_dict={'a': 0}, + ) + bstr = d.to_base64(protocol=protocol, compress=compress) + d2 = MyComplexDoc.from_base64(bstr, protocol=protocol, compress=compress) + assert d2.aux_dict == {'a': 0} + assert len(d2.content_dict_doclist['test1']) == 2 + assert d2.content_dict_doclist['test1'][0].title == '123' + assert d2.content_dict_doclist['test1'][1].title == '456' + assert len(d2.content_dict_list['test1']) == 2 + assert d2.content_dict_list['test1'][0].title == '123' + assert d2.content_dict_list['test1'][1].title == '456' diff --git a/tests/units/document/test_text_document.py b/tests/units/document/test_text_document.py new file mode 100644 index 00000000000..153e2922ead --- /dev/null +++ b/tests/units/document/test_text_document.py @@ -0,0 +1,30 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from docarray.documents import TextDoc + + +def test_text_document_init(): + text = TextDoc('hello world') + assert text.text == 'hello world' + assert text == 'hello world' + + text = TextDoc(text='hello world') + assert text.text == 'hello world' + assert text == 'hello world' + + text = TextDoc() + assert text is not None + assert text.text is None diff --git a/tests/units/document/test_to_schema.py b/tests/units/document/test_to_schema.py new file mode 100644 index 00000000000..ad0b7444acd --- /dev/null +++ b/tests/units/document/test_to_schema.py @@ -0,0 +1,92 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytest + +from docarray import BaseDoc +from docarray.base_doc.io.json import orjson_dumps +from docarray.typing import NdArray, TorchTensor + + +class NpDoc(BaseDoc): + embedding: NdArray[3, 4] + embedding_no_shape: NdArray + + +class TorchDoc(BaseDoc): + embedding: TorchTensor[3, 4] + embedding_no_shape: TorchTensor + + +def test_np_schema(): + schema = NpDoc.schema() + assert schema['properties']['embedding']['tensor/array shape'] == '[3, 4]' + assert schema['properties']['embedding']['type'] == 'array' + assert schema['properties']['embedding']['items']['type'] == 'number' + assert ( + schema['properties']['embedding']['example'] + == orjson_dumps(np.zeros([3, 4])).decode() + ) + + assert ( + schema['properties']['embedding_no_shape']['tensor/array shape'] + == 'not specified' + ) + assert schema['properties']['embedding_no_shape']['type'] == 'array' + assert schema['properties']['embedding']['items']['type'] == 'number' + + +def test_torch_schema(): + schema = TorchDoc.schema() + assert schema['properties']['embedding']['tensor/array shape'] == '[3, 4]' + assert schema['properties']['embedding']['type'] == 'array' + assert schema['properties']['embedding']['items']['type'] == 'number' + assert ( + schema['properties']['embedding']['example'] + == orjson_dumps(np.zeros([3, 4])).decode() + ) + + assert ( + schema['properties']['embedding_no_shape']['tensor/array shape'] + == 'not specified' + ) + assert schema['properties']['embedding_no_shape']['type'] == 'array' + assert schema['properties']['embedding']['items']['type'] == 'number' + + +@pytest.mark.tensorflow +def test_tensorflow_schema(): + from docarray.typing import TensorFlowTensor + + class TensorflowDoc(BaseDoc): + embedding: TensorFlowTensor[3, 4] + embedding_no_shape: TensorFlowTensor + + schema = TensorflowDoc.schema() + assert schema['properties']['embedding']['tensor/array shape'] == '[3, 4]' + assert schema['properties']['embedding']['type'] == 'array' + assert schema['properties']['embedding']['items']['type'] == 'number' + assert ( + schema['properties']['embedding']['example'] + == orjson_dumps(np.zeros([3, 4])).decode() + ) + + assert ( + schema['properties']['embedding_no_shape']['tensor/array shape'] + == 'not specified' + ) + assert schema['properties']['embedding_no_shape']['type'] == 'array' + assert schema['properties']['embedding']['items']['type'] == 'number' diff --git a/tests/units/document/test_update.py b/tests/units/document/test_update.py new file mode 100644 index 00000000000..5e76caa0dc2 --- /dev/null +++ b/tests/units/document/test_update.py @@ -0,0 +1,104 @@ +from typing import Dict, List, Optional, Set + +import pytest + +from docarray import BaseDoc, DocList +from docarray.documents import ImageDoc + + +class InnerDoc(BaseDoc): + integer: int + inner_list: List + + +class MMDoc(BaseDoc): + text: str = '' + price: int = 0 + categories: Optional[List[str]] = None + image: Optional[ImageDoc] = None + matches: Optional[DocList] = None + matches_with_same_id: Optional[DocList] = None + opt_int: Optional[int] = None + test_set: Optional[Set] = None + inner_doc: Optional[InnerDoc] = None + test_dict: Optional[Dict] = None + + +@pytest.fixture +def doc1(): + return MMDoc( + text='hey here', + categories=['a', 'b', 'c'], + price=10, + matches=DocList[MMDoc]([MMDoc()]), + matches_with_same_id=DocList[MMDoc]( + [MMDoc(id='a', matches=DocList[MMDoc]([MMDoc()]))] + ), + test_set={'a', 'a'}, + inner_doc=InnerDoc(integer=2, inner_list=['c', 'd']), + test_dict={'a': 0, 'b': 2, 'd': 4, 'z': 3}, + ) + + +@pytest.fixture +def doc2(doc1): + return MMDoc( + id=doc1.id, + text='hey here 2', + categories=['d', 'e', 'f'], + price=5, + opt_int=5, + matches=DocList[MMDoc]([MMDoc()]), + matches_with_same_id=DocList[MMDoc]( + [MMDoc(id='a', matches=DocList[MMDoc]([MMDoc()]))] + ), + test_set={'a', 'b'}, + inner_doc=InnerDoc(integer=3, inner_list=['a', 'b']), + test_dict={'a': 10, 'b': 10, 'c': 3, 'z': None}, + ) + + +def test_update_complex(doc1, doc2): + doc1.update(doc2) + # doc1 is changed in place (no extra memory) + assert doc1.text == 'hey here 2' + assert doc1.categories == ['a', 'b', 'c', 'd', 'e', 'f'] + assert len(doc1.matches) == 2 + assert doc1.opt_int == 5 + assert doc1.price == 5 + assert doc1.test_set == {'a', 'b'} + assert len(doc1.matches_with_same_id) == 1 + assert len(doc1.matches_with_same_id[0].matches) == 2 + assert doc1.inner_doc.integer == 3 + assert doc1.inner_doc.inner_list == ['c', 'd', 'a', 'b'] + assert doc1.test_dict == {'a': 10, 'b': 10, 'c': 3, 'd': 4, 'z': None} + + +def test_update_simple(): + class MyDocument(BaseDoc): + content: str + title: Optional[str] = None + tags_: List + + my_doc1 = MyDocument( + content='Core content of the document', title='Title', tags_=['python', 'AI'] + ) + my_doc2 = MyDocument(content='Core content updated', tags_=['docarray']) + + my_doc1.update(my_doc2) + assert my_doc1.content == 'Core content updated' + assert my_doc1.title == 'Title' + assert my_doc1.tags_ == ['python', 'AI', 'docarray'] + + +def test_update_different_schema_fails(): + class DocA(BaseDoc): + content: str + + class DocB(BaseDoc): + image: Optional[ImageDoc] = None + + docA = DocA(content='haha') + docB = DocB() + with pytest.raises(Exception): + docA.update(docB) diff --git a/tests/units/document/test_view.py b/tests/units/document/test_view.py new file mode 100644 index 00000000000..ecd53a918fa --- /dev/null +++ b/tests/units/document/test_view.py @@ -0,0 +1,52 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np + +from docarray import BaseDoc +from docarray.array import DocVec +from docarray.array.doc_vec.column_storage import ColumnStorageView +from docarray.typing import AnyTensor + + +def test_document_view(): + class MyDoc(BaseDoc): + tensor: AnyTensor + name: str + + docs = [MyDoc(tensor=np.zeros((10, 10)), name='hello', id=str(i)) for i in range(4)] + + doc_vec = DocVec[MyDoc](docs) + storage = doc_vec._storage + + result = str(doc_vec[0]) + assert 'MyDoc' in result + assert 'id' in result + assert 'tensor' in result + assert 'name' in result + + doc = MyDoc.from_view(ColumnStorageView(0, storage)) + assert doc.is_view() + assert doc.id == '0' + assert (doc.tensor == np.zeros(10)).all() + assert doc.name == 'hello' + + storage.columns['id'][0] = '12345' + storage.columns['tensor'][0] = np.ones(10) + storage.columns['name'][0] = 'byebye' + + assert doc.id == '12345' + assert (doc.tensor == np.ones(10)).all() + assert doc.name == 'byebye' diff --git a/tests/units/test_helper.py b/tests/units/test_helper.py new file mode 100644 index 00000000000..0c68fe9884d --- /dev/null +++ b/tests/units/test_helper.py @@ -0,0 +1,181 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Optional + +import pytest + +from docarray import BaseDoc, DocList +from docarray.documents import ImageDoc +from docarray.helper import ( + _access_path_dict_to_nested_dict, + _access_path_to_dict, + _dict_to_access_paths, + _is_access_path_valid, + _update_nested_dicts, + get_paths, +) + + +@pytest.fixture() +def nested_doc(): + class Inner(BaseDoc): + img: Optional[ImageDoc] + + class Middle(BaseDoc): + img: Optional[ImageDoc] + inner: Optional[Inner] + + class Outer(BaseDoc): + img: Optional[ImageDoc] + middle: Optional[Middle] + da: DocList[Inner] + + doc = Outer( + img=ImageDoc(), + middle=Middle(img=ImageDoc(), inner=Inner(img=ImageDoc())), + da=DocList[Inner]([Inner(img=ImageDoc(url='test.png'))]), + ) + return doc + + +def test_is_access_path_valid(nested_doc): + assert _is_access_path_valid(nested_doc.__class__, 'img') + assert _is_access_path_valid(nested_doc.__class__, 'middle__img') + assert _is_access_path_valid(nested_doc.__class__, 'middle__inner__img') + assert _is_access_path_valid(nested_doc.__class__, 'middle') + assert _is_access_path_valid(nested_doc.__class__, 'da__img__url') + + +def test_is_access_path_not_valid(nested_doc): + assert not _is_access_path_valid(nested_doc.__class__, 'inner') + assert not _is_access_path_valid(nested_doc.__class__, 'some__other__path') + assert not _is_access_path_valid(nested_doc.__class__, 'middle.inner') + + +def test_get_access_paths(): + class Painting(BaseDoc): + title: str + img: ImageDoc + + access_paths = Painting._get_access_paths() + assert access_paths == [ + 'id', + 'title', + 'img__id', + 'img__url', + 'img__tensor', + 'img__embedding', + 'img__bytes_', + ] + + +def test_dict_to_access_paths(): + d = { + 'a0': {'b0': {'c0': 0}, 'b1': {'c0': 1}}, + 'a1': {'b0': {'c0': 2, 'c1': 3}, 'b1': 4}, + } + casted = _dict_to_access_paths(d) + assert casted == { + 'a0__b0__c0': 0, + 'a0__b1__c0': 1, + 'a1__b0__c0': 2, + 'a1__b0__c1': 3, + 'a1__b1': 4, + } + + +def test_access_path_to_dict(): + access_path = 'a__b__c__d__e' + value = 1 + result = {'a': {'b': {'c': {'d': {'e': value}}}}} + assert _access_path_to_dict(access_path, value) == result + + +def test_access_path_dict_to_nested_dict(): + d = { + 'a0__b0__c0': 0, + 'a0__b1__c0': 1, + 'a1__b0__c0': 2, + 'a1__b0__c1': 3, + 'a1__b1': 4, + } + casted = _access_path_dict_to_nested_dict(d) + assert casted == { + 'a0': {'b0': {'c0': 0}, 'b1': {'c0': 1}}, + 'a1': {'b0': {'c0': 2, 'c1': 3}, 'b1': 4}, + } + + +def test_update_nested_dict(): + d1 = {'text': 'hello', 'image': {'tensor': None}} + d2 = {'image': {'url': 'some.png'}} + + _update_nested_dicts(d1, d2) + assert d1 == {'text': 'hello', 'image': {'tensor': None, 'url': 'some.png'}} + + +def test_get_paths(): + paths = list(get_paths(patterns='*.py')) + for path in paths: + assert path.endswith('.py') + + +def test_get_paths_recursive(): + paths_rec = list(get_paths(patterns='**', recursive=True)) + paths_not_rec = list(get_paths(patterns='**', recursive=False)) + + assert len(paths_rec) > len(paths_not_rec) + + +def test_get_paths_exclude(): + paths = list(get_paths(patterns='*.py')) + paths_wo_init = list(get_paths(patterns='*.py', exclude_regex='__init__.[a-z]*')) + + assert len(paths_wo_init) <= len(paths) + assert '__init__.py' not in paths_wo_init + + +def test_shallow_copy(): + from torch import rand + + from docarray import BaseDoc + from docarray.helper import _shallow_copy_doc + from docarray.typing import TorchTensor, VideoUrl + + class VideoDoc(BaseDoc): + url: VideoUrl + tensor_video: TorchTensor + + class MyDoc(BaseDoc): + docs: DocList[VideoDoc] + tensor: TorchTensor + + doc_ori = MyDoc( + docs=DocList[VideoDoc]( + [ + VideoDoc( + url=f'http://example.ai/videos/{i}', + tensor_video=rand(256), + ) + for i in range(10) + ] + ), + tensor=rand(256), + ) + + doc_copy = _shallow_copy_doc(doc_ori) + + assert doc_copy == doc_ori diff --git a/tests/units/typing/__init__.py b/tests/units/typing/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/typing/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/typing/da/__init__.py b/tests/units/typing/da/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/typing/da/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/typing/da/test_relations.py b/tests/units/typing/da/test_relations.py new file mode 100644 index 00000000000..cadac712f5a --- /dev/null +++ b/tests/units/typing/da/test_relations.py @@ -0,0 +1,59 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from docarray import BaseDoc, DocList +from docarray.utils._internal.pydantic import is_pydantic_v2 + + +@pytest.mark.skipif( + is_pydantic_v2, + reason="Subscripted generics cannot be used with class and instance checks", +) +def test_instance_and_equivalence(): + class MyDoc(BaseDoc): + text: str + + docs = DocList[MyDoc]([MyDoc(text='hello')]) + + assert issubclass(DocList[MyDoc], DocList[MyDoc]) + assert issubclass(docs.__class__, DocList[MyDoc]) + + assert isinstance(docs, DocList[MyDoc]) + + +@pytest.mark.skipif( + is_pydantic_v2, + reason="Subscripted generics cannot be used with class and instance checks", +) +def test_subclassing(): + class MyDoc(BaseDoc): + text: str + + class MyDocList(DocList[MyDoc]): + pass + + docs = MyDocList([MyDoc(text='hello')]) + + assert issubclass(MyDocList, DocList[MyDoc]) + assert issubclass(docs.__class__, DocList[MyDoc]) + + assert isinstance(docs, MyDocList) + assert isinstance(docs, DocList[MyDoc]) + + assert issubclass(MyDoc, BaseDoc) + assert not issubclass(DocList[MyDoc], DocList[BaseDoc]) + assert not issubclass(MyDocList, DocList[BaseDoc]) diff --git a/tests/units/typing/tensor/__init__.py b/tests/units/typing/tensor/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/typing/tensor/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/typing/tensor/test_audio_tensor.py b/tests/units/typing/tensor/test_audio_tensor.py new file mode 100644 index 00000000000..45b54caf654 --- /dev/null +++ b/tests/units/typing/tensor/test_audio_tensor.py @@ -0,0 +1,182 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import numpy as np +import pytest +import torch +from pydantic import parse_obj_as + +from docarray import BaseDoc +from docarray.typing import AudioTensor +from docarray.typing.bytes.audio_bytes import AudioBytes +from docarray.typing.tensor.audio.audio_ndarray import AudioNdArray +from docarray.typing.tensor.audio.audio_torch_tensor import AudioTorchTensor +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + import tensorflow._api.v2.experimental.numpy as tnp + + from docarray.typing.tensor.audio import AudioTensorFlowTensor + + +@pytest.mark.parametrize( + 'tensor,cls_audio_tensor,cls_tensor', + [ + (torch.zeros(1000, 2), AudioTorchTensor, torch.Tensor), + (np.zeros((1000, 2)), AudioNdArray, np.ndarray), + ], +) +def test_set_audio_tensor(tensor, cls_audio_tensor, cls_tensor): + class MyAudioDoc(BaseDoc): + tensor: cls_audio_tensor + + doc = MyAudioDoc(tensor=tensor) + assert isinstance(doc.tensor, cls_audio_tensor) + assert isinstance(doc.tensor, cls_tensor) + assert (doc.tensor == tensor).all() + + +@pytest.mark.tensorflow +def test_set_audio_tensorflow_tensor(): + class MyAudioDoc(BaseDoc): + tensor: AudioTensorFlowTensor + + doc = MyAudioDoc(tensor=tf.zeros((1000, 2))) + assert isinstance(doc.tensor, AudioTensorFlowTensor) + assert isinstance(doc.tensor.tensor, tf.Tensor) + assert tnp.allclose(doc.tensor.tensor, tf.zeros((1000, 2))) + + +@pytest.mark.parametrize( + 'cls_tensor,tensor', + [ + (AudioNdArray, np.zeros((1000, 2))), + (AudioTorchTensor, torch.zeros(1000, 2)), + (AudioTorchTensor, np.zeros((1000, 2))), + ], +) +def test_validation(cls_tensor, tensor): + arr = parse_obj_as(cls_tensor, tensor) + assert isinstance(arr, cls_tensor) + + +@pytest.mark.tensorflow +def test_validation_tensorflow(): + arr = parse_obj_as(AudioTensorFlowTensor, tf.zeros((1000, 2))) + assert isinstance(arr, AudioTensorFlowTensor) + + +@pytest.mark.parametrize( + 'cls_tensor,tensor,expect_error', + [ + (AudioNdArray, torch.zeros(1000, 2), False), + (AudioNdArray, 'hello', True), + (AudioTorchTensor, 'hello', True), + ], +) +def test_illegal_validation(cls_tensor, tensor, expect_error): + if expect_error: + with pytest.raises(ValueError): + parse_obj_as(cls_tensor, tensor) + else: + parse_obj_as(cls_tensor, tensor) + + +@pytest.mark.proto +@pytest.mark.parametrize( + 'cls_tensor,tensor,proto_key', + [ + (AudioTorchTensor, torch.zeros(1000, 2), AudioTorchTensor._proto_type_name), + (AudioNdArray, np.zeros((1000, 2)), AudioNdArray._proto_type_name), + ], +) +def test_proto_tensor(cls_tensor, tensor, proto_key): + tensor = parse_obj_as(cls_tensor, tensor) + proto = tensor._to_node_protobuf() + assert proto_key in str(proto) + + +@pytest.mark.tensorflow +def test_proto_tensor_tensorflow(): + tensor = parse_obj_as(AudioTensorFlowTensor, tf.zeros((1000, 2))) + proto = tensor._to_node_protobuf() + assert AudioTensorFlowTensor._proto_type_name in str(proto) + + +@pytest.mark.parametrize( + 'cls_tensor,tensor', + [ + (AudioTorchTensor, torch.zeros(1000, 2)), + (AudioNdArray, np.zeros((1000, 2))), + ], +) +def test_save_audio_tensor_to_wav_file(cls_tensor, tensor, tmpdir): + tmp_file = str(tmpdir / 'tmp.wav') + audio_tensor = parse_obj_as(cls_tensor, tensor) + audio_tensor.save(tmp_file) + assert os.path.isfile(tmp_file) + + +@pytest.mark.tensorflow +def test_save_audio_tensorflow_tensor_to_wav_file(tmpdir): + tmp_file = str(tmpdir / 'tmp.wav') + audio_tensor = parse_obj_as(AudioTensorFlowTensor, tf.zeros((1000, 2))) + audio_tensor.save(tmp_file) + assert os.path.isfile(tmp_file) + + +@pytest.mark.parametrize( + 'audio_tensor', + [ + parse_obj_as(AudioTorchTensor, torch.zeros(1000, 2)), + parse_obj_as(AudioNdArray, np.zeros((1000, 2))), + ], +) +def test_save_audio_tensor_to_bytes(audio_tensor): + b = audio_tensor.to_bytes() + isinstance(b, bytes) + isinstance(b, AudioBytes) + + +@pytest.mark.parametrize( + 'tensor,cls_audio_tensor,cls_tensor', + [ + (torch.zeros(1000, 2), AudioTorchTensor, torch.Tensor), + (np.zeros((1000, 2)), AudioNdArray, np.ndarray), + ], +) +def test_torch_ndarray_to_audio_tensor(tensor, cls_audio_tensor, cls_tensor): + class MyAudioDoc(BaseDoc): + tensor: AudioTensor + + doc = MyAudioDoc(tensor=tensor) + assert isinstance(doc.tensor, cls_audio_tensor) + assert isinstance(doc.tensor, cls_tensor) + assert (doc.tensor == tensor).all() + + +@pytest.mark.tensorflow +def test_tensorflow_to_audio_tensor(): + class MyAudioDoc(BaseDoc): + tensor: AudioTensor + + doc = MyAudioDoc(tensor=tf.zeros((1000, 2))) + assert isinstance(doc.tensor, AudioTensorFlowTensor) + assert isinstance(doc.tensor.tensor, tf.Tensor) + assert tnp.allclose(doc.tensor.tensor, tf.zeros((1000, 2))) diff --git a/tests/units/typing/tensor/test_cross_backend.py b/tests/units/typing/tensor/test_cross_backend.py new file mode 100644 index 00000000000..cd5403c49c7 --- /dev/null +++ b/tests/units/typing/tensor/test_cross_backend.py @@ -0,0 +1,44 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytest +from pydantic import parse_obj_as + +from docarray.typing import NdArray, TorchTensor + +try: + from docarray.typing import TensorFlowTensor +except (ImportError, TypeError): + pass + + +@pytest.mark.tensorflow +def test_coercion_behavior(): + t_np = parse_obj_as(NdArray[128], np.zeros(128)) + t_th = parse_obj_as(TorchTensor[128], np.zeros(128)) + t_tf = parse_obj_as(TensorFlowTensor[128], np.zeros(128)) + + assert isinstance(t_np, NdArray[128]) + assert not isinstance(t_np, TensorFlowTensor[128]) + assert not isinstance(t_np, TorchTensor[128]) + + assert isinstance(t_th, TorchTensor[128]) + assert not isinstance(t_th, NdArray[128]) + assert not isinstance(t_th, TensorFlowTensor[128]) + + assert isinstance(t_tf, TensorFlowTensor[128]) + assert not isinstance(t_tf, TorchTensor[128]) + assert not isinstance(t_tf, NdArray[128]) diff --git a/tests/units/typing/tensor/test_embedding.py b/tests/units/typing/tensor/test_embedding.py new file mode 100644 index 00000000000..078cb7fddbb --- /dev/null +++ b/tests/units/typing/tensor/test_embedding.py @@ -0,0 +1,75 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytest +import torch +from pydantic.tools import parse_obj_as, schema_json_of + +from docarray import BaseDoc +from docarray.base_doc.io.json import orjson_dumps +from docarray.typing import AnyEmbedding, NdArrayEmbedding, TorchEmbedding +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + + from docarray.computation.tensorflow_backend import tnp + from docarray.typing.tensor.embedding import TensorFlowEmbedding + + +@pytest.mark.proto +def test_proto_embedding(): + embedding = parse_obj_as(AnyEmbedding, np.zeros((3, 224, 224))) + + embedding._to_node_protobuf() + + +def test_json_schema(): + schema_json_of(AnyEmbedding) + + +def test_dump_json(): + tensor = parse_obj_as(AnyEmbedding, np.zeros((3, 224, 224))) + orjson_dumps(tensor) + + +@pytest.mark.parametrize( + 'tensor,cls_audio_tensor,cls_tensor', + [ + (torch.zeros(1000, 2), TorchEmbedding, torch.Tensor), + (np.zeros((1000, 2)), NdArrayEmbedding, np.ndarray), + ], +) +def test_torch_ndarray_to_any_embedding(tensor, cls_audio_tensor, cls_tensor): + class MyEmbeddingDoc(BaseDoc): + tensor: AnyEmbedding + + doc = MyEmbeddingDoc(tensor=tensor) + assert isinstance(doc.tensor, cls_audio_tensor) + assert isinstance(doc.tensor, cls_tensor) + assert (doc.tensor == tensor).all() + + +@pytest.mark.tensorflow +def test_tensorflow_to_any_embedding(): + class MyEmbeddingDoc(BaseDoc): + tensor: AnyEmbedding + + doc = MyEmbeddingDoc(tensor=tf.zeros((1000, 2))) + assert isinstance(doc.tensor, TensorFlowEmbedding) + assert isinstance(doc.tensor.tensor, tf.Tensor) + assert tnp.allclose(doc.tensor.tensor, tf.zeros((1000, 2))) diff --git a/tests/units/typing/tensor/test_image_tensor.py b/tests/units/typing/tensor/test_image_tensor.py new file mode 100644 index 00000000000..b05a71403a2 --- /dev/null +++ b/tests/units/typing/tensor/test_image_tensor.py @@ -0,0 +1,95 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import numpy as np +import pytest +import torch +from pydantic import parse_obj_as + +from docarray import BaseDoc +from docarray.typing import ImageBytes, ImageNdArray, ImageTensor, ImageTorchTensor +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + + from docarray.computation.tensorflow_backend import tnp + from docarray.typing.tensor.image import ImageTensorFlowTensor + + +@pytest.mark.parametrize( + 'cls_tensor,tensor', + [ + (ImageTorchTensor, torch.zeros((224, 224, 3))), + (ImageNdArray, np.zeros((224, 224, 3))), + ], +) +def test_save_image_tensor_to_file(cls_tensor, tensor, tmpdir): + tmp_file = str(tmpdir / 'tmp.jpg') + image_tensor = parse_obj_as(cls_tensor, tensor) + image_tensor.save(tmp_file) + assert os.path.isfile(tmp_file) + + +@pytest.mark.tensorflow +def test_save_image_tensorflow_tensor_to_file(tmpdir): + tmp_file = str(tmpdir / 'tmp.jpg') + image_tensor = parse_obj_as(ImageTensorFlowTensor, tf.zeros((224, 224, 3))) + image_tensor.save(tmp_file) + assert os.path.isfile(tmp_file) + + +@pytest.mark.parametrize( + 'image_tensor', + [ + parse_obj_as(ImageTorchTensor, torch.zeros(224, 224, 3)), + parse_obj_as(ImageNdArray, np.zeros((224, 224, 3))), + ], +) +def test_save_image_tensor_to_bytes(image_tensor): + b = image_tensor.to_bytes() + isinstance(b, bytes) + isinstance(b, ImageBytes) + + +@pytest.mark.parametrize( + 'tensor,cls_audio_tensor,cls_tensor', + [ + (torch.zeros(1000, 2), ImageTorchTensor, torch.Tensor), + (np.zeros((1000, 2)), ImageNdArray, np.ndarray), + ], +) +def test_torch_ndarray_to_image_tensor(tensor, cls_audio_tensor, cls_tensor): + class MyImageDoc(BaseDoc): + tensor: ImageTensor + + doc = MyImageDoc(tensor=tensor) + assert isinstance(doc.tensor, cls_audio_tensor) + assert isinstance(doc.tensor, cls_tensor) + assert (doc.tensor == tensor).all() + + +@pytest.mark.tensorflow +def test_tensorflow_to_image_tensor(): + class MyImageDoc(BaseDoc): + tensor: ImageTensor + + doc = MyImageDoc(tensor=tf.zeros((1000, 2))) + assert isinstance(doc.tensor, ImageTensorFlowTensor) + assert isinstance(doc.tensor.tensor, tf.Tensor) + assert tnp.allclose(doc.tensor.tensor, tf.zeros((1000, 2))) diff --git a/tests/units/typing/tensor/test_jax_array.py b/tests/units/typing/tensor/test_jax_array.py new file mode 100644 index 00000000000..34b4c979dfc --- /dev/null +++ b/tests/units/typing/tensor/test_jax_array.py @@ -0,0 +1,201 @@ +import numpy as np +import pytest +from pydantic import schema_json_of +from pydantic.tools import parse_obj_as + +from docarray.base_doc.io.json import orjson_dumps +from docarray.utils._internal.misc import is_jax_available + +jax_available = is_jax_available() +if jax_available: + import jax.numpy as jnp + from jax._src.core import InconclusiveDimensionOperation + + from docarray.typing import JaxArray + + +@pytest.mark.jax +def test_proto_tensor(): + from docarray.proto.pb2.docarray_pb2 import NdArrayProto + + tensor = parse_obj_as(JaxArray, jnp.zeros((3, 224, 224))) + proto = tensor.to_protobuf() + assert isinstance(proto, NdArrayProto) + + from_proto = JaxArray.from_protobuf(proto) + assert isinstance(from_proto, JaxArray) + assert jnp.allclose(tensor.tensor, from_proto.tensor) + + +@pytest.mark.jax +def test_json_schema(): + schema_json_of(JaxArray) + + +@pytest.mark.jax +@pytest.mark.skip +def test_dump_json(): + tensor = parse_obj_as(JaxArray, jnp.zeros((2, 56, 56))) + orjson_dumps(tensor) + + +@pytest.mark.jax +def test_unwrap(): + tf_tensor = parse_obj_as(JaxArray, jnp.zeros((3, 224, 224))) + unwrapped = tf_tensor.unwrap() + + assert not isinstance(unwrapped, JaxArray) + assert isinstance(tf_tensor, JaxArray) + assert isinstance(unwrapped, jnp.ndarray) + + assert np.allclose(unwrapped, np.zeros((3, 224, 224))) + + +@pytest.mark.jax +def test_from_ndarray(): + nd = np.array([1, 2, 3]) + tensor = JaxArray.from_ndarray(nd) + assert isinstance(tensor, JaxArray) + assert isinstance(tensor.tensor, jnp.ndarray) + + +@pytest.mark.jax +def test_ellipsis_in_shape(): + # ellipsis in the end, two extra dimensions needed + tf_tensor = parse_obj_as(JaxArray[3, ...], jnp.zeros((3, 128, 224))) + assert isinstance(tf_tensor, JaxArray) + assert isinstance(tf_tensor.tensor, jnp.ndarray) + assert tf_tensor.tensor.shape == (3, 128, 224) + + # ellipsis in the beginning, two extra dimensions needed + tf_tensor = parse_obj_as(JaxArray[..., 224], jnp.zeros((3, 128, 224))) + assert isinstance(tf_tensor, JaxArray) + assert isinstance(tf_tensor.tensor, jnp.ndarray) + assert tf_tensor.tensor.shape == (3, 128, 224) + + # more than one ellipsis in the shape + with pytest.raises(ValueError): + parse_obj_as(JaxArray[3, ..., 128, ...], jnp.zeros((3, 128, 224))) + + # wrong shape + with pytest.raises(ValueError): + parse_obj_as(JaxArray[3, 224, ...], jnp.zeros((3, 128, 224))) + + +@pytest.mark.jax +def test_parametrized(): + # correct shape, single axis + tf_tensor = parse_obj_as(JaxArray[128], jnp.zeros(128)) + assert isinstance(tf_tensor, JaxArray) + assert isinstance(tf_tensor.tensor, jnp.ndarray) + assert tf_tensor.tensor.shape == (128,) + + # correct shape, multiple axis + tf_tensor = parse_obj_as(JaxArray[3, 224, 224], jnp.zeros((3, 224, 224))) + assert isinstance(tf_tensor, JaxArray) + assert isinstance(tf_tensor.tensor, jnp.ndarray) + assert tf_tensor.tensor.shape == (3, 224, 224) + + # wrong but reshapable shape + tf_tensor = parse_obj_as(JaxArray[3, 224, 224], jnp.zeros((224, 3, 224))) + assert isinstance(tf_tensor, JaxArray) + assert isinstance(tf_tensor.tensor, jnp.ndarray) + assert tf_tensor.tensor.shape == (3, 224, 224) + + # wrong and not reshapable shape + with pytest.raises(InconclusiveDimensionOperation): + parse_obj_as(JaxArray[3, 224, 224], jnp.zeros((224, 224))) + + +@pytest.mark.jax +def test_parametrized_with_str(): + # test independent variable dimensions + tf_tensor = parse_obj_as(JaxArray[3, 'x', 'y'], jnp.zeros((3, 224, 224))) + assert isinstance(tf_tensor, JaxArray) + assert isinstance(tf_tensor.tensor, jnp.ndarray) + assert tf_tensor.tensor.shape == (3, 224, 224) + + tf_tensor = parse_obj_as(JaxArray[3, 'x', 'y'], jnp.zeros((3, 60, 128))) + assert isinstance(tf_tensor, JaxArray) + assert isinstance(tf_tensor.tensor, jnp.ndarray) + assert tf_tensor.tensor.shape == (3, 60, 128) + + with pytest.raises(ValueError): + parse_obj_as(JaxArray[3, 'x', 'y'], jnp.zeros((4, 224, 224))) + + with pytest.raises(ValueError): + parse_obj_as(JaxArray[3, 'x', 'y'], jnp.zeros((100, 1))) + + # test dependent variable dimensions + tf_tensor = parse_obj_as(JaxArray[3, 'x', 'x'], jnp.zeros((3, 224, 224))) + assert isinstance(tf_tensor, JaxArray) + assert isinstance(tf_tensor.tensor, jnp.ndarray) + assert tf_tensor.tensor.shape == (3, 224, 224) + + with pytest.raises(ValueError): + _ = parse_obj_as(JaxArray[3, 'x', 'x'], jnp.zeros((3, 60, 128))) + + with pytest.raises(ValueError): + _ = parse_obj_as(JaxArray[3, 'x', 'x'], jnp.zeros((3, 60))) + + +@pytest.mark.jax +@pytest.mark.parametrize('shape', [(3, 224, 224), (224, 224, 3)]) +def test_parameterized_tensor_class_name(shape): + MyTFT = JaxArray[3, 224, 224] + tensor = parse_obj_as(MyTFT, jnp.zeros(shape)) + + assert MyTFT.__name__ == 'JaxArray[3, 224, 224]' + assert MyTFT.__qualname__ == 'JaxArray[3, 224, 224]' + + assert tensor.__class__.__name__ == 'JaxArray' + assert tensor.__class__.__qualname__ == 'JaxArray' + assert f'{tensor.tensor[0][0][0]}' == '0.0' + + +@pytest.mark.jax +def test_parametrized_subclass(): + c1 = JaxArray[128] + c2 = JaxArray[128] + assert issubclass(c1, c2) + assert issubclass(c1, JaxArray) + + assert not issubclass(c1, JaxArray[256]) + + +@pytest.mark.jax +def test_parametrized_instance(): + t = parse_obj_as(JaxArray[128], jnp.zeros((128,))) + assert isinstance(t, JaxArray[128]) + assert isinstance(t, JaxArray) + # assert isinstance(t, jnp.ndarray) + + assert not isinstance(t, JaxArray[256]) + assert not isinstance(t, JaxArray[2, 128]) + assert not isinstance(t, JaxArray[2, 2, 64]) + + +@pytest.mark.jax +def test_parametrized_equality(): + t1 = parse_obj_as(JaxArray[128], jnp.zeros((128,))) + t2 = parse_obj_as(JaxArray[128], jnp.zeros((128,))) + assert jnp.allclose(t1.tensor, t2.tensor) + + +@pytest.mark.jax +def test_parametrized_operations(): + t1 = parse_obj_as(JaxArray[128], jnp.zeros((128,))) + t2 = parse_obj_as(JaxArray[128], jnp.zeros((128,))) + t_result = t1.tensor + t2.tensor + assert isinstance(t_result, jnp.ndarray) + assert not isinstance(t_result, JaxArray) + assert not isinstance(t_result, JaxArray[128]) + + +@pytest.mark.jax +def test_set_item(): + t = JaxArray(tensor=jnp.zeros((3, 224, 224))) + t[0] = jnp.ones((1, 224, 224)) + assert jnp.allclose(t.tensor[0], jnp.ones((1, 224, 224))) + assert jnp.allclose(t.tensor[1], jnp.zeros((1, 224, 224))) + assert jnp.allclose(t.tensor[2], jnp.zeros((1, 224, 224))) diff --git a/tests/units/typing/tensor/test_ndarray.py b/tests/units/typing/tensor/test_ndarray.py new file mode 100644 index 00000000000..93ed58b3824 --- /dev/null +++ b/tests/units/typing/tensor/test_ndarray.py @@ -0,0 +1,279 @@ +import numpy as np +import orjson +import pytest +import torch +from pydantic.tools import parse_obj_as, schema_json_of + +from docarray import BaseDoc +from docarray.base_doc.io.json import orjson_dumps +from docarray.typing import AudioNdArray, NdArray, TorchTensor +from docarray.typing.tensor import NdArrayEmbedding +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + + +@pytest.mark.proto +def test_proto_tensor(): + tensor = parse_obj_as(NdArray, np.zeros((3, 224, 224))) + + tensor._to_node_protobuf() + + +def test_from_list(): + tensor = parse_obj_as(NdArray, [[0.0, 0.0], [0.0, 0.0]]) + + assert (tensor == np.zeros((2, 2))).all() + + +def test_json_schema(): + schema_json_of(NdArray) + + +def test_dump_json(): + tensor = parse_obj_as(NdArray, np.zeros((3, 224, 224))) + orjson_dumps(tensor) + + +def test_load_json(): + tensor = parse_obj_as(NdArray, np.zeros((2, 2))) + + json = orjson_dumps(tensor) + print(json) + print(type(json)) + new_tensor = orjson.loads(json) + + assert (new_tensor == tensor).all() + + +def test_unwrap(): + tensor = parse_obj_as(NdArray, np.zeros((3, 224, 224))) + ndarray = tensor.unwrap() + + assert not isinstance(ndarray, NdArray) + assert isinstance(ndarray, np.ndarray) + assert isinstance(tensor, NdArray) + assert (ndarray == np.zeros((3, 224, 224))).all() + + +@pytest.mark.parametrize( + 'tensor_class, tensor_type, tensor_fn', + [(NdArray, np.ndarray, np.zeros), (TorchTensor, torch.Tensor, torch.zeros)], +) +def test_ellipsis_in_shape(tensor_class, tensor_type, tensor_fn): + # ellipsis in the end, two extra dimensions needed + tensor = parse_obj_as(tensor_class[3, ...], tensor_fn((3, 128, 224))) + assert isinstance(tensor, tensor_class) + assert isinstance(tensor, tensor_type) + assert tensor.shape == (3, 128, 224) + + # ellipsis in the middle, one extra dimension needed + tensor = parse_obj_as(tensor_class[3, ..., 224], tensor_fn((3, 128, 224))) + assert isinstance(tensor, tensor_class) + assert isinstance(tensor, tensor_type) + assert tensor.shape == (3, 128, 224) + + # ellipsis in the beginning, two extra dimensions needed + tensor = parse_obj_as(tensor_class[..., 224], tensor_fn((3, 128, 224))) + assert isinstance(tensor, tensor_class) + assert isinstance(tensor, tensor_type) + assert tensor.shape == (3, 128, 224) + + # more than one ellipsis in the shape + with pytest.raises(ValueError): + parse_obj_as(tensor_class[3, ..., 128, ...], tensor_fn((3, 128, 224))) + + # bigger dimension than expected + with pytest.raises(ValueError): + parse_obj_as(tensor_class[3, 128, 224, ...], tensor_fn((3, 128))) + + # no extra dimension needed + with pytest.raises(ValueError): + parse_obj_as(tensor_class[3, 128, 224, ...], tensor_fn((3, 128, 224))) + + # wrong shape + with pytest.raises(ValueError): + parse_obj_as(tensor_class[3, 224, ...], tensor_fn((3, 128, 224))) + + # passing only ellipsis as a shape + with pytest.raises(TypeError): + parse_obj_as(tensor_class[...], tensor_fn((3, 128, 224))) + + +@pytest.mark.parametrize( + 'tensor_class, tensor_type, tensor_fn', + [(NdArray, np.ndarray, np.zeros), (TorchTensor, torch.Tensor, torch.zeros)], +) +def test_parametrized(tensor_class, tensor_type, tensor_fn): + # correct shape, single axis + tensor = parse_obj_as(tensor_class[128], tensor_fn(128)) + assert isinstance(tensor, tensor_class) + assert isinstance(tensor, tensor_type) + assert tensor.shape == (128,) + + # correct shape, multiple axis + tensor = parse_obj_as(tensor_class[3, 224, 224], tensor_fn((3, 224, 224))) + assert isinstance(tensor, tensor_class) + assert isinstance(tensor, tensor_type) + assert tensor.shape == (3, 224, 224) + + # wrong but reshapable shape + tensor = parse_obj_as(tensor_class[3, 224, 224], tensor_fn((3, 224, 224))) + assert isinstance(tensor, tensor_class) + assert isinstance(tensor, tensor_type) + assert tensor.shape == (3, 224, 224) + + # wrong and not reshapable shape + with pytest.raises(ValueError): + parse_obj_as(tensor_class[3, 224, 224], tensor_fn((224, 224))) + + # test independent variable dimensions + tensor = parse_obj_as(tensor_class[3, 'x', 'y'], tensor_fn((3, 224, 224))) + assert isinstance(tensor, tensor_class) + assert isinstance(tensor, tensor_type) + assert tensor.shape == (3, 224, 224) + + tensor = parse_obj_as(tensor_class[3, 'x', 'y'], tensor_fn((3, 60, 128))) + assert isinstance(tensor, tensor_class) + assert isinstance(tensor, tensor_type) + assert tensor.shape == (3, 60, 128) + + with pytest.raises(ValueError): + parse_obj_as(tensor_class[3, 'x', 'y'], tensor_fn((4, 224, 224))) + + with pytest.raises(ValueError): + parse_obj_as(tensor_class[3, 'x', 'y'], tensor_fn((100, 1))) + + # test dependent variable dimensions + tensor = parse_obj_as(tensor_class[3, 'x', 'x'], tensor_fn((3, 224, 224))) + assert isinstance(tensor, tensor_class) + assert isinstance(tensor, tensor_type) + assert tensor.shape == (3, 224, 224) + + with pytest.raises(ValueError): + tensor = parse_obj_as(tensor_class[3, 'x', 'x'], tensor_fn((3, 60, 128))) + + with pytest.raises(ValueError): + tensor = parse_obj_as(tensor_class[3, 'x', 'x'], tensor_fn((3, 60))) + + +def test_np_embedding(): + # correct shape + tensor = parse_obj_as(NdArrayEmbedding[128], np.zeros((128,))) + assert isinstance(tensor, NdArrayEmbedding) + assert isinstance(tensor, NdArray) + assert isinstance(tensor, np.ndarray) + assert tensor.shape == (128,) + + # wrong shape at data setting time + with pytest.raises(ValueError): + parse_obj_as(NdArrayEmbedding[128], np.zeros((256,))) + + # illegal shape at class creation time + with pytest.raises(ValueError): + parse_obj_as(NdArrayEmbedding[128, 128], np.zeros((128, 128))) + + +def test_parametrized_subclass(): + c1 = NdArray[128] + c2 = NdArray[128] + assert issubclass(c1, c2) + assert issubclass(c1, NdArray) + assert issubclass(c1, np.ndarray) + + assert not issubclass(c1, NdArray[256]) + + +def test_parametrized_instance(): + t = parse_obj_as(NdArray[128], np.zeros(128)) + assert isinstance(t, NdArray[128]) + assert isinstance(t, NdArray) + assert isinstance(t, np.ndarray) + + assert not isinstance(t, NdArray[256]) + assert not isinstance(t, NdArray[2, 64]) + assert not isinstance(t, NdArray[2, 2, 32]) + + +def test_parametrized_equality(): + t1 = parse_obj_as(NdArray[128], np.zeros(128)) + t2 = parse_obj_as(NdArray[128], np.zeros(128)) + t3 = parse_obj_as(NdArray[128], np.ones(128)) + assert (t1 == t2).all() + assert not (t1 == t3).any() + + +def test_parametrized_operations(): + t1 = parse_obj_as(NdArray[128], np.zeros(128)) + t2 = parse_obj_as(NdArray[128], np.zeros(128)) + t_result = t1 + t2 + assert isinstance(t_result, np.ndarray) + assert isinstance(t_result, NdArray) + assert isinstance(t_result, NdArray[128]) + + +def test_class_equality(): + assert NdArray == NdArray + assert NdArray[128] == NdArray[128] + assert NdArray[128] != NdArray[256] + assert NdArray[128] != NdArray[2, 64] + assert not NdArray[128] == NdArray[2, 64] + + assert NdArrayEmbedding == NdArrayEmbedding + assert NdArrayEmbedding[128] == NdArrayEmbedding[128] + assert NdArrayEmbedding[128] != NdArrayEmbedding[256] + + assert AudioNdArray == AudioNdArray + assert AudioNdArray[128] == AudioNdArray[128] + assert AudioNdArray[128] != AudioNdArray[256] + + +def test_class_hash(): + assert hash(NdArray) == hash(NdArray) + assert hash(NdArray[128]) == hash(NdArray[128]) + assert hash(NdArray[128]) != hash(NdArray[256]) + assert hash(NdArray[128]) != hash(NdArray[2, 64]) + assert not hash(NdArray[128]) == hash(NdArray[2, 64]) + + assert hash(NdArrayEmbedding) == hash(NdArrayEmbedding) + assert hash(NdArrayEmbedding[128]) == hash(NdArrayEmbedding[128]) + assert hash(NdArrayEmbedding[128]) != hash(NdArrayEmbedding[256]) + + assert hash(AudioNdArray) == hash(AudioNdArray) + assert hash(AudioNdArray[128]) == hash(AudioNdArray[128]) + assert hash(AudioNdArray[128]) != hash(AudioNdArray[256]) + + +@pytest.mark.parametrize( + 'tensor', + [ + torch.zeros(10), + TorchTensor(torch.zeros(10)), + np.zeros(10), + ], +) +def test_torch_numpy_to_ndarray(tensor): + class MyAudioDoc(BaseDoc): + tensor: NdArray + + doc = MyAudioDoc(tensor=tensor) + assert isinstance(doc.tensor, np.ndarray) + assert isinstance(doc.tensor, NdArray) + assert isinstance(doc.tensor, NdArray[10]) + + +@pytest.mark.tensorflow +def test_tensorflow_to_ndarray(): + class MyAudioDoc(BaseDoc): + tensor: NdArray + + doc = MyAudioDoc( + tensor=tf.zeros( + 10, + ) + ) + assert isinstance(doc.tensor, np.ndarray) + assert isinstance(doc.tensor, NdArray) + assert isinstance(doc.tensor, NdArray[10]) diff --git a/tests/units/typing/tensor/test_np_ops.py b/tests/units/typing/tensor/test_np_ops.py new file mode 100644 index 00000000000..27da03c5aee --- /dev/null +++ b/tests/units/typing/tensor/test_np_ops.py @@ -0,0 +1,38 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np + +from docarray import BaseDoc +from docarray.typing import NdArray + + +def test_tensor_ops(): + class A(BaseDoc): + tensor: NdArray[3, 224, 224] + + class B(BaseDoc): + tensor: NdArray[3, 112, 224] + + tensor = A(tensor=np.ones((3, 224, 224))).tensor + tensord = A(tensor=np.ones((3, 224, 224))).tensor + tensorn = np.zeros((3, 224, 224)) + tensorhalf = B(tensor=np.ones((3, 112, 224))).tensor + tensorfull = np.concatenate([tensorhalf, tensorhalf], axis=1) + + assert type(tensor) == NdArray + assert type(tensor + tensord) == NdArray + assert type(tensor + tensorn) == NdArray + assert type(tensor + tensorfull) == NdArray diff --git a/tests/units/typing/tensor/test_tensor.py b/tests/units/typing/tensor/test_tensor.py new file mode 100644 index 00000000000..6d506201cbc --- /dev/null +++ b/tests/units/typing/tensor/test_tensor.py @@ -0,0 +1,48 @@ +import numpy as np +import pytest +import torch + +from docarray import BaseDoc +from docarray.typing import AnyTensor, NdArray, TorchTensor +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + + from docarray.computation.tensorflow_backend import tnp + from docarray.typing import TensorFlowTensor + + +@pytest.mark.parametrize( + 'tensor,cls_audio_tensor,cls_tensor', + [ + (torch.zeros(1000, 2), TorchTensor, torch.Tensor), + (np.zeros((1000, 2)), NdArray, np.ndarray), + ], +) +def test_torch_ndarray_to_any_tensor(tensor, cls_audio_tensor, cls_tensor): + class MyTensorDoc(BaseDoc): + tensor: AnyTensor + + doc = MyTensorDoc(tensor=tensor) + assert isinstance(doc.tensor, cls_audio_tensor) + assert isinstance(doc.tensor, cls_tensor) + assert doc.tensor.shape == (1000, 2) + assert (doc.tensor == tensor).all() + + +@pytest.mark.tensorflow +def test_tensorflow_to_any_tensor(): + class MyTensorDoc(BaseDoc): + tensor: AnyTensor + + doc = MyTensorDoc(tensor=tf.zeros((1000, 2))) + assert isinstance(doc.tensor, TensorFlowTensor) + assert isinstance(doc.tensor.tensor, tf.Tensor) + assert tnp.allclose(doc.tensor.tensor, tf.zeros((1000, 2))) + + +def test_equals_type(): + # see https://github.com/docarray/docarray/pull/1739 + assert not (TorchTensor == type) diff --git a/tests/units/typing/tensor/test_tensor_coercion.py b/tests/units/typing/tensor/test_tensor_coercion.py new file mode 100644 index 00000000000..e358e0eb7ee --- /dev/null +++ b/tests/units/typing/tensor/test_tensor_coercion.py @@ -0,0 +1,50 @@ +import numpy as np +import pytest +import torch +from pydantic import parse_obj_as + +from docarray.typing import NdArray, TorchTensor +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + + from docarray.typing import TensorFlowTensor +else: + + ### This is needed to fake the import of tensorflow when it is not installed + class TfNotInstalled: + def zeros(self, *args, **kwargs): + return 0 + + class TensorFlowTensor: + def _docarray_from_native(self, *args, **kwargs): + return 0 + + tf = TfNotInstalled() + + +pure_tensor_to_test = [ + np.zeros((3, 224, 224)), + torch.zeros(3, 224, 224), + tf.zeros((3, 224, 224)), +] + +docarray_tensor_to_test = [ + NdArray._docarray_from_native(np.zeros((3, 224, 224))), + TorchTensor._docarray_from_native(torch.zeros(3, 224, 224)), + TensorFlowTensor._docarray_from_native(tf.zeros((3, 224, 224))), +] + + +@pytest.mark.tensorflow +@pytest.mark.parametrize('tensor', pure_tensor_to_test + docarray_tensor_to_test) +@pytest.mark.parametrize('tensor_cls', [NdArray, TorchTensor, TensorFlowTensor]) +def test_torch_tensor_coerse(tensor_cls, tensor): + t = parse_obj_as(tensor_cls, tensor) + assert isinstance(t, tensor_cls) + + t_numpy = t._docarray_to_ndarray() + assert t_numpy.shape == (3, 224, 224) + assert (t_numpy == np.zeros((3, 224, 224))).all() diff --git a/tests/units/typing/tensor/test_tensor_flow_tensor.py b/tests/units/typing/tensor/test_tensor_flow_tensor.py new file mode 100644 index 00000000000..3a51f5a95aa --- /dev/null +++ b/tests/units/typing/tensor/test_tensor_flow_tensor.py @@ -0,0 +1,201 @@ +import numpy as np +import pytest +from pydantic import schema_json_of +from pydantic.tools import parse_obj_as + +from docarray.base_doc.io.json import orjson_dumps +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + import tensorflow._api.v2.experimental.numpy as tnp # type: ignore + from tensorflow.python.framework.errors_impl import InvalidArgumentError + + from docarray.typing import TensorFlowTensor + + +@pytest.mark.tensorflow +def test_proto_tensor(): + from docarray.proto.pb2.docarray_pb2 import NdArrayProto + + tensor = parse_obj_as(TensorFlowTensor, tf.zeros((3, 224, 224))) + proto = tensor.to_protobuf() + assert isinstance(proto, NdArrayProto) + + from_proto = TensorFlowTensor.from_protobuf(proto) + assert isinstance(from_proto, TensorFlowTensor) + assert tnp.allclose(tensor.tensor, from_proto.tensor) + + +@pytest.mark.tensorflow +def test_json_schema(): + schema_json_of(TensorFlowTensor) + + +@pytest.mark.tensorflow +def test_dump_json(): + tensor = parse_obj_as(TensorFlowTensor, tf.zeros((3, 224, 224))) + orjson_dumps(tensor) + + +@pytest.mark.tensorflow +def test_unwrap(): + tf_tensor = parse_obj_as(TensorFlowTensor, tf.zeros((3, 224, 224))) + unwrapped = tf_tensor.unwrap() + + assert not isinstance(unwrapped, TensorFlowTensor) + assert isinstance(tf_tensor, TensorFlowTensor) + assert isinstance(unwrapped, tf.Tensor) + + assert np.allclose(unwrapped, np.zeros((3, 224, 224))) + + +@pytest.mark.tensorflow +def test_from_ndarray(): + nd = np.array([1, 2, 3]) + tensor = TensorFlowTensor.from_ndarray(nd) + assert isinstance(tensor, TensorFlowTensor) + assert isinstance(tensor.tensor, tf.Tensor) + + +@pytest.mark.tensorflow +def test_ellipsis_in_shape(): + # ellipsis in the end, two extra dimensions needed + tf_tensor = parse_obj_as(TensorFlowTensor[3, ...], tf.zeros((3, 128, 224))) + assert isinstance(tf_tensor, TensorFlowTensor) + assert isinstance(tf_tensor.tensor, tf.Tensor) + assert tf_tensor.tensor.shape == (3, 128, 224) + + # ellipsis in the beginning, two extra dimensions needed + tf_tensor = parse_obj_as(TensorFlowTensor[..., 224], tf.zeros((3, 128, 224))) + assert isinstance(tf_tensor, TensorFlowTensor) + assert isinstance(tf_tensor.tensor, tf.Tensor) + assert tf_tensor.tensor.shape == (3, 128, 224) + + # more than one ellipsis in the shape + with pytest.raises(ValueError): + parse_obj_as(TensorFlowTensor[3, ..., 128, ...], tf.zeros((3, 128, 224))) + + # wrong shape + with pytest.raises(ValueError): + parse_obj_as(TensorFlowTensor[3, 224, ...], tf.zeros((3, 128, 224))) + + +@pytest.mark.tensorflow +def test_parametrized(): + # correct shape, single axis + tf_tensor = parse_obj_as(TensorFlowTensor[128], tf.zeros(128)) + assert isinstance(tf_tensor, TensorFlowTensor) + assert isinstance(tf_tensor.tensor, tf.Tensor) + assert tf_tensor.tensor.shape == (128,) + + # correct shape, multiple axis + tf_tensor = parse_obj_as(TensorFlowTensor[3, 224, 224], tf.zeros((3, 224, 224))) + assert isinstance(tf_tensor, TensorFlowTensor) + assert isinstance(tf_tensor.tensor, tf.Tensor) + assert tf_tensor.tensor.shape == (3, 224, 224) + + # wrong but reshapable shape + tf_tensor = parse_obj_as(TensorFlowTensor[3, 224, 224], tf.zeros((224, 3, 224))) + assert isinstance(tf_tensor, TensorFlowTensor) + assert isinstance(tf_tensor.tensor, tf.Tensor) + assert tf_tensor.tensor.shape == (3, 224, 224) + + # wrong and not reshapable shape + with pytest.raises(InvalidArgumentError): + parse_obj_as(TensorFlowTensor[3, 224, 224], tf.zeros((224, 224))) + + +@pytest.mark.tensorflow +def test_parametrized_with_str(): + # test independent variable dimensions + tf_tensor = parse_obj_as(TensorFlowTensor[3, 'x', 'y'], tf.zeros((3, 224, 224))) + assert isinstance(tf_tensor, TensorFlowTensor) + assert isinstance(tf_tensor.tensor, tf.Tensor) + assert tf_tensor.tensor.shape == (3, 224, 224) + + tf_tensor = parse_obj_as(TensorFlowTensor[3, 'x', 'y'], tf.zeros((3, 60, 128))) + assert isinstance(tf_tensor, TensorFlowTensor) + assert isinstance(tf_tensor.tensor, tf.Tensor) + assert tf_tensor.tensor.shape == (3, 60, 128) + + with pytest.raises(ValueError): + parse_obj_as(TensorFlowTensor[3, 'x', 'y'], tf.zeros((4, 224, 224))) + + with pytest.raises(ValueError): + parse_obj_as(TensorFlowTensor[3, 'x', 'y'], tf.zeros((100, 1))) + + # test dependent variable dimensions + tf_tensor = parse_obj_as(TensorFlowTensor[3, 'x', 'x'], tf.zeros((3, 224, 224))) + assert isinstance(tf_tensor, TensorFlowTensor) + assert isinstance(tf_tensor.tensor, tf.Tensor) + assert tf_tensor.tensor.shape == (3, 224, 224) + + with pytest.raises(ValueError): + _ = parse_obj_as(TensorFlowTensor[3, 'x', 'x'], tf.zeros((3, 60, 128))) + + with pytest.raises(ValueError): + _ = parse_obj_as(TensorFlowTensor[3, 'x', 'x'], tf.zeros((3, 60))) + + +@pytest.mark.tensorflow +@pytest.mark.parametrize('shape', [(3, 224, 224), (224, 224, 3)]) +def test_parameterized_tensor_class_name(shape): + MyTFT = TensorFlowTensor[3, 224, 224] + tensor = parse_obj_as(MyTFT, tf.zeros(shape)) + + assert MyTFT.__name__ == 'TensorFlowTensor[3, 224, 224]' + assert MyTFT.__qualname__ == 'TensorFlowTensor[3, 224, 224]' + + assert tensor.__class__.__name__ == 'TensorFlowTensor' + assert tensor.__class__.__qualname__ == 'TensorFlowTensor' + assert f'{tensor.tensor[0][0][0]}' == '0.0' + + +@pytest.mark.tensorflow +def test_parametrized_subclass(): + c1 = TensorFlowTensor[128] + c2 = TensorFlowTensor[128] + assert issubclass(c1, c2) + assert issubclass(c1, TensorFlowTensor) + + assert not issubclass(c1, TensorFlowTensor[256]) + + +@pytest.mark.tensorflow +def test_parametrized_instance(): + t = parse_obj_as(TensorFlowTensor[128], tf.zeros((128,))) + assert isinstance(t, TensorFlowTensor[128]) + assert isinstance(t, TensorFlowTensor) + # assert isinstance(t, tf.Tensor) + + assert not isinstance(t, TensorFlowTensor[256]) + assert not isinstance(t, TensorFlowTensor[2, 128]) + assert not isinstance(t, TensorFlowTensor[2, 2, 64]) + + +@pytest.mark.tensorflow +def test_parametrized_equality(): + t1 = parse_obj_as(TensorFlowTensor[128], tf.zeros((128,))) + t2 = parse_obj_as(TensorFlowTensor[128], tf.zeros((128,))) + assert tf.experimental.numpy.allclose(t1.tensor, t2.tensor) + + +@pytest.mark.tensorflow +def test_parametrized_operations(): + t1 = parse_obj_as(TensorFlowTensor[128], tf.zeros((128,))) + t2 = parse_obj_as(TensorFlowTensor[128], tf.zeros((128,))) + t_result = t1.tensor + t2.tensor + assert isinstance(t_result, tf.Tensor) + assert not isinstance(t_result, TensorFlowTensor) + assert not isinstance(t_result, TensorFlowTensor[128]) + + +@pytest.mark.tensorflow +def test_set_item(): + t = TensorFlowTensor(tensor=tf.zeros((3, 224, 224))) + t[0] = tf.ones((1, 224, 224)) + assert tnp.allclose(t.tensor[0], tf.ones((1, 224, 224))) + assert tnp.allclose(t.tensor[1], tf.zeros((1, 224, 224))) + assert tnp.allclose(t.tensor[2], tf.zeros((1, 224, 224))) diff --git a/tests/units/typing/tensor/test_torch_ops.py b/tests/units/typing/tensor/test_torch_ops.py new file mode 100644 index 00000000000..7e6e4a54f96 --- /dev/null +++ b/tests/units/typing/tensor/test_torch_ops.py @@ -0,0 +1,38 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch + +from docarray import BaseDoc +from docarray.typing import TorchTensor + + +def test_tensor_ops(): + class A(BaseDoc): + tensor: TorchTensor[3, 224, 224] + + class B(BaseDoc): + tensor: TorchTensor[3, 112, 224] + + tensor = A(tensor=torch.ones(3, 224, 224)).tensor + tensord = A(tensor=torch.ones(3, 224, 224)).tensor + tensorn = torch.zeros(3, 224, 224) + tensorhalf = B(tensor=torch.ones(3, 112, 224)).tensor + tensorfull = torch.cat([tensorhalf, tensorhalf], dim=1) + + assert type(tensor) == TorchTensor + assert type(tensor + tensord) == TorchTensor + assert type(tensor + tensorn) == TorchTensor + assert type(tensor + tensorfull) == TorchTensor diff --git a/tests/units/typing/tensor/test_torch_tensor.py b/tests/units/typing/tensor/test_torch_tensor.py new file mode 100644 index 00000000000..dbe8b58a8e5 --- /dev/null +++ b/tests/units/typing/tensor/test_torch_tensor.py @@ -0,0 +1,250 @@ +import pytest +import torch +from pydantic.tools import parse_obj_as, schema_json_of + +from docarray import BaseDoc +from docarray.base_doc.io.json import orjson_dumps +from docarray.proto import DocProto +from docarray.typing import TorchEmbedding, TorchTensor + + +class MyDoc(BaseDoc): + tens: TorchTensor + + +@pytest.mark.proto +def test_proto_tensor(): + tensor = parse_obj_as(TorchTensor, torch.zeros(3, 224, 224)) + + tensor._to_node_protobuf() + + +def test_json_schema(): + schema_json_of(TorchTensor) + + +def test_dump_json(): + tensor = parse_obj_as(TorchTensor, torch.zeros(3, 224, 224)) + orjson_dumps(tensor) + + +def test_unwrap(): + tensor = parse_obj_as(TorchTensor, torch.zeros(3, 224, 224)) + ndarray = tensor.unwrap() + + assert not isinstance(ndarray, TorchTensor) + assert isinstance(tensor, TorchTensor) + assert isinstance(ndarray, torch.Tensor) + + assert tensor.data_ptr() == ndarray.data_ptr() + + assert (ndarray == torch.zeros(3, 224, 224)).all() + + +def test_parametrized_correct_axis_shape(): + # correct shape, single axis + tensor = parse_obj_as(TorchTensor[128], torch.zeros(128)) + assert isinstance(tensor, TorchTensor) + assert isinstance(tensor, torch.Tensor) + assert tensor.shape == (128,) + + +def test_correct_shape_multiple_axis(): + # correct shape, multiple axis + tensor = parse_obj_as(TorchTensor[3, 224, 224], torch.zeros(3, 224, 224)) + assert isinstance(tensor, TorchTensor) + assert isinstance(tensor, torch.Tensor) + assert tensor.shape == (3, 224, 224) + + +def test_wrong_but_reshapable(): + # wrong but reshapable shape + tensor = parse_obj_as(TorchTensor[3, 224, 224], torch.zeros(224, 3, 224)) + assert isinstance(tensor, TorchTensor) + assert isinstance(tensor, torch.Tensor) + assert tensor.shape == (3, 224, 224) + + # wrong and not reshapable shape + with pytest.raises(ValueError): + parse_obj_as(TorchTensor[3, 224, 224], torch.zeros(224, 224)) + + +def test_independent_variable_dim(): + # test independent variable dimensions + tensor = parse_obj_as(TorchTensor[3, 'x', 'y'], torch.zeros(3, 224, 224)) + assert isinstance(tensor, TorchTensor) + assert isinstance(tensor, torch.Tensor) + assert tensor.shape == (3, 224, 224) + + +def test_param(): + tensor = parse_obj_as(TorchTensor[3, 'x', 'y'], torch.zeros(3, 60, 128)) + assert isinstance(tensor, TorchTensor) + assert isinstance(tensor, torch.Tensor) + assert tensor.shape == (3, 60, 128) + + with pytest.raises(ValueError): + parse_obj_as(TorchTensor[3, 'x', 'y'], torch.zeros(4, 224, 224)) + + with pytest.raises(ValueError): + parse_obj_as(TorchTensor[3, 'x', 'y'], torch.zeros(100, 1)) + + +def test_dependent_variable_dim(): + # test dependent variable dimensions + tensor = parse_obj_as(TorchTensor[3, 'x', 'x'], torch.zeros(3, 224, 224)) + assert isinstance(tensor, TorchTensor) + assert isinstance(tensor, torch.Tensor) + assert tensor.shape == (3, 224, 224) + + with pytest.raises(ValueError): + _ = parse_obj_as(TorchTensor[3, 'x', 'x'], torch.zeros(3, 60, 128)) + + with pytest.raises(ValueError): + _ = parse_obj_as(TorchTensor[3, 'x', 'x'], torch.zeros(3, 60)) + + +@pytest.mark.parametrize('shape', [(3, 224, 224), (224, 224, 3)]) +def test_parameterized_tensor_class_name(shape): + MyTT = TorchTensor[3, 224, 224] + tensor = parse_obj_as(MyTT, torch.zeros(shape)) + + assert MyTT.__name__ == 'TorchTensor[3, 224, 224]' + assert MyTT.__qualname__ == 'TorchTensor[3, 224, 224]' + + assert tensor.__class__.__name__ == 'TorchTensor' + assert tensor.__class__.__qualname__ == 'TorchTensor' + assert f'{tensor[0][0][0]}' == 'TorchTensor(0.)' + + +def test_torch_embedding(): + # correct shape + tensor = parse_obj_as(TorchEmbedding[128], torch.zeros(128)) + assert isinstance(tensor, TorchEmbedding) + assert isinstance(tensor, torch.Tensor) + assert tensor.shape == (128,) + + # wrong shape at data setting time + with pytest.raises(ValueError): + parse_obj_as(TorchEmbedding[128], torch.zeros(256)) + + # illegal shape at class creation time + with pytest.raises(ValueError): + parse_obj_as(TorchEmbedding[128, 128], torch.zeros(128, 128)) + + +def test_parametrized_subclass(): + c1 = TorchTensor[128] + c2 = TorchTensor[128] + assert issubclass(c1, c2) + assert issubclass(c1, TorchTensor) + assert issubclass(c1, torch.Tensor) + + assert not issubclass(c1, TorchTensor[256]) + + +def test_parametrized_instance(): + t = parse_obj_as(TorchTensor[128], torch.zeros(128)) + assert isinstance(t, TorchTensor[128]) + assert isinstance(t, TorchTensor) + assert isinstance(t, torch.Tensor) + + assert not isinstance(t, TorchTensor[256]) + assert not isinstance(t, TorchTensor[2, 128]) + assert not isinstance(t, TorchTensor[2, 2, 64]) + + +def test_parametrized_equality(): + t1 = parse_obj_as(TorchTensor[128], torch.zeros(128)) + t2 = parse_obj_as(TorchTensor[128], torch.zeros(128)) + assert (t1 == t2).all() + + +def test_parametrized_operations(): + t1 = parse_obj_as(TorchTensor[128], torch.zeros(128)) + t2 = parse_obj_as(TorchTensor[128], torch.zeros(128)) + t_result = t1 + t2 + assert isinstance(t_result, torch.Tensor) + assert isinstance(t_result, TorchTensor) + assert isinstance(t_result, TorchTensor[128]) + + +def test_deepcopy(): + from docarray import BaseDoc + + class MMdoc(BaseDoc): + embedding: TorchEmbedding + + doc = MMdoc(embedding=torch.randn(32)) + doc_copy = doc.copy(deep=True) + + assert doc.embedding.data_ptr() != doc_copy.embedding.data_ptr() + assert (doc.embedding == doc_copy.embedding).all() + + doc_copy.embedding = torch.randn(32) + assert not (doc.embedding == doc_copy.embedding).all() + + +def test_deepcopy_tensor(): + from docarray import BaseDoc + + class MMdoc(BaseDoc): + embedding: TorchTensor + + doc = MMdoc(embedding=torch.randn(32)) + doc_copy = doc.copy(deep=True) + + assert doc.embedding.data_ptr() != doc_copy.embedding.data_ptr() + assert (doc.embedding == doc_copy.embedding).all() + + doc_copy.embedding = torch.randn(32) + assert not (doc.embedding == doc_copy.embedding).all() + + +@pytest.mark.parametrize('requires_grad', [True]) # , False]) +def test_json_serialization(requires_grad: bool): + orig_doc = MyDoc(tens=torch.rand(10, requires_grad=requires_grad)) + serialized_doc = orig_doc.to_json() + assert serialized_doc + assert isinstance(serialized_doc, str) + + new_doc = MyDoc.from_json(serialized_doc) + assert len(new_doc.tens) == 10 + + +@pytest.mark.parametrize('protocol', ['pickle', 'protobuf']) +@pytest.mark.parametrize('requires_grad', [True, False]) +def test_bytes_serialization(requires_grad, protocol): + orig_doc = MyDoc(tens=torch.rand(10, requires_grad=requires_grad)) + serialized_doc = orig_doc.to_bytes(protocol=protocol) + assert serialized_doc + assert isinstance(serialized_doc, bytes) + + conv_doc = MyDoc.from_bytes(serialized_doc, protocol=protocol) + assert isinstance(conv_doc.tens, TorchTensor) + assert conv_doc.tens.shape == (10,) + + +@pytest.mark.parametrize('protocol', ['pickle', 'protobuf']) +@pytest.mark.parametrize('requires_grad', [True, False]) +def test_base64_serialization(requires_grad, protocol): + orig_doc = MyDoc(tens=torch.rand(10, requires_grad=requires_grad)) + serialized_doc = orig_doc.to_base64(protocol=protocol) + assert serialized_doc + assert isinstance(serialized_doc, str) + + conv_doc = MyDoc.from_base64(serialized_doc, protocol=protocol) + assert isinstance(conv_doc.tens, TorchTensor) + assert conv_doc.tens.shape == (10,) + + +@pytest.mark.parametrize('requires_grad', [True, False]) +def test_protobuf_serialization(requires_grad: bool): + orig_doc = MyDoc(tens=torch.rand(10, requires_grad=requires_grad)) + serialized_doc = orig_doc.to_protobuf() + assert serialized_doc + assert isinstance(serialized_doc, DocProto) + + conv_doc = MyDoc.from_protobuf(serialized_doc) + assert isinstance(conv_doc.tens, TorchTensor) + assert conv_doc.tens.shape == (10,) diff --git a/tests/units/typing/tensor/test_video_tensor.py b/tests/units/typing/tensor/test_video_tensor.py new file mode 100644 index 00000000000..7cd44537d18 --- /dev/null +++ b/tests/units/typing/tensor/test_video_tensor.py @@ -0,0 +1,219 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import numpy as np +import pytest +import torch +from pydantic.tools import parse_obj_as + +from docarray import BaseDoc +from docarray.typing import ( + AudioNdArray, + AudioTorchTensor, + VideoBytes, + VideoNdArray, + VideoTensor, + VideoTorchTensor, +) +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + import tensorflow._api.v2.experimental.numpy as tnp + + from docarray.typing.tensor.video import VideoTensorFlowTensor + + +@pytest.mark.parametrize( + 'tensor,cls_video_tensor,cls_tensor', + [ + (torch.zeros(1, 224, 224, 3), VideoTorchTensor, torch.Tensor), + (np.zeros((1, 224, 224, 3)), VideoNdArray, np.ndarray), + ], +) +def test_set_video_tensor(tensor, cls_video_tensor, cls_tensor): + class MyVideoDoc(BaseDoc): + tensor: cls_video_tensor + + doc = MyVideoDoc(tensor=tensor) + + assert isinstance(doc.tensor, cls_video_tensor) + assert isinstance(doc.tensor, cls_tensor) + assert (doc.tensor == tensor).all() + + +@pytest.mark.tensorflow +def test_set_video_tensor_tensorflow(): + class MyVideoDoc(BaseDoc): + tensor: VideoTensorFlowTensor + + doc = MyVideoDoc(tensor=tf.zeros((1, 224, 224, 3))) + + assert isinstance(doc.tensor, VideoTensorFlowTensor) + assert isinstance(doc.tensor.tensor, tf.Tensor) + assert tnp.allclose(doc.tensor.tensor, tf.zeros((1, 224, 224, 3))) + + +@pytest.mark.parametrize( + 'cls_tensor,tensor', + [ + (VideoNdArray, np.zeros((1, 224, 224, 3))), + (VideoTorchTensor, torch.zeros(1, 224, 224, 3)), + (VideoTorchTensor, np.zeros((1, 224, 224, 3))), + ], +) +def test_validation(cls_tensor, tensor): + arr = parse_obj_as(cls_tensor, tensor) + assert isinstance(arr, cls_tensor) + + +@pytest.mark.tensorflow +def test_validation_tensorflow(): + arr = parse_obj_as(VideoTensorFlowTensor, np.zeros((1, 224, 224, 3))) + assert isinstance(arr, VideoTensorFlowTensor) + + arr = parse_obj_as(VideoTensorFlowTensor, tf.zeros((1, 224, 224, 3))) + assert isinstance(arr, VideoTensorFlowTensor) + + arr = parse_obj_as(VideoTensorFlowTensor, torch.zeros((1, 224, 224, 3))) + assert isinstance(arr, VideoTensorFlowTensor) + + +@pytest.mark.parametrize( + 'cls_tensor,tensor,expect_error', + [ + (VideoNdArray, torch.zeros(1, 224, 224, 3), False), + (VideoNdArray, torch.zeros(1, 224, 224, 100), True), + (VideoTorchTensor, torch.zeros(1, 224, 224, 3), False), + (VideoTorchTensor, torch.zeros(1, 224, 224, 100), True), + (VideoNdArray, 'hello', True), + (VideoTorchTensor, 'hello', True), + ], +) +def test_illegal_validation(cls_tensor, tensor, expect_error): + if expect_error: + with pytest.raises(ValueError): + parse_obj_as(cls_tensor, tensor) + else: + parse_obj_as(cls_tensor, tensor) + + +@pytest.mark.parametrize( + 'cls_tensor,tensor,proto_key', + [ + ( + VideoTorchTensor, + torch.zeros(1, 224, 224, 3), + VideoTorchTensor._proto_type_name, + ), + (VideoNdArray, np.zeros((1, 224, 224, 3)), VideoNdArray._proto_type_name), + ], +) +def test_proto_tensor(cls_tensor, tensor, proto_key): + tensor = parse_obj_as(cls_tensor, tensor) + proto = tensor._to_node_protobuf() + assert proto_key in str(proto) + + +@pytest.mark.tensorflow +def test_proto_tensor_tensorflow(): + tensor = parse_obj_as(VideoTensorFlowTensor, tf.zeros((1, 224, 224, 3))) + proto = tensor._to_node_protobuf() + assert VideoTensorFlowTensor._proto_type_name in str(proto) + + +@pytest.mark.parametrize( + 'video_tensor', + [ + parse_obj_as(VideoTorchTensor, torch.zeros(1, 224, 224, 3)), + parse_obj_as(VideoNdArray, np.zeros((1, 224, 224, 3))), + ], +) +def test_save_video_tensor_to_file(video_tensor, tmpdir): + tmp_file = str(tmpdir / 'tmp.mp4') + video_tensor.save(tmp_file) + assert os.path.isfile(tmp_file) + + +@pytest.mark.parametrize( + 'video_tensor', + [ + parse_obj_as(VideoTorchTensor, torch.zeros(1, 224, 224, 3)), + parse_obj_as(VideoNdArray, np.zeros((1, 224, 224, 3))), + ], +) +def test_save_video_tensor_to_bytes(video_tensor): + b = video_tensor.to_bytes() + isinstance(b, bytes) + isinstance(b, VideoBytes) + + +@pytest.mark.tensorflow +def test_save_video_tensorflow_tensor_to_file(tmpdir): + tmp_file = str(tmpdir / 'tmp.mp4') + video_tensor = parse_obj_as(VideoTensorFlowTensor, tf.zeros((1, 224, 224, 3))) + video_tensor.save(tmp_file) + assert os.path.isfile(tmp_file) + + +@pytest.mark.parametrize( + 'video_tensor', + [ + parse_obj_as(VideoTorchTensor, torch.zeros(1, 224, 224, 3)), + parse_obj_as(VideoNdArray, np.zeros((1, 224, 224, 3))), + ], +) +@pytest.mark.parametrize( + 'audio_tensor', + [ + parse_obj_as(AudioTorchTensor, torch.randn(100, 1, 1024).to(torch.float32)), + parse_obj_as(AudioNdArray, np.random.randn(100, 1, 1024).astype('float32')), + ], +) +def test_save_video_tensor_to_file_including_audio(video_tensor, audio_tensor, tmpdir): + tmp_file = str(tmpdir / 'tmp.mp4') + video_tensor.save(tmp_file, audio_tensor=audio_tensor) + assert os.path.isfile(tmp_file) + + +@pytest.mark.parametrize( + 'tensor,cls_audio_tensor,cls_tensor', + [ + (torch.zeros(2, 10, 10, 3), VideoTorchTensor, torch.Tensor), + (np.zeros((2, 10, 10, 3)), VideoNdArray, np.ndarray), + ], +) +def test_torch_ndarray_to_video_tensor(tensor, cls_audio_tensor, cls_tensor): + class MyAudioDoc(BaseDoc): + tensor: VideoTensor + + doc = MyAudioDoc(tensor=tensor) + assert isinstance(doc.tensor, cls_audio_tensor) + assert isinstance(doc.tensor, cls_tensor) + assert (doc.tensor == tensor).all() + + +@pytest.mark.tensorflow +def test_tensorflow_to_video_tensor(): + class MyAudioDoc(BaseDoc): + tensor: VideoTensor + + doc = MyAudioDoc(tensor=tf.zeros((2, 10, 10, 3))) + assert isinstance(doc.tensor, VideoTensorFlowTensor) + assert isinstance(doc.tensor.tensor, tf.Tensor) + assert tnp.allclose(doc.tensor.tensor, tf.zeros((2, 10, 10, 3))) diff --git a/tests/units/typing/test_bytes.py b/tests/units/typing/test_bytes.py new file mode 100644 index 00000000000..4415f809db5 --- /dev/null +++ b/tests/units/typing/test_bytes.py @@ -0,0 +1,40 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +from pydantic import parse_obj_as + +from docarray.typing import ImageBytes, ImageTensor, ImageUrl + +CUR_DIR = os.path.dirname(os.path.abspath(__file__)) +PATH_TO_IMAGE_DATA = os.path.join(CUR_DIR, '..', '..', 'toydata', 'image-data') +IMAGE_PATHS = { + 'png': os.path.join(PATH_TO_IMAGE_DATA, 'so_good.png'), + 'jpg': os.path.join(PATH_TO_IMAGE_DATA, '05984.jpg'), + 'jpeg': os.path.join(PATH_TO_IMAGE_DATA, '05984-2.jpeg'), +} + + +def test_bytes_load(): + url = parse_obj_as(ImageUrl, IMAGE_PATHS['png']) + + tensor = parse_obj_as(ImageTensor, url.load()) + + bytes_ = parse_obj_as(ImageBytes, tensor.to_bytes()) + + tensor_new = parse_obj_as(ImageTensor, bytes_.load()) + + assert (tensor_new == tensor).all() diff --git a/tests/units/typing/test_id.py b/tests/units/typing/test_id.py new file mode 100644 index 00000000000..10eb46694b4 --- /dev/null +++ b/tests/units/typing/test_id.py @@ -0,0 +1,52 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from uuid import UUID + +import pytest +from pydantic import schema_json_of +from pydantic.tools import parse_obj_as + +from docarray.base_doc.io.json import orjson_dumps +from docarray.typing import ID + + +@pytest.mark.parametrize( + 'id', ['1234', 1234, UUID('cf57432e-809e-4353-adbd-9d5c0d733868')] +) +def test_id_validation(id): + parsed_id = parse_obj_as(ID, id) + + assert parsed_id == str(id) + + +def test_json_schema(): + schema_json_of(ID) + + +def test_dump_json(): + id = parse_obj_as(ID, 1234) + orjson_dumps(id) + + +@pytest.mark.parametrize( + 'id', ['1234', 1234, UUID('cf57432e-809e-4353-adbd-9d5c0d733868')] +) +def test_operators(id): + parsed_id = parse_obj_as(ID, id) + assert parsed_id == str(id) + assert parsed_id != 'aljdñjd' + assert str(id)[0:1] in parsed_id + assert 'docarray' not in parsed_id diff --git a/tests/units/typing/url/__init__.py b/tests/units/typing/url/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/typing/url/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/typing/url/test_any_url.py b/tests/units/typing/url/test_any_url.py new file mode 100644 index 00000000000..d6633f1fe8a --- /dev/null +++ b/tests/units/typing/url/test_any_url.py @@ -0,0 +1,59 @@ +import pytest +from pydantic.tools import parse_obj_as, schema_json_of + +from docarray.base_doc.io.json import orjson_dumps +from docarray.typing import AnyUrl + + +@pytest.mark.proto +def test_proto_any_url(): + uri = parse_obj_as(AnyUrl, 'http://jina.ai/img.png') + + uri._to_node_protobuf() + + +def test_json_schema(): + schema_json_of(AnyUrl) + + +def test_dump_json(): + url = parse_obj_as(AnyUrl, 'http://jina.ai/img.png') + orjson_dumps(url) + + +@pytest.mark.parametrize( + 'relative_path', + [ + 'data/05978.jpg', + '../../data/05978.jpg', + ], +) +def test_relative_path(relative_path): + # see issue: https://github.com/docarray/docarray/issues/978 + url = parse_obj_as(AnyUrl, relative_path) + assert url == relative_path + + +def test_operators(): + url = parse_obj_as(AnyUrl, 'data/05978.jpg') + assert url == 'data/05978.jpg' + assert url != 'aljdñjd' + assert 'data' in url + assert 'docarray' not in url + + +def test_get_url_extension(): + # Test with a URL with extension + assert AnyUrl._get_url_extension('https://jina.ai/hey.md?model=gpt-4') == 'md' + assert AnyUrl._get_url_extension('https://jina.ai/text.txt') == 'txt' + assert AnyUrl._get_url_extension('bla.jpg') == 'jpg' + + # Test with a URL without extension + assert not AnyUrl._get_url_extension('https://jina.ai') + assert not AnyUrl._get_url_extension('https://jina.ai/?model=gpt-4') + + # Test with a text without extension + assert not AnyUrl._get_url_extension('some_text') + + # Test with empty input + assert not AnyUrl._get_url_extension('') diff --git a/tests/units/typing/url/test_audio_url.py b/tests/units/typing/url/test_audio_url.py new file mode 100644 index 00000000000..a787847abb0 --- /dev/null +++ b/tests/units/typing/url/test_audio_url.py @@ -0,0 +1,155 @@ +import os +from typing import Optional + +import numpy as np +import pytest +import torch +from pydantic.tools import parse_obj_as, schema_json_of + +from docarray import BaseDoc +from docarray.base_doc.io.json import orjson_dumps +from docarray.typing import AudioBytes, AudioTorchTensor, AudioUrl +from docarray.typing.url.mimetypes import ( + AUDIO_MIMETYPE, + IMAGE_MIMETYPE, + OBJ_MIMETYPE, + TEXT_MIMETYPE, + VIDEO_MIMETYPE, +) +from docarray.utils._internal.misc import is_tf_available +from tests import TOYDATA_DIR + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + + from docarray.typing.tensor.audio import AudioTensorFlowTensor + +AUDIO_FILES = [ + str(TOYDATA_DIR / 'hello.wav'), + str(TOYDATA_DIR / 'olleh.wav'), +] +REMOTE_AUDIO_FILE = 'https://github.com/docarray/docarray/blob/main/tests/toydata/olleh.wav?raw=true' # noqa: E501 + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_url', + [*AUDIO_FILES, REMOTE_AUDIO_FILE], +) +def test_audio_url(file_url): + uri = parse_obj_as(AudioUrl, file_url) + tensor, _ = uri.load() + assert isinstance(tensor, np.ndarray) + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_url', + [*AUDIO_FILES, REMOTE_AUDIO_FILE], +) +def test_load_audio_url_to_audio_torch_tensor_field(file_url): + class MyAudioDoc(BaseDoc): + audio_url: AudioUrl + tensor: Optional[AudioTorchTensor] = None + + doc = MyAudioDoc(audio_url=file_url) + doc.tensor, _ = doc.audio_url.load() + + assert isinstance(doc.tensor, torch.Tensor) + assert isinstance(doc.tensor, AudioTorchTensor) + + +@pytest.mark.tensorflow +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_url', + [*AUDIO_FILES, REMOTE_AUDIO_FILE], +) +def test_load_audio_url_to_audio_tensorflow_tensor_field(file_url): + class MyAudioDoc(BaseDoc): + audio_url: AudioUrl + tensor: Optional[AudioTensorFlowTensor] = None + + doc = MyAudioDoc(audio_url=file_url) + doc.tensor, _ = doc.audio_url.load() + + assert isinstance(doc.tensor, AudioTensorFlowTensor) + assert isinstance(doc.tensor.tensor, tf.Tensor) + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_url', + [*AUDIO_FILES, REMOTE_AUDIO_FILE], +) +def test_load(file_url): + url = parse_obj_as(AudioUrl, file_url) + tensor, _ = url.load() + assert isinstance(tensor, np.ndarray) + + +def test_json_schema(): + schema_json_of(AudioUrl) + + +def test_dump_json(): + url = parse_obj_as(AudioUrl, REMOTE_AUDIO_FILE) + orjson_dumps(url) + + +@pytest.mark.parametrize( + 'path_to_file', + [*AUDIO_FILES, REMOTE_AUDIO_FILE], +) +def test_validation(path_to_file): + url = parse_obj_as(AudioUrl, path_to_file) + assert isinstance(url, AudioUrl) + assert isinstance(url, str) + + +@pytest.mark.proto +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_url', + [*AUDIO_FILES, REMOTE_AUDIO_FILE], +) +def test_proto_audio_url(file_url): + uri = parse_obj_as(AudioUrl, file_url) + proto = uri._to_node_protobuf() + assert 'audio_url' in str(proto) + + +def test_load_bytes(): + uri = parse_obj_as(AudioUrl, REMOTE_AUDIO_FILE) + audio_bytes = uri.load_bytes() + assert isinstance(audio_bytes, bytes) + assert isinstance(audio_bytes, AudioBytes) + assert len(audio_bytes) > 0 + + +@pytest.mark.parametrize( + 'file_type, file_source', + [ + (AUDIO_MIMETYPE, AUDIO_FILES[0]), + (AUDIO_MIMETYPE, AUDIO_FILES[1]), + (AUDIO_MIMETYPE, REMOTE_AUDIO_FILE), + (IMAGE_MIMETYPE, os.path.join(TOYDATA_DIR, 'test.png')), + (VIDEO_MIMETYPE, os.path.join(TOYDATA_DIR, 'mov_bbb.mp4')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'test' 'test.html')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'test' 'test.md')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'penal_colony.txt')), + (OBJ_MIMETYPE, os.path.join(TOYDATA_DIR, 'test.glb')), + ], +) +def test_file_validation(file_type, file_source): + if file_type != AudioUrl.mime_type(): + with pytest.raises(ValueError): + parse_obj_as(AudioUrl, file_source) + else: + parse_obj_as(AudioUrl, file_source) diff --git a/tests/units/typing/url/test_image_url.py b/tests/units/typing/url/test_image_url.py new file mode 100644 index 00000000000..e5cc246da55 --- /dev/null +++ b/tests/units/typing/url/test_image_url.py @@ -0,0 +1,223 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import urllib + +import numpy as np +import PIL +import pytest +from PIL import Image +from pydantic.tools import parse_obj_as, schema_json_of + +from docarray.base_doc.io.json import orjson_dumps +from docarray.typing import ImageUrl +from docarray.typing.url.mimetypes import ( + OBJ_MIMETYPE, + AUDIO_MIMETYPE, + VIDEO_MIMETYPE, + IMAGE_MIMETYPE, + TEXT_MIMETYPE, +) +from tests import TOYDATA_DIR + +CUR_DIR = os.path.dirname(os.path.abspath(__file__)) +PATH_TO_IMAGE_DATA = os.path.join(CUR_DIR, '..', '..', '..', 'toydata', 'image-data') +IMAGE_PATHS = { + 'png': os.path.join(PATH_TO_IMAGE_DATA, 'so_good.png'), + 'jpg': os.path.join(PATH_TO_IMAGE_DATA, '05984.jpg'), + 'jpeg': os.path.join(PATH_TO_IMAGE_DATA, '05984-2.jpeg'), +} +REMOTE_JPG = ( + 'https://upload.wikimedia.org/wikipedia/commons/8/80/' + 'Dag_Sebastian_Ahlander_at_G%C3%B6teborg_Book_Fair_2012b.jpg' +) + + +@pytest.mark.slow +@pytest.mark.internet +def test_image_url(): + uri = parse_obj_as(ImageUrl, REMOTE_JPG) + + tensor = uri.load() + + assert isinstance(tensor, np.ndarray) + + +@pytest.mark.proto +def test_proto_image_url(): + uri = parse_obj_as(ImageUrl, REMOTE_JPG) + + uri._to_node_protobuf() + + +def test_json_schema(): + schema_json_of(ImageUrl) + + +def test_dump_json(): + url = parse_obj_as(ImageUrl, 'http://jina.ai/img.png') + orjson_dumps(url) + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'image_format,path_to_img', + [ + ('png', IMAGE_PATHS['png']), + ('jpg', IMAGE_PATHS['jpg']), + ('jpeg', IMAGE_PATHS['jpeg']), + ('remote-jpg', REMOTE_JPG), + ], +) +def test_load(image_format, path_to_img): + url = parse_obj_as(ImageUrl, path_to_img) + tensor = url.load() + assert isinstance(tensor, np.ndarray) + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'image_format,path_to_img', + [ + ('png', IMAGE_PATHS['png']), + ('jpg', IMAGE_PATHS['jpg']), + ('jpeg', IMAGE_PATHS['jpeg']), + ('remote-jpg', REMOTE_JPG), + ], +) +def test_load_pil(image_format, path_to_img): + url = parse_obj_as(ImageUrl, path_to_img) + img = url.load_pil() + assert isinstance(img, PIL.Image.Image) + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'image_format,path_to_img', + [ + ('png', IMAGE_PATHS['png']), + ('jpg', IMAGE_PATHS['jpg']), + ('jpeg', IMAGE_PATHS['jpeg']), + ('remote-jpg', REMOTE_JPG), + ], +) +@pytest.mark.parametrize('width,height', [(224, None), (None, 224), (224, 224)]) +def test_load_width_height(image_format, path_to_img, width, height): + url = parse_obj_as(ImageUrl, path_to_img) + tensor = url.load(width=width, height=height) + assert isinstance(tensor, np.ndarray) + + shape = tensor.shape + if width: + assert shape[1] == width + if height: + assert shape[0] == height + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'image_format,path_to_img', + [ + ('png', IMAGE_PATHS['png']), + ('jpg', IMAGE_PATHS['jpg']), + ('jpeg', IMAGE_PATHS['jpeg']), + ('remote-jpg', REMOTE_JPG), + ], +) +@pytest.mark.parametrize( + 'axis_layout', + [ + ('H', 'W', 'C'), + ('H', 'C', 'W'), + ('C', 'H', 'W'), + ('C', 'W', 'H'), + ('W', 'C', 'H'), + ('W', 'H', 'C'), + ], +) +def test_load_channel_axis(image_format, path_to_img, axis_layout): + sizes = {'H': 100, 'W': 200, 'C': 3} + url = parse_obj_as(ImageUrl, path_to_img) + tensor = url.load(axis_layout=axis_layout, height=sizes['H'], width=sizes['W']) + assert isinstance(tensor, np.ndarray) + + shape = tensor.shape + for axis, axis_name in enumerate(axis_layout): + assert shape[axis] == sizes[axis_name] + + +@pytest.mark.internet +def test_load_timeout(): + url = parse_obj_as(ImageUrl, REMOTE_JPG) + with pytest.raises(urllib.error.URLError): + _ = url.load(timeout=0.001) + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'image_format,path_to_img', + [ + ('png', IMAGE_PATHS['png']), + ('jpg', IMAGE_PATHS['jpg']), + ('jpeg', IMAGE_PATHS['jpeg']), + ('jpg', REMOTE_JPG), + ], +) +def test_load_to_bytes(image_format, path_to_img): + url = parse_obj_as(ImageUrl, path_to_img) + _bytes = url.load_bytes() + assert isinstance(_bytes, bytes) + img = Image.frombytes(mode='1', size=(224, 224), data=_bytes) + assert isinstance(img, Image.Image) + + +@pytest.mark.parametrize( + 'path_to_img', + [*IMAGE_PATHS.values(), REMOTE_JPG], +) +def test_validation(path_to_img): + url = parse_obj_as(ImageUrl, path_to_img) + assert isinstance(url, ImageUrl) + assert isinstance(url, str) + + +@pytest.mark.parametrize( + 'file_type, file_source', + [ + (IMAGE_MIMETYPE, IMAGE_PATHS['png']), + (IMAGE_MIMETYPE, IMAGE_PATHS['jpg']), + (IMAGE_MIMETYPE, IMAGE_PATHS['jpeg']), + (IMAGE_MIMETYPE, REMOTE_JPG), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.mp3')), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.wav')), + (VIDEO_MIMETYPE, os.path.join(TOYDATA_DIR, 'mov_bbb.mp4')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'test' 'test.html')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'test' 'test.md')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'penal_colony.txt')), + (OBJ_MIMETYPE, os.path.join(TOYDATA_DIR, 'test.glb')), + ], +) +def test_file_validation(file_type, file_source): + if file_type != ImageUrl.mime_type(): + with pytest.raises(ValueError): + parse_obj_as(ImageUrl, file_source) + else: + parse_obj_as(ImageUrl, file_source) diff --git a/tests/units/typing/url/test_mesh_url.py b/tests/units/typing/url/test_mesh_url.py new file mode 100644 index 00000000000..df807ffa501 --- /dev/null +++ b/tests/units/typing/url/test_mesh_url.py @@ -0,0 +1,126 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import numpy as np +import pytest +from pydantic.tools import parse_obj_as, schema_json_of + +from docarray.base_doc.io.json import orjson_dumps +from docarray.typing import Mesh3DUrl, NdArray +from docarray.typing.url.mimetypes import ( + OBJ_MIMETYPE, + AUDIO_MIMETYPE, + VIDEO_MIMETYPE, + IMAGE_MIMETYPE, + TEXT_MIMETYPE, +) +from tests import TOYDATA_DIR + +MESH_FILES = { + 'obj': str(TOYDATA_DIR / 'tetrahedron.obj'), + 'glb': str(TOYDATA_DIR / 'test.glb'), + 'ply': str(TOYDATA_DIR / 'cube.ply'), +} +REMOTE_OBJ_FILE = 'https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj' + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_format, file_path', + [ + ('obj', MESH_FILES['obj']), + ('glb', MESH_FILES['glb']), + ('ply', MESH_FILES['ply']), + ('remote-obj', REMOTE_OBJ_FILE), + ], +) +def test_load(file_format, file_path): + url = parse_obj_as(Mesh3DUrl, file_path) + tensors = url.load() + + assert isinstance(tensors.vertices, np.ndarray) + assert isinstance(tensors.vertices, NdArray) + assert isinstance(tensors.faces, np.ndarray) + assert isinstance(tensors.faces, NdArray) + assert tensors.vertices.shape[1] == 3 + assert tensors.faces.shape[1] == 3 + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_path', + [*MESH_FILES.values(), REMOTE_OBJ_FILE], +) +@pytest.mark.parametrize('field', ['vertices', 'faces']) +def test_load_one_of_fields(file_path, field): + url = parse_obj_as(Mesh3DUrl, file_path) + field = getattr(url.load(), field) + + assert isinstance(field, np.ndarray) + assert isinstance(field, NdArray) + + +def test_json_schema(): + schema_json_of(Mesh3DUrl) + + +def test_dump_json(): + url = parse_obj_as(Mesh3DUrl, REMOTE_OBJ_FILE) + orjson_dumps(url) + + +@pytest.mark.parametrize( + 'path_to_file', + [*MESH_FILES.values(), REMOTE_OBJ_FILE], +) +def test_validation(path_to_file): + url = parse_obj_as(Mesh3DUrl, path_to_file) + assert isinstance(url, Mesh3DUrl) + assert isinstance(url, str) + + +@pytest.mark.proto +def test_proto_mesh_url(): + uri = parse_obj_as(Mesh3DUrl, REMOTE_OBJ_FILE) + uri._to_node_protobuf() + + +@pytest.mark.parametrize( + 'file_type, file_source', + [ + (OBJ_MIMETYPE, MESH_FILES['obj']), + (OBJ_MIMETYPE, MESH_FILES['glb']), + (OBJ_MIMETYPE, MESH_FILES['ply']), + (OBJ_MIMETYPE, REMOTE_OBJ_FILE), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.aac')), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.mp3')), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.ogg')), + (VIDEO_MIMETYPE, os.path.join(TOYDATA_DIR, 'mov_bbb.mp4')), + (IMAGE_MIMETYPE, os.path.join(TOYDATA_DIR, 'test.png')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'test' 'test.html')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'test' 'test.md')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'penal_colony.txt')), + ], +) +def test_file_validation(file_type, file_source): + if file_type != Mesh3DUrl.mime_type(): + with pytest.raises(ValueError): + parse_obj_as(Mesh3DUrl, file_source) + else: + parse_obj_as(Mesh3DUrl, file_source) diff --git a/tests/units/typing/url/test_point_cloud_url.py b/tests/units/typing/url/test_point_cloud_url.py new file mode 100644 index 00000000000..3deb3e5779a --- /dev/null +++ b/tests/units/typing/url/test_point_cloud_url.py @@ -0,0 +1,130 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os + +import numpy as np +import pytest +from pydantic.tools import parse_obj_as, schema_json_of + +from docarray.base_doc.io.json import orjson_dumps +from docarray.typing import NdArray, PointCloud3DUrl +from docarray.typing.url.mimetypes import ( + OBJ_MIMETYPE, + AUDIO_MIMETYPE, + VIDEO_MIMETYPE, + IMAGE_MIMETYPE, + TEXT_MIMETYPE, +) +from tests import TOYDATA_DIR + +MESH_FILES = { + 'obj': str(TOYDATA_DIR / 'tetrahedron.obj'), + 'glb': str(TOYDATA_DIR / 'test.glb'), + 'ply': str(TOYDATA_DIR / 'cube.ply'), +} +REMOTE_OBJ_FILE = 'https://people.sc.fsu.edu/~jburkardt/data/obj/al.obj' + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_format, file_path', + [ + ('obj', MESH_FILES['obj']), + ('glb', MESH_FILES['glb']), + ('ply', MESH_FILES['ply']), + ('remote-obj', REMOTE_OBJ_FILE), + ], +) +def test_load(file_format, file_path): + n_samples = 100 + url = parse_obj_as(PointCloud3DUrl, file_path) + tensors = url.load(samples=n_samples) + + assert isinstance(tensors.points, np.ndarray) + assert isinstance(tensors.points, NdArray) + assert tensors.points.shape == (n_samples, 3) + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_format, file_path', + [ + ('obj', MESH_FILES['obj']), + ('glb', MESH_FILES['glb']), + ('ply', MESH_FILES['ply']), + ('remote-obj', REMOTE_OBJ_FILE), + ], +) +def test_load_with_multiple_geometries_true(file_format, file_path): + n_samples = 100 + url = parse_obj_as(PointCloud3DUrl, file_path) + tensors = url.load(samples=n_samples, multiple_geometries=True) + + assert isinstance(tensors.points, np.ndarray) + assert len(tensors.points.shape) == 3 + assert tensors.points.shape[1:] == (100, 3) + + +def test_json_schema(): + schema_json_of(PointCloud3DUrl) + + +def test_dump_json(): + url = parse_obj_as(PointCloud3DUrl, REMOTE_OBJ_FILE) + orjson_dumps(url) + + +@pytest.mark.parametrize( + 'path_to_file', + [*MESH_FILES.values(), REMOTE_OBJ_FILE], +) +def test_validation(path_to_file): + url = parse_obj_as(PointCloud3DUrl, path_to_file) + assert isinstance(url, PointCloud3DUrl) + assert isinstance(url, str) + + +@pytest.mark.proto +def test_proto_point_cloud_url(): + uri = parse_obj_as(PointCloud3DUrl, REMOTE_OBJ_FILE) + uri._to_node_protobuf() + + +@pytest.mark.parametrize( + 'file_type, file_source', + [ + (OBJ_MIMETYPE, MESH_FILES['obj']), + (OBJ_MIMETYPE, MESH_FILES['glb']), + (OBJ_MIMETYPE, MESH_FILES['ply']), + (OBJ_MIMETYPE, REMOTE_OBJ_FILE), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.aac')), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.mp3')), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.ogg')), + (VIDEO_MIMETYPE, os.path.join(TOYDATA_DIR, 'mov_bbb.mp4')), + (IMAGE_MIMETYPE, os.path.join(TOYDATA_DIR, 'test.png')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'test' 'test.html')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'test' 'test.md')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'penal_colony.txt')), + ], +) +def test_file_validation(file_type, file_source): + if file_type != PointCloud3DUrl.mime_type(): + with pytest.raises(ValueError): + parse_obj_as(PointCloud3DUrl, file_source) + else: + parse_obj_as(PointCloud3DUrl, file_source) diff --git a/tests/units/typing/url/test_text_url.py b/tests/units/typing/url/test_text_url.py new file mode 100644 index 00000000000..a755344f394 --- /dev/null +++ b/tests/units/typing/url/test_text_url.py @@ -0,0 +1,119 @@ +import os +import urllib + +import pytest +from pydantic import parse_obj_as, schema_json_of + +from docarray.base_doc.io.json import orjson_dumps +from docarray.typing import TextUrl +from docarray.typing.url.mimetypes import ( + OBJ_MIMETYPE, + AUDIO_MIMETYPE, + VIDEO_MIMETYPE, + IMAGE_MIMETYPE, + TEXT_MIMETYPE, +) +from tests import TOYDATA_DIR + +REMOTE_TEXT_FILE = 'https://de.wikipedia.org/wiki/Brixen' +CUR_DIR = os.path.dirname(os.path.abspath(__file__)) +LOCAL_TEXT_FILES = [ + str(TOYDATA_DIR / 'penal_colony.txt'), + str(TOYDATA_DIR / 'test.md'), + str(TOYDATA_DIR / 'test.html'), + str(TOYDATA_DIR / 'test.css'), + str(TOYDATA_DIR / 'test.csv'), + str(TOYDATA_DIR / 'test.log'), +] +LOCAL_TEXT_FILES_AND_BEGINNING = [ + (str(TOYDATA_DIR / 'penal_colony.txt'), '“It’s a peculiar apparatus,”'), + (str(TOYDATA_DIR / 'test.md'), "# Hello"), + (str(TOYDATA_DIR / 'test.html'), ""), + (str(TOYDATA_DIR / 'test.css'), 'body {'), + (str(TOYDATA_DIR / 'test.csv'), "John,Doe"), + (str(TOYDATA_DIR / 'test.log'), "2022-11-25 12:34:56 INFO: Program started"), +] + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'url,expected_beginning', + [(REMOTE_TEXT_FILE, ''), *LOCAL_TEXT_FILES_AND_BEGINNING], +) +def test_load(url, expected_beginning): + uri = parse_obj_as(TextUrl, url) + + txt = uri.load() + assert txt.startswith(expected_beginning) + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize('url', [REMOTE_TEXT_FILE, *LOCAL_TEXT_FILES]) +def test_load_to_bytes(url): + uri = parse_obj_as(TextUrl, url) + + txt_bytes = uri.load_bytes() + assert isinstance(txt_bytes, bytes) + + +@pytest.mark.proto +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize('url', [REMOTE_TEXT_FILE, *LOCAL_TEXT_FILES]) +def test_proto_text_url(url): + uri = parse_obj_as(TextUrl, url) + + proto = uri._to_node_protobuf() + assert 'text_url' in str(proto) + + +@pytest.mark.internet +def test_load_timeout(): + url = parse_obj_as(TextUrl, REMOTE_TEXT_FILE) + with pytest.raises(urllib.error.URLError): + _ = url.load(timeout=0.001) + with pytest.raises(urllib.error.URLError): + _ = url.load_bytes(timeout=0.001) + + +def test_json_schema(): + schema_json_of(TextUrl) + + +@pytest.mark.internet +def test_dump_json(): + url = parse_obj_as(TextUrl, REMOTE_TEXT_FILE) + orjson_dumps(url) + + +@pytest.mark.parametrize( + 'path_to_file', + [REMOTE_TEXT_FILE, *LOCAL_TEXT_FILES], +) +def test_validation(path_to_file): + url = parse_obj_as(TextUrl, path_to_file) + assert isinstance(url, TextUrl) + assert isinstance(url, str) + + +@pytest.mark.parametrize( + 'file_type, file_source', + [ + *[(TEXT_MIMETYPE, file) for file in LOCAL_TEXT_FILES], + (TEXT_MIMETYPE, REMOTE_TEXT_FILE), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.aac')), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.mp3')), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.ogg')), + (IMAGE_MIMETYPE, os.path.join(TOYDATA_DIR, 'test.png')), + (VIDEO_MIMETYPE, os.path.join(TOYDATA_DIR, 'mov_bbb.mp4')), + (OBJ_MIMETYPE, os.path.join(TOYDATA_DIR, 'test.glb')), + ], +) +def test_file_validation(file_type, file_source): + if file_type != TextUrl.mime_type(): + with pytest.raises(ValueError): + parse_obj_as(TextUrl, file_source) + else: + parse_obj_as(TextUrl, file_source) diff --git a/tests/units/typing/url/test_video_url.py b/tests/units/typing/url/test_video_url.py new file mode 100644 index 00000000000..0bd889f37bf --- /dev/null +++ b/tests/units/typing/url/test_video_url.py @@ -0,0 +1,179 @@ +import os +from typing import Optional + +import numpy as np +import pytest +import torch +from pydantic.tools import parse_obj_as, schema_json_of + +from docarray import BaseDoc +from docarray.base_doc.io.json import orjson_dumps +from docarray.typing import ( + AudioNdArray, + NdArray, + VideoBytes, + VideoNdArray, + VideoTorchTensor, + VideoUrl, +) +from docarray.typing.url.mimetypes import ( + AUDIO_MIMETYPE, + IMAGE_MIMETYPE, + OBJ_MIMETYPE, + TEXT_MIMETYPE, + VIDEO_MIMETYPE, +) +from docarray.utils._internal.misc import is_tf_available +from tests import TOYDATA_DIR + +tf_available = is_tf_available() +if tf_available: + import tensorflow as tf + + from docarray.typing.tensor.video import VideoTensorFlowTensor + +LOCAL_VIDEO_FILE = str(TOYDATA_DIR / 'mov_bbb.mp4') +REMOTE_VIDEO_FILE = 'https://github.com/docarray/docarray/blob/main/tests/toydata/mov_bbb.mp4?raw=true' # noqa: E501 + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_url', + [LOCAL_VIDEO_FILE, REMOTE_VIDEO_FILE], +) +def test_load(file_url): + url = parse_obj_as(VideoUrl, file_url) + video, audio, indices = url.load() + + assert isinstance(audio, np.ndarray) + assert isinstance(audio, AudioNdArray) + + assert isinstance(video, np.ndarray) + assert isinstance(video, VideoNdArray) + + assert isinstance(indices, np.ndarray) + assert isinstance(indices, NdArray) + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_url', + [LOCAL_VIDEO_FILE, REMOTE_VIDEO_FILE], +) +@pytest.mark.parametrize( + 'field, attr_cls', + [ + ('video', VideoNdArray), + ('audio', AudioNdArray), + ('key_frame_indices', NdArray), + ], +) +def test_load_one_of_named_tuple_results(file_url, field, attr_cls): + url = parse_obj_as(VideoUrl, file_url) + result = getattr(url.load(), field) + + assert isinstance(result, np.ndarray) + assert isinstance(result, attr_cls) + + +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_url', + [LOCAL_VIDEO_FILE, REMOTE_VIDEO_FILE], +) +def test_load_video_url_to_video_torch_tensor_field(file_url): + class MyVideoDoc(BaseDoc): + video_url: VideoUrl + tensor: Optional[VideoTorchTensor] = None + + doc = MyVideoDoc(video_url=file_url) + doc.tensor = doc.video_url.load().video + + assert isinstance(doc.tensor, torch.Tensor) + assert isinstance(doc.tensor, VideoTorchTensor) + + +@pytest.mark.tensorflow +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_url', + [LOCAL_VIDEO_FILE, REMOTE_VIDEO_FILE], +) +def test_load_video_url_to_video_tensorflow_tensor_field(file_url): + class MyVideoDoc(BaseDoc): + video_url: VideoUrl + tensor: Optional[VideoTensorFlowTensor] = None + + doc = MyVideoDoc(video_url=file_url) + doc.tensor = doc.video_url.load().video + + assert isinstance(doc.tensor, VideoTensorFlowTensor) + assert isinstance(doc.tensor.tensor, tf.Tensor) + + +def test_json_schema(): + schema_json_of(VideoUrl) + + +def test_dump_json(): + url = parse_obj_as(VideoUrl, REMOTE_VIDEO_FILE) + orjson_dumps(url) + + +@pytest.mark.parametrize( + 'path_to_file', + [LOCAL_VIDEO_FILE, REMOTE_VIDEO_FILE], +) +def test_validation(path_to_file): + url = parse_obj_as(VideoUrl, path_to_file) + assert isinstance(url, VideoUrl) + assert isinstance(url, str) + + +@pytest.mark.proto +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'file_url', + [LOCAL_VIDEO_FILE, REMOTE_VIDEO_FILE], +) +def test_proto_video_url(file_url): + uri = parse_obj_as(VideoUrl, file_url) + proto = uri._to_node_protobuf() + assert 'video_url' in str(proto) + + +def test_load_bytes(): + file_url = LOCAL_VIDEO_FILE + uri = parse_obj_as(VideoUrl, file_url) + video_bytes = uri.load_bytes() + assert isinstance(video_bytes, bytes) + assert isinstance(video_bytes, VideoBytes) + assert len(video_bytes) > 0 + + +@pytest.mark.parametrize( + 'file_type, file_source', + [ + (VIDEO_MIMETYPE, LOCAL_VIDEO_FILE), + (VIDEO_MIMETYPE, REMOTE_VIDEO_FILE), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.aac')), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.mp3')), + (AUDIO_MIMETYPE, os.path.join(TOYDATA_DIR, 'hello.ogg')), + (IMAGE_MIMETYPE, os.path.join(TOYDATA_DIR, 'test.png')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'test' 'test.html')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'test' 'test.md')), + (TEXT_MIMETYPE, os.path.join(TOYDATA_DIR, 'penal_colony.txt')), + (OBJ_MIMETYPE, os.path.join(TOYDATA_DIR, 'test.glb')), + ], +) +def test_file_validation(file_type, file_source): + if file_type != VideoUrl.mime_type(): + with pytest.raises(ValueError): + parse_obj_as(VideoUrl, file_source) + else: + parse_obj_as(VideoUrl, file_source) diff --git a/tests/units/util/__init__.py b/tests/units/util/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/util/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/util/query_language/__init__.py b/tests/units/util/query_language/__init__.py new file mode 100644 index 00000000000..74f8f7582cd --- /dev/null +++ b/tests/units/util/query_language/__init__.py @@ -0,0 +1,15 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/units/util/query_language/test_lookup.py b/tests/units/util/query_language/test_lookup.py new file mode 100644 index 00000000000..844f5475b9e --- /dev/null +++ b/tests/units/util/query_language/test_lookup.py @@ -0,0 +1,68 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import pytest + +from docarray.utils._internal.query_language.lookup import dunder_get, lookup + + +class A: + class B: + c = 0 + d = 'docarray' + e = [0, 1] + f = {} + + b: B = B() + + +@pytest.mark.parametrize('input', [A(), {'b': {'c': 0, 'd': 'docarray'}}]) +def test_dunder_get(input): + assert dunder_get(input, 'b__c') == 0 + + expected_exception = KeyError if isinstance(input, dict) else AttributeError + with pytest.raises(expected_exception): + _ = dunder_get(input, 'z') + + with pytest.raises(expected_exception): + _ = dunder_get(input, 'b__z') + + +@pytest.mark.parametrize( + 'input', [A(), {'b': {'c': 0, 'd': 'docarray', 'e': [0, 1], 'f': {}}}] +) +def test_lookup(input): + assert lookup('b__c.exact', 0, input) + assert not lookup('b__c.gt', 0, input) + assert lookup('b__c.gte', 0, input) + assert not lookup('b__c.lt', 0, input) + assert lookup('b__c.lte', 0, input) + assert lookup('b__d.regex', 'array*', input) + assert lookup('b__d.contains', 'array', input) + assert lookup('b__d.icontains', 'Array', input) + assert lookup('b__d.in', ['a', 'docarray'], input) + assert lookup('b__d.nin', ['a', 'b'], input) + assert lookup('b__d.startswith', 'doc', input) + assert lookup('b__d.istartswith', 'Doc', input) + assert lookup('b__d.endswith', 'array', input) + assert lookup('b__d.iendswith', 'Array', input) + assert lookup('b__e.size', 2, input) + assert not lookup('b__e.size', 3, input) + assert lookup('b__d.size', len('docarray'), input) + assert not lookup('b__e.size', len('docarray') + 1, input) + assert not lookup('b__z.exists', True, input) + assert lookup('b__z.exists', False, input) + assert not lookup('b__f__z.exists', True, input) + assert lookup('b__f__z.exists', False, input) diff --git a/tests/units/util/test_create_dynamic_code_class.py b/tests/units/util/test_create_dynamic_code_class.py new file mode 100644 index 00000000000..b7df497816d --- /dev/null +++ b/tests/units/util/test_create_dynamic_code_class.py @@ -0,0 +1,358 @@ +from typing import Any, Dict, List, Optional, Union, ClassVar + +import numpy as np +import pytest +from pydantic import Field + +from docarray import BaseDoc, DocList +from docarray.documents import TextDoc +from docarray.typing import AnyTensor, ImageUrl +from docarray.utils.create_dynamic_doc_class import ( + create_base_doc_from_schema, + create_pure_python_type_model, +) +from docarray.utils._internal.pydantic import is_pydantic_v2 + + +@pytest.mark.parametrize('transformation', ['proto', 'json']) +def test_create_pydantic_model_from_schema(transformation): + class Nested2Doc(BaseDoc): + value: str + classvar: ClassVar[str] = 'classvar2' + + class Nested1Doc(BaseDoc): + nested: Nested2Doc + classvar: ClassVar[str] = 'classvar1' + + class CustomDoc(BaseDoc): + tensor: Optional[AnyTensor] = None + url: ImageUrl + num: float = 0.5 + num_num: List[float] = [1.5, 2.5] + lll: List[List[List[int]]] = [[[5]]] + fff: List[List[List[float]]] = [[[5.2]]] + single_text: TextDoc + texts: DocList[TextDoc] + d: Dict[str, str] = {'a': 'b'} + di: Optional[Dict[str, int]] = None + u: Union[str, int] + lu: List[Union[str, int]] = [0, 1, 2] + tags: Optional[Dict[str, Any]] = None + nested: Nested1Doc + classvar: ClassVar[str] = 'classvar' + + CustomDocCopy = create_pure_python_type_model(CustomDoc) + new_custom_doc_model = create_base_doc_from_schema( + CustomDocCopy.schema(), 'CustomDoc', {} + ) + print(f'new_custom_doc_model {new_custom_doc_model.schema()}') + + original_custom_docs = DocList[CustomDoc]( + [ + CustomDoc( + num=3.5, + num_num=[4.5, 5.5], + url='photo.jpg', + lll=[[[40]]], + fff=[[[40.2]]], + d={'b': 'a'}, + texts=DocList[TextDoc]([TextDoc(text='hey ha', embedding=np.zeros(3))]), + single_text=TextDoc(text='single hey ha', embedding=np.zeros(2)), + u='a', + lu=[3, 4], + nested=Nested1Doc(nested=Nested2Doc(value='hello world')), + ) + ] + ) + for doc in original_custom_docs: + doc.tensor = np.zeros((10, 10, 10)) + doc.di = {'a': 2} + + if transformation == 'proto': + custom_partial_da = DocList[new_custom_doc_model].from_protobuf( + original_custom_docs.to_protobuf() + ) + original_back = DocList[CustomDoc].from_protobuf( + custom_partial_da.to_protobuf() + ) + elif transformation == 'json': + custom_partial_da = DocList[new_custom_doc_model].from_json( + original_custom_docs.to_json() + ) + original_back = DocList[CustomDoc].from_json(custom_partial_da.to_json()) + + assert len(custom_partial_da) == 1 + assert custom_partial_da[0].url == 'photo.jpg' + assert custom_partial_da[0].num == 3.5 + assert custom_partial_da[0].num_num == [4.5, 5.5] + assert custom_partial_da[0].lll == [[[40]]] + if is_pydantic_v2: + assert custom_partial_da[0].lu == [3, 4] + else: + assert custom_partial_da[0].lu == ['3', '4'] # Union validates back to string + assert custom_partial_da[0].fff == [[[40.2]]] + assert custom_partial_da[0].di == {'a': 2} + assert custom_partial_da[0].d == {'b': 'a'} + assert len(custom_partial_da[0].texts) == 1 + assert custom_partial_da[0].texts[0].text == 'hey ha' + assert custom_partial_da[0].texts[0].embedding.shape == (3,) + assert custom_partial_da[0].tensor.shape == (10, 10, 10) + assert custom_partial_da[0].u == 'a' + assert custom_partial_da[0].single_text.text == 'single hey ha' + assert custom_partial_da[0].single_text.embedding.shape == (2,) + assert original_back[0].nested.nested.value == 'hello world' + assert original_back[0].num == 3.5 + assert original_back[0].num_num == [4.5, 5.5] + assert original_back[0].classvar == 'classvar' + assert original_back[0].nested.classvar == 'classvar1' + assert original_back[0].nested.nested.classvar == 'classvar2' + + assert len(original_back) == 1 + assert original_back[0].url == 'photo.jpg' + assert original_back[0].lll == [[[40]]] + if is_pydantic_v2: + assert original_back[0].lu == [3, 4] # Union validates back to string + else: + assert original_back[0].lu == ['3', '4'] # Union validates back to string + assert original_back[0].fff == [[[40.2]]] + assert original_back[0].di == {'a': 2} + assert original_back[0].d == {'b': 'a'} + assert len(original_back[0].texts) == 1 + assert original_back[0].texts[0].text == 'hey ha' + assert original_back[0].texts[0].embedding.shape == (3,) + assert original_back[0].tensor.shape == (10, 10, 10) + assert original_back[0].u == 'a' + assert original_back[0].single_text.text == 'single hey ha' + assert original_back[0].single_text.embedding.shape == (2,) + + class TextDocWithId(BaseDoc): + ia: str + + TextDocWithIdCopy = create_pure_python_type_model(TextDocWithId) + new_textdoc_with_id_model = create_base_doc_from_schema( + TextDocWithIdCopy.schema(), 'TextDocWithId', {} + ) + print(f'new_textdoc_with_id_model {new_textdoc_with_id_model.schema()}') + + original_text_doc_with_id = DocList[TextDocWithId]( + [TextDocWithId(ia=f'ID {i}') for i in range(10)] + ) + if transformation == 'proto': + custom_da = DocList[new_textdoc_with_id_model].from_protobuf( + original_text_doc_with_id.to_protobuf() + ) + original_back = DocList[TextDocWithId].from_protobuf(custom_da.to_protobuf()) + elif transformation == 'json': + custom_da = DocList[new_textdoc_with_id_model].from_json( + original_text_doc_with_id.to_json() + ) + original_back = DocList[TextDocWithId].from_json(custom_da.to_json()) + + assert len(custom_da) == 10 + for i, doc in enumerate(custom_da): + assert doc.ia == f'ID {i}' + + assert len(original_back) == 10 + for i, doc in enumerate(original_back): + assert doc.ia == f'ID {i}' + + class ResultTestDoc(BaseDoc): + matches: DocList[TextDocWithId] + + ResultTestDocCopy = create_pure_python_type_model(ResultTestDoc) + new_result_test_doc_with_id_model = create_base_doc_from_schema( + ResultTestDocCopy.schema(), 'ResultTestDoc', {} + ) + result_test_docs = DocList[ResultTestDoc]( + [ResultTestDoc(matches=original_text_doc_with_id)] + ) + + if transformation == 'proto': + custom_da = DocList[new_result_test_doc_with_id_model].from_protobuf( + result_test_docs.to_protobuf() + ) + original_back = DocList[ResultTestDoc].from_protobuf(custom_da.to_protobuf()) + elif transformation == 'json': + custom_da = DocList[new_result_test_doc_with_id_model].from_json( + result_test_docs.to_json() + ) + original_back = DocList[ResultTestDoc].from_json(custom_da.to_json()) + + assert len(custom_da) == 1 + assert len(custom_da[0].matches) == 10 + for i, doc in enumerate(custom_da[0].matches): + assert doc.ia == f'ID {i}' + + assert len(original_back) == 1 + assert len(original_back[0].matches) == 10 + for i, doc in enumerate(original_back[0].matches): + assert doc.ia == f'ID {i}' + + +@pytest.mark.parametrize('transformation', ['proto', 'json']) +def test_create_empty_doc_list_from_schema(transformation): + class CustomDoc(BaseDoc): + tensor: Optional[AnyTensor] + url: ImageUrl + lll: List[List[List[int]]] = [[[5]]] + fff: List[List[List[float]]] = [[[5.2]]] + single_text: TextDoc + texts: DocList[TextDoc] + d: Dict[str, str] = {'a': 'b'} + di: Optional[Dict[str, int]] = None + u: Union[str, int] + lu: List[Union[str, int]] = [0, 1, 2] + tags: Optional[Dict[str, Any]] = None + lf: List[float] = [3.0, 4.1] + + CustomDocCopy = create_pure_python_type_model(CustomDoc) + new_custom_doc_model = create_base_doc_from_schema( + CustomDocCopy.schema(), 'CustomDoc' + ) + print(f'new_custom_doc_model {new_custom_doc_model.schema()}') + + original_custom_docs = DocList[CustomDoc]() + if transformation == 'proto': + custom_partial_da = DocList[new_custom_doc_model].from_protobuf( + original_custom_docs.to_protobuf() + ) + original_back = DocList[CustomDoc].from_protobuf( + custom_partial_da.to_protobuf() + ) + elif transformation == 'json': + custom_partial_da = DocList[new_custom_doc_model].from_json( + original_custom_docs.to_json() + ) + original_back = DocList[CustomDoc].from_json(custom_partial_da.to_json()) + + assert len(custom_partial_da) == 0 + assert len(original_back) == 0 + + class TextDocWithId(BaseDoc): + ia: str + + TextDocWithIdCopy = create_pure_python_type_model(TextDocWithId) + new_textdoc_with_id_model = create_base_doc_from_schema( + TextDocWithIdCopy.schema(), 'TextDocWithId', {} + ) + print(f'new_textdoc_with_id_model {new_textdoc_with_id_model.schema()}') + + original_text_doc_with_id = DocList[TextDocWithId]() + if transformation == 'proto': + custom_da = DocList[new_textdoc_with_id_model].from_protobuf( + original_text_doc_with_id.to_protobuf() + ) + original_back = DocList[TextDocWithId].from_protobuf(custom_da.to_protobuf()) + elif transformation == 'json': + custom_da = DocList[new_textdoc_with_id_model].from_json( + original_text_doc_with_id.to_json() + ) + original_back = DocList[TextDocWithId].from_json(custom_da.to_json()) + + assert len(original_back) == 0 + assert len(custom_da) == 0 + + class ResultTestDoc(BaseDoc): + matches: DocList[TextDocWithId] + + ResultTestDocCopy = create_pure_python_type_model(ResultTestDoc) + new_result_test_doc_with_id_model = create_base_doc_from_schema( + ResultTestDocCopy.schema(), 'ResultTestDoc', {} + ) + print( + f'new_result_test_doc_with_id_model {new_result_test_doc_with_id_model.schema()}' + ) + result_test_docs = DocList[ResultTestDoc]() + + if transformation == 'proto': + custom_da = DocList[new_result_test_doc_with_id_model].from_protobuf( + result_test_docs.to_protobuf() + ) + original_back = DocList[ResultTestDoc].from_protobuf(custom_da.to_protobuf()) + elif transformation == 'json': + custom_da = DocList[new_result_test_doc_with_id_model].from_json( + result_test_docs.to_json() + ) + original_back = DocList[ResultTestDoc].from_json(custom_da.to_json()) + + assert len(original_back) == 0 + assert len(custom_da) == 0 + + +def test_create_with_field_info(): + class CustomDoc(BaseDoc): + """Here I have the description of the class""" + + a: str = Field(examples=['Example here'], another_extra='I am another extra') + + CustomDocCopy = create_pure_python_type_model(CustomDoc) + new_custom_doc_model = create_base_doc_from_schema( + CustomDocCopy.schema(), 'CustomDoc' + ) + assert new_custom_doc_model.schema().get('properties')['a']['examples'] == [ + 'Example here' + ] + assert ( + new_custom_doc_model.schema().get('properties')['a']['another_extra'] + == 'I am another extra' + ) + assert ( + new_custom_doc_model.schema().get('description') + == 'Here I have the description of the class' + ) + + +def test_dynamic_class_creation_multiple_doclist_nested(): + from docarray import BaseDoc, DocList + + class MyTextDoc(BaseDoc): + text: str + + class QuoteFile(BaseDoc): + texts: DocList[MyTextDoc] + + class SearchResult(BaseDoc): + results: DocList[QuoteFile] = None + + models_created_by_name = {} + SearchResult_aux = create_pure_python_type_model(SearchResult) + m = create_base_doc_from_schema( + SearchResult_aux.schema(), 'SearchResult', models_created_by_name + ) + print(f'm {m.schema()}') + QuoteFile_reconstructed_in_gateway_from_Search_results = models_created_by_name[ + 'QuoteFile' + ] + textlist = DocList[models_created_by_name['MyTextDoc']]( + [models_created_by_name['MyTextDoc'](id='11', text='hey')] + ) + + reconstructed_in_gateway_from_Search_results = ( + QuoteFile_reconstructed_in_gateway_from_Search_results(id='0', texts=textlist) + ) + assert reconstructed_in_gateway_from_Search_results.texts[0].text == 'hey' + + +def test_id_optional(): + from docarray import BaseDoc + import json + + class MyTextDoc(BaseDoc): + text: str + opt: Optional[str] = None + + MyTextDoc_aux = create_pure_python_type_model(MyTextDoc) + td = create_base_doc_from_schema(MyTextDoc_aux.schema(), 'MyTextDoc') + print(f'{td.schema()}') + direct = MyTextDoc.from_json(json.dumps({"text": "text"})) + aux = MyTextDoc_aux.from_json(json.dumps({"text": "text"})) + indirect = td.from_json(json.dumps({"text": "text"})) + assert direct.text == 'text' + assert aux.text == 'text' + assert indirect.text == 'text' + direct = MyTextDoc(text='hey') + aux = MyTextDoc_aux(text='hey') + indirect = td(text='hey') + assert direct.text == 'hey' + assert aux.text == 'hey' + assert indirect.text == 'hey' diff --git a/tests/units/util/test_filter.py b/tests/units/util/test_filter.py new file mode 100644 index 00000000000..417bde4232e --- /dev/null +++ b/tests/units/util/test_filter.py @@ -0,0 +1,289 @@ +import json +from typing import Any, Dict, List, Optional + +import pytest + +from docarray import BaseDoc, DocList +from docarray.documents import ImageDoc, TextDoc +from docarray.utils.filter import filter_docs + + +class MMDoc(BaseDoc): + text_doc: TextDoc + text: str = '' + image: Optional[ImageDoc] = None + price: int = 0 + optional_num: Optional[int] = None + boolean: bool = False + categories: Optional[List[str]] = None + sub_docs: Optional[List[TextDoc]] = None + dictionary: Optional[Dict[str, Any]] = None + + +@pytest.fixture +def docs(): + mmdoc1 = MMDoc( + text_doc=TextDoc(text='Text Doc of Document 1'), + text='Text of Document 1', + sub_docs=[TextDoc(text='subtext1'), TextDoc(text='subtext2')], + dictionary={}, + ) + mmdoc2 = MMDoc( + text_doc=TextDoc(text='Text Doc of Document 2'), + text='Text of Document 2', + image=ImageDoc(url='exampleimage.jpg'), + price=3, + dictionary={'a': 0, 'b': 1, 'c': 2, 'd': {'e': 3}}, + ) + mmdoc3 = MMDoc( + text_doc=TextDoc(text='Text Doc of Document 3'), + text='Text of Document 3', + price=1000, + boolean=True, + categories=['cat1', 'cat2'], + sub_docs=[TextDoc(text='subtext1'), TextDoc(text='subtext2')], + optional_num=30, + dictionary={'a': 0, 'b': 1}, + ) + docs = DocList[MMDoc]([mmdoc1, mmdoc2, mmdoc3]) + + return docs + + +@pytest.mark.parametrize('dict_api', [True, False]) +def test_empty_filter(docs, dict_api): + q = {} if dict_api else '{}' + result = filter_docs(docs, q) + assert len(result) == len(docs) + + +@pytest.mark.parametrize('dict_api', [True, False]) +def test_simple_filter(docs, dict_api): + if dict_api: + method = lambda query: filter_docs(docs, query) # noqa: E731 + else: + method = lambda query: filter_docs(docs, json.dumps(query)) # noqa: E731 + + result = method({'text': {'$eq': 'Text of Document 1'}}) + assert len(result) == 1 + assert result[0].text == 'Text of Document 1' + + result = method({'text': {'$neq': 'Text of Document 1'}}) + assert len(result) == 2 + + result = method({'text_doc': {'$eq': 'Text Doc of Document 1'}}) + assert len(result) == 1 + assert result[0].text_doc == 'Text Doc of Document 1' + + result = method({'text_doc': {'$neq': 'Text Doc of Document 1'}}) + assert len(result) == 2 + + result = method({'text': {'$regex': 'Text*'}}) + assert len(result) == 3 + + result = method({'text': {'$regex': 'TeAxt*'}}) + assert len(result) == 0 + + result = method({'text_doc': {'$regex': 'Text*'}}) + assert len(result) == 3 + + result = method({'text_doc': {'$regex': 'TeAxt*'}}) + assert len(result) == 0 + + result = method({'price': {'$gte': 500}}) + assert len(result) == 1 + + result = method({'price': {'$lte': 500}}) + assert len(result) == 2 + + result = method({'dictionary': {'$eq': {}}}) + assert len(result) == 1 + assert result[0].dictionary == {} + + result = method({'dictionary': {'$eq': {'a': 0, 'b': 1}}}) + assert len(result) == 1 + assert result[0].dictionary == {'a': 0, 'b': 1} + + result = method({'text': {'$neq': 'Text of Document 1'}}) + assert len(result) == 2 + + # EXISTS DOES NOT SEEM TO WORK + result = method({'optional_num': {'$exists': True}}) + assert len(result) == 3 + result = method({'optional_num': {'$exists': False}}) + assert len(result) == 0 + + result = method({'price': {'$exists': True}}) + assert len(result) == 3 + result = method({'price': {'$exists': False}}) + assert len(result) == 0 + + # DOES NOT SEEM TO WORK WITH OPTIONAL NUMBERS + result = method({'optional_num': {'$gte': 20}}) + assert len(result) == 1 + + result = method({'optional_num': {'$lte': 20}}) + assert len(result) == 0 + + +@pytest.mark.parametrize('dict_api', [True, False]) +def test_nested_filter(docs, dict_api): + if dict_api: + method = lambda query: filter_docs(docs, query) # noqa: E731 + else: + method = lambda query: filter_docs(docs, json.dumps(query)) # noqa: E731 + + result = method({'dictionary__a': {'$eq': 0}}) + assert len(result) == 2 + for res in result: + assert res.dictionary['a'] == 0 + + result = method({'dictionary__c': {'$exists': True}}) + assert len(result) == 1 + assert result[0].dictionary['c'] == 2 + + result = method({'dictionary__d__e': {'$exists': True}}) + assert len(result) == 1 + assert result[0].dictionary['d'] == {'e': 3} + + result = method({'dictionary__d__e': {'$eq': 3}}) + assert len(result) == 1 + assert result[0].dictionary['d'] == {'e': 3} + + result = method({'image__url': {'$eq': 'exampleimage.jpg'}}) + assert len(result) == 1 + assert result[0].image.url == 'exampleimage.jpg' + + +@pytest.mark.parametrize('dict_api', [True, False]) +def test_array_simple_filters(docs, dict_api): + if dict_api: + method = lambda query: filter_docs(docs, query) # noqa: E731 + else: + method = lambda query: filter_docs(docs, json.dumps(query)) # noqa: E731 + + # SIZE DOES NOT SEEM TO WORK + result = method({'sub_docs': {'$size': 2}}) + assert len(result) == 2 + + result = method({'categories': {'$size': 2}}) + assert len(result) == 1 + + +@pytest.mark.parametrize('dict_api', [True, False]) +def test_placehold_filter(dict_api): + docs = DocList[MMDoc]( + [ + MMDoc(text='A', text_doc=TextDoc(text='A')), + MMDoc(text='A', text_doc=TextDoc(text='B')), + ] + ) + + if dict_api: + method = lambda query: filter_docs(docs, query) # noqa: E731 + else: + method = lambda query: filter_docs(docs, json.dumps(query)) # noqa: E731 + + # DOES NOT SEEM TO WORK + result = method({'text': {'$eq': '{text_doc}'}}) + assert len(result) == 1 + + result = method({'text_doc': {'$eq': '{text}'}}) + assert len(result) == 1 + + +@pytest.mark.parametrize('dict_api', [True, False]) +def test_logic_filter(docs, dict_api): + if dict_api: + method = lambda query: filter_docs(docs, query) # noqa: E731 + else: + method = lambda query: filter_docs(docs, json.dumps(query)) # noqa: E731 + result = method( + { + '$or': { + 'text': {'$eq': 'Text of Document 1'}, + 'text_doc': {'$eq': 'Text Doc of Document 2'}, + } + } + ) + assert len(result) == 2 + + result = method( + { + '$not': { + '$or': { + 'text': {'$eq': 'Text of Document 1'}, + 'text_doc': {'$eq': 'Text Doc of Document 2'}, + } + } + } + ) + assert len(result) == 1 + + result = method( + { + '$and': { + 'text': {'$eq': 'Text of Document 1'}, + 'text_doc': {'$eq': 'Text Doc of Document 2'}, + } + } + ) + assert len(result) == 0 + + result = method( + { + '$not': { + '$and': { + 'text': {'$eq': 'Text of Document 1'}, + 'text_doc': {'$eq': 'Text Doc of Document 2'}, + } + } + } + ) + assert len(result) == 3 + + +@pytest.mark.parametrize('dict_api', [True, False]) +def test_from_docstring(dict_api): + class MyDocument(BaseDoc): + caption: TextDoc + image: ImageDoc + price: int + + docs = DocList[MyDocument]( + [ + MyDocument( + caption='A tiger in the jungle', + image=ImageDoc(url='tigerphoto.png'), + price=100, + ), + MyDocument( + caption='A swimming turtle', + image=ImageDoc(url='turtlepic.png'), + price=50, + ), + MyDocument( + caption='A couple birdwatching with binoculars', + image=ImageDoc(url='binocularsphoto.png'), + price=30, + ), + ] + ) + + query = { + '$and': { + 'image__url': {'$regex': 'photo'}, + 'price': {'$lte': 50}, + } + } + + if dict_api: + method = lambda query: filter_docs(docs, query) # noqa: E731 + else: + method = lambda query: filter_docs(docs, json.dumps(query)) # noqa: E731 + + results = method(query) + assert len(results) == 1 + assert results[0].price == 30 + assert results[0].caption == 'A couple birdwatching with binoculars' + assert results[0].image.url == 'binocularsphoto.png' diff --git a/tests/units/util/test_find.py b/tests/units/util/test_find.py new file mode 100644 index 00000000000..ca7cbe7160a --- /dev/null +++ b/tests/units/util/test_find.py @@ -0,0 +1,413 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Optional, Union + +import numpy as np +import pytest +import torch + +from docarray import BaseDoc, DocList +from docarray.typing import NdArray, TorchTensor +from docarray.utils.find import find, find_batched + + +class TorchDoc(BaseDoc): + tensor: TorchTensor + + +class NdDoc(BaseDoc): + tensor: NdArray + + +@pytest.fixture() +def random_torch_query(): + return TorchDoc(tensor=torch.rand(128)) + + +@pytest.fixture() +def random_torch_batch_query(): + return DocList[TorchDoc]([TorchDoc(tensor=torch.rand(128)) for _ in range(5)]) + + +@pytest.fixture() +def random_nd_query(): + return NdDoc(tensor=np.random.rand(128)) + + +@pytest.fixture() +def random_nd_batch_query(): + return DocList[NdDoc]([NdDoc(tensor=np.random.rand(128)) for _ in range(5)]) + + +@pytest.fixture() +def random_torch_index(): + return DocList[TorchDoc](TorchDoc(tensor=torch.rand(128)) for _ in range(10)) + + +@pytest.fixture() +def random_nd_index(): + return DocList[NdDoc](NdDoc(tensor=np.random.rand(128)) for _ in range(10)) + + +@pytest.mark.parametrize('metric', ['cosine_sim', 'euclidean_dist', 'sqeuclidean_dist']) +def test_find_torch(random_torch_query, random_torch_index, metric): + top_k, scores = find( + random_torch_index, + random_torch_query, + search_field='tensor', + limit=7, + metric=metric, + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert top_k.doc_type == random_torch_index.doc_type + + if metric.endswith('_dist'): + assert (torch.stack(sorted(scores)) == scores).all() + else: + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() + + +def test_find_torch_tensor_query(random_torch_query, random_torch_index): + query = random_torch_query.tensor + top_k, scores = find( + random_torch_index, + query, + search_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() + + +def test_find_torch_stacked(random_torch_query, random_torch_index): + random_torch_index = random_torch_index.to_doc_vec() + top_k, scores = find( + random_torch_index, + random_torch_query, + search_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() + + +@pytest.mark.parametrize('metric', ['cosine_sim', 'euclidean_dist', 'sqeuclidean_dist']) +def test_find_np(random_nd_query, random_nd_index, metric): + top_k, scores = find( + random_nd_index, + random_nd_query, + search_field='tensor', + limit=7, + metric=metric, + ) + assert len(top_k) == 7 + assert len(scores) == 7 + if metric.endswith('_dist'): + assert (sorted(scores) == scores).all() + else: + assert (sorted(scores, reverse=True) == scores).all() + + +def test_find_np_tensor_query(random_nd_query, random_nd_index): + query = random_nd_query.tensor + top_k, scores = find( + random_nd_index, + query, + search_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (sorted(scores, reverse=True) == scores).all() + + +def test_find_np_stacked(random_nd_query, random_nd_index): + random_nd_index = random_nd_index.to_doc_vec() + top_k, scores = find( + random_nd_index, + random_nd_query, + search_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (sorted(scores, reverse=True) == scores).all() + + +@pytest.mark.parametrize('metric', ['cosine_sim', 'euclidean_dist', 'sqeuclidean_dist']) +def test_find_batched_torch(random_torch_batch_query, random_torch_index, metric): + documents, scores = find_batched( + random_torch_index, + random_torch_batch_query, + search_field='tensor', + limit=7, + metric=metric, + ) + assert len(documents) == len(random_torch_batch_query) + assert len(scores) == len(random_torch_batch_query) + for top_k, top_scores in zip(documents, scores): + assert len(top_k) == 7 + assert len(top_scores) == 7 + assert top_k.doc_type == random_torch_index.doc_type + + for sc in scores: + if metric.endswith('_dist'): + assert (torch.stack(sorted(sc)) == sc).all() + else: + assert (torch.stack(sorted(sc, reverse=True)) == sc).all() + + +def test_find_batched_torch_tensor_query(random_torch_batch_query, random_torch_index): + query = torch.stack(random_torch_batch_query.tensor) + documents, scores = find_batched( + random_torch_index, + query, + search_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(documents) == len(random_torch_batch_query) + assert len(scores) == len(random_torch_batch_query) + for top_k, top_scores in zip(documents, scores): + assert len(top_k) == 7 + assert len(top_scores) == 7 + for sc in scores: + assert (torch.stack(sorted(sc, reverse=True)) == sc).all() + + +@pytest.mark.parametrize('stack_what', ['index', 'query', 'both']) +def test_find_batched_torch_stacked( + random_torch_batch_query, random_torch_index, stack_what +): + if stack_what in ('index', 'both'): + random_torch_index = random_torch_index.to_doc_vec() + if stack_what in ('query', 'both'): + random_torch_batch_query = random_torch_batch_query.to_doc_vec() + + documents, scores = find_batched( + random_torch_index, + random_torch_batch_query, + search_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(documents) == len(random_torch_batch_query) + assert len(scores) == len(random_torch_batch_query) + for top_k, top_scores in zip(documents, scores): + assert len(top_k) == 7 + assert len(top_scores) == 7 + for sc in scores: + assert (torch.stack(sorted(sc, reverse=True)) == sc).all() + + +@pytest.mark.parametrize('metric', ['cosine_sim', 'euclidean_dist', 'sqeuclidean_dist']) +def test_find_batched_np(random_nd_batch_query, random_nd_index, metric): + documents, scores = find_batched( + random_nd_index, + random_nd_batch_query, + search_field='tensor', + limit=7, + metric=metric, + ) + assert len(documents) == len(random_nd_batch_query) + assert len(scores) == len(random_nd_batch_query) + for top_k, top_scores in zip(documents, scores): + assert len(top_k) == 7 + assert len(top_scores) == 7 + for sc in scores: + if metric.endswith('_dist'): + assert (sorted(sc) == sc).all() + else: + assert (sorted(sc, reverse=True) == sc).all() + + +def test_find_batched_np_tensor_query(random_nd_batch_query, random_nd_index): + query = np.stack(random_nd_batch_query.tensor) + documents, scores = find_batched( + random_nd_index, + query, + search_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(documents) == len(random_nd_batch_query) + assert len(scores) == len(random_nd_batch_query) + for top_k, top_scores in zip(documents, scores): + assert len(top_k) == 7 + assert len(top_scores) == 7 + for sc in scores: + assert (sorted(sc, reverse=True) == sc).all() + + +@pytest.mark.parametrize('stack_what', ['index', 'query', 'both']) +def test_find_batched_np_stacked(random_nd_batch_query, random_nd_index, stack_what): + if stack_what in ('index', 'both'): + random_nd_index = random_nd_index.to_doc_vec() + if stack_what in ('query', 'both'): + random_nd_batch_query = random_nd_batch_query.to_doc_vec() + documents, scores = find_batched( + random_nd_index, + random_nd_batch_query, + search_field='tensor', + limit=7, + metric='cosine_sim', + ) + assert len(documents) == len(random_nd_batch_query) + assert len(scores) == len(random_nd_batch_query) + for top_k, top_scores in zip(documents, scores): + assert len(top_k) == 7 + assert len(top_scores) == 7 + for sc in scores: + assert (sorted(sc, reverse=True) == sc).all() + + +def test_find_optional(): + class MyDoc(BaseDoc): + embedding: Optional[TorchTensor] + + query = MyDoc(embedding=torch.rand(10)) + index = DocList[MyDoc]([MyDoc(embedding=torch.rand(10)) for _ in range(10)]) + + top_k, scores = find( + index, + query, + search_field='embedding', + limit=7, + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() + + +def test_find_union(): + class MyDoc(BaseDoc): + embedding: Union[TorchTensor, NdArray] + + query = MyDoc(embedding=torch.rand(10)) + index = DocList[MyDoc]([MyDoc(embedding=torch.rand(10)) for _ in range(10)]) + + top_k, scores = find( + index, + query, + search_field='embedding', + limit=7.0, + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() + + +@pytest.mark.parametrize('stack', [False, True]) +def test_find_nested(stack): + class InnerDoc(BaseDoc): + title: str + embedding: TorchTensor + + class MyDoc(BaseDoc): + inner: InnerDoc + + query = MyDoc(inner=InnerDoc(title='query', embedding=torch.rand(2))) + index = DocList[MyDoc]( + [ + MyDoc(inner=InnerDoc(title=f'doc {i}', embedding=torch.rand(2))) + for i in range(10) + ] + ) + if stack: + index = index.to_doc_vec() + + top_k, scores = find( + index, + query, + search_field='inner__embedding', + limit=7, + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() + + +def test_find_nested_union_optional(): + class MyDoc(BaseDoc): + embedding: Union[Optional[TorchTensor], Optional[NdArray]] + embedding2: Optional[Union[TorchTensor, NdArray]] + embedding3: Optional[Optional[TorchTensor]] + embedding4: Union[Optional[Union[TorchTensor, NdArray]], TorchTensor] + + query = MyDoc( + embedding=torch.rand(10), + embedding2=torch.rand(10), + embedding3=torch.rand(10), + embedding4=torch.rand(10), + ) + index = DocList[MyDoc]( + [ + MyDoc( + embedding=torch.rand(10), + embedding2=torch.rand(10), + embedding3=torch.rand(10), + embedding4=torch.rand(10), + ) + for _ in range(10) + ] + ) + + top_k, scores = find( + index, + query, + search_field='embedding', + limit=7, + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() + + top_k, scores = find( + index, + query, + search_field='embedding2', + limit=7.0, + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() + + top_k, scores = find( + index, + query, + search_field='embedding3', + limit=7, + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() + + top_k, scores = find( + index, + query, + search_field='embedding4', + limit=7, + ) + assert len(top_k) == 7 + assert len(scores) == 7 + assert (torch.stack(sorted(scores, reverse=True)) == scores).all() diff --git a/tests/units/util/test_map.py b/tests/units/util/test_map.py new file mode 100644 index 00000000000..65dd3c17389 --- /dev/null +++ b/tests/units/util/test_map.py @@ -0,0 +1,101 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Generator, Optional + +import pytest + +from docarray import BaseDoc, DocList +from docarray.documents import ImageDoc +from docarray.typing import ImageUrl, NdArray +from docarray.utils.map import map_docs, map_docs_batched +from tests.units.typing.test_bytes import IMAGE_PATHS + +N_DOCS = 2 + + +def load_from_doc(d: ImageDoc) -> ImageDoc: + if d.url is not None: + d.tensor = d.url.load() + return d + + +@pytest.fixture() +def da(): + da = DocList[ImageDoc]([ImageDoc(url=IMAGE_PATHS['png']) for _ in range(N_DOCS)]) + return da + + +@pytest.mark.parametrize('backend', ['thread', 'process']) +def test_map(da, backend): + for tensor in da.tensor: + assert tensor is None + + docs = list(map_docs(docs=da, func=load_from_doc, backend=backend)) + + assert len(docs) == N_DOCS + for doc in docs: + assert doc.tensor is not None + + +def test_map_multiprocessing_lambda_func_raise_exception(da): + with pytest.raises(ValueError, match='Multiprocessing does not allow'): + list(map_docs(docs=da, func=lambda x: x, backend='process')) + + +def test_map_multiprocessing_local_func_raise_exception(da): + def local_func(x): + return x + + with pytest.raises(ValueError, match='Multiprocessing does not allow'): + list(map_docs(docs=da, func=local_func, backend='process')) + + +@pytest.mark.parametrize('backend', ['thread', 'process']) +def test_check_order(backend): + da = DocList[ImageDoc]([ImageDoc(id=str(i)) for i in range(N_DOCS)]) + + docs = list(map_docs(docs=da, func=load_from_doc, backend=backend)) + + assert len(docs) == N_DOCS + for i, doc in enumerate(docs): + assert doc.id == str(i) + + +def load_from_da(da: DocList) -> DocList: + for doc in da: + doc.tensor = doc.url.load() + return da + + +class MyImage(BaseDoc): + tensor: Optional[NdArray] = None + url: ImageUrl + + +@pytest.mark.slow +@pytest.mark.parametrize('n_docs,batch_size', [(10, 5), (10, 8)]) +@pytest.mark.parametrize('backend', ['thread', 'process']) +def test_map_docs_batched(n_docs, batch_size, backend): + da = DocList[MyImage]([MyImage(url=IMAGE_PATHS['png']) for _ in range(n_docs)]) + it = map_docs_batched( + docs=da, func=load_from_da, batch_size=batch_size, backend=backend + ) + assert isinstance(it, Generator) + + for batch in it: + assert isinstance(batch, DocList) + for d in batch: + assert isinstance(d, MyImage) diff --git a/tests/units/util/test_reduce.py b/tests/units/util/test_reduce.py new file mode 100644 index 00000000000..816796831fc --- /dev/null +++ b/tests/units/util/test_reduce.py @@ -0,0 +1,139 @@ +from typing import Dict, List, Optional, Set + +import pytest + +from docarray import BaseDoc, DocList +from docarray.documents import ImageDoc +from docarray.utils.reduce import reduce, reduce_all + + +class InnerDoc(BaseDoc): + integer: int + inner_list: List + + +class MMDoc(BaseDoc): + text: str = '' + price: int = 0 + categories: Optional[List[str]] = None + image: Optional[ImageDoc] = None + matches: Optional[DocList] = None + matches_with_same_id: Optional[DocList] = None + opt_int: Optional[int] = None + test_set: Optional[Set] = None + inner_doc: Optional[InnerDoc] = None + test_dict: Optional[Dict] = None + + +@pytest.fixture +def doc1(): + return MMDoc( + text='hey here', + categories=['a', 'b', 'c'], + price=10, + matches=DocList[MMDoc]([MMDoc()]), + matches_with_same_id=DocList[MMDoc]( + [MMDoc(id='a', matches=DocList[MMDoc]([MMDoc()]))] + ), + test_set={'a', 'a'}, + inner_doc=InnerDoc(integer=2, inner_list=['c', 'd']), + test_dict={'a': 0, 'b': 2, 'd': 4, 'z': 3}, + ) + + +@pytest.fixture +def doc2(doc1): + return MMDoc( + id=doc1.id, + text='hey here 2', + categories=['d', 'e', 'f'], + price=5, + opt_int=5, + matches=DocList[MMDoc]([MMDoc()]), + matches_with_same_id=DocList[MMDoc]( + [MMDoc(id='a', matches=DocList[MMDoc]([MMDoc()]))] + ), + test_set={'a', 'b'}, + inner_doc=InnerDoc(integer=3, inner_list=['a', 'b']), + test_dict={'a': 10, 'b': 10, 'c': 3, 'z': None}, + ) + + +def test_reduce_different_ids(): + da1 = DocList[MMDoc]([MMDoc() for _ in range(10)]) + da2 = DocList[MMDoc]([MMDoc() for _ in range(10)]) + result = reduce(da1, da2) + assert len(result) == 20 + # da1 is changed in place (no extra memory) + assert len(da1) == 20 + + +def test_reduce(doc1, doc2): + da1 = DocList[MMDoc]([doc1, MMDoc()]) + da2 = DocList[MMDoc]([MMDoc(), doc2]) + result = reduce(da1, da2) + assert len(result) == 3 + # da1 is changed in place (no extra memory) + assert len(da1) == 3 + merged_doc = result[0] + assert merged_doc.text == 'hey here 2' + assert merged_doc.categories == ['a', 'b', 'c', 'd', 'e', 'f'] + assert len(merged_doc.matches) == 2 + assert merged_doc.opt_int == 5 + assert merged_doc.price == 5 + assert merged_doc.test_set == {'a', 'b'} + assert len(merged_doc.matches_with_same_id) == 1 + assert len(merged_doc.matches_with_same_id[0].matches) == 2 + assert merged_doc.inner_doc.integer == 3 + assert merged_doc.inner_doc.inner_list == ['c', 'd', 'a', 'b'] + + +def test_reduce_all(doc1, doc2): + da1 = DocList[MMDoc]([doc1, MMDoc()]) + da2 = DocList[MMDoc]([MMDoc(), doc2]) + da3 = DocList[MMDoc]([MMDoc(), MMDoc(), doc1]) + result = reduce_all([da1, da2, da3]) + assert len(result) == 5 + # da1 is changed in place (no extra memory) + assert len(da1) == 5 + merged_doc = result[0] + assert merged_doc.text == 'hey here 2' + assert merged_doc.categories == [ + 'a', + 'b', + 'c', + 'd', + 'e', + 'f', + 'a', + 'b', + 'c', + 'd', + 'e', + 'f', + ] + assert len(merged_doc.matches) == 2 + assert merged_doc.opt_int == 5 + assert merged_doc.price == 5 + assert merged_doc.test_set == {'a', 'b'} + assert len(merged_doc.matches_with_same_id) == 1 + assert len(merged_doc.matches_with_same_id[0].matches) == 2 + assert merged_doc.inner_doc.integer == 3 + assert merged_doc.inner_doc.inner_list == ['c', 'd', 'a', 'b', 'c', 'd', 'a', 'b'] + + +def test_update_ndarray(): + from docarray.typing import NdArray + + import numpy as np + + class MyDoc(BaseDoc): + embedding: NdArray[128] + + embedding1 = np.random.rand(128) + embedding2 = np.random.rand(128) + + doc1 = MyDoc(id='0', embedding=embedding1) + doc2 = MyDoc(id='0', embedding=embedding2) + doc1.update(doc2) + assert (doc1.embedding == embedding2).all() diff --git a/tests/units/util/test_typing.py b/tests/units/util/test_typing.py new file mode 100644 index 00000000000..f40fde4ab21 --- /dev/null +++ b/tests/units/util/test_typing.py @@ -0,0 +1,108 @@ +# Licensed to the LF AI & Data foundation under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Dict, List, Optional, Set, Tuple, Union + +import pytest + +from docarray.typing import NdArray, TorchTensor +from docarray.typing.tensor.abstract_tensor import AbstractTensor +from docarray.utils._internal._typing import ( + is_tensor_union, + is_type_tensor, + safe_issubclass, +) +from docarray.utils._internal.misc import is_tf_available + +tf_available = is_tf_available() +if tf_available: + from docarray.typing import TensorFlowTensor +else: + TensorFlowTensor = None + + +@pytest.mark.parametrize( + 'type_, is_tensor', + [ + (int, False), + (TorchTensor, True), + (NdArray, True), + (AbstractTensor, True), + (Optional[TorchTensor], False), + (Union[TorchTensor, NdArray], False), + (None, False), + (Dict, False), + ], +) +def test_is_type_tensor(type_, is_tensor): + assert is_type_tensor(type_) == is_tensor + + +@pytest.mark.tensorflow +@pytest.mark.parametrize( + 'type_, is_tensor', + [ + (TensorFlowTensor, True), + (Optional[TensorFlowTensor], False), + ], +) +def test_is_type_tensor_with_tf(type_, is_tensor): + assert is_type_tensor(type_) == is_tensor + + +@pytest.mark.parametrize( + 'type_, is_union_tensor', + [ + (int, False), + (TorchTensor, False), + (NdArray, False), + (Optional[TorchTensor], True), + (Optional[NdArray], True), + (Union[NdArray, TorchTensor], True), + (Union[NdArray, TorchTensor, AbstractTensor], True), + (Union[NdArray, TorchTensor, Optional[TorchTensor]], True), + (Union[NdArray, TorchTensor, None], True), + ], +) +def test_is_union_type_tensor(type_, is_union_tensor): + assert is_tensor_union(type_) == is_union_tensor + + +@pytest.mark.tensorflow +@pytest.mark.parametrize( + 'type_, is_union_tensor', + [ + (TensorFlowTensor, False), + (Optional[TensorFlowTensor], True), + (Union[NdArray, TorchTensor, TensorFlowTensor], True), + (Union[NdArray, TorchTensor, Optional[TensorFlowTensor]], True), + ], +) +def test_is_union_type_tensor_with_tf(type_, is_union_tensor): + assert is_tensor_union(type_) == is_union_tensor + + +@pytest.mark.parametrize( + 'type_, cls, is_subclass', + [ + (List[str], object, False), + (List[List[int]], object, False), + (Set[str], object, False), + (Dict, object, False), + (Tuple[int, int], object, False), + ], +) +def test_safe_issubclass(type_, cls, is_subclass): + assert safe_issubclass(type_, cls) == is_subclass