fix: datafusion aggregations in parity tests

tokoko · Sep 5, 2024 · ac5c557 · ac5c557
1 parent c6e10bd
commit ac5c557
Showing 1 changed file with 26 additions and 14 deletions.
diff --git a/ibis_substrait/tests/compiler/test_parity.py b/ibis_substrait/tests/compiler/test_parity.py
@@ -158,26 +158,37 @@ def test_left_join(consumer: str, request):
 
 @pytest.mark.parametrize(
     "consumer",
-    [
-        "acero_consumer",
-        pytest.param(
-            "datafusion_consumer",
-            marks=[pytest.mark.xfail(Exception, reason="")],
-        ),
-    ],
+    ["acero_consumer", "datafusion_consumer"],
 )
 def test_filter_groupby(consumer: str, request):
     filter_table = orders.join(
         stores, orders["fk_store_id"] == stores["store_id"]
     ).filter(lambda t: t.order_total > 30)
 
-    expr = filter_table.group_by("city").aggregate(
-        sales=filter_table["order_id"].count()
+    expr = (
+        filter_table.group_by("city")
+        .aggregate(sales=filter_table["order_id"].count())
+        .filter(ibis.literal(True))
     )
 
     run_parity_test(request.getfixturevalue(consumer), expr)
 
 
+@pytest.mark.parametrize(
+    "consumer",
+    [
+        pytest.param(
+            "datafusion_consumer",
+            marks=[pytest.mark.xfail(Exception, reason="")],
+        ),
+    ],
+)
+def test_groupby_datafusion(consumer: str, request):
+    expr = orders.group_by("fk_store_id").aggregate(sales=orders["order_id"].count())
+
+    run_parity_test(request.getfixturevalue(consumer), expr)
+
+
 @pytest.mark.parametrize(
     "consumer",
     [
@@ -187,18 +198,19 @@ def test_filter_groupby(consumer: str, request):
                 pytest.mark.xfail(pa.ArrowNotImplementedError, reason="Unimplemented")
             ],
         ),
-        pytest.param(
-            "datafusion_consumer",
-            marks=[pytest.mark.xfail(Exception, reason="")],
-        ),
+        "datafusion_consumer",
     ],
 )
 def test_filter_groupby_count_distinct(consumer: str, request):
     filter_table = orders.join(
         stores, orders["fk_store_id"] == stores["store_id"]
     ).filter(lambda t: t.order_total > 30)
 
-    expr = filter_table.group_by("city").aggregate(sales=filter_table["city"].nunique())
+    expr = (
+        filter_table.group_by("city")
+        .aggregate(sales=filter_table["city"].nunique())
+        .filter(ibis.literal(True))
+    )
 
     run_parity_test(request.getfixturevalue(consumer), expr)