fix(hogql): combinate all the aggregations (#41515)

The CH docs are pretty loosy goosy about which of these are actually valid and what ordering works, so I've both assumed a lot and also decided to be more strict in some ways (technically `-If` can preceded `-ForEach`, for example, but I don't allow it here, and I have no idea if stuff like `avgIfForEachIf` would be permitted, but I also disallow that). 

This was sparked by wanting `avgForEach` and it not being available, and my gut saying we should just throw the whole possible space at the wall here rather than continuing to add them one by one.

Author: oliverb123

posthog/hogql/functions/aggregations.py

@@ -1,11 +1,83 @@
+from collections.abc import Callable
+
 from posthog.hogql.ast import ArrayType, BooleanType, StringType
 from posthog.hogql.base import UnknownType
 
 from .core import HogQLFunctionMeta
 
+COMBINATORS = {
+    "If": {"allowedSuffixes": [], "argMap": lambda min, max: [min + 1, max + 1]},
+    "Array": {"allowedSuffixes": ["If", "OrDefault", "OrNull"], "argMap": lambda min, max: [min, max]},
+    "Map": {"allowedSuffixes": ["If", "OrDefault", "OrNull"], "argMap": lambda min, max: [min, max]},
+    "State": {"allowedSuffixes": ["If", "OrDefault", "OrNull"], "argMap": lambda min, max: [min, max]},
+    "Merge": {"allowedSuffixes": ["If", "OrDefault", "OrNull"], "argMap": lambda min, max: [min, max]},
+    "ForEach": {"allowedSuffixes": ["If", "OrDefault", "OrNull"], "argMap": lambda min, max: [min, max]},
+    "OrDefault": {"allowedSuffixes": ["If"], "argMap": lambda min, max: [min, max]},
+    "OrNull": {"allowedSuffixes": ["If"], "argMap": lambda min, max: [min, max]},
+    "ArgMin": {"allowedSuffixes": ["If", "OrDefault", "OrNull"], "argMap": lambda min, max: [min + 1, max + 1]},
+    "ArgMax": {"allowedSuffixes": ["If", "OrDefault", "OrNull"], "argMap": lambda min, max: [min + 1, max + 1]},
+}
+
+COMBINATOR_AGGREGATIONS = {
+    "avg": HogQLFunctionMeta("avg", 1, 1, aggregate=True),
+    "sum": HogQLFunctionMeta("sum", 1, 1, aggregate=True),
+    "min": HogQLFunctionMeta("min", 1, 1, aggregate=True),
+    "max": HogQLFunctionMeta("max", 1, 1, aggregate=True),
+    "count": HogQLFunctionMeta("count", 0, 1, aggregate=True),
+    "countDistinct": HogQLFunctionMeta("countDistinct", 1, 1, aggregate=True),
+    "median": HogQLFunctionMeta("median", 1, 1, aggregate=True),
+}
+
+
+def _generate_suffix_combinations(
+    base_name: str, base_meta: HogQLFunctionMeta, current_suffixes: list[str] | None = None
+):
+    result = {}
+
+    if current_suffixes is None:
+        current_suffixes = []
+
+    if current_suffixes:
+        func_name = base_name + "".join(current_suffixes)
+        # Calculate new parameter ranges based on suffix rules
+        min_params, max_params = base_meta.min_args, base_meta.max_args
+        for suffix in current_suffixes:
+            if suffix in COMBINATORS:
+                arg_map: Callable[[int, int | None], list[int]] = COMBINATORS[suffix]["argMap"]  # type: ignore
+                min_params, max_params = arg_map(min_params, max_params)
+
+        result[func_name] = HogQLFunctionMeta(func_name, min_params, max_params, aggregate=True)
+
+    if not current_suffixes:
+        available_suffixes = list(COMBINATORS.keys())
+    else:
+        last_suffix = current_suffixes[-1]
+        allowed_suffixes: list[str] = COMBINATORS.get(last_suffix, {}).get("allowedSuffixes", [])  # type: ignore
+        available_suffixes = allowed_suffixes
+
+    for suffix in available_suffixes:
+        if suffix not in current_suffixes:
+            nested_result = _generate_suffix_combinations(base_name, base_meta, [*current_suffixes, suffix])
+            result.update(nested_result)
+
+    return result
+
+
+def generate_combinator_suffix_combinations():
+    result = {}
+
+    for base_name, base_meta in COMBINATOR_AGGREGATIONS.items():
+        combinations = _generate_suffix_combinations(base_name, base_meta)
+        result.update(combinations)
+
+    return result
+
+
 # Permitted HogQL aggregations
 # Keep in sync with the posthog.com repository: contents/docs/sql/aggregations.mdx
 HOGQL_AGGREGATIONS: dict[str, HogQLFunctionMeta] = {
+    # Generated combinator functions
+    **generate_combinator_suffix_combinations(),
     # Standard aggregate functions
     "count": HogQLFunctionMeta("count", 0, 1, aggregate=True, case_sensitive=False),
     "countIf": HogQLFunctionMeta("countIf", 1, 2, aggregate=True),