feat: sortBy now works with compound columns

Author: mchav

src/DataFrame/Operations/Permutation.hs

@@ -22,8 +22,9 @@ import Data.Vector.Internal.Check (HasCallStack)
 import DataFrame.Errors (DataFrameException (..))
 import DataFrame.Internal.Column (Column (..), Columnable, atIndicesStable)
 import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn)
-import DataFrame.Internal.Expression (Expr (Col))
+import DataFrame.Internal.Expression (Expr (Col), getColumns)
 import DataFrame.Operations.Core (columnNames, dimensions)
+import DataFrame.Operations.Transformations (derive)
 import System.Random (Random (randomR), RandomGen)
 import Type.Reflection (typeRep)
 
@@ -38,15 +39,40 @@ instance Eq SortOrder where
     (==) (Desc _) (Desc _) = True
     (==) _ _ = False
 
-getSortColumnName :: SortOrder -> T.Text
-getSortColumnName (Asc (Col n)) = n
-getSortColumnName (Desc (Col n)) = n
-getSortColumnName _ = error "Sorting on compound column"
+sortOrderColumns :: SortOrder -> [T.Text]
+sortOrderColumns (Asc e) = getColumns e
+sortOrderColumns (Desc e) = getColumns e
 
 mustFlipCompare :: SortOrder -> Bool
 mustFlipCompare (Asc _) = True
 mustFlipCompare (Desc _) = False
 
+{- | Materialize any compound sort expressions into synthetic columns on
+a working dataframe, returning rewritten 'SortOrder's that reference
+those columns by name.
+-}
+prepareSortColumns :: [SortOrder] -> DataFrame -> ([SortOrder], DataFrame)
+prepareSortColumns = go 0
+  where
+    go _ [] acc = ([], acc)
+    go i (ord : rest) acc =
+        let (ord', acc') = materializeSortOrder i ord acc
+            (rest', acc'') = go (i + 1) rest acc'
+         in (ord' : rest', acc'')
+
+materializeSortOrder :: Int -> SortOrder -> DataFrame -> (SortOrder, DataFrame)
+materializeSortOrder _ ord@(Asc (Col _)) df = (ord, df)
+materializeSortOrder _ ord@(Desc (Col _)) df = (ord, df)
+materializeSortOrder i (Asc (e :: Expr a)) df =
+    let name = syntheticName i
+     in (Asc (Col name :: Expr a), derive name e df)
+materializeSortOrder i (Desc (e :: Expr a)) df =
+    let name = syntheticName i
+     in (Desc (Col name :: Expr a), derive name e df)
+
+syntheticName :: Int -> T.Text
+syntheticName i = "__sortBy_synthetic_" <> T.pack (show i) <> "__"
+
 {- | O(k log n) Sorts the dataframe by a given row.
 
 > sortBy Ascending ["Age"] df
@@ -56,22 +82,24 @@ sortBy ::
     DataFrame ->
     DataFrame
 sortBy sortOrds df
-    | any (`notElem` columnNames df) names =
+    | not (null missing) =
         throw $
             ColumnsNotFoundException
-                (names L.\\ columnNames df)
+                missing
                 "sortBy"
                 (columnNames df)
     | otherwise =
         let
-            comparators = map (`sortOrderComparator` df) sortOrds
+            (sortOrds', df') = prepareSortColumns sortOrds df
+            comparators = map (`sortOrderComparator` df') sortOrds'
             compositeCompare i j = mconcat [c i j | c <- comparators]
-            nRows = fst (dataframeDimensions df)
+            nRows = fst (dataframeDimensions df')
             indexes = sortIndices compositeCompare nRows
          in
             df{columns = V.map (atIndicesStable indexes) (columns df)}
   where
-    names = map getSortColumnName sortOrds
+    referenced = L.nub (concatMap sortOrderColumns sortOrds)
+    missing = referenced L.\\ columnNames df
 
 {- | Build a row-index comparator from a SortOrder and a DataFrame.
 The Ord dictionary is recovered from the SortOrder GADT.

tests/Operations/Sort.hs

@@ -89,11 +89,83 @@ sortByColumnDoesNotExist =
             (print $ D.sortBy [D.Asc (F.col @Int "test0")] testData)
         )
 
+compoundTestData :: D.DataFrame
+compoundTestData =
+    D.fromNamedColumns
+        [ ("a", DI.fromList ([3, 1, 4, 1, 5] :: [Int]))
+        , ("b", DI.fromList ([10, 40, 20, 30, 50] :: [Int]))
+        ]
+
+sortByCompoundExpression :: Test
+sortByCompoundExpression =
+    TestCase
+        ( assertEqual
+            "Sorting by Asc (a + b) orders rows by the sum without leaking a synthetic column"
+            ( D.fromNamedColumns
+                [ ("a", DI.fromList ([3, 4, 1, 1, 5] :: [Int]))
+                , ("b", DI.fromList ([10, 20, 30, 40, 50] :: [Int]))
+                ]
+            )
+            (D.sortBy [D.Asc (F.col @Int "a" + F.col @Int "b")] compoundTestData)
+        )
+
+sortByCompoundExpressionDescending :: Test
+sortByCompoundExpressionDescending =
+    TestCase
+        ( assertEqual
+            "Sorting by Desc (b - a) orders rows by descending difference"
+            ( D.fromNamedColumns
+                [ ("a", DI.fromList ([5, 1, 1, 4, 3] :: [Int]))
+                , ("b", DI.fromList ([50, 40, 30, 20, 10] :: [Int]))
+                ]
+            )
+            (D.sortBy [D.Desc (F.col @Int "b" - F.col @Int "a")] compoundTestData)
+        )
+
+sortByCompoundMixedWithBareColumn :: Test
+sortByCompoundMixedWithBareColumn =
+    TestCase
+        ( assertEqual
+            "Mixing a compound Asc key with a bare Desc tie-breaker works"
+            ( D.fromNamedColumns
+                [ ("a", DI.fromList ([1, 1, 3, 4, 5] :: [Int]))
+                , ("b", DI.fromList ([40, 30, 10, 20, 50] :: [Int]))
+                ]
+            )
+            ( D.sortBy
+                [D.Asc (F.col @Int "a" * 2), D.Desc (F.col @Int "b")]
+                compoundTestData
+            )
+        )
+
+sortByCompoundMissingColumn :: Test
+sortByCompoundMissingColumn =
+    TestCase
+        ( assertExpectException
+            "[Error Case]"
+            ( D.columnsNotFound
+                ["nope"]
+                "sortBy"
+                (D.columnNames compoundTestData)
+            )
+            ( print $
+                D.sortBy
+                    [D.Asc (F.col @Int "nope" + F.col @Int "a")]
+                    compoundTestData
+            )
+        )
+
 tests :: [Test]
 tests =
     [ TestLabel "sortByAscendingWAI" sortByAscendingWAI
     , TestLabel "sortByDescendingWAI" sortByDescendingWAI
     , TestLabel "sortByColumnDoesNotExist" sortByColumnDoesNotExist
     , TestLabel "sortByTwoColumns" sortByTwoColumns
     , TestLabel "sortByOneColumnAscOneColumnDesc" sortByOneColumnAscOneColumnDesc
+    , TestLabel "sortByCompoundExpression" sortByCompoundExpression
+    , TestLabel
+        "sortByCompoundExpressionDescending"
+        sortByCompoundExpressionDescending
+    , TestLabel "sortByCompoundMixedWithBareColumn" sortByCompoundMixedWithBareColumn
+    , TestLabel "sortByCompoundMissingColumn" sortByCompoundMissingColumn
     ]