Patch fix to how COUNT(*) calls above joins get split (#437)

Patching the logic of how `COUNT(*)` calls are handled in the aggregation splitting logic to avoid bugs elsewhere int he codebase. Also added several simplification rules exposed by the changes:
- `INTEGER(true)` -> `1`
- `INTEGER(false)` -> `0`
- `INTEGER(int_literal)` -> `int_literal`
- `1 * x` -> `x`
- `x * 1` -> `x`
- `0 * x` -> `0`
- `x * 0` -> `0`

Author: knassre-bodo

pydough/conversion/agg_split.py

@@ -375,11 +375,6 @@ def attempt_join_aggregate_transpose(
         # perform the transpose.
         return node, True
 
-    if need_count_aggs and not (can_push_left and can_push_right):
-        # If we need to push down COUNT(*) aggregates, but cannot push into
-        # both sides of the join, we cannot perform the transpose.
-        return node, True
-
     # Parse the join condition to identify the lists of equi-join keys
     # from the LHS and RHS, and verify that all of the columns used by
     # the condition are in those lists.
@@ -406,21 +401,31 @@ def attempt_join_aggregate_transpose(
 
     # If we need count aggregates, add one to each side of the join.
     if need_count_aggs:
-        assert len(count_aggs) > 0
-        lhs_aggs.append(count_aggs.pop())
-        new_agg_name: str
+        lhs_agg_name: str
+        rhs_agg_name: str
         idx: int = 0
         while True:
-            new_agg_name = f"agg_{idx}"
-            if new_agg_name not in node.columns:
+            lhs_agg_name = f"agg_{idx}"
+            idx += 1
+            if lhs_agg_name not in node.columns and lhs_agg_name not in join.columns:
                 break
+        while True:
+            rhs_agg_name = f"agg_{idx}"
             idx += 1
-        node.aggregations[new_agg_name] = CallExpression(
+            if rhs_agg_name not in node.columns and rhs_agg_name not in join.columns:
+                break
+        node.aggregations[lhs_agg_name] = CallExpression(
             pydop.COUNT,
             NumericType(),
             [],
         )
-        rhs_aggs.append(new_agg_name)
+        node.aggregations[rhs_agg_name] = CallExpression(
+            pydop.COUNT,
+            NumericType(),
+            [],
+        )
+        lhs_aggs.append(lhs_agg_name)
+        rhs_aggs.append(rhs_agg_name)
 
     # Loop over both inputs and perform the pushdown into whichever one(s)
     # will allow an aggregate to be pushed into them.
@@ -445,20 +450,32 @@ def attempt_join_aggregate_transpose(
         if side_count_ref is not None:
             count_refs.append(side_count_ref)
 
-    # For each COUNT(*) aggregate, replace with the product of the COUNT(*)
-    # calls that were pushed into each side of the join.
+    # For each COUNT(*) aggregate, replace with the product of the calls that
+    # were pushed into each side of the join.
     for count_call_name in count_aggs:
-        assert len(count_refs) > 1
-        product: RelationalExpression = CallExpression(
-            pydop.MUL, NumericType(), count_refs
-        )
-        projection_columns[count_call_name] = product
-        need_projection = True
+        if len(count_refs) > 1:
+            product: RelationalExpression = CallExpression(
+                pydop.MUL, NumericType(), count_refs
+            )
+            product_sum: CallExpression = CallExpression(
+                pydop.SUM, NumericType(), [product]
+            )
+            node.aggregations[count_call_name] = product_sum
+            node.columns[count_call_name] = product_sum
+        elif len(count_refs) == 1:
+            regular_sum: CallExpression = CallExpression(
+                pydop.SUM, NumericType(), [count_refs[0]]
+            )
+            node.aggregations[count_call_name] = regular_sum
+            node.columns[count_call_name] = regular_sum
 
     # If the node requires projection at the end, create a new Project node on
     # top of the top aggregate.
     if need_projection:
-        return Project(node, projection_columns), True
+        new_node: RelationalNode = node.copy(
+            inputs=[split_partial_aggregates(input, config) for input in node.inputs]
+        )
+        return Project(new_node, projection_columns), False
     else:
         return node, True
 

pydough/conversion/relational_simplification.py

@@ -884,6 +884,16 @@ def simplify_function_call(
                     not_null=True, not_negative=True
                 )
 
+            # INTEGER(x) -> x if x is a literal integer. Also simplify for
+            # booleans.
+            case pydop.INTEGER:
+                if isinstance(expr.inputs[0], LiteralExpression) and isinstance(
+                    expr.inputs[0].value, (int, bool)
+                ):
+                    output_expr = LiteralExpression(
+                        int(expr.inputs[0].value), expr.data_type
+                    )
+
             # The result of addition is non-negative or positive if all the
             # operands are. It is also positive if all the operands are
             # non-negative and at least one of them is positive.
@@ -895,11 +905,35 @@ def simplify_function_call(
                     output_predicates.positive = True
 
             # The result of multiplication is non-negative or positive if all
-            # the operands are.
+            # the operands are. Also, simplify when any argument is 0 to the
+            # output being 0, and remove any arguments that are 1.
             case pydop.MUL:
                 output_predicates |= intersect_set & PredicateSet(
                     not_negative=True, positive=True
                 )
+                remaining_args: list[RelationalExpression] = [
+                    arg
+                    for arg in expr.inputs
+                    if not (
+                        isinstance(arg, LiteralExpression)
+                        and arg.value in (1, 1.0, True)
+                    )
+                ]
+                if len(remaining_args) == 0:
+                    output_expr = expr.inputs[0]
+                elif len(remaining_args) == 1:
+                    output_expr = remaining_args[0]
+                elif len(remaining_args) < len(expr.inputs):
+                    output_expr = CallExpression(
+                        pydop.MUL, expr.data_type, remaining_args
+                    )
+                for arg in expr.inputs:
+                    if isinstance(arg, LiteralExpression) and arg.value in (
+                        0,
+                        0.0,
+                        False,
+                    ):
+                        output_expr = LiteralExpression(0, expr.data_type)
 
             # The result of division is non-negative or positive if all the
             # operands are, and is also non-null if both operands are non-null

tests/test_pipeline_custom_datasets.py

@@ -30,9 +30,7 @@
                 "synthea",
                 lambda: pd.DataFrame(
                     {
-                        "condition_description": [
-                            "Escherichia coli urinary tract infection"
-                        ],
+                        "condition_description": ["Normal pregnancy"],
                     }
                 ),
                 "synthea_most_common_conditions",

tests/test_pipeline_defog.py

@@ -1503,5 +1503,5 @@ def test_defog_e2e(
     same database connector. Run on the defog.ai queries.
     """
     defog_pipeline_test_data.run_e2e_test(
-        defog_graphs, sqlite_defog_connection, defog_config
+        defog_graphs, sqlite_defog_connection, defog_config, coerce_types=True
     )

tests/test_pipeline_defog_custom.py

@@ -2304,7 +2304,7 @@ def test_pipeline_e2e_defog_custom(
     schemas.
     """
     defog_custom_pipeline_test_data.run_e2e_test(
-        defog_graphs, sqlite_defog_connection, config=defog_config
+        defog_graphs, sqlite_defog_connection, config=defog_config, coerce_types=True
     )
 
 

tests/test_pipeline_mysql.py

@@ -554,4 +554,5 @@ def test_pipeline_e2e_mysql_defog(
         defog_config,
         reference_database=sqlite_defog_connection,
         coerce_types=True,
+        rtol=1e4,
     )

tests/test_plan_refsols/common_prefix_b.txt

@@ -1,8 +1,8 @@
 ROOT(columns=[('name', r_name), ('n_nations', n_rows), ('n_customers', sum_sum_expr), ('n_suppliers', sum_n_rows)], orderings=[(r_name):asc_first])
  JOIN(condition=t0.r_regionkey == t1.n_regionkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_rows': t1.n_rows, 'r_name': t0.r_name, 'sum_n_rows': t1.sum_n_rows, 'sum_sum_expr': t1.sum_sum_expr})
   SCAN(table=tpch.REGION, columns={'r_name': r_name, 'r_regionkey': r_regionkey})
-  AGGREGATE(keys={'n_regionkey': n_regionkey}, aggregations={'n_rows': COUNT(), 'sum_n_rows': SUM(n_rows), 'sum_sum_expr': SUM(sum_expr)})
-   JOIN(condition=t0.n_nationkey == t1.s_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_regionkey': t0.n_regionkey, 'n_rows': t1.n_rows, 'sum_expr': t0.n_rows})
+  AGGREGATE(keys={'n_regionkey': n_regionkey}, aggregations={'n_rows': COUNT(), 'sum_n_rows': SUM(n_rows), 'sum_sum_expr': SUM(sum_expr_4)})
+   JOIN(condition=t0.n_nationkey == t1.s_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_regionkey': t0.n_regionkey, 'n_rows': t1.n_rows, 'sum_expr_4': t0.n_rows})
     JOIN(condition=t0.n_nationkey == t1.c_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_nationkey': t0.n_nationkey, 'n_regionkey': t0.n_regionkey, 'n_rows': t1.n_rows})
      SCAN(table=tpch.NATION, columns={'n_nationkey': n_nationkey, 'n_regionkey': n_regionkey})
      AGGREGATE(keys={'c_nationkey': c_nationkey}, aggregations={'n_rows': COUNT()})

tests/test_plan_refsols/common_prefix_c.txt

@@ -1,7 +1,7 @@
-ROOT(columns=[('name', r_name), ('n_nations', n_rows), ('n_customers', sum_sum_expr), ('n_suppliers', sum_n_rows), ('n_orders', DEFAULT_TO(sum_sum_expr_15, 0:numeric)), ('n_parts', sum_sum_n_rows)], orderings=[(r_name):asc_first])
- JOIN(condition=t0.r_regionkey == t1.n_regionkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_rows': t1.n_rows, 'r_name': t0.r_name, 'sum_n_rows': t1.sum_sum_n_rows_0, 'sum_sum_expr': t1.sum_n_rows, 'sum_sum_expr_15': t1.sum_sum_expr_15, 'sum_sum_n_rows': t1.sum_sum_n_rows})
+ROOT(columns=[('name', r_name), ('n_nations', n_rows), ('n_customers', sum_expr), ('n_suppliers', sum_n_rows), ('n_orders', DEFAULT_TO(sum_sum_expr, 0:numeric)), ('n_parts', sum_sum_n_rows)], orderings=[(r_name):asc_first])
+ JOIN(condition=t0.r_regionkey == t1.n_regionkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_rows': t1.n_rows, 'r_name': t0.r_name, 'sum_expr': t1.sum_n_rows, 'sum_n_rows': t1.sum_sum_n_rows_0, 'sum_sum_expr': t1.sum_sum_expr, 'sum_sum_n_rows': t1.sum_sum_n_rows})
   SCAN(table=tpch.REGION, columns={'r_name': r_name, 'r_regionkey': r_regionkey})
-  AGGREGATE(keys={'n_regionkey': n_regionkey}, aggregations={'n_rows': COUNT(), 'sum_n_rows': SUM(n_rows), 'sum_sum_expr_15': SUM(sum_expr_15), 'sum_sum_n_rows': SUM(sum_n_rows), 'sum_sum_n_rows_0': SUM(sum_n_rows_0)})
+  AGGREGATE(keys={'n_regionkey': n_regionkey}, aggregations={'n_rows': COUNT(), 'sum_n_rows': SUM(n_rows), 'sum_sum_expr': SUM(sum_expr_15), 'sum_sum_n_rows': SUM(sum_n_rows), 'sum_sum_n_rows_0': SUM(sum_n_rows_0)})
    JOIN(condition=t0.n_nationkey == t1.s_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_regionkey': t0.n_regionkey, 'n_rows': t0.n_rows, 'sum_expr_15': t0.sum_n_rows, 'sum_n_rows': t1.sum_n_rows, 'sum_n_rows_0': t1.n_rows})
     JOIN(condition=t0.n_nationkey == t1.c_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_nationkey': t0.n_nationkey, 'n_regionkey': t0.n_regionkey, 'n_rows': t1.n_rows, 'sum_n_rows': t1.sum_n_rows})
      SCAN(table=tpch.NATION, columns={'n_nationkey': n_nationkey, 'n_regionkey': n_regionkey})

tests/test_plan_refsols/common_prefix_d.txt

@@ -1,24 +1,23 @@
-ROOT(columns=[('name', r_name), ('n_nations', n_rows), ('n_customers', sum_sum_expr), ('n_suppliers', sum_n_rows), ('n_orders_94', DEFAULT_TO(sum_sum_sum_expr, 0:numeric)), ('n_orders_95', DEFAULT_TO(sum_sum_expr_10, 0:numeric)), ('n_orders_96', DEFAULT_TO(sum_sum_n_rows, 0:numeric))], orderings=[(r_name):asc_first])
- JOIN(condition=t0.r_regionkey == t1.n_regionkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_rows': t1.n_rows, 'r_name': t0.r_name, 'sum_n_rows': t1.sum_sum_n_rows_0, 'sum_sum_expr': t1.sum_n_rows, 'sum_sum_expr_10': t1.sum_sum_expr_10, 'sum_sum_n_rows': t1.sum_sum_n_rows, 'sum_sum_sum_expr': t1.sum_sum_sum_expr})
+ROOT(columns=[('name', r_name), ('n_nations', n_rows), ('n_customers', sum_expr), ('n_suppliers', sum_n_rows), ('n_orders_94', DEFAULT_TO(sum_sum_expr, 0:numeric)), ('n_orders_95', DEFAULT_TO(sum_sum_expr_10, 0:numeric)), ('n_orders_96', DEFAULT_TO(sum_sum_n_rows, 0:numeric))], orderings=[(r_name):asc_first])
+ JOIN(condition=t0.r_regionkey == t1.n_regionkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_rows': t1.n_rows, 'r_name': t0.r_name, 'sum_expr': t1.sum_n_rows, 'sum_n_rows': t1.sum_sum_n_rows_0, 'sum_sum_expr': t1.sum_sum_expr, 'sum_sum_expr_10': t1.sum_sum_expr_10, 'sum_sum_n_rows': t1.sum_sum_n_rows})
   SCAN(table=tpch.REGION, columns={'r_name': r_name, 'r_regionkey': r_regionkey})
-  AGGREGATE(keys={'n_regionkey': n_regionkey}, aggregations={'n_rows': COUNT(), 'sum_n_rows': SUM(n_rows), 'sum_sum_expr_10': SUM(sum_expr_10), 'sum_sum_n_rows': SUM(sum_n_rows), 'sum_sum_n_rows_0': SUM(sum_n_rows_0), 'sum_sum_sum_expr': SUM(sum_sum_expr)})
-   JOIN(condition=t0.n_nationkey == t1.s_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_regionkey': t0.n_regionkey, 'n_rows': t0.n_rows, 'sum_expr_10': t0.sum_expr_10, 'sum_n_rows': t0.sum_n_rows, 'sum_n_rows_0': t1.n_rows, 'sum_sum_expr': t0.sum_sum_expr})
-    PROJECT(columns={'n_nationkey': n_nationkey, 'n_regionkey': n_regionkey, 'n_rows': n_rows, 'sum_expr_10': sum_expr_10, 'sum_n_rows': sum_n_rows, 'sum_sum_expr': sum_expr})
-     JOIN(condition=t0.n_nationkey == t1.c_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_nationkey': t0.n_nationkey, 'n_regionkey': t0.n_regionkey, 'n_rows': t1.n_rows, 'sum_expr': t1.sum_expr_7, 'sum_expr_10': t1.sum_expr_10, 'sum_n_rows': t1.sum_n_rows})
-      SCAN(table=tpch.NATION, columns={'n_nationkey': n_nationkey, 'n_regionkey': n_regionkey})
-      AGGREGATE(keys={'c_nationkey': c_nationkey}, aggregations={'n_rows': COUNT(), 'sum_expr_10': SUM(expr_10), 'sum_expr_7': SUM(expr_7), 'sum_n_rows': SUM(n_rows)})
-       JOIN(condition=t0.c_custkey == t1.o_custkey, type=LEFT, cardinality=SINGULAR_FILTER, columns={'c_nationkey': t0.c_nationkey, 'expr_10': t0.n_rows, 'expr_7': t0.expr_7, 'n_rows': t1.n_rows})
-        JOIN(condition=t0.c_custkey == t1.o_custkey, type=LEFT, cardinality=SINGULAR_FILTER, columns={'c_custkey': t0.c_custkey, 'c_nationkey': t0.c_nationkey, 'expr_7': t0.n_rows, 'n_rows': t1.n_rows})
-         JOIN(condition=t0.c_custkey == t1.o_custkey, type=LEFT, cardinality=SINGULAR_FILTER, columns={'c_custkey': t0.c_custkey, 'c_nationkey': t0.c_nationkey, 'n_rows': t1.n_rows})
-          SCAN(table=tpch.CUSTOMER, columns={'c_custkey': c_custkey, 'c_nationkey': c_nationkey})
-          AGGREGATE(keys={'o_custkey': o_custkey}, aggregations={'n_rows': COUNT()})
-           FILTER(condition=YEAR(o_orderdate) == 1994:numeric, columns={'o_custkey': o_custkey})
-            SCAN(table=tpch.ORDERS, columns={'o_custkey': o_custkey, 'o_orderdate': o_orderdate})
+  AGGREGATE(keys={'n_regionkey': n_regionkey}, aggregations={'n_rows': COUNT(), 'sum_n_rows': SUM(n_rows), 'sum_sum_expr': SUM(sum_expr), 'sum_sum_expr_10': SUM(sum_expr_10), 'sum_sum_n_rows': SUM(sum_n_rows), 'sum_sum_n_rows_0': SUM(sum_n_rows_0)})
+   JOIN(condition=t0.n_nationkey == t1.s_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_regionkey': t0.n_regionkey, 'n_rows': t0.n_rows, 'sum_expr': t0.sum_expr, 'sum_expr_10': t0.sum_expr_10, 'sum_n_rows': t0.sum_n_rows, 'sum_n_rows_0': t1.n_rows})
+    JOIN(condition=t0.n_nationkey == t1.c_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_nationkey': t0.n_nationkey, 'n_regionkey': t0.n_regionkey, 'n_rows': t1.n_rows, 'sum_expr': t1.sum_expr_7, 'sum_expr_10': t1.sum_expr_10, 'sum_n_rows': t1.sum_n_rows})
+     SCAN(table=tpch.NATION, columns={'n_nationkey': n_nationkey, 'n_regionkey': n_regionkey})
+     AGGREGATE(keys={'c_nationkey': c_nationkey}, aggregations={'n_rows': COUNT(), 'sum_expr_10': SUM(expr_10), 'sum_expr_7': SUM(expr_7), 'sum_n_rows': SUM(n_rows)})
+      JOIN(condition=t0.c_custkey == t1.o_custkey, type=LEFT, cardinality=SINGULAR_FILTER, columns={'c_nationkey': t0.c_nationkey, 'expr_10': t0.n_rows, 'expr_7': t0.expr_7, 'n_rows': t1.n_rows})
+       JOIN(condition=t0.c_custkey == t1.o_custkey, type=LEFT, cardinality=SINGULAR_FILTER, columns={'c_custkey': t0.c_custkey, 'c_nationkey': t0.c_nationkey, 'expr_7': t0.n_rows, 'n_rows': t1.n_rows})
+        JOIN(condition=t0.c_custkey == t1.o_custkey, type=LEFT, cardinality=SINGULAR_FILTER, columns={'c_custkey': t0.c_custkey, 'c_nationkey': t0.c_nationkey, 'n_rows': t1.n_rows})
+         SCAN(table=tpch.CUSTOMER, columns={'c_custkey': c_custkey, 'c_nationkey': c_nationkey})
          AGGREGATE(keys={'o_custkey': o_custkey}, aggregations={'n_rows': COUNT()})
-          FILTER(condition=YEAR(o_orderdate) == 1995:numeric, columns={'o_custkey': o_custkey})
+          FILTER(condition=YEAR(o_orderdate) == 1994:numeric, columns={'o_custkey': o_custkey})
            SCAN(table=tpch.ORDERS, columns={'o_custkey': o_custkey, 'o_orderdate': o_orderdate})
         AGGREGATE(keys={'o_custkey': o_custkey}, aggregations={'n_rows': COUNT()})
-         FILTER(condition=YEAR(o_orderdate) == 1996:numeric, columns={'o_custkey': o_custkey})
+         FILTER(condition=YEAR(o_orderdate) == 1995:numeric, columns={'o_custkey': o_custkey})
           SCAN(table=tpch.ORDERS, columns={'o_custkey': o_custkey, 'o_orderdate': o_orderdate})
+       AGGREGATE(keys={'o_custkey': o_custkey}, aggregations={'n_rows': COUNT()})
+        FILTER(condition=YEAR(o_orderdate) == 1996:numeric, columns={'o_custkey': o_custkey})
+         SCAN(table=tpch.ORDERS, columns={'o_custkey': o_custkey, 'o_orderdate': o_orderdate})
     AGGREGATE(keys={'s_nationkey': s_nationkey}, aggregations={'n_rows': COUNT()})
      SCAN(table=tpch.SUPPLIER, columns={'s_nationkey': s_nationkey})

tests/test_plan_refsols/common_prefix_f.txt

@@ -1,8 +1,8 @@
 ROOT(columns=[('name', r_name), ('n_customers', sum_sum_expr), ('n_nations', n_rows), ('n_suppliers', sum_n_rows)], orderings=[(r_name):asc_first])
  JOIN(condition=t0.r_regionkey == t1.n_regionkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_rows': t1.n_rows, 'r_name': t0.r_name, 'sum_n_rows': t1.sum_n_rows, 'sum_sum_expr': t1.sum_sum_expr})
   SCAN(table=tpch.REGION, columns={'r_name': r_name, 'r_regionkey': r_regionkey})
-  AGGREGATE(keys={'n_regionkey': n_regionkey}, aggregations={'n_rows': COUNT(), 'sum_n_rows': SUM(n_rows), 'sum_sum_expr': SUM(sum_expr)})
-   JOIN(condition=t0.n_nationkey == t1.s_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_regionkey': t0.n_regionkey, 'n_rows': t1.n_rows, 'sum_expr': t0.n_rows})
+  AGGREGATE(keys={'n_regionkey': n_regionkey}, aggregations={'n_rows': COUNT(), 'sum_n_rows': SUM(n_rows), 'sum_sum_expr': SUM(sum_expr_4)})
+   JOIN(condition=t0.n_nationkey == t1.s_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_regionkey': t0.n_regionkey, 'n_rows': t1.n_rows, 'sum_expr_4': t0.n_rows})
     JOIN(condition=t0.n_nationkey == t1.c_nationkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_nationkey': t0.n_nationkey, 'n_regionkey': t0.n_regionkey, 'n_rows': t1.n_rows})
      SCAN(table=tpch.NATION, columns={'n_nationkey': n_nationkey, 'n_regionkey': n_regionkey})
      AGGREGATE(keys={'c_nationkey': c_nationkey}, aggregations={'n_rows': COUNT()})