add docstr for filter function (#177)

numb3r3 · web-flow · commit a045b376f038 · 2022-03-08T09:45:20.000+01:00
* docs(filter): init docstr

* fix: add regex filter test
diff --git a/docarray/array/mixins/find.py b/docarray/array/mixins/find.py
@@ -167,6 +167,35 @@ def _filter(
         limit: Optional[int] = None,
         only_id: bool = False,
     ) -> 'DocumentArray':
+        """Returns a subset of documents by filtering by the given query.
+        The query language we provide now is following the
+        [MongoDB](https://docs.mongodb.com/manual/reference/operator/query/) query language. For example::
+
+            >>> docs._filter({'text': {'$eq': 'hello'}})
+
+            The above will return a `DocumentArray` in which each document has doc.text == 'hello'. And we also support
+            placeholder format by using the following syntax::
+
+            >>> docs._filter({'text': {'$eq': '{tags__name}'}})
+
+            will return a `DocumentArray` in which each document has doc.text == doc.tags['name'].
+
+        Now, only the subset of MongoDB's query operators are supported:
+            - `$eq` - Equal to (number, string)
+            - `$ne` - Not equal to (number, string)
+            - `$gt` - Greater than (number)
+            - `$gte` - Greater than or equal to (number)
+            - `$lt` - Less than (number)
+            - `$lte` - Less than or equal to (number)
+            - `$in` - Included in an array
+            - `$nin` - Not included in an array
+            - `$regex` - Match a specified regular expression
+
+        :param query: the input query dictionary.
+        :param limit: the maximum number of matches, when not given defaults to 20.
+        :param only_id: if set, then returning documents will only contain ``id``
+        :return: a `DocumentArray` containing the `Document` objects for matching with the query.
+        """
         from ... import DocumentArray
         from ..queryset import QueryParser
 
diff --git a/tests/unit/array/mixins/test_filter.py b/tests/unit/array/mixins/test_filter.py
@@ -22,7 +22,7 @@ def test_empty_filter(docs):
     assert len(result) == 5
 
 
-def test_sample_filter(docs):
+def test_simple_filter(docs):
     result = docs.find({'text': {'$eq': 'hello'}})
     assert len(result) == 1
     assert result[0].text == 'hello'
@@ -31,6 +31,14 @@ def test_sample_filter(docs):
     assert len(result) == 1
     assert result[0].tags['x'] == 0.8
 
+    result = docs.find({'tags__name': {'$regex': '^h'}})
+    assert len(result) == 2
+    assert result[1].id == docs[1].id
+
+    result = docs.find({'text': {'$regex': '^h'}})
+    assert len(result) == 1
+    assert result[0].id == docs[0].id
+
 
 def test_logic_filter(docs):
     result = docs.find({'$or': {'tags__x': {'$gte': 0.1}, 'tags__y': {'$gte': 0.5}}})
diff --git a/tests/unit/array/test_lookup.py b/tests/unit/array/test_lookup.py
@@ -31,6 +31,9 @@ def test_lookup_ops(doc):
     assert lookup('tags__x__lte', 0.1, doc)
     assert not lookup('tags__y__lt', 1.5, doc)
 
+    assert lookup('text__regex', '^test', doc)
+    assert not lookup('text__regex', '^est', doc)
+
 
 def test_lookup_pl(doc):
     from docarray.array.queryset.lookup import lookup
