INTPYTHON-792 Add support for window functions

Author: timgraham

.github/workflows/runtests.py

@@ -47,6 +47,7 @@
     "empty_models",
     "expressions",
     "expressions_case",
+    "expressions_window",
     "field_defaults",
     "file_storage",
     "file_uploads",

django_mongodb_backend/compiler.py

@@ -24,6 +24,7 @@
     Subquery,
     Value,
     When,
+    Window,
 )
 from django.db.models.sql import Query
 
@@ -157,6 +158,16 @@ def query(self, compiler, connection, get_wrapping_pipeline=None, as_expr=False)
         else:
             if subquery.aggregation_pipeline is None:
                 subquery.aggregation_pipeline = []
+            # Window stages must precede the wrapping pipeline so the $group
+            # can reference the window function outputs.
+            if subquery.window_pipeline:
+                subquery.aggregation_pipeline.extend(subquery.window_pipeline)
+                subquery.window_pipeline = None
+            # Qualify ($match on window results) must come after window stages
+            # but before the wrapping pipeline's $group stage.
+            if subquery.qualify_mql:
+                subquery.aggregation_pipeline.append({"$match": subquery.qualify_mql})
+                subquery.qualify_mql = None
             subquery.aggregation_pipeline.extend(wrapping_result_pipeline)
         # Erase project_fields since the required value is projected above.
         subquery.project_fields = None
@@ -209,6 +220,10 @@ def when(self, compiler, connection):
     return self.condition.as_mql(compiler, connection, as_expr=True)
 
 
+def window(self, compiler, connection):  # noqa: ARG001
+    raise NotSupportedError("Window expressions must be used as annotations.")
+
+
 def value(self, compiler, connection, as_expr=False):  # noqa: ARG001
     value = self.value
     if isinstance(value, (list, int, str, dict, tuple)) and as_expr:
@@ -258,3 +273,4 @@ def register_expressions():
     Subquery.as_mql_path = partialmethod(subquery, as_expr=False)
     When.as_mql_expr = when
     Value.as_mql = value
+    Window.as_mql_expr = window

django_mongodb_backend/expressions/builtins.py

@@ -33,11 +33,13 @@ class DatabaseFeatures(GISFeatures, BaseDatabaseFeatures):
     supports_expression_defaults = False
     supports_expression_indexes = False
     supports_foreign_keys = False
+    supports_frame_range_fixed_distance = True
     supports_ignore_conflicts = False
     supports_json_field_contains = False
     # BSON Date type doesn't support microsecond precision.
     supports_microsecond_precision = False
     supports_nulls_distinct_unique_constraints = True
+    supports_over_clause = True
     supports_paramstyle_pyformat = False
     supports_sequence_reset = False
     supports_slicing_ordering_in_compound = True
@@ -527,6 +529,10 @@ def django_test_skips(self):
             "prefetch_related.tests.LookupOrderingTest.test_order",
             "prefetch_related.tests.MultiDbTests.test_using_is_honored_m2m",
             "prefetch_related.tests.MultiTableInheritanceTest.test_m2m_to_inheriting_model",
+            "prefetch_related.tests.PrefetchLimitTests.test_empty_order",
+            "prefetch_related.tests.PrefetchLimitTests.test_m2m_forward",
+            "prefetch_related.tests.PrefetchLimitTests.test_m2m_reverse",
+            "prefetch_related.tests.PrefetchLimitTests.test_reverse_ordering",
             "prefetch_related.tests.PrefetchRelatedMTICacheTests.test_add_clears_prefetched_objects_in_grandparent",
             "prefetch_related.tests.PrefetchRelatedMTICacheTests.test_add_clears_prefetched_objects_in_parent",
             "prefetch_related.tests.PrefetchRelatedMTICacheTests.test_grandparent_m2m_available_in_child",

django_mongodb_backend/features.py

@@ -4,6 +4,7 @@
 from .json import register_json
 from .math import register_math
 from .text import register_text
+from .window import register_window
 
 
 def register_functions():
@@ -13,3 +14,4 @@ def register_functions():
     register_json()
     register_math()
     register_text()
+    register_window()

django_mongodb_backend/functions/__init__.py

