Adding support for overlap

Author: WaVEV

django_mongodb_backend/fields/array.py

@@ -326,9 +326,7 @@ def get_subquery_wrapping_pipeline(self, compiler, connection, field_name, expr)
     def as_mql(self, compiler, connection):
         lhs_mql = process_lhs(self, compiler, connection)
         value = process_rhs(self, compiler, connection)
-        return {
-            "$and": [{"$ne": [lhs_mql, None]}, {"$size": {"$setIntersection": [value, lhs_mql]}}]
-        }
+        return {"$and": [{"$isArray": lhs_mql}, {"$size": {"$setIntersection": [value, lhs_mql]}}]}
 
 
 @ArrayField.register_lookup

django_mongodb_backend/fields/embedded_model_array.py

@@ -57,11 +57,20 @@ def get_transform(self, name):
         transform = super().get_transform(name)
         if transform:
             return transform
-        return KeyTransformFactory(name, self.base_field)
+        return KeyTransformFactory(name, self)
+
+
+class ProcessRHSMixin:
+    def process_rhs(self, compiler, connection):
+        if isinstance(self.lhs, KeyTransform):
+            get_db_prep_value = self.lhs._lhs.output_field.get_db_prep_value
+        else:
+            get_db_prep_value = self.lhs.output_field.get_db_prep_value
+        return None, [get_db_prep_value(v, connection, prepared=True) for v in self.rhs]
 
 
 @EmbeddedModelArrayField.register_lookup
-class EMFArrayExact(EMFExact):
+class EMFArrayExact(EMFExact, ProcessRHSMixin):
     def as_mql(self, compiler, connection):
         lhs_mql = process_lhs(self, compiler, connection)
         value = process_rhs(self, compiler, connection)
@@ -103,15 +112,61 @@ def as_mql(self, compiler, connection):
         }
 
 
+@EmbeddedModelArrayField.register_lookup
+class ArrayOverlap(EMFExact, ProcessRHSMixin):
+    lookup_name = "overlap"
+
+    def as_mql(self, compiler, connection):
+        lhs_mql = process_lhs(self, compiler, connection)
+        values = process_rhs(self, compiler, connection)
+        if isinstance(self.lhs, KeyTransform):
+            lhs_mql, inner_lhs_mql = lhs_mql
+            return {
+                "$anyElementTrue": {
+                    "$ifNull": [
+                        {
+                            "$map": {
+                                "input": lhs_mql,
+                                "as": "item",
+                                "in": {"$in": [inner_lhs_mql, values]},
+                            }
+                        },
+                        [],
+                    ]
+                }
+            }
+        conditions = []
+        inner_lhs_mql = "$$item"
+        for value in values:
+            if isinstance(value, models.Model):
+                value, emf_data = self.model_to_dict(value)
+                # Get conditions for any nested EmbeddedModelFields.
+                conditions.append({"$and": self.get_conditions({inner_lhs_mql: (value, emf_data)})})
+        return {
+            "$anyElementTrue": {
+                "$ifNull": [
+                    {
+                        "$map": {
+                            "input": lhs_mql,
+                            "as": "item",
+                            "in": {"$or": conditions},
+                        }
+                    },
+                    [],
+                ]
+            }
+        }
+
+
 class KeyTransform(Transform):
     # it should be different class than EMF keytransform even most of the methods are equal.
-    def __init__(self, key_name, base_field, *args, **kwargs):
+    def __init__(self, key_name, array_field, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.base_field = base_field
+        self.array_field = array_field
         self.key_name = key_name
         # The iteration items begins from the base_field, a virtual column with
         # base field output type is created.
-        column_target = base_field.clone()
+        column_target = array_field.base_field.embedded_model._meta.get_field(key_name).clone()
         column_name = f"$item.{key_name}"
         column_target.db_column = column_name
         column_target.set_attributes_from_name(column_name)
@@ -134,7 +189,7 @@ def _get_missing_field_or_lookup_exception(self, lhs, name):
             suggestion = "."
         raise FieldDoesNotExist(
             f"Unsupported lookup '{name}' for "
-            f"{self.base_field.__class__.__name__} '{self.base_field.name}'"
+            f"{self.array_field.base_field.__class__.__name__} '{self.array_field.base_field.name}'"
             f"{suggestion}"
         )
 
@@ -147,7 +202,9 @@ def get_transform(self, name):
         transform = (
             self._lhs.get_transform(name)
             if isinstance(self._lhs, Transform)
-            else self.base_field.embedded_model._meta.get_field(self.key_name).get_transform(name)
+            else self.array_field.base_field.embedded_model._meta.get_field(
+                self.key_name
+            ).get_transform(name)
         )
         if transform:
             self._sub_transform = transform
@@ -163,7 +220,7 @@ def as_mql(self, compiler, connection):
 
     @property
     def output_field(self):
-        return EmbeddedModelArrayField(self.base_field)
+        return self.array_field
 
 
 class KeyTransformFactory:

tests/model_fields_/test_embedded_model.py

@@ -282,6 +282,55 @@ def test_len(self):
             MuseumExhibit.objects.filter(sections__1__artifacts__len=1), [self.wonders]
         )
 
+    def test_overlap_simplefield(self):
+        self.assertSequenceEqual(
+            MuseumExhibit.objects.filter(sections__section_number__overlap=[10]), []
+        )
+        self.assertSequenceEqual(
+            MuseumExhibit.objects.filter(sections__section_number__overlap=[1]),
+            [self.egypt, self.wonders, self.new_descoveries],
+        )
+        self.assertSequenceEqual(
+            MuseumExhibit.objects.filter(sections__section_number__overlap=[2]), [self.wonders]
+        )
+
+    def test_overlap_emf(self):
+        self.assertSequenceEqual(
+            Movie.objects.filter(reviews__overlap=[Review(title="The best", rating=10)]),
+            [self.clouds],
+        )
+
+    """
+    def test_overlap_charfield_including_expression(self):
+        obj_1 = CharArrayModel.objects.create(field=["TEXT", "lower text"])
+        obj_2 = CharArrayModel.objects.create(field=["lower text", "TEXT"])
+        CharArrayModel.objects.create(field=["lower text", "text"])
+        self.assertSequenceEqual(
+            CharArrayModel.objects.filter(
+                field__overlap=[
+                    Upper(Value("text")),
+                    "other",
+                ]
+            ),
+            [obj_1, obj_2],
+        )
+
+    def test_overlap_values(self):
+        qs = NullableIntegerArrayModel.objects.filter(order__lt=3)
+        self.assertCountEqual(
+            NullableIntegerArrayModel.objects.filter(
+                field__overlap=qs.values_list("field"),
+            ),
+            self.objs[:3],
+        )
+        self.assertCountEqual(
+            NullableIntegerArrayModel.objects.filter(
+                field__overlap=qs.values("field"),
+            ),
+            self.objs[:3],
+        )
+    """
+
 
 class QueryingTests(TestCase):
     @classmethod