add argMax aggregation (#134)

Author: gambuzzi

CHANGELOG.md

@@ -1,6 +1,7 @@
 ### 1.5.0
 - feat: #140: Adding pre-commit with ruff as linter and code formatter. 
 - fix: #139: Fix to replicas query when using default as cluster name
+- feat: #134: add `argMax` aggregation https://clickhouse.com/docs/sql-reference/aggregate-functions/reference/argmax
 - feat: #133: Fix simultaneous queries error when iteration is interrupted
 - feat: #130: Add `distributed_migrations` database setting to support distributed migration queries.
 - feat: #129: Add `toYearWeek` datetime functionality

clickhouse_backend/models/aggregates.py

@@ -11,6 +11,7 @@
     "uniqHLL12",
     "uniqTheta",
     "anyLast",
+    "argMax",
 ]
 
 
@@ -67,3 +68,10 @@ class uniqTheta(uniq):
 
 class anyLast(Aggregate):
     pass
+
+
+class argMax(Aggregate):
+    arity = 2
+
+    def _resolve_output_field(self):
+        return self.get_source_fields()[0]

tests/aggregates/tests.py

@@ -1,4 +1,4 @@
-from django.db.models import Q
+from django.db.models import Q, F
 from django.test import TestCase
 
 from clickhouse_backend.models import (
@@ -11,6 +11,8 @@
     uniqTheta,
 )
 
+from clickhouse_backend.models.aggregates import argMax
+
 from .models import WatchSeries
 
 
@@ -170,6 +172,36 @@ def setUpTestData(cls):
         # Use bulk_create to insert the list of objects in a single query
         WatchSeries.objects.bulk_create(watch_series_list)
 
+    def test_argMax(self):
+        result = (
+            WatchSeries.objects.values("show")
+            .annotate(episode=argMax("episode", "date_id"))
+            .order_by("show")
+        )
+
+        expected_result = [
+            {"show": "Bridgerton", "episode": "S1E1"},
+            {"show": "Game of Thrones", "episode": "S1E1"},
+        ]
+
+        self.assertQuerysetEqual(result, expected_result, transform=dict)
+
+    def test_argMax_output_field(self):
+        result = (
+            WatchSeries.objects.values("show")
+            .annotate(date=argMax("date_id", "episode"))
+            .annotate(week=F("date__week"), week_day=F("date__week_day"))
+            .values("show", "week", "week_day")
+            .order_by("show")
+        )
+
+        expected_result = [
+            {"show": "Bridgerton", "week": 20, "week_day": 6},
+            {"show": "Game of Thrones", "week": 20, "week_day": 5},
+        ]
+
+        self.assertQuerysetEqual(result, expected_result, transform=dict)
+
     def _test_uniq(self, cls_uniq):
         result = (
             WatchSeries.objects.values("show", "episode")