InsertColumn.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

module Operations.InsertColumn where

import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import qualified DataFrame as D
import qualified DataFrame.Internal.Column as DI
import qualified DataFrame.Internal.DataFrame as DI

import Test.HUnit

testData :: D.DataFrame
testData =
    D.fromNamedColumns
        [ ("test1", DI.fromList ([1 .. 26] :: [Int]))
        , ("test2", DI.fromList ['a' .. 'z'])
        , ("test3", DI.fromList ([1 .. 26] :: [Int]))
        , ("test4", DI.fromList ['a' .. 'z'])
        , ("test5", DI.fromList ([1 .. 26] :: [Int]))
        , ("test6", DI.fromList ['a' .. 'z'])
        , ("test7", DI.fromList ([1 .. 26] :: [Int]))
        , ("test8", DI.fromList ['a' .. 'z'])
        ]

-- Adding a boxed vector to an empty dataframe creates a new column boxed containing the vector elements.
addBoxedColumn :: Test
addBoxedColumn =
    TestCase
        ( assertEqual
            "Two columns should be equal"
            (Just $ DI.BoxedColumn (V.fromList ["Thuba" :: T.Text, "Zodwa", "Themba"]))
            (DI.getColumn "new" $ D.insertVector "new" (V.fromList ["Thuba" :: T.Text, "Zodwa", "Themba"]) D.empty)
        )

addBoxedColumn' :: Test
addBoxedColumn' =
    TestCase
        ( assertEqual
            "Two columns should be equal"
            (Just $ DI.fromList ["Thuba" :: T.Text, "Zodwa", "Themba"])
            (DI.getColumn "new" $ D.insertColumn "new" (DI.fromList ["Thuba" :: T.Text, "Zodwa", "Themba"]) D.empty)
        )

-- Adding an boxed vector with an unboxable type (Int/Double) to an empty dataframe creates a new column boxed containing the vector elements.
addUnboxedColumn :: Test
addUnboxedColumn =
    TestCase
        ( assertEqual
            "Value should be boxed"
            (Just $ DI.UnboxedColumn (VU.fromList [1 :: Int, 2, 3]))
            (DI.getColumn "new" $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3]) D.empty)
        )

addUnboxedColumn' :: Test
addUnboxedColumn' =
    TestCase
        ( assertEqual
            "Value should be boxed"
            (Just $ DI.fromList [1 :: Int, 2, 3])
            (DI.getColumn "new" $ D.insertColumn "new" (DI.fromList [1 :: Int, 2, 3]) D.empty)
        )

-- Adding a column with less values than the current DF dimensions adds column with optionals.
addSmallerColumnBoxed :: Test
addSmallerColumnBoxed =
    TestCase
        ( assertEqual
            "Missing values should be replaced with Nothing"
            (Just $ DI.OptionalColumn (V.fromList [Just "a" :: Maybe T.Text, Just "b", Just "c", Nothing, Nothing]))
            (DI.getColumn "newer" $ D.insertVector "newer" (V.fromList ["a" :: T.Text, "b", "c"]) $ D.insertVector "new" (V.fromList ["a" :: T.Text, "b", "c", "d", "e"]) D.empty)
        )

addSmallerColumnUnboxed :: Test
addSmallerColumnUnboxed =
    TestCase
        ( assertEqual
            "Missing values should be replaced with Nothing"
            (Just $ DI.OptionalColumn (V.fromList [Just 1 :: Maybe Int, Just 2, Just 3, Nothing, Nothing]))
            (DI.getColumn "newer" $ D.insertVector "newer" (V.fromList [1 :: Int, 2, 3]) $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3, 4, 5]) D.empty)
        )

insertColumnWithDefaultFillsWithDefault :: Test
insertColumnWithDefaultFillsWithDefault =
    TestCase
        ( assertEqual
            "Missing values should be replaced with Nothing"
            (Just $ DI.UnboxedColumn (VU.fromList [1 :: Int, 2, 3, 0, 0]))
            (DI.getColumn "newer" $ D.insertVectorWithDefault 0 "newer" (V.fromList [1 :: Int, 2, 3]) $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3, 4, 5]) D.empty)
        )

