Grouping no longer maintains order.

Test fix to not rely on ordering: apache/arrow#36709. Distinct aggregation supports aliasing.
coady · Jul 18, 2023 · 3cc9a6b · 3cc9a6b
1 parent 412cc0a
commit 3cc9a6b
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 8 deletions.
diff --git a/graphique/interface.py b/graphique/interface.py
@@ -295,8 +295,8 @@ def group(
             for agg in aggs['list']:
                 columns[agg.name] = ListChunk.inner_flatten(*columns[agg.name].chunks)
             for agg in distinct:
-                column = columns[agg.name]
-                columns[agg.name] = ListChunk.map_list(column, agg.distinct, **agg.options)
+                column = columns[agg.alias]
+                columns[agg.alias] = ListChunk.map_list(column, agg.distinct, **agg.options)
         return type(self)(self.apply_list(pa.table(columns), list_func))
 
     @strawberry.field(deprecation_reason="use `group(by: [<fragment key>, ...])`")

diff --git a/tests/test_core.py b/tests/test_core.py
@@ -24,9 +24,7 @@ def test_chunks():
     array = pa.chunked_array([list('aba'), list('bcb')])
     table = pa.table({'col': array})
     groups = T.group(table, 'col', counts='counts')
-    assert groups['col'].to_pylist() == list('abc')
-    assert groups['counts'].to_pylist() == [2, 3, 1]
-    assert table['col'].unique() == pa.array('abc')
+    assert dict(zip(*groups.to_pydict().values())) == {'a': 2, 'b': 3, 'c': 1}
     array = pa.chunked_array([pa.array(list(chunk)).dictionary_encode() for chunk in ('aba', 'ca')])
     assert C.index(array, 'a') == 0
     assert C.index(array, 'c') == 3

diff --git a/tests/test_dataset.py b/tests/test_dataset.py
@@ -122,11 +122,11 @@ def test_list(partclient):
     )
     assert data['group']['tables'][0] == {'length': 1, 'row': {'state': 'PR', 'county': 'Adjuntas'}}
     data = partclient.execute(
-        '''{ group(by: "state", aggregate: {distinct: {name: "county"}}) {
-        tables { row { state } columns { county { length } } } } }'''
+        '''{ group(by: "state", aggregate: {distinct: {alias: "counties", name: "county"}}) {
+        tables { row { state } column(name: "counties") { length } } } } '''
     )
     table = data['group']['tables'][0]
-    assert table == {'row': {'state': 'PR'}, 'columns': {'county': {'length': 78}}}
+    assert table == {'row': {'state': 'PR'}, 'column': {'length': 78}}
     data = partclient.execute(
         '''{ group(by: "north", aggregate: {distinct: {name: "west"}}) {
         tables { row { north } columns { west { length } } } } }'''