dataframe-0.3.1.1: src/DataFrame/Operations/Core.hs
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
module DataFrame.Operations.Core 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
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.Either
import Data.Function (on, (&))
import Data.Maybe
import Data.Type.Equality (TestEquality (..))
import DataFrame.Errors
import DataFrame.Internal.Column (Column (..), Columnable, columnLength, columnTypeString, expandColumn, fromList, fromVector)
import DataFrame.Internal.DataFrame (DataFrame (..), empty, getColumn)
import DataFrame.Internal.Parsing (isNullish)
import Type.Reflection
import Prelude hiding (null)
{- | O(1) Get DataFrame dimensions i.e. (rows, columns)
==== __Example__
@
ghci> D.dimensions df
(100, 3)
@
-}
dimensions :: DataFrame -> (Int, Int)
dimensions = dataframeDimensions
{-# INLINE dimensions #-}
{- | O(k) Get column names of the DataFrame in order of insertion.
==== __Example__
@
ghci> D.columnNames df
["col_a", "col_b", "col_c"]
@
-}
columnNames :: DataFrame -> [T.Text]
columnNames = map fst . L.sortBy (compare `on` snd) . M.toList . columnIndices
{-# INLINE columnNames #-}
{- | Adds a vector to the dataframe. If the vector has less elements than the dataframe and the dataframe is not empty
the vector is converted to type `Maybe a` filled with `Nothing` to match the size of the dataframe. Similarly,
if the vector has more elements than what's currently in the dataframe, the other columns in the dataframe are
change to `Maybe <Type>` and filled with `Nothing`.
==== __Example__
@
ghci> import qualified Data.Vector as V
ghci> D.insertVector "numbers" (V.fromList [1..10]) D.empty
---------------
index | numbers
------|--------
Int | Int
------|--------
0 | 1
1 | 2
2 | 3
3 | 4
4 | 5
5 | 6
6 | 7
7 | 8
8 | 9
9 | 10
@
-}
insertVector ::
forall a.
(Columnable a) =>
-- | Column Name
T.Text ->
-- | Vector to add to column
V.Vector a ->
-- | DataFrame to add column to
DataFrame ->
DataFrame
insertVector name xs = insertColumn name (fromVector xs)
{-# INLINE insertVector #-}
{- | /O(k)/ Add a column to the dataframe providing a default.
This constructs a new vector and also may convert it
to an unboxed vector if necessary. Since columns are usually
large the runtime is dominated by the length of the list, k.
-}
insertVectorWithDefault ::
forall a.
(Columnable a) =>
-- | Default Value
a ->
-- | Column name
T.Text ->
-- | Data to add to column
V.Vector a ->
-- | DataFrame to add the column to
DataFrame ->
DataFrame
insertVectorWithDefault defaultValue name xs d =
let (rows, _) = dataframeDimensions d
values = xs V.++ V.replicate (rows - V.length xs) defaultValue
in insertColumn name (fromVector values) d
{- | /O(n)/ Adds an unboxed vector to the dataframe.
Same as insertVector but takes an unboxed vector. If you insert a vector of numbers through insertVector it will either way be converted
into an unboxed vector so this function saves that extra work/conversion.
-}
insertUnboxedVector ::
forall a.
(Columnable a, VU.Unbox a) =>
-- | Column Name
T.Text ->
-- | Unboxed vector to add to column
VU.Vector a ->
-- | DataFrame to add the column to
DataFrame ->
DataFrame
insertUnboxedVector name xs = insertColumn name (UnboxedColumn xs)
{- | /O(n)/ Add a column to the dataframe.
==== __Example__
@
ghci> D.insertColumn "numbers" (D.fromList [1..10]) D.empty
---------------
index | numbers
------|--------
Int | Int
------|--------
0 | 1
1 | 2
2 | 3
3 | 4
4 | 5
5 | 6
6 | 7
7 | 8
8 | 9
9 | 10
@
-}
insertColumn ::
-- | Column Name
T.Text ->
-- | Column to add
Column ->
-- | DataFrame to add the column to
DataFrame ->
DataFrame
insertColumn name column d =
let
(r, c) = dataframeDimensions d
n = max (columnLength column) r
in
case M.lookup name (columnIndices d) of
Just i -> DataFrame (V.map (expandColumn n) (columns d V.// [(i, column)])) (columnIndices d) (n, c)
Nothing -> DataFrame (V.map (expandColumn n) (columns d `V.snoc` column)) (M.insert name c (columnIndices d)) (n, c + 1)
{- | /O(n)/ Clones a column and places it under a new name in the dataframe.
==== __Example__
@
ghci> import qualified Data.Vector as V
ghci> df = insertVector "numbers" (V.fromList [1..10]) D.empty
ghci> D.cloneColumn "numbers" "others" df
------------------------
index | numbers | others
------|---------|-------
Int | Int | Int
------|---------|-------
0 | 1 | 1
1 | 2 | 2
2 | 3 | 3
3 | 4 | 4
4 | 5 | 5
5 | 6 | 6
6 | 7 | 7
7 | 8 | 8
8 | 9 | 9
9 | 10 | 10
@
-}
cloneColumn :: T.Text -> T.Text -> DataFrame -> DataFrame
cloneColumn original new df = fromMaybe (throw $ ColumnNotFoundException original "cloneColumn" (M.keys $ columnIndices df)) $ do
column <- getColumn original df
return $ insertColumn new column df
{- | /O(n)/ Renames a single column.
==== __Example__
@
ghci> import qualified Data.Vector as V
ghci> df = insertVector "numbers" (V.fromList [1..10]) D.empty
ghci> D.rename "numbers" "others" df
--------------
index | others
------|-------
Int | Int
------|-------
0 | 1
1 | 2
2 | 3
3 | 4
4 | 5
5 | 6
6 | 7
7 | 8
8 | 9
9 | 10
@
-}
rename :: T.Text -> T.Text -> DataFrame -> DataFrame
rename orig new df = either throw id (renameSafe orig new df)
{- | /O(n)/ Renames many columns.
==== __Example__
@
ghci> import qualified Data.Vector as V
ghci> df = D.insertVector "others" (V.fromList [11..20]) (D.insertVector "numbers" (V.fromList [1..10]) D.empty)
ghci> df
------------------------
index | numbers | others
------|---------|-------
Int | Int | Int
------|---------|-------
0 | 1 | 11
1 | 2 | 12
2 | 3 | 13
3 | 4 | 14
4 | 5 | 15
5 | 6 | 16
6 | 7 | 17
7 | 8 | 18
8 | 9 | 19
9 | 10 | 20
ghci> D.renameMany [("numbers", "first_10"), ("others", "next_10")] df
--------------------------
index | first_10 | next_10
------|----------|--------
Int | Int | Int
------|----------|--------
0 | 1 | 11
1 | 2 | 12
2 | 3 | 13
3 | 4 | 14
4 | 5 | 15
5 | 6 | 16
6 | 7 | 17
7 | 8 | 18
8 | 9 | 19
9 | 10 | 20
@
-}
renameMany :: [(T.Text, T.Text)] -> DataFrame -> DataFrame
renameMany replacements df = fold (uncurry rename) replacements df
renameSafe :: T.Text -> T.Text -> DataFrame -> Either DataFrameException DataFrame
renameSafe orig new df = fromMaybe (Left $ ColumnNotFoundException orig "rename" (M.keys $ columnIndices df)) $ do
columnIndex <- M.lookup orig (columnIndices df)
let origRemoved = M.delete orig (columnIndices df)
let newAdded = M.insert new columnIndex origRemoved
return (Right df{columnIndices = newAdded})
data ColumnInfo = ColumnInfo
{ nameOfColumn :: !T.Text
, nonNullValues :: !Int
, nullValues :: !Int
, partiallyParsedValues :: !Int
, uniqueValues :: !Int
, typeOfColumn :: !T.Text
}
{- | O(n * k ^ 2) Returns the number of non-null columns in the dataframe and the type associated with each column.
==== __Example__
@
ghci> import qualified Data.Vector as V
ghci> df = D.insertVector "others" (V.fromList [11..20]) (D.insertVector "numbers" (V.fromList [1..10]) D.empty)
ghci> D.describeColumns df
-----------------------------------------------------------------------------------------------------
index | Column Name | # Non-null Values | # Null Values | # Partially parsed | # Unique Values | Type
------|-------------|-------------------|---------------|--------------------|-----------------|-----
Int | Text | Int | Int | Int | Int | Text
------|-------------|-------------------|---------------|--------------------|-----------------|-----
0 | others | 10 | 0 | 0 | 10 | Int
1 | numbers | 10 | 0 | 0 | 10 | Int
@
-}
describeColumns :: DataFrame -> DataFrame
describeColumns df =
empty
& insertColumn "Column Name" (fromList (map nameOfColumn infos))
& insertColumn "# Non-null Values" (fromList (map nonNullValues infos))
& insertColumn "# Null Values" (fromList (map nullValues infos))
& insertColumn "# Partially parsed" (fromList (map partiallyParsedValues infos))
& insertColumn "# Unique Values" (fromList (map uniqueValues infos))
& insertColumn "Type" (fromList (map typeOfColumn infos))
where
infos = L.sortBy (compare `on` nonNullValues) (V.ifoldl' go [] (columns df)) :: [ColumnInfo]
indexMap = M.fromList (map (\(a, b) -> (b, a)) $ M.toList (columnIndices df))
columnName i = M.lookup i indexMap
go acc i col@(OptionalColumn (c :: V.Vector a)) =
let
cname = columnName i
countNulls = nulls col
countPartial = partiallyParsed col
columnType = T.pack $ show $ typeRep @a
unique = S.size $ VG.foldr S.insert S.empty c
in
if isNothing cname then acc else ColumnInfo (fromMaybe "" cname) (columnLength col - countNulls) countNulls countPartial unique columnType : acc
go acc i col@(BoxedColumn (c :: V.Vector a)) =
let
cname = columnName i
countPartial = partiallyParsed col
columnType = T.pack $ show $ typeRep @a
unique = S.size $ VG.foldr S.insert S.empty c
in
if isNothing cname then acc else ColumnInfo (fromMaybe "" cname) (columnLength col) 0 countPartial unique columnType : acc
go acc i col@(UnboxedColumn c) =
let
cname = columnName i
columnType = T.pack $ columnTypeString col
unique = S.size $ VG.foldr S.insert S.empty c
in
-- Unboxed columns cannot have nulls since Maybe
-- is not an instance of Unbox a
if isNothing cname then acc else ColumnInfo (fromMaybe "" cname) (columnLength col) 0 0 unique columnType : acc
nulls :: Column -> Int
nulls (OptionalColumn xs) = VG.length $ VG.filter isNothing xs
nulls (BoxedColumn (xs :: V.Vector a)) = case testEquality (typeRep @a) (typeRep @T.Text) of
Just Refl -> VG.length $ VG.filter isNullish xs
Nothing -> case testEquality (typeRep @a) (typeRep @String) of
Just Refl -> VG.length $ VG.filter (isNullish . T.pack) xs
Nothing -> case typeRep @a of
App t1 t2 -> case eqTypeRep t1 (typeRep @Maybe) of
Just HRefl -> VG.length $ VG.filter isNothing xs
Nothing -> 0
_ -> 0
nulls _ = 0
partiallyParsed :: Column -> Int
partiallyParsed (BoxedColumn (xs :: V.Vector a)) =
case typeRep @a of
App (App tycon t1) t2 -> case eqTypeRep tycon (typeRep @Either) of
Just HRefl -> VG.length $ VG.filter isLeft xs
Nothing -> 0
_ -> 0
partiallyParsed _ = 0
{- | Creates a dataframe from a list of tuples with name and column.
==== __Example__
@
ghci> df = D.fromNamedColumns [("numbers", D.fromList [1..10]), ("others", D.fromList [11..20])]
ghci> df
------------------------
index | numbers | others
------|---------|-------
Int | Int | Int
------|---------|-------
0 | 1 | 11
1 | 2 | 12
2 | 3 | 13
3 | 4 | 14
4 | 5 | 15
5 | 6 | 16
6 | 7 | 17
7 | 8 | 18
8 | 9 | 19
9 | 10 | 20
@
-}
fromNamedColumns :: [(T.Text, Column)] -> DataFrame
fromNamedColumns = L.foldl' (\df (name, column) -> insertColumn name column df) empty
{- | Create a dataframe from a list of columns. The column names are "0", "1"... etc.
Useful for quick exploration but you should probably alwyas rename the columns after
or drop the ones you don't want.
==== __Example__
@
ghci> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]]
ghci> df
-----------------
index | 0 | 1
------|-----|----
Int | Int | Int
------|-----|----
0 | 1 | 11
1 | 2 | 12
2 | 3 | 13
3 | 4 | 14
4 | 5 | 15
5 | 6 | 16
6 | 7 | 17
7 | 8 | 18
8 | 9 | 19
9 | 10 | 20
@
-}
fromUnnamedColumns :: [Column] -> DataFrame
fromUnnamedColumns = fromNamedColumns . zip (map (T.pack . show) [0 ..])
{- | O (k * n) Counts the occurences of each value in a given column.
==== __Example__
@
ghci> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]]
ghci> D.valueCounts @Int "0" df
[(1,1),(2,1),(3,1),(4,1),(5,1),(6,1),(7,1),(8,1),(9,1),(10,1)]
@
-}
valueCounts :: forall a. (Columnable a) => T.Text -> DataFrame -> [(a, Int)]
valueCounts columnName df = case getColumn columnName df of
Nothing -> throw $ ColumnNotFoundException columnName "valueCounts" (M.keys $ columnIndices df)
Just (BoxedColumn (column' :: V.Vector c)) ->
let
column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column'
in
case (typeRep @a) `testEquality` (typeRep @c) of
Nothing ->
throw $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right $ typeRep @a
, expectedType = Right $ typeRep @c
, errorColumnName = Just (T.unpack columnName)
, callingFunctionName = Just "valueCounts"
}
)
Just Refl -> M.toAscList column
Just (OptionalColumn (column' :: V.Vector c)) ->
let
column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column'
in
case (typeRep @a) `testEquality` (typeRep @c) of
Nothing ->
throw $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right $ typeRep @a
, expectedType = Right $ typeRep @c
, errorColumnName = Just (T.unpack columnName)
, callingFunctionName = Just "valueCounts"
}
)
Just Refl -> M.toAscList column
Just (UnboxedColumn (column' :: VU.Vector c)) ->
let
column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty (V.convert column')
in
case (typeRep @a) `testEquality` (typeRep @c) of
Nothing ->
throw $
TypeMismatchException
( MkTypeErrorContext
{ userType = Right $ typeRep @a
, expectedType = Right $ typeRep @c
, errorColumnName = Just (T.unpack columnName)
, callingFunctionName = Just "valueCounts"
}
)
Just Refl -> M.toAscList column
{- | A left fold for dataframes that takes the dataframe as the last object.
This makes it easier to chain operations.
==== __Example__
@
ghci> D.fold (const id) [1..5] df
-----------------
index | 0 | 1
------|-----|----
Int | Int | Int
------|-----|----
0 | 1 | 11
1 | 2 | 12
2 | 3 | 13
3 | 4 | 14
4 | 5 | 15
5 | 6 | 16
6 | 7 | 17
7 | 8 | 18
8 | 9 | 19
9 | 10 | 20
@
-}
fold :: (a -> DataFrame -> DataFrame) -> [a] -> DataFrame -> DataFrame
fold f xs acc = L.foldl' (flip f) acc xs