Add sortby tests and cleanup test structure.

Author: mchav

dataframe.cabal

@@ -106,13 +106,14 @@ benchmark dataframe-benchmark
 
 test-suite tests
     type: exitcode-stdio-1.0
-    main-is: DataFrameTests.hs
+    main-is: Main.hs
     other-modules: Assertions,
-                   AddColumn,
-                   Apply,
-                   Filter,
-                   Sort,
-                   Take
+                   Operations.AddColumn,
+                   Operations.Apply,
+                   Operations.Filter,
+                   Operations.GroupBy,
+                   Operations.Sort,
+                   Operations.Take
     build-depends: base >= 4.17.2.0 && < 4.21,
                    HUnit ^>= 1.6,
                    random >= 1,

src/Data/DataFrame/Internal.hs

@@ -11,6 +11,7 @@
 
 module Data.DataFrame.Internal where
 
+import qualified Data.List as L
 import qualified Data.Map as M
 import qualified Data.Map.Strict as MS
 import qualified Data.Set as S
@@ -183,14 +184,16 @@ instance Eq Column where
     case testEquality (typeRep @t1) (typeRep @t2) of
       Nothing -> False
       Just Refl -> a == b
+  -- Note: comparing grouped columns is expensive. We do this for stable tests
+  -- but also you should probably aggregate grouped columns soon after creating them.
   (==) (GroupedBoxedColumn (a :: V.Vector t1)) (GroupedBoxedColumn (b :: V.Vector t2)) =
     case testEquality (typeRep @t1) (typeRep @t2) of
       Nothing -> False
-      Just Refl -> a == b
+      Just Refl -> V.map (L.sort . VG.toList) a == V.map (L.sort . VG.toList) b
   (==) (GroupedUnboxedColumn (a :: V.Vector t1)) (GroupedUnboxedColumn (b :: V.Vector t2)) =
     case testEquality (typeRep @t1) (typeRep @t2) of
       Nothing -> False
-      Just Refl -> a == b
+      Just Refl -> V.map (L.sort . VG.toList) a == V.map (L.sort . VG.toList) b
   (==) _ _ = False
 
 -- Traversing columns.

tests/DataFrameTests.hs tests/Main.hstests/DataFrameTests.hs renamed to tests/Main.hs

@@ -17,11 +17,12 @@ import Test.HUnit
 
 import Assertions
 
