diff --git a/.github/workflows/runtests.py b/.github/workflows/runtests.py index 1ef0108c0..0f4a325a2 100755 --- a/.github/workflows/runtests.py +++ b/.github/workflows/runtests.py @@ -47,6 +47,7 @@ "empty_models", "expressions", "expressions_case", + "expressions_window", "field_defaults", "file_storage", "file_uploads", diff --git a/django_mongodb_backend/compiler.py b/django_mongodb_backend/compiler.py index 39aa1214d..4acce1969 100644 --- a/django_mongodb_backend/compiler.py +++ b/django_mongodb_backend/compiler.py @@ -1,4 +1,5 @@ import itertools +import json from collections import defaultdict from bson import SON, ObjectId, json_util @@ -6,7 +7,7 @@ from django.db import IntegrityError, NotSupportedError from django.db.models import Count from django.db.models.aggregates import Aggregate, Variance -from django.db.models.expressions import Case, Col, OrderBy, Ref, Value, When +from django.db.models.expressions import Case, Col, OrderBy, Ref, RowRange, Value, When, Window from django.db.models.functions.comparison import Coalesce from django.db.models.functions.math import Power from django.db.models.lookups import IsNull @@ -38,6 +39,8 @@ def __init__(self, *args, **kwargs): self.subqueries = [] # Search stage. self.search_pipeline = [] + # Window function pipeline stages ($setWindowFields + post-$addFields). + self.window_pipeline = None # Does the aggregation have no GROUP BY fields and need wrapping? self.needs_wrap_aggregation = False # The MQL equivalent to a SQL HAVING clause. @@ -212,6 +215,11 @@ def _get_group_expressions(self, order_by): having_group_by = self.having.get_group_by_cols() if self.having else () for expr in having_group_by: expressions.add(expr) + # Include partition/sort columns from window annotations so that + # $setWindowFields can partition the grouped data. + for expr in self.query.annotations.values(): + if expr.contains_over_clause: + expressions |= set(expr.get_group_by_cols()) return expressions def _get_group_id_expressions(self, order_by): @@ -308,7 +316,8 @@ def pre_sql_setup(self, with_col_aliases=False): extra_select, order_by, group_by = super().pre_sql_setup(with_col_aliases=with_col_aliases) search_replacements = self._prepare_search_query_for_aggregation_pipeline(order_by) group, group_replacements = self._prepare_annotations_for_aggregation_pipeline(order_by) - all_replacements = {**search_replacements, **group_replacements} + window_stages, window_replacements = self._prepare_window_annotations_for_pipeline() + all_replacements = {**search_replacements, **group_replacements, **window_replacements} self.search_pipeline = self._compound_searches_queries(search_replacements) # query.group_by is either: # - None: no GROUP BY @@ -329,6 +338,7 @@ def pre_sql_setup(self, with_col_aliases=False): self.subqueries = [] pipeline.append({"$match": having}) self.aggregation_pipeline = pipeline + self.window_pipeline = window_stages self.annotations = { target: expr.replace_expressions(all_replacements) for target, expr in self.query.annotation_select.items() @@ -336,6 +346,10 @@ def pre_sql_setup(self, with_col_aliases=False): self.order_by_objs = [expr.replace_expressions(all_replacements) for expr, _ in order_by] if (where := self.get_where()) and search_replacements: self.set_where(where.replace_expressions(search_replacements)) + # Apply window/annotation replacements to the qualify clause so that + # Window expressions are replaced by their pre-computed field Refs. + if getattr(self, "qualify", None) and window_replacements: + self.qualify = self.qualify.replace_expressions(all_replacements) return extra_select, order_by, group_by def get_project_columns(self, columns): @@ -493,6 +507,9 @@ def build_query(self, columns=None): if extra_fields: query.extra_fields = self.get_project_fields(extra_fields, force_expression=True) query.subqueries = self.subqueries + query.window_pipeline = self.window_pipeline + if getattr(self, "qualify", None): + query.qualify_mql = self.qualify.as_mql(self, self.connection) return query @cached_property @@ -800,6 +817,9 @@ def _get_aggregate_expressions(self, expr): def _get_search_expressions(self, expr): return self._get_all_expressions_of_type(expr, SearchExpression) + def _get_window_expressions(self, expr): + return self._get_all_expressions_of_type(expr, Window) + def _get_all_expressions_of_type(self, expr, target_type): stack = [expr] while stack: @@ -809,6 +829,248 @@ def _get_all_expressions_of_type(self, expr, target_type): elif hasattr(expr, "get_source_expressions"): stack.extend(expr.get_source_expressions()) + def _compile_window_partition_by(self, partition_by): + """Compile Window.partition_by to a MongoDB partitionBy expression.""" + if partition_by is None: + return None + exprs = partition_by.get_source_expressions() + if len(exprs) == 1: + return exprs[0].as_mql(self, self.connection, as_expr=True) + return {str(i): e.as_mql(self, self.connection, as_expr=True) for i, e in enumerate(exprs)} + + def _compile_window_sort_by(self, order_by_list, sort_fields): + """ + Compile Window.order_by to a MongoDB sortBy document and any pre-sort + $addFields. + + sort_fields is a shared dict {expr_mql_key: field_name} across all + windows in the same query, so that identical expressions reuse the same + field name and different expressions get distinct names. + """ + if order_by_list is None: + return {}, {} + exprs = order_by_list.get_source_expressions() + sort_doc = {} + pre_add_fields = {} + for item in exprs: + if isinstance(item, OrderBy): + expr = item.expression + direction = DESCENDING if item.descending else ASCENDING + else: + # item an an expression without an OrderBy wrapper. + expr = item + direction = ASCENDING + if isinstance(expr, (Col, Ref)): + field_name = expr.as_mql(self, self.connection, as_expr=False) + val_mql = expr.as_mql(self, self.connection, as_expr=True) + else: + val_mql = expr.as_mql(self, self.connection, as_expr=True) + expr_key = json.dumps(val_mql, sort_keys=True) + if expr_key not in sort_fields: + sort_fields[expr_key] = f"__wsort{len(sort_fields) + 1}" + field_name = sort_fields[expr_key] + pre_add_fields[field_name] = val_mql + if isinstance(item, OrderBy) and (item.nulls_first or item.nulls_last): + # Encode null position into a sub-document so that a single + # sortBy field suffices ($rank etc. require exactly one field). + # MongoDB compares documents field-by-field in insertion order, + # so {ind: 0, val: v} vs {ind: 1, val: null} is decided by + # "ind" before "val", giving clean null positioning. + # null_ind=1 puts null AFTER non-null in ASC / BEFORE in DESC. + # nulls_last != descending means nulls will be last in ASC or + # first in DESC — both require the higher ind value (1). + null_ind = int(bool(item.nulls_last) != bool(item.descending)) + null_field = f"__wsort{len(sort_fields) + 1}" + sort_fields[f"__null_{null_field}"] = null_field + pre_add_fields[null_field] = { + "ind": {"$cond": [{"$eq": [val_mql, None]}, null_ind, 1 - null_ind]}, + "val": val_mql, + } + sort_doc[null_field] = direction + else: + sort_doc[field_name] = direction + return sort_doc, pre_add_fields + + @staticmethod + def _sort_subset_of_partition(sort_doc, pre_add_fields, partition_mql): + """ + Return True if all sort fields are derived from the partition + expressions. + + When sort == partition, all rows in the partition are tied in sort + order, so the correct SQL RANGE window is the full partition (unbounded + both ways). + """ + if not partition_mql or not sort_doc: + return False + + def dump(expr): + return json.dumps(expr, sort_keys=True) + + if isinstance(partition_mql, dict) and not any(k.startswith("$") for k in partition_mql): + partition_set = {dump(v) for v in partition_mql.values()} + else: + partition_set = {dump(partition_mql)} + sort_set = set() + for field in sort_doc: + if field in pre_add_fields: + sort_set.add(dump(pre_add_fields[field])) + else: + sort_set.add(dump(f"${field}")) + return sort_set.issubset(partition_set) + + def _window_function_output(self, alias, window, idx, use_full_partition=False): + """ + Build $setWindowFields output entries and post-$addFields for a Window + expression. + + Return (output_entries, add_fields) where output_entries goes into + $setWindowFields.output and $addFields uses those results for complex + functions. + """ + expr = window.source_expression + frame = window.frame + has_order_by = bool(window.order_by and window.order_by.get_source_expressions()) + + def frame_mql(): + if frame is None: + return None + if frame.exclusion: + raise NotSupportedError("This backend does not support window frame exclusions.") + + def boundary(v): + if v is None: + return "unbounded" + if v == 0: + return "current" + return v + + bounds = [boundary(frame.start.value), boundary(frame.end.value)] + return {"documents": bounds} if isinstance(frame, RowRange) else {"range": bounds} + + def default_frame(): + if frame: + return frame_mql() + if has_order_by and not use_full_partition: + return {"documents": ["unbounded", "current"]} + return {"documents": ["unbounded", "unbounded"]} + + return expr.get_window_mql(self, self.connection, alias, idx, default_frame) + + def _prepare_window_annotations_for_pipeline(self): + """ + Build $setWindowFields pipeline stages for window annotations. + + Return (stages, replacements) where stages is a list of pipeline dicts + and replacements maps Window sub-expressions and full over-clause + annotations to Ref objects for use in $project. + """ + # Step 1: Collect all Window sub-expressions and their aliases. Process + # all annotations and any Window expressions used directly in the + # qualify (filter) clause. + seen_ids = set() + window_list = [] # list of (alias, window_expr) + idx = itertools.count(start=1) + window_replacements = {} # {Window instance: Ref} + + def _collect_windows_from_expr(expr, target=None): + for window in self._get_window_expressions(expr): + if id(window) in seen_ids: + continue + seen_ids.add(id(window)) + alias = target if (target and window is expr) else f"__window{next(idx)}" + window_list.append((alias, window)) + window_replacements[window] = Ref(alias, window) + + for target, expr in self.query.annotations.items(): + if not expr.contains_over_clause: + continue + _collect_windows_from_expr(expr, target) + # Scan the qualify clause for Windows used directly in filters (e.g., + # Employee.objects.filter(Exact(Window(...), 1))). + qualify = getattr(self, "qualify", None) + if qualify: + _collect_windows_from_expr(qualify) + if not window_list: + return [], {} + # Step 2: Group windows by (partition_key, sort_key) for + # $setWindowFields. + groups = {} + pre_add_fields = {} + sort_fields = {} # shared across all windows, expr_mql_key: field_name + for alias, window in window_list: + partition_mql = self._compile_window_partition_by(window.partition_by) + sort_doc, pre_add_fields_ = self._compile_window_sort_by(window.order_by, sort_fields) + pre_add_fields.update(pre_add_fields_) + partition_key = json.dumps(partition_mql) + sort_key = json.dumps(sort_doc) + group_key = (partition_key, sort_key) + if group_key not in groups: + groups[group_key] = {"partition": partition_mql, "sort": sort_doc, "windows": []} + groups[group_key]["windows"].append((alias, window)) + # Step 3: Build $setWindowFields stages with all outputs. + stages = [] + if pre_add_fields: + stages.append({"$addFields": pre_add_fields}) + post_add_fields = {} + for _group_key, group_info in groups.items(): + set_window_fields_output = {} + # Determine if sort fields are covered by partition fields (all + # rows tied). In that case, use the full-partition window instead + # of current-row window. + use_full_partition = self._sort_subset_of_partition( + group_info["sort"], + pre_add_fields, + group_info["partition"], + ) + for alias, window in group_info["windows"]: + output_entries, post_entries = self._window_function_output( + alias, window, idx, use_full_partition + ) + set_window_fields_output.update(output_entries) + post_add_fields.update(post_entries) + # $documentNumber, $rank, and $denseRank each require exactly one + # sortBy field; add a default when none is specified. + # $rank/$denseRank with no ORDER BY means all rows are tied (rank + # 1), so sort by a constant. $documentNumber needs unique row + # numbers, so sort by _id. + sort_doc = group_info["sort"] + if not sort_doc: + output_ops = { + op for v in set_window_fields_output.values() if isinstance(v, dict) for op in v + } + if output_ops & {"$rank", "$denseRank"}: + const_field = f"__wsort{len(sort_fields) + 1}" + pre_add_fields[const_field] = {"$literal": 1} + sort_doc = {const_field: 1} + elif "$documentNumber" in output_ops: + sort_doc = {"_id": 1} + set_window_fields = {"output": set_window_fields_output} + if group_info["partition"]: + set_window_fields["partitionBy"] = group_info["partition"] + if sort_doc: + set_window_fields["sortBy"] = sort_doc + stages.append({"$setWindowFields": set_window_fields}) + if post_add_fields: + stages.append({"$addFields": post_add_fields}) + # Step 4: Compute all over-clause annotations in an intermediate + # $addFields so they're available for the qualify $match. + intermediate_add_fields = {} + annotation_replacements = {} + for target, expr in self.query.annotations.items(): + if not expr.contains_over_clause: + continue + replaced_expr = expr.replace_expressions(window_replacements) + annotation_replacements[expr] = Ref(target, expr) + if isinstance(replaced_expr, Ref) and replaced_expr.refs == target: + continue # $setWindowFields already wrote this field. + intermediate_add_fields[target] = replaced_expr.as_mql( + self, self.connection, as_expr=True + ) + if intermediate_add_fields: + stages.append({"$addFields": intermediate_add_fields}) + return stages, {**window_replacements, **annotation_replacements} + def get_project_fields(self, columns=None, ordering=None, force_expression=False): if not columns: return {} @@ -922,6 +1184,11 @@ def execute_sql(self, returning_fields=None): f"You can't set {field.name} (a non-nullable field) to None." ) if hasattr(value, "as_mql"): + if getattr(value, "contains_over_clause", False): + raise FieldError( + f"Window expressions are not allowed in this query " + f"({field.name}={value!r})." + ) raise NotSupportedError( f"MongoDB does not support creating models with " f"expressions: got {value} for field {field.name}." @@ -988,7 +1255,8 @@ def execute_sql(self, result_type): ) if value.contains_over_clause: raise FieldError( - f"Window expressions are not allowed in this query ({field.name}={value})." + "Window expressions are not allowed in this query " + f"({field.name}={value!r})." ) elif hasattr(value, "prepare_database_save"): if field.remote_field: diff --git a/django_mongodb_backend/expressions/builtins.py b/django_mongodb_backend/expressions/builtins.py index c96c35595..ce23a341f 100644 --- a/django_mongodb_backend/expressions/builtins.py +++ b/django_mongodb_backend/expressions/builtins.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 diff --git a/django_mongodb_backend/features.py b/django_mongodb_backend/features.py index 70a24afbe..24c60e685 100644 --- a/django_mongodb_backend/features.py +++ b/django_mongodb_backend/features.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", diff --git a/django_mongodb_backend/functions/__init__.py b/django_mongodb_backend/functions/__init__.py index 2304c127d..c93657ebd 100644 --- a/django_mongodb_backend/functions/__init__.py +++ b/django_mongodb_backend/functions/__init__.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() diff --git a/django_mongodb_backend/functions/window.py b/django_mongodb_backend/functions/window.py new file mode 100644 index 000000000..f6d616ace --- /dev/null +++ b/django_mongodb_backend/functions/window.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 diff --git a/django_mongodb_backend/query.py b/django_mongodb_backend/query.py index dcde93ab1..8beb32688 100644 --- a/django_mongodb_backend/query.py +++ b/django_mongodb_backend/query.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 diff --git a/docs/releases/6.0.x.rst b/docs/releases/6.0.x.rst index dc8f71691..46fef9bb3 100644 --- a/docs/releases/6.0.x.rst +++ b/docs/releases/6.0.x.rst @@ -10,6 +10,8 @@ Django MongoDB Backend 6.0.x New features ------------ +- Added support for :ref:`window functions `. + - Added support for the following database functions: - :class:`~django.db.models.functions.ExtractQuarter` diff --git a/tests/expressions_/models.py b/tests/expressions_/models.py new file mode 100644 index 000000000..ea36bd3b2 --- /dev/null +++ b/tests/expressions_/models.py @@ -0,0 +1,5 @@ +from django.db import models + + +class Number(models.Model): + value = models.IntegerField() diff --git a/tests/expressions_/test_window.py b/tests/expressions_/test_window.py new file mode 100644 index 000000000..5f8e293b0 --- /dev/null +++ b/tests/expressions_/test_window.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())))