Sorting.hs

{-# LANGUAGE OverloadedStrings #-}
module DataFrame.Operations.Sorting where

import qualified Data.List as L
import qualified Data.Text as T
import qualified Data.Vector as V

import Control.Exception (throw)
import DataFrame.Errors (DataFrameException(..))
import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame (DataFrame(..), getColumn)
import DataFrame.Internal.Row
import DataFrame.Operations.Core

-- | Sort order taken as a parameter by the sortby function.
data SortOrder = Ascending | Descending deriving (Eq)

-- | O(k log n) Sorts the dataframe by a given row.
--
-- > sortBy "Age" df
sortBy ::
  SortOrder ->
  [T.Text] ->
  DataFrame ->
  DataFrame
sortBy order names df
  | any (`notElem` columnNames df) names = throw $ ColumnNotFoundException (T.pack $ show $ names L.\\ columnNames df) "sortBy" (columnNames df)
  | otherwise = let
      -- TODO: Remove the SortOrder defintion from operations so we can share it between here and internal and
      -- we don't have to do this Bool mapping.
      indexes = sortedIndexes' (order == Ascending) (toRowVector names df)
      pick idxs col = atIndicesStable idxs col
    in df {columns = V.map (pick indexes) (columns df)}