-import qualified AddColumn
-import qualified Apply
-import qualified Filter
-import qualified Sort
-import qualified Take
+import qualified Operations.AddColumn
+import qualified Operations.Apply
+import qualified Operations.Filter
+import qualified Operations.GroupBy
+import qualified Operations.Sort
+import qualified Operations.Take
 
 testData :: D.DataFrame
 testData = D.fromList [ ("test1", DI.toColumn ([1..26] :: [Int]))
@@ -68,11 +69,12 @@ parseTests = [
 
 tests :: Test
 tests = TestList $ dimensionsTest
-                ++ AddColumn.tests
-                ++ Apply.tests
-                ++ Filter.tests
-                ++ Sort.tests
-                ++ Take.tests
+                ++ Operations.AddColumn.tests
+                ++ Operations.Apply.tests
+                ++ Operations.Filter.tests
+                ++ Operations.GroupBy.tests
+                ++ Operations.Sort.tests
+                ++ Operations.Take.tests
                 ++ parseTests
 
 main :: IO ()

tests/AddColumn.hs tests/Operations/AddColumn.hstests/AddColumn.hs renamed to tests/Operations/AddColumn.hs

@@ -1,6 +1,6 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE OverloadedStrings #-}
-module AddColumn where
+module Operations.AddColumn where
 
 import qualified Data.DataFrame as D
 import qualified Data.DataFrame.Internal as DI

tests/Apply.hs tests/Operations/Apply.hstests/Apply.hs renamed to tests/Operations/Apply.hs

@@ -1,7 +1,7 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE TupleSections #-}
-module Apply where
+module Operations.Apply where
 
 import qualified Data.DataFrame as D
 import qualified Data.DataFrame.Internal as DI

tests/Filter.hs tests/Operations/Filter.hstests/Filter.hs renamed to tests/Operations/Filter.hs

@@ -1,6 +1,6 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE OverloadedStrings #-}
-module Filter where
+module Operations.Filter where
 
 import qualified Data.DataFrame as D
 import qualified Data.DataFrame.Internal as DI

tests/Operations/GroupBy.hs

@@ -0,0 +1,40 @@
+{-# LANGUAGE OverloadedStrings #-}
+module Operations.GroupBy where
+
+import qualified Data.DataFrame as D
+import qualified Data.DataFrame.Internal as DI
+import qualified Data.DataFrame.Errors as DE
+import qualified Data.DataFrame.Operations as DO
+import qualified Data.DataFrame.Util as DU
+import qualified Data.Text as T
+import qualified Data.Vector as V
+import qualified Data.Vector.Unboxed as VU
+
+import Assertions
+import Test.HUnit
+
+values :: [(T.Text, DI.Column)]
+values = [ ("test1", DI.toColumn (concatMap (replicate 10) [1 :: Int, 2, 3, 4]))
+         , ("test2", DI.toColumn (take 40 $ cycle [1 :: Int,2]))
+         , ("test3", DI.toColumn [(1 :: Int)..40])
+         , ("test4", DI.toColumn (reverse [(1 :: Int)..40]))
+         ]
+
+testData :: D.DataFrame
+testData = D.fromList values
+
+groupBySingleRowWAI :: Test
+groupBySingleRowWAI = TestCase (assertEqual "Groups by single column"
+                (D.fromList [("test1", DI.toColumn [(1::Int)..4]),
+                             -- This just makes rows with [1, 2] for every unique test1 row
+                             ("test2", DI.GroupedUnboxedColumn (V.replicate 4 $ VU.fromList (take 10 $ cycle [1 :: Int, 2]))),
+                             ("test3", DI.GroupedUnboxedColumn (V.generate 4 (\i -> VU.fromList [(i * 10 + 1)..((i + 1) * 10)]))),
+                             ("test4", DI.GroupedUnboxedColumn (V.generate 4 (\i -> VU.fromList [(((3 - i) + 1) * 10),(((3 - i) + 1) * 10 - 1)..((3 - i) * 10 + 1)])))
+                            ])
+                (D.groupBy ["test1"] testData D.|> D.sortBy D.Ascending "test1"))
+
+tests :: [Test]
+tests = [ TestLabel "groupBySingleRowWAI" groupBySingleRowWAI
+        , TestLabel "groupBySingleRowWAI" groupBySingleRowWAI
+        ]
+

tests/Sort.hs tests/Operations/Sort.hstests/Sort.hs renamed to tests/Operations/Sort.hs

@@ -1,6 +1,5 @@
-{-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE OverloadedStrings #-}
-module Sort where
+module Operations.Sort where
 
 import qualified Data.DataFrame as D
 import qualified Data.DataFrame.Internal as DI
@@ -11,7 +10,6 @@ import qualified Data.Vector as V
 import qualified Data.Vector.Unboxed as VU
 
 import Assertions
-
 import Control.Monad
 import Data.Char
 import System.Random
@@ -29,13 +27,13 @@ testData :: D.DataFrame
 testData = D.fromList values
 
 sortByAscendingWAI :: Test
-sortByAscendingWAI = TestCase (assertEqual "Non existent filter value returns no rows"
+sortByAscendingWAI = TestCase (assertEqual "Sorting rows by ascending works as intended"
                     (D.fromList [("test1", DI.toColumn [(1::Int)..26]),
                                  ("test2", DI.toColumn ['a'..'z'])])
                     (D.sortBy D.Ascending "test1" testData))
 
 sortByDescendingWAI :: Test
-sortByDescendingWAI = TestCase (assertEqual "Non existent filter value returns no rows"
+sortByDescendingWAI = TestCase (assertEqual "Sorting rows by descending works as intended"
                     (D.fromList [("test1", DI.toColumn $ reverse [(1::Int)..26]),
                                  ("test2", DI.toColumn $ reverse ['a'..'z'])])
                     (D.sortBy D.Descending "test1" testData))

tests/Take.hs tests/Operations/Take.hstests/Take.hs renamed to tests/Operations/Take.hs

@@ -1,5 +1,5 @@
 {-# LANGUAGE OverloadedStrings #-}
-module Take where
+module Operations.Take where
 
 import qualified Data.DataFrame as D
 import qualified Data.DataFrame.Internal as DI