GroupBy.rowNumber() #16

* testing combination of sort and row number
dflib · Mar 31, 2019 · 6c952dc · 6c952dc
1 parent 28f090f
commit 6c952dc
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/dflib/src/main/java/com/nhl/dflib/GroupBy.java b/dflib/src/main/java/com/nhl/dflib/GroupBy.java
@@ -92,6 +92,31 @@ public <V extends Comparable<? super V>> GroupBy sort(RowToValueMapper<V> sortKe
         return new GroupBy(ungrouped, sorted);
     }
 
+    public GroupBy sort(String column, boolean ascending) {
+
+        Comparator<RowProxy> comparator = Sorters.sorter(ungrouped.getColumnsIndex(), column, ascending);
+        Map<Object, Series<Integer>> sorted = new LinkedHashMap<>((int) (groupsIndex.size() / 0.75));
+
+        for (Map.Entry<Object, Series<Integer>> e : groupsIndex.entrySet()) {
+            Series<Integer> sortedGroup = new IndexSorter(ungrouped, e.getValue()).sortIndex(comparator);
+            sorted.put(e.getKey(), sortedGroup);
+        }
+
+        return new GroupBy(ungrouped, sorted);
+    }
+
+    public GroupBy sort(int column, boolean ascending) {
+        Comparator<RowProxy> comparator = Sorters.sorter(ungrouped.getColumnsIndex(), column, ascending);
+        Map<Object, Series<Integer>> sorted = new LinkedHashMap<>((int) (groupsIndex.size() / 0.75));
+
+        for (Map.Entry<Object, Series<Integer>> e : groupsIndex.entrySet()) {
+            Series<Integer> sortedGroup = new IndexSorter(ungrouped, e.getValue()).sortIndex(comparator);
+            sorted.put(e.getKey(), sortedGroup);
+        }
+
+        return new GroupBy(ungrouped, sorted);
+    }
+
     public GroupBy sort(String[] columns, boolean[] ascending) {
         if (columns.length == 0) {
             return this;

diff --git a/dflib/src/test/java/com/nhl/dflib/DataFrame_GroupBy_WindowFuncsTest.java b/dflib/src/test/java/com/nhl/dflib/DataFrame_GroupBy_WindowFuncsTest.java
@@ -46,4 +46,18 @@ public void testGroupBy_RowNumbers2() {
         Series<Integer> rn = df.group("a").rowNumbers();
         new SeriesAsserts(rn).expectData(0, 0, 1, 2, 0);
     }
+
+    @Test
+    public void testGroupBy_RowNumbers_Sort() {
+        Index i = Index.forLabels("a", "b", "c");
+        DataFrame df = createDf(i,
+                3, "x", "m",
+                0, "y", "n",
+                3, "z", "k",
+                3, "a", "f",
+                1, "x", "s");
+
+        Series<Integer> rn = df.group("a").sort("b", true).rowNumbers();
+        new SeriesAsserts(rn).expectData(1, 0, 2, 0, 0);
+    }
 }