Transformations.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module DataFrame.Operations.Transformations where

import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Unboxed as VU

import Control.Exception (throw)
import Data.Maybe
import DataFrame.Errors (DataFrameException (..), TypeErrorContext (..))
import DataFrame.Internal.Column (Column (..), Columnable, TypedColumn (TColumn), columnTypeString, ifoldrColumn, imapColumn, mapColumn, unwrapTypedColumn)
import DataFrame.Internal.DataFrame (DataFrame (..), getColumn)
import DataFrame.Internal.Expression
import DataFrame.Internal.Row (Any, mkRowFromArgs, toAny)
import DataFrame.Operations.Core
import Type.Reflection (TypeRep, typeOf, typeRep)

-- | O(k) Apply a function to a given column in a dataframe.
apply ::
    forall b c.
    (Columnable b, Columnable c) =>
    -- | function to apply
    (b -> c) ->
    -- | Column name
    T.Text ->
    -- | DataFrame to apply operation to
    DataFrame ->
    DataFrame
apply f columnName d = case safeApply f columnName d of
    Left exception -> throw exception
    Right df -> df

-- | O(k) Safe version of the apply function. Returns (instead of throwing) the error.
safeApply ::
    forall b c.
    (Columnable b, Columnable c) =>
    -- | function to apply
    (b -> c) ->
    -- | Column name
    T.Text ->
    -- | DataFrame to apply operation to
    DataFrame ->
    Either DataFrameException DataFrame
safeApply f columnName d = case getColumn columnName d of
    Nothing -> Left $ ColumnNotFoundException columnName "apply" (map fst $ M.toList $ columnIndices d)
    Just column -> case mapColumn f column of
        Nothing ->
            Left $
                TypeMismatchException
                    ( MkTypeErrorContext
                        { userType = Right $ typeRep @b
                        , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
                        , errorColumnName = Just (T.unpack columnName)
                        , callingFunctionName = Just "apply"
                        }
                    )
        Just column' -> Right $ insertColumn columnName column' d

{- | O(k) Apply a function to a combination of columns in a dataframe and
add the result into `alias` column.
-}
derive :: forall a. (Columnable a) => T.Text -> Expr a -> DataFrame -> DataFrame
derive name expr df =
    let
        value = interpret @a df expr
     in
        insertColumn name (unwrapTypedColumn value) df

-- | O(k * n) Apply a function to given column names in a dataframe.
applyMany ::
    (Columnable b, Columnable c) =>
    (b -> c) ->
    [T.Text] ->
    DataFrame ->
    DataFrame
applyMany f names df = L.foldl' (flip (apply f)) df names

-- | O(k) Convenience function that applies to an int column.
applyInt ::
    (Columnable b) =>
    {- | Column name
    | function to apply
    -}
    (Int -> b) ->
    T.Text ->
    -- | DataFrame to apply operation to
    DataFrame ->
    DataFrame
applyInt = apply

-- | O(k) Convenience function that applies to an double column.
applyDouble ::
    (Columnable b) =>
    {- | Column name
    | function to apply
    -}
    (Double -> b) ->
    T.Text ->
    -- | DataFrame to apply operation to
    DataFrame ->
    DataFrame
applyDouble = apply

{- | O(k * n) Apply a function to a column only if there is another column
value that matches the given criterion.

> applyWhere "Age" (<20) "Generation" (const "Gen-Z")
-}
applyWhere ::
    forall a b.
    (Columnable a, Columnable b) =>
    (a -> Bool) -> -- Filter condition
    T.Text -> -- Criterion Column
    (b -> b) -> -- function to apply
    T.Text -> -- Column name
    DataFrame -> -- DataFrame to apply operation to
    DataFrame
applyWhere condition filterColumnName f columnName df = case getColumn filterColumnName df of
    Nothing -> throw $ ColumnNotFoundException filterColumnName "applyWhere" (map fst $ M.toList $ columnIndices df)
    Just column -> case ifoldrColumn (\i val acc -> if condition val then V.cons i acc else acc) V.empty column of
        Nothing ->
            throw $
                TypeMismatchException
                    ( MkTypeErrorContext
                        { userType = Right $ typeRep @a
                        , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
                        , errorColumnName = Just (T.unpack columnName)
                        , callingFunctionName = Just "applyWhere"
                        }
                    )
        Just indexes ->
            if V.null indexes
                then df
                else L.foldl' (\d i -> applyAtIndex i f columnName d) df indexes

-- | O(k) Apply a function to the column at a given index.
applyAtIndex ::
    forall a.
    (Columnable a) =>
    -- | Index
    Int ->
    -- | function to apply
    (a -> a) ->
    -- | Column name
    T.Text ->
    -- | DataFrame to apply operation to
    DataFrame ->
    DataFrame
applyAtIndex i f columnName df = case getColumn columnName df of
    Nothing -> throw $ ColumnNotFoundException columnName "applyAtIndex" (map fst $ M.toList $ columnIndices df)
    Just column -> case imapColumn (\index value -> if index == i then f value else value) column of
        Nothing ->
            throw $
                TypeMismatchException
                    ( MkTypeErrorContext
                        { userType = Right $ typeRep @a
                        , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
                        , errorColumnName = Just (T.unpack columnName)
                        , callingFunctionName = Just "applyAtIndex"
                        }
                    )
        Just column' -> insertColumn columnName column' df

-- | Replace all instances of `Nothing` in a column with the given value.
impute ::
    forall b.
    (Columnable b) =>
    T.Text ->
    b ->
    DataFrame ->
    DataFrame
impute columnName value df = case getColumn columnName df of
    Nothing -> throw $ ColumnNotFoundException columnName "impute" (map fst $ M.toList $ columnIndices df)
    Just (OptionalColumn _) -> case safeApply (fromMaybe value) columnName df of
        Left (TypeMismatchException context) -> throw $ TypeMismatchException (context{callingFunctionName = Just "impute"})
        Left exception -> throw exception
        Right res -> res
    _ -> error "Cannot impute to a non-Empty column"