@@ -0,0 +1,240 @@
+from django.db.models.aggregates import Aggregate, Count, StdDev, Variance
+from django.db.models.functions.window import (
+    CumeDist,
+    DenseRank,
+    FirstValue,
+    Lag,
+    LastValue,
+    Lead,
+    NthValue,
+    Ntile,
+    PercentRank,
+    Rank,
+    RowNumber,
+)
+
+
+def aggregate(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    agg_mql = self.as_mql_expr(compiler, connection)
+    return {alias: {**agg_mql, "window": default_frame()}}, {}
+
+
+def count(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    # Count.as_mql_expr() returns {"$sum": lhs_mql} for non-distinct counts.
+    # MongoDB's $count is a pipeline stage (not a window accumulator), so use
+    # $sum instead.
+    agg_mql = self.as_mql_expr(compiler, connection)
+    return {alias: {**agg_mql, "window": default_frame()}}, {}
+
+
+def cume_dist(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    # CumeDist = (rank + group_size - 1) / total. Use $rank for the base rank
+    # (1-based, same value for ties), range:[0, 0] for group_size (count of all
+    # rows with identical sort key value), and full-partition sum for total.
+    rank_alias = f"__wtemp{next(idx)}"
+    group_size_alias = f"__wtemp{next(idx)}"
+    total_alias = f"__wtemp{next(idx)}"
+    output = {
+        rank_alias: {"$rank": {}},
+        group_size_alias: {"$sum": 1, "window": {"range": [0, 0]}},
+        total_alias: {"$sum": 1, "window": {"documents": ["unbounded", "unbounded"]}},
+    }
+    add_fields = {
+        alias: {
+            "$divide": [
+                {"$subtract": [{"$add": [f"${rank_alias}", f"${group_size_alias}"]}, 1]},
+                f"${total_alias}",
+            ]
+        }
+    }
+    return output, add_fields
+
+
+def dense_rank(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    return {alias: {"$denseRank": {}}}, {}
+
+
+def first_value(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    expr = self.get_source_expressions()[0]
+    mql = expr.as_mql(compiler, connection, as_expr=True)
+    return {alias: {"$first": mql, "window": default_frame()}}, {}
+
+
+def lag(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    exprs = self.get_source_expressions()
+    expr_mql = exprs[0].as_mql(compiler, connection, as_expr=True)
+    offset = exprs[1].value
+    mql = {"output": expr_mql, "by": -offset}
+    if len(exprs) > 2:
+        mql["default"] = exprs[2].as_mql(compiler, connection, as_expr=True)
+    return {alias: {"$shift": mql}}, {}
+
+
+def last_value(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    expr = self.get_source_expressions()[0]
+    mql = expr.as_mql(compiler, connection, as_expr=True)
+    return {alias: {"$last": mql, "window": default_frame()}}, {}
+
+
+def lead(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    exprs = self.get_source_expressions()
+    expr_mql = exprs[0].as_mql(compiler, connection, as_expr=True)
+    offset = exprs[1].value
+    mql = {"output": expr_mql, "by": offset}
+    if len(exprs) > 2:
+        mql["default"] = exprs[2].as_mql(compiler, connection, as_expr=True)
+    return {alias: {"$shift": mql}}, {}
+
+
+def nth_value(self, compiler, connection, alias, idx, default_frame):
+    expr, nth_expr = self.get_source_expressions()
+    mql = expr.as_mql(compiler, connection, as_expr=True)
+    nth = nth_expr.value
+    push_alias = f"__wtemp{next(idx)}"
+    output = {push_alias: {"$push": mql, "window": default_frame()}}
+    add_fields = {
+        alias: {
+            "$cond": {
+                "if": {"$gte": [{"$size": f"${push_alias}"}, nth]},
+                "then": {"$arrayElemAt": [f"${push_alias}", nth - 1]},
+                "else": None,
+            }
+        }
+    }
+    return output, add_fields
+
+
+def ntile(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    num_buckets = self.get_source_expressions()[0].value
+    doc_num_alias = f"__wtemp{next(idx)}"
+    total_alias = f"__wtemp{next(idx)}"
+    output = {
+        doc_num_alias: {"$documentNumber": {}},
+        total_alias: {"$sum": 1, "window": {"documents": ["unbounded", "unbounded"]}},
+    }
+    # SQL NTILE: rows 1..extra*(per+1) go to the first `extra` buckets (each
+    # gets per+1 rows); remaining rows are distributed evenly across the rest.
+    add_fields = {
+        alias: {
+            "$let": {
+                "vars": {
+                    "per": {"$floor": {"$divide": [f"${total_alias}", num_buckets]}},
+                    "xtra": {"$mod": [f"${total_alias}", num_buckets]},
+                },
+                "in": {
+                    "$let": {
+                        "vars": {
+                            "thr": {"$multiply": ["$$xtra", {"$add": ["$$per", 1]}]},
+                        },
+                        "in": {
+                            "$cond": {
+                                "if": {"$lte": [f"${doc_num_alias}", "$$thr"]},
+                                "then": {
+                                    "$add": [
+                                        {
+                                            "$floor": {
+                                                "$divide": [
+                                                    {"$subtract": [f"${doc_num_alias}", 1]},
+                                                    {"$add": ["$$per", 1]},
+                                                ]
+                                            }
+                                        },
+                                        1,
+                                    ]
+                                },
+                                "else": {
+                                    "$add": [
+                                        "$$xtra",
+                                        {
+                                            "$floor": {
+                                                "$divide": [
+                                                    {
+                                                        "$subtract": [
+                                                            f"${doc_num_alias}",
+                                                            {"$add": ["$$thr", 1]},
+                                                        ]
+                                                    },
+                                                    # Guard against division by
+                                                    # zero when num_buckets >
+                                                    # total (per==0).
+                                                    {"$max": ["$$per", 1]},
+                                                ]
+                                            }
+                                        },
+                                        1,
+                                    ]
+                                },
+                            }
+                        },
+                    }
+                },
+            }
+        }
+    }
+    return output, add_fields
+
+
+def percent_rank(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    # Use document-position rank ($sum:1 docs unbounded-to-current) instead of
+    # $rank, which MongoDB limits to a single sort field.
+    row_num_alias = f"__wtemp{next(idx)}"
+    count_alias = f"__wtemp{next(idx)}"
+    output = {
+        row_num_alias: {"$sum": 1, "window": {"documents": ["unbounded", "current"]}},
+        count_alias: {"$sum": 1, "window": {"documents": ["unbounded", "unbounded"]}},
+    }
+    add_fields = {
+        alias: {
+            "$cond": {
+                "if": {"$lte": [f"${count_alias}", 1]},
+                "then": {"$literal": 0.0},
+                "else": {
+                    "$divide": [
+                        {"$subtract": [f"${row_num_alias}", 1]},
+                        {"$subtract": [f"${count_alias}", 1]},
+                    ]
+                },
+            }
+        }
+    }
+    return output, add_fields
+
+
+def rank(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    return {alias: {"$rank": {}}}, {}
+
+
+def row_number(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    return {alias: {"$documentNumber": {}}}, {}
+
+
+def stddev(self, compiler, connection, alias, idx, default_frame):  # noqa: ARG001
+    agg_mql = self.as_mql_expr(compiler, connection)
+    return {alias: {**agg_mql, "window": default_frame()}}, {}
+
+
+def variance(self, compiler, connection, alias, idx, default_frame):
+    # MongoDB has no $varPop/$varSamp window accumulator; compute as stdDev².
+    agg_mql = self.as_mql_expr(compiler, connection)
+    stddev_alias = f"__wtemp{next(idx)}"
+    output = {stddev_alias: {**agg_mql, "window": default_frame()}}
+    add_fields = {alias: {"$pow": [f"${stddev_alias}", 2]}}
+    return output, add_fields
+
+
+def register_window():
+    Aggregate.get_window_mql = aggregate
+    Count.get_window_mql = count
+    CumeDist.get_window_mql = cume_dist
+    DenseRank.get_window_mql = dense_rank
+    FirstValue.get_window_mql = first_value
+    Lag.get_window_mql = lag
+    LastValue.get_window_mql = last_value
+    Lead.get_window_mql = lead
+    NthValue.get_window_mql = nth_value
+    Ntile.get_window_mql = ntile
+    PercentRank.get_window_mql = percent_rank
+    Rank.get_window_mql = rank
+    RowNumber.get_window_mql = row_number
+    StdDev.get_window_mql = stddev
+    Variance.get_window_mql = variance

django_mongodb_backend/functions/window.py

@@ -52,6 +52,8 @@ def __init__(self, compiler):
         self.project_fields = None
         self.aggregation_pipeline = compiler.aggregation_pipeline
         self.search_pipeline = compiler.search_pipeline
+        self.window_pipeline = None
+        self.qualify_mql = None
         self.extra_fields = None
         self.combinator_pipeline = None
         # $lookup stage that encapsulates the pipeline for performing a nested
@@ -93,6 +95,10 @@ def get_pipeline(self):
             pipeline.append({"$match": self.match_mql})
         if self.aggregation_pipeline:
             pipeline.extend(self.aggregation_pipeline)
+        if self.window_pipeline:
+            pipeline.extend(self.window_pipeline)
+        if self.qualify_mql:
+            pipeline.append({"$match": self.qualify_mql})
         if self.needs_wrap_aggregation:
             if self.compiler.connection.client_encryption:
                 # Automatic encryption doesn't support $unionWith, so use

django_mongodb_backend/query.py

@@ -10,6 +10,8 @@ Django MongoDB Backend 6.0.x
 New features
 ------------
 
+- Added support for :ref:`window functions <django:window-functions>`.
+
 - Added support for the following database functions:
 
   - :class:`~django.db.models.functions.ExtractQuarter`

docs/releases/6.0.x.rst

@@ -0,0 +1,5 @@
+from django.db import models
+
+
+class Number(models.Model):
+    value = models.IntegerField()

tests/expressions_/models.py

@@ -0,0 +1,13 @@
+from django.db import NotSupportedError
+from django.db.models.expressions import Window
+from django.db.models.functions.window import Rank
+from django.test import TestCase
+
+from .models import Number
+
+
+class WindowTests(TestCase):
+    def test_window_order_by(self):
+        msg = "Window expressions must be used as annotations."
+        with self.assertRaisesMessage(NotSupportedError, msg):
+            list(Number.objects.order_by(Window(Rank())))