Add support for count_if in decomposing to measures (#1347)

* Add support for count_if, max_by, and min_by in decomposing to measures

* Fix gradle wrapper

* Only support count_if

* Fix

* Test

* Fix error in countif post-agg

* Fix test for countif post-agg

Author: shangyian

.github/workflows/test.yml

@@ -150,16 +150,13 @@ jobs:
       run:
         working-directory: ./datajunction-clients/java
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up JDK 17
-        uses: actions/setup-java@v3
+        uses: actions/setup-java@v4
         with:
           java-version: '17'
           distribution: 'temurin'
-      - name: Official Gradle Wrapper Validation Action
-        uses: gradle/actions/wrapper-validation@v3
+      - name: Setup Gradle
+        uses: gradle/actions/setup-gradle@v4
       - name: Build with Gradle
-        uses: gradle/actions/setup-gradle@v3
-        with:
-          arguments: build -x test
-          build-root-directory: ./datajunction-clients/java
+        run: ./gradlew build -x test

datajunction-server/datajunction_server/sql/decompose.py

@@ -22,8 +22,11 @@ def __init__(self, query_ast: ast.Query):
         self.handlers = {
             dj_functions.Sum: self._simple_associative_agg,
             dj_functions.Count: self._simple_associative_agg,
+            dj_functions.CountIf: self._simple_associative_agg,
             dj_functions.Max: self._simple_associative_agg,
+            dj_functions.MaxBy: self._simple_associative_agg,
             dj_functions.Min: self._simple_associative_agg,
+            dj_functions.MinBy: self._simple_associative_agg,
             dj_functions.Avg: self._avg,
             dj_functions.AnyValue: self._simple_associative_agg,
         }
@@ -158,7 +161,7 @@ def update_ast(func, measures: list[Measure]):
                 ),
             )
         elif (
-            func.function() == dj_functions.Count
+            func.function() in (dj_functions.Count, dj_functions.CountIf)
             and func.quantifier != ast.SetQuantifier.Distinct
         ):
             func.name.name = "SUM"

datajunction-server/datajunction_server/sql/functions.py

@@ -1222,6 +1222,8 @@ class CountIf(Function):
     count_if(expr) - Returns the number of true values in expr.
     """
 
+    is_aggregation = True
+
 
 @CountIf.register  # type: ignore
 def infer_type(arg: ct.BooleanType) -> ct.IntegerType:
@@ -3244,6 +3246,8 @@ class MaxBy(Function):
     max_by(val, key) - Returns the value of val corresponding to the maximum value of key.
     """
 
+    is_aggregation = True
+
 
 @MaxBy.register  # type: ignore
 def infer_type(val: ct.ColumnType, key: ct.ColumnType) -> ct.ColumnType:
@@ -3334,6 +3338,8 @@ class MinBy(Function):
     min_by(val, key) - Returns the value of val corresponding to the minimum value of key.
     """
 
+    is_aggregation = True
+
 
 @MinBy.register  # type: ignore
 def infer_type(val: ct.ColumnType, key: ct.ColumnType) -> ct.ColumnType:

datajunction-server/tests/sql/decompose_test.py

@@ -618,6 +618,38 @@ def test_unsupported_aggregation_function():
     )
 
 
+def test_count_if():
+    """
+    Test decomposition for count_if.
+    """
+    extractor = MeasureExtractor.from_query_string(
+        "SELECT CAST(COUNT_IF(ARRAY_CONTAINS(field_a, 'xyz')) AS FLOAT) / COUNT(*) "
+        "FROM parent_node",
+    )
+    measures, derived_sql = extractor.extract()
+    expected_measures = [
+        Measure(
+            name="field_a_count_if_c1f2ed10",
+            expression="ARRAY_CONTAINS(field_a, 'xyz')",
+            aggregation="COUNT_IF",
+            rule=AggregationRule(type=Aggregability.FULL),
+        ),
+        Measure(
+            name="count_3389dae3",
+            expression="*",
+            aggregation="COUNT",
+            rule=AggregationRule(type=Aggregability.FULL),
+        ),
+    ]
+    assert measures == expected_measures
+    assert str(derived_sql) == str(
+        parse(
+            "SELECT  CAST(SUM(field_a_count_if_c1f2ed10) AS FLOAT) / SUM(count_3389dae3) "
+            "FROM parent_node",
+        ),
+    )
+
+
 def test_metric_query_with_aliases():
     """
     Test behavior when the query contains unsupported aggregation functions. We just return an