insertColumnWithDefaultFillsLargerNoop :: Test
insertColumnWithDefaultFillsLargerNoop =
    TestCase
        ( assertEqual
            "Lists should be the same size"
            (Just $ DI.UnboxedColumn (VU.fromList [(6 :: Int) .. 10]))
            (DI.getColumn "newer" $ D.insertVectorWithDefault 0 "newer" (V.fromList [(6 :: Int) .. 10]) $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3, 4, 5]) D.empty)
        )

addLargerColumnBoxed :: Test
addLargerColumnBoxed =
    TestCase
        ( assertEqual
            "Smaller lists should grow and contain optionals"
            ( D.fromNamedColumns
                [ ("new", D.fromList [Just "a" :: Maybe T.Text, Just "b", Just "c", Nothing, Nothing])
                , ("newer", D.fromList ["a" :: T.Text, "b", "c", "d", "e"])
                ]
            )
            ( D.insertVector "newer" (V.fromList ["a" :: T.Text, "b", "c", "d", "e"]) $
                D.insertVector "new" (V.fromList ["a" :: T.Text, "b", "c"]) D.empty
            )
        )

addLargerColumnUnboxed :: Test
addLargerColumnUnboxed =
    TestCase
        ( assertEqual
            "Smaller lists should grow and contain optionals"
            ( D.fromNamedColumns
                [ ("old", D.fromList [Just 1 :: Maybe Int, Just 2, Nothing, Nothing, Nothing])
                , ("new", D.fromList [Just 1 :: Maybe Int, Just 2, Just 3, Nothing, Nothing])
                , ("newer", D.fromList [1 :: Int, 2, 3, 4, 5])
                ]
            )
            ( D.insertVector "newer" (V.fromList [1 :: Int, 2, 3, 4, 5]) $
                D.insertVector "new" (V.fromList [1 :: Int, 2, 3]) $
                    D.insertVector "old" (V.fromList [1 :: Int, 2]) D.empty
            )
        )

dimensionsChangeAfterAdd :: Test
dimensionsChangeAfterAdd =
    TestCase
        ( assertEqual
            "should be (26, 3)"
            (26, 9)
            (D.dimensions $ D.insertVector @Int "new" (V.fromList [1 .. 26]) testData)
        )

dimensionsNotChangedAfterDuplicate :: Test
dimensionsNotChangedAfterDuplicate =
    TestCase
        ( assertEqual
            "should be (26, 3)"
            (26, 9)
            ( D.dimensions $
                D.insertVector @Int "new" (V.fromList [1 .. 26]) $
                    D.insertVector @Int "new" (V.fromList [1 .. 26]) testData
            )
        )

tests :: [Test]
tests =
    [ TestLabel "dimensionsChangeAfterAdd" dimensionsChangeAfterAdd
    , TestLabel "dimensionsNotChangedAfterDuplicate" dimensionsNotChangedAfterDuplicate
    , TestLabel "addBoxedColunmToEmpty" addBoxedColumn
    , TestLabel "addBoxedColumnAutoUnboxes" addBoxedColumn
    , TestLabel "addSmallerColumnBoxed" addSmallerColumnBoxed
    , TestLabel "addSmallerColumnUnboxed" addSmallerColumnUnboxed
    , TestLabel "addLargerColumnBoxed" addLargerColumnBoxed
    , TestLabel "addLargerColumnUnboxed" addLargerColumnUnboxed
    , TestLabel "insertColumnWithDefaultFillsWithDefault" insertColumnWithDefaultFillsWithDefault
    , TestLabel "insertColumnWithDefaultFillsLargerNoop" insertColumnWithDefaultFillsLargerNoop
    ]