{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module DataFrame.Internal.DataFrame where

import qualified Data.Map as M
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU

import Control.Exception (throw)
import Data.Function (on)
import Data.List (sortBy, (\\))
import Data.Maybe (fromMaybe)
import Data.Type.Equality (
    TestEquality (testEquality),
    type (:~:) (Refl),
    type (:~~:) (HRefl),
 )
import DataFrame.Display.Terminal.PrettyPrint
import DataFrame.Errors
import DataFrame.Internal.Column
import DataFrame.Internal.Expression
import DataFrame.Internal.PackedText (packedIndexText)
import Text.Printf
import Type.Reflection (Typeable, eqTypeRep, typeRep, pattern App)
import Prelude hiding (null)

data DataFrame = DataFrame
    { DataFrame -> Vector Column
columns :: V.Vector Column
    -- ^ Column-oriented storage: the frame is a vector of columns.
    , DataFrame -> Map Text Int
columnIndices :: M.Map T.Text Int
    -- ^ Keeps the column names in the order they were inserted in.
    , DataFrame -> (Int, Int)
dataframeDimensions :: (Int, Int)
    -- ^ (rows, columns)
    , DataFrame -> Map Text UExpr
derivingExpressions :: M.Map T.Text UExpr
    }

{- | Force evaluation of all columns in a DataFrame. Replacement for the removed
@instance NFData DataFrame@; used by the IO and lazy-executor strict paths.
-}
forceDataFrame :: DataFrame -> DataFrame
forceDataFrame :: DataFrame -> DataFrame
forceDataFrame df :: DataFrame
df@(DataFrame Vector Column
cols Map Text Int
idx (Int, Int)
dims Map Text UExpr
_exprs) =
    (() -> Column -> ()) -> () -> Vector Column -> ()
forall a b. (a -> b -> a) -> a -> Vector b -> a
V.foldl' (\() Column
c -> Column -> ()
forceColumn Column
c) () Vector Column
cols () -> DataFrame -> DataFrame
forall a b. a -> b -> b
`seq` Map Text Int
idx Map Text Int -> DataFrame -> DataFrame
forall a b. a -> b -> b
`seq` (Int, Int)
dims (Int, Int) -> DataFrame -> DataFrame
forall a b. a -> b -> b
`seq` DataFrame
df

{- | A record that contains information about how and what
rows are grouped in the dataframe. This can only be used with
`aggregate`.
-}
data GroupedDataFrame = Grouped
    { GroupedDataFrame -> DataFrame
fullDataframe :: DataFrame
    , GroupedDataFrame -> [Text]
groupedColumns :: [T.Text]
    , GroupedDataFrame -> Vector Int
valueIndices :: VU.Vector Int
    , GroupedDataFrame -> Vector Int
offsets :: VU.Vector Int
    , GroupedDataFrame -> Vector Int
rowToGroup :: VU.Vector Int
    {- ^ rowToGroup[i] = group index for row i.  Length n (one per row).
    Built once in 'groupBy'; reused by every aggregation.
    -}
    }

instance Show GroupedDataFrame where
    show :: GroupedDataFrame -> String
show (Grouped DataFrame
df [Text]
cols Vector Int
_indices Vector Int
_os Vector Int
_rtg) =
        String -> String -> ShowS
forall r. PrintfType r => String -> r
printf
            String
"{ keyColumns: %s groupedColumns: %s }"
            ([Text] -> String
forall a. Show a => a -> String
show [Text]
cols)
            ([Text] -> String
forall a. Show a => a -> String
show (Map Text Int -> [Text]
forall k a. Map k a -> [k]
M.keys (DataFrame -> Map Text Int
columnIndices DataFrame
df) [Text] -> [Text] -> [Text]
forall a. Eq a => [a] -> [a] -> [a]
\\ [Text]
cols))

instance Eq GroupedDataFrame where
    == :: GroupedDataFrame -> GroupedDataFrame -> Bool
(==) (Grouped DataFrame
df [Text]
cols Vector Int
_indices Vector Int
_os Vector Int
_rtg) (Grouped DataFrame
df' [Text]
cols' Vector Int
_indices' Vector Int
_os' Vector Int
_rtg') = (DataFrame
df DataFrame -> DataFrame -> Bool
forall a. Eq a => a -> a -> Bool
== DataFrame
df') Bool -> Bool -> Bool
&& ([Text]
cols [Text] -> [Text] -> Bool
forall a. Eq a => a -> a -> Bool
== [Text]
cols')

instance Eq DataFrame where
    (==) :: DataFrame -> DataFrame -> Bool
    DataFrame
a == :: DataFrame -> DataFrame -> Bool
== DataFrame
b =
        Map Text Int -> [Text]
forall k a. Map k a -> [k]
M.keys (DataFrame -> Map Text Int
columnIndices DataFrame
a) [Text] -> [Text] -> Bool
forall a. Eq a => a -> a -> Bool
== Map Text Int -> [Text]
forall k a. Map k a -> [k]
M.keys (DataFrame -> Map Text Int
columnIndices DataFrame
b)
            Bool -> Bool -> Bool
&& ((Text, Int) -> Bool -> Bool) -> Bool -> [(Text, Int)] -> Bool
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
                ( \(Text
name, Int
index) Bool
acc -> Bool
acc Bool -> Bool -> Bool
&& (DataFrame -> Vector Column
columns DataFrame
a Vector Column -> Int -> Maybe Column
forall a. Vector a -> Int -> Maybe a
V.!? Int
index Maybe Column -> Maybe Column -> Bool
forall a. Eq a => a -> a -> Bool
== (DataFrame -> Vector Column
columns DataFrame
b Vector Column -> Int -> Maybe Column
forall a. Vector a -> Int -> Maybe a
V.!? (DataFrame -> Map Text Int
columnIndices DataFrame
b Map Text Int -> Text -> Int
forall k a. Ord k => Map k a -> k -> a
M.! Text
name)))
                )
                Bool
True
                (Map Text Int -> [(Text, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Text Int -> [(Text, Int)]) -> Map Text Int -> [(Text, Int)]
forall a b. (a -> b) -> a -> b
$ DataFrame -> Map Text Int
columnIndices DataFrame
a)

instance Show DataFrame where
    show :: DataFrame -> String
    show :: DataFrame -> String
show DataFrame
d =
        let (Int
r, Int
_) = DataFrame -> (Int, Int)
dataframeDimensions DataFrame
d
            cfg :: TruncateConfig
cfg = TruncateConfig
defaultTruncateConfig
            shown :: Int
shown = if TruncateConfig -> Int
maxRows TruncateConfig
cfg Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 then Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (TruncateConfig -> Int
maxRows TruncateConfig
cfg) Int
r else Int
r
            body :: Text
body = RenderFormat -> Maybe TruncateConfig -> DataFrame -> Text
asTextWith RenderFormat
Plain (TruncateConfig -> Maybe TruncateConfig
forall a. a -> Maybe a
Just TruncateConfig
cfg) DataFrame
d
            footer :: Text
footer
                | Int
shown Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r =
                    Text
"\nShowing "
                        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show Int
shown)
                        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" rows out of "
                        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show Int
r)
                | Bool
otherwise = Text
T.empty
         in Text -> String
T.unpack (Text
body Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
footer)

{- | Configures how a 'DataFrame' is rendered as text: 'maxRows' caps rendered
rows, 'maxColumns' collapses middle columns past the limit into an ellipsis, and
'maxCellWidth' truncates long cells. A non-positive field means \"no limit\".
-}
data TruncateConfig = TruncateConfig
    { TruncateConfig -> Int
maxRows :: Int
    , TruncateConfig -> Int
maxColumns :: Int
    , TruncateConfig -> Int
maxCellWidth :: Int
    }
    deriving (Int -> TruncateConfig -> ShowS
[TruncateConfig] -> ShowS
TruncateConfig -> String
(Int -> TruncateConfig -> ShowS)
-> (TruncateConfig -> String)
-> ([TruncateConfig] -> ShowS)
-> Show TruncateConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TruncateConfig -> ShowS
showsPrec :: Int -> TruncateConfig -> ShowS
$cshow :: TruncateConfig -> String
show :: TruncateConfig -> String
$cshowList :: [TruncateConfig] -> ShowS
showList :: [TruncateConfig] -> ShowS
Show, TruncateConfig -> TruncateConfig -> Bool
(TruncateConfig -> TruncateConfig -> Bool)
-> (TruncateConfig -> TruncateConfig -> Bool) -> Eq TruncateConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TruncateConfig -> TruncateConfig -> Bool
== :: TruncateConfig -> TruncateConfig -> Bool
$c/= :: TruncateConfig -> TruncateConfig -> Bool
/= :: TruncateConfig -> TruncateConfig -> Bool
Eq)

-- | Sensible defaults for GHCi: 20 rows, 10 columns, 30 characters per cell.
defaultTruncateConfig :: TruncateConfig
defaultTruncateConfig :: TruncateConfig
defaultTruncateConfig =
    TruncateConfig{maxRows :: Int
maxRows = Int
20, maxColumns :: Int
maxColumns = Int
10, maxCellWidth :: Int
maxCellWidth = Int
30}

-- | Ellipsis character used to mark elided columns and clipped cells.
ellipsisText :: T.Text
ellipsisText :: Text
ellipsisText = Text
"\x2026"

-- | For showing the dataframe as markdown in notebooks.
toMarkdown :: DataFrame -> T.Text
toMarkdown :: DataFrame -> Text
toMarkdown = RenderFormat -> DataFrame -> Text
asText RenderFormat
Markdown

-- | For showing the dataframe as a string markdown in notebooks.
toMarkdown' :: DataFrame -> String
toMarkdown' :: DataFrame -> String
toMarkdown' = Text -> String
T.unpack (Text -> String) -> (DataFrame -> Text) -> DataFrame -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame -> Text
toMarkdown

asText :: RenderFormat -> DataFrame -> T.Text
asText :: RenderFormat -> DataFrame -> Text
asText RenderFormat
fmt = RenderFormat -> Maybe TruncateConfig -> DataFrame -> Text
asTextWith RenderFormat
fmt Maybe TruncateConfig
forall a. Maybe a
Nothing

asTextWith :: RenderFormat -> Maybe TruncateConfig -> DataFrame -> T.Text
asTextWith :: RenderFormat -> Maybe TruncateConfig -> DataFrame -> Text
asTextWith RenderFormat
fmt Maybe TruncateConfig
mTrunc DataFrame
d =
    let allHeaders :: [Text]
allHeaders =
            ((Text, Int) -> Text) -> [(Text, Int)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Int) -> Text
forall a b. (a, b) -> a
fst (((Text, Int) -> (Text, Int) -> Ordering)
-> [(Text, Int)] -> [(Text, Int)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> ((Text, Int) -> Int) -> (Text, Int) -> (Text, Int) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (Text, Int) -> Int
forall a b. (a, b) -> b
snd) (Map Text Int -> [(Text, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (DataFrame -> Map Text Int
columnIndices DataFrame
d)))
        nCols :: Int
nCols = [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
allHeaders
        (Int
totalRows, Int
_) = DataFrame -> (Int, Int)
dataframeDimensions DataFrame
d

        rowCap :: Int
rowCap = case Maybe TruncateConfig
mTrunc of
            Just TruncateConfig
cfg | TruncateConfig -> Int
maxRows TruncateConfig
cfg Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
totalRows (TruncateConfig -> Int
maxRows TruncateConfig
cfg)
            Maybe TruncateConfig
_ -> Int
totalRows

        ([Text]
visibleHeaders, Maybe Int
ellipsisAt) = Maybe TruncateConfig -> Int -> [Text] -> ([Text], Maybe Int)
forall a. Maybe TruncateConfig -> Int -> [a] -> ([a], Maybe Int)
pickColumns Maybe TruncateConfig
mTrunc Int
nCols [Text]
allHeaders

        lookupCol :: Text -> Maybe Column
lookupCol Text
name =
            (Column -> Column) -> Maybe Column -> Maybe Column
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
                (Int -> Column -> Column
takeColumn Int
rowCap)
                (Vector Column -> Int -> Maybe Column
forall a. Vector a -> Int -> Maybe a
(V.!?) (DataFrame -> Vector Column
columns DataFrame
d) (Map Text Int -> Text -> Int
forall k a. Ord k => Map k a -> k -> a
(M.!) (DataFrame -> Map Text Int
columnIndices DataFrame
d) Text
name))
        survivingCols :: [Maybe Column]
survivingCols = (Text -> Maybe Column) -> [Text] -> [Maybe Column]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Maybe Column
lookupCol [Text]
visibleHeaders
        survivingTypes :: [Text]
survivingTypes = (Maybe Column -> Text) -> [Maybe Column] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> (Column -> Text) -> Maybe Column -> Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Text
"" Column -> Text
getType) [Maybe Column]
survivingCols
        survivingData :: [Vector Text]
survivingData = (Maybe Column -> Vector Text) -> [Maybe Column] -> [Vector Text]
forall a b. (a -> b) -> [a] -> [b]
map Maybe Column -> Vector Text
get [Maybe Column]
survivingCols

        clipCell :: Text -> Text
clipCell = case Maybe TruncateConfig
mTrunc of
            Just TruncateConfig
cfg | TruncateConfig -> Int
maxCellWidth TruncateConfig
cfg Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> Int -> Text -> Text
truncateCell (TruncateConfig -> Int
maxCellWidth TruncateConfig
cfg)
            Maybe TruncateConfig
_ -> Text -> Text
forall a. a -> a
id

        ([Text]
finalHeaders, [Text]
finalTypes, [Vector Text]
finalCols) = case Maybe Int
ellipsisAt of
            Maybe Int
Nothing -> ([Text]
visibleHeaders, [Text]
survivingTypes, [Vector Text]
survivingData)
            Just Int
i ->
                let ellipsisCol :: Vector Text
ellipsisCol = Int -> Text -> Vector Text
forall a. Int -> a -> Vector a
V.replicate Int
rowCap Text
ellipsisText
                 in ( Int -> Text -> [Text] -> [Text]
forall a. Int -> a -> [a] -> [a]
insertAt Int
i Text
ellipsisText [Text]
visibleHeaders
                    , Int -> Text -> [Text] -> [Text]
forall a. Int -> a -> [a] -> [a]
insertAt Int
i Text
ellipsisText [Text]
survivingTypes
                    , Int -> Vector Text -> [Vector Text] -> [Vector Text]
forall a. Int -> a -> [a] -> [a]
insertAt Int
i Vector Text
ellipsisCol [Vector Text]
survivingData
                    )

        getType :: Column -> T.Text
        showMaybeType :: forall a. (Typeable a) => String
        showMaybeType :: forall a. Typeable a => String
showMaybeType =
            let s :: String
s = TypeRep a -> String
forall a. Show a => a -> String
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a)
             in String
"Maybe " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> if Char
' ' Char -> String -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` String
s then String
"(" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
s String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
")" else String
s
        getType :: Column -> Text
getType (BoxedColumn Maybe Bitmap
Nothing (Vector a
_ :: V.Vector a)) = String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ TypeRep a -> String
forall a. Show a => a -> String
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a)
        getType (BoxedColumn (Just Bitmap
_) (Vector a
_ :: V.Vector a)) = String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ forall a. Typeable a => String
showMaybeType @a
        getType (UnboxedColumn Maybe Bitmap
Nothing (Vector a
_ :: VU.Vector a)) = String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ TypeRep a -> String
forall a. Show a => a -> String
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a)
        getType (UnboxedColumn (Just Bitmap
_) (Vector a
_ :: VU.Vector a)) = String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ forall a. Typeable a => String
showMaybeType @a
        getType (PackedText Maybe Bitmap
Nothing PackedTextData
_) = String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ TypeRep Text -> String
forall a. Show a => a -> String
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @T.Text)
        getType (PackedText (Just Bitmap
_) PackedTextData
_) = String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ forall a. Typeable a => String
showMaybeType @T.Text
        getType c :: Column
c@(MergedColumn Column
_ Column
_) = Column -> Text
getType (Column -> Column
mergedHead Column
c)

        get :: Maybe Column -> V.Vector T.Text
        get :: Maybe Column -> Vector Text
get (Just c :: Column
c@(MergedColumn Column
_ Column
_)) = Maybe Column -> Vector Text
get (Column -> Maybe Column
forall a. a -> Maybe a
Just (Column -> Column
materializeMerged Column
c))
        get (Just (BoxedColumn (Just Bitmap
bm) (Vector a
column :: V.Vector a))) =
            Int -> (Int -> Text) -> Vector Text
forall a. Int -> (Int -> a) -> Vector a
V.generate (Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
column) ((Int -> Text) -> Vector Text) -> (Int -> Text) -> Vector Text
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i
                    then String -> Text
T.pack (Maybe a -> String
forall a. Show a => a -> String
show (a -> Maybe a
forall a. a -> Maybe a
Just (Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex Vector a
column Int
i)))
                    else Text
"Nothing"
        get (Just (BoxedColumn Maybe Bitmap
Nothing (Vector a
column :: V.Vector a))) =
            case TypeRep a -> TypeRep Text -> Maybe (a :~: Text)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @T.Text) of
                Just a :~: Text
Refl -> Vector a
Vector Text
column
                Maybe (a :~: Text)
Nothing -> case TypeRep a -> TypeRep String -> Maybe (a :~: String)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @String) of
                    Just a :~: String
Refl -> (String -> Text) -> Vector String -> Vector Text
forall a b. (a -> b) -> Vector a -> Vector b
V.map String -> Text
T.pack Vector a
Vector String
column
                    Maybe (a :~: String)
Nothing -> (a -> Text) -> Vector a -> Vector Text
forall a b. (a -> b) -> Vector a -> Vector b
V.map (String -> Text
T.pack (String -> Text) -> (a -> String) -> a -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> String
forall a. Show a => a -> String
show) Vector a
column
        get (Just (UnboxedColumn (Just Bitmap
bm) Vector a
column)) =
            Int -> (Int -> Text) -> Vector Text
forall a. Int -> (Int -> a) -> Vector a
V.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
column) ((Int -> Text) -> Vector Text) -> (Int -> Text) -> Vector Text
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                if Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i
                    then String -> Text
T.pack (Maybe a -> String
forall a. Show a => a -> String
show (a -> Maybe a
forall a. a -> Maybe a
Just (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
column Int
i)))
                    else Text
"Nothing"
        get (Just (UnboxedColumn Maybe Bitmap
Nothing Vector a
column)) =
            Int -> (Int -> Text) -> Vector Text
forall a. Int -> (Int -> a) -> Vector a
V.generate (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
column) (String -> Text
T.pack (String -> Text) -> (Int -> String) -> Int -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> String
forall a. Show a => a -> String
show (a -> String) -> (Int -> a) -> Int -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
column)
        get (Just c :: Column
c@(PackedText Maybe Bitmap
_ PackedTextData
_)) = Maybe Column -> Vector Text
get (Column -> Maybe Column
forall a. a -> Maybe a
Just (Column -> Column
materializePacked Column
c))
        get Maybe Column
Nothing = Vector Text
forall a. Vector a
V.empty
     in RenderFormat -> [Text] -> [Text] -> [Vector Text] -> Text
showTable
            RenderFormat
fmt
            ((Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Text
clipCell [Text]
finalHeaders)
            ((Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Text
clipCell [Text]
finalTypes)
            ((Vector Text -> Vector Text) -> [Vector Text] -> [Vector Text]
forall a b. (a -> b) -> [a] -> [b]
map ((Text -> Text) -> Vector Text -> Vector Text
forall a b. (a -> b) -> Vector a -> Vector b
V.map Text -> Text
clipCell) [Vector Text]
finalCols)

{- | Decide which columns survive horizontal truncation and where to splice the
ellipsis column. Splits with the extra column on the left for odd 'maxColumns';
inserts the ellipsis only when it actually saves space.
-}
pickColumns ::
    Maybe TruncateConfig ->
    Int ->
    [a] ->
    ([a], Maybe Int)
pickColumns :: forall a. Maybe TruncateConfig -> Int -> [a] -> ([a], Maybe Int)
pickColumns Maybe TruncateConfig
mTrunc Int
nCols [a]
xs = case Maybe TruncateConfig
mTrunc of
    Just TruncateConfig
cfg
        | let c :: Int
c = TruncateConfig -> Int
maxColumns TruncateConfig
cfg
        , Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
        , Int
nCols Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 ->
            let leftN :: Int
leftN = (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
                rightN :: Int
rightN = Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
leftN
             in ( Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
Prelude.take Int
leftN [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
Prelude.drop (Int
nCols Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
rightN) [a]
xs
                , Int -> Maybe Int
forall a. a -> Maybe a
Just Int
leftN
                )
    Maybe TruncateConfig
_ -> ([a]
xs, Maybe Int
forall a. Maybe a
Nothing)

-- | Splice @x@ into @xs@ at index @i@ (0-based), shifting later elements right.
insertAt :: Int -> a -> [a] -> [a]
insertAt :: forall a. Int -> a -> [a] -> [a]
insertAt Int
i a
x [a]
xs = let ([a]
l, [a]
r) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
i [a]
xs in [a]
l [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
r

-- | Cap a single cell's rendered length, appending an ellipsis when shortened.
truncateCell :: Int -> T.Text -> T.Text
truncateCell :: Int -> Text -> Text
truncateCell Int
n Text
t
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Text
t
    | Text -> Int -> Ordering
T.compareLength Text
t Int
n Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
/= Ordering
GT = Text
t
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Text
ellipsisText
    | Bool
otherwise = Int -> Text -> Text
T.take (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Text
t Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
ellipsisText

-- | O(1) Creates an empty dataframe
empty :: DataFrame
empty :: DataFrame
empty =
    DataFrame
        { columns :: Vector Column
columns = Vector Column
forall a. Vector a
V.empty
        , columnIndices :: Map Text Int
columnIndices = Map Text Int
forall k a. Map k a
M.empty
        , dataframeDimensions :: (Int, Int)
dataframeDimensions = (Int
0, Int
0)
        , derivingExpressions :: Map Text UExpr
derivingExpressions = Map Text UExpr
forall k a. Map k a
M.empty
        }

-- | O(k) Get column names of the DataFrame in order of insertion.
columnNames :: DataFrame -> [T.Text]
columnNames :: DataFrame -> [Text]
columnNames = ((Text, Int) -> Text) -> [(Text, Int)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Int) -> Text
forall a b. (a, b) -> a
fst ([(Text, Int)] -> [Text])
-> (DataFrame -> [(Text, Int)]) -> DataFrame -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Text, Int) -> (Text, Int) -> Ordering)
-> [(Text, Int)] -> [(Text, Int)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> ((Text, Int) -> Int) -> (Text, Int) -> (Text, Int) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (Text, Int) -> Int
forall a b. (a, b) -> b
snd) ([(Text, Int)] -> [(Text, Int)])
-> (DataFrame -> [(Text, Int)]) -> DataFrame -> [(Text, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Text Int -> [(Text, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Text Int -> [(Text, Int)])
-> (DataFrame -> Map Text Int) -> DataFrame -> [(Text, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame -> Map Text Int
columnIndices
{-# INLINE columnNames #-}

{- | Insert a column into a DataFrame. If a column with the same name already
exists it is replaced in-place; otherwise the column is appended at the end.
Other columns are expanded (padded with nulls) to match the new row count.
-}
insertColumn :: T.Text -> Column -> DataFrame -> DataFrame
insertColumn :: Text -> Column -> DataFrame -> DataFrame
insertColumn Text
name Column
column DataFrame
d =
    let
        (Int
r, Int
c) = DataFrame -> (Int, Int)
dataframeDimensions DataFrame
d
        n :: Int
n = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Column -> Int
columnLength Column
column) Int
r
        exprs :: Map Text UExpr
exprs = Text -> Map Text UExpr -> Map Text UExpr
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Text
name (DataFrame -> Map Text UExpr
derivingExpressions DataFrame
d)
     in
        case Text -> Map Text Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Text
name (DataFrame -> Map Text Int
columnIndices DataFrame
d) of
            Just Int
i ->
                Vector Column
-> Map Text Int -> (Int, Int) -> Map Text UExpr -> DataFrame
DataFrame
                    ((Column -> Column) -> Vector Column -> Vector Column
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Int -> Column -> Column
expandColumn Int
n) (DataFrame -> Vector Column
columns DataFrame
d Vector Column -> [(Int, Column)] -> Vector Column
forall a. Vector a -> [(Int, a)] -> Vector a
V.// [(Int
i, Column
column)]))
                    (DataFrame -> Map Text Int
columnIndices DataFrame
d)
                    (Int
n, Int
c)
                    Map Text UExpr
exprs
            Maybe Int
Nothing ->
                Vector Column
-> Map Text Int -> (Int, Int) -> Map Text UExpr -> DataFrame
DataFrame
                    ((Column -> Column) -> Vector Column -> Vector Column
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Int -> Column -> Column
expandColumn Int
n) (DataFrame -> Vector Column
columns DataFrame
d Vector Column -> Column -> Vector Column
forall a. Vector a -> a -> Vector a
`V.snoc` Column
column))
                    (Text -> Int -> Map Text Int -> Map Text Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Text
name Int
c (DataFrame -> Map Text Int
columnIndices DataFrame
d))
                    (Int
n, Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    Map Text UExpr
exprs

-- | Build a DataFrame from a list of @(name, column)@ pairs using 'insertColumn'.
fromNamedColumns :: [(T.Text, Column)] -> DataFrame
fromNamedColumns :: [(Text, Column)] -> DataFrame
fromNamedColumns = (DataFrame -> (Text, Column) -> DataFrame)
-> DataFrame -> [(Text, Column)] -> DataFrame
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\DataFrame
df (Text
name, Column
column) -> Text -> Column -> DataFrame -> DataFrame
insertColumn Text
name Column
column DataFrame
df) DataFrame
empty

{- | Safely retrieves a column by name from the dataframe.

Returns 'Nothing' if the column does not exist.

==== __Examples__

>>> getColumn "age" df
Just (UnboxedColumn ...)

>>> getColumn "nonexistent" df
Nothing
-}
getColumn :: T.Text -> DataFrame -> Maybe Column
getColumn :: Text -> DataFrame -> Maybe Column
getColumn Text
name DataFrame
df
    | DataFrame -> Bool
null DataFrame
df = Maybe Column
forall a. Maybe a
Nothing
    | Bool
otherwise = do
        Int
i <- DataFrame -> Map Text Int
columnIndices DataFrame
df Map Text Int -> Text -> Maybe Int
forall k a. Ord k => Map k a -> k -> Maybe a
M.!? Text
name
        DataFrame -> Vector Column
columns DataFrame
df Vector Column -> Int -> Maybe Column
forall a. Vector a -> Int -> Maybe a
V.!? Int
i

{- | Retrieve a column by name, throwing 'ColumnsNotFoundException' if it does not
exist. Unsafe version of 'getColumn'; use when the column is certain to exist.
-}
unsafeGetColumn :: T.Text -> DataFrame -> Column
unsafeGetColumn :: Text -> DataFrame -> Column
unsafeGetColumn Text
name DataFrame
df = case Text -> DataFrame -> Maybe Column
getColumn Text
name DataFrame
df of
    Maybe Column
Nothing -> DataFrameException -> Column
forall a e. Exception e => e -> a
throw (DataFrameException -> Column) -> DataFrameException -> Column
forall a b. (a -> b) -> a -> b
$ [Text] -> Text -> [Text] -> DataFrameException
ColumnsNotFoundException [Text
name] Text
"" (Map Text Int -> [Text]
forall k a. Map k a -> [k]
M.keys (Map Text Int -> [Text]) -> Map Text Int -> [Text]
forall a b. (a -> b) -> a -> b
$ DataFrame -> Map Text Int
columnIndices DataFrame
df)
    Just Column
col -> Column
col

{- | Checks if the dataframe is empty (has no columns).

Returns 'True' if the dataframe has no columns, 'False' otherwise.
Note that a dataframe with columns but no rows is not considered null.
-}
null :: DataFrame -> Bool
null :: DataFrame -> Bool
null DataFrame
df = Vector Column -> Bool
forall a. Vector a -> Bool
V.null (DataFrame -> Vector Column
columns DataFrame
df)

-- | Convert a DataFrame to a CSV (comma-separated) text.
toCsv :: DataFrame -> T.Text
toCsv :: DataFrame -> Text
toCsv = Char -> DataFrame -> Text
toSeparated Char
','

-- | Convert a DataFrame to a CSV (comma-separated) string.
toCsv' :: DataFrame -> String
toCsv' :: DataFrame -> String
toCsv' = Text -> String
T.unpack (Text -> String) -> (DataFrame -> Text) -> DataFrame -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> DataFrame -> Text
toSeparated Char
','

-- | Convert a DataFrame to a text representation with a custom separator.
toSeparated :: Char -> DataFrame -> T.Text
toSeparated :: Char -> DataFrame -> Text
toSeparated Char
sep DataFrame
df
    | DataFrame -> Bool
null DataFrame
df = Text
T.empty
    | Bool
otherwise =
        let (Int
rows, Int
_) = DataFrame -> (Int, Int)
dataframeDimensions DataFrame
df
            headers :: [Text]
headers = ((Text, Int) -> Text) -> [(Text, Int)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Int) -> Text
forall a b. (a, b) -> a
fst (((Text, Int) -> (Text, Int) -> Ordering)
-> [(Text, Int)] -> [(Text, Int)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> ((Text, Int) -> Int) -> (Text, Int) -> (Text, Int) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (Text, Int) -> Int
forall a b. (a, b) -> b
snd) (Map Text Int -> [(Text, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (DataFrame -> Map Text Int
columnIndices DataFrame
df)))
            sepText :: Text
sepText = Char -> Text
T.singleton Char
sep
            headerLine :: Text
headerLine = Text -> [Text] -> Text
T.intercalate Text
sepText [Text]
headers
            dataLines :: [Text]
dataLines = (Int -> Text) -> [Int] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> [Text] -> Text
T.intercalate Text
sepText ([Text] -> Text) -> (Int -> [Text]) -> Int -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame -> Int -> [Text]
getRowAsText DataFrame
df) [Int
0 .. Int
rows Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
         in [Text] -> Text
T.unlines (Text
headerLine Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: [Text]
dataLines)

getRowAsText :: DataFrame -> Int -> [T.Text]
getRowAsText :: DataFrame -> Int -> [Text]
getRowAsText DataFrame
df Int
i = (Column -> Text) -> [Column] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Column -> Int -> Text
`showElement` Int
i) (Vector Column -> [Column]
forall a. Vector a -> [a]
V.toList (DataFrame -> Vector Column
columns DataFrame
df))

showElement :: Column -> Int -> T.Text
showElement :: Column -> Int -> Text
showElement (MergedColumn Column
a Column
b) Int
i =
    Column -> Int -> Text
showElement
        (Column -> Column
materializeMerged (Column -> Column -> Column
MergedColumn (Int -> Int -> Column -> Column
sliceColumn Int
i Int
1 Column
a) (Int -> Int -> Column -> Column
sliceColumn Int
i Int
1 Column
b)))
        Int
0
showElement (BoxedColumn Maybe Bitmap
bm (Vector a
c :: V.Vector a)) Int
i = case Maybe Bitmap
bm of
    Just Bitmap
b | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
b Int
i) -> Text
"null"
    Maybe Bitmap
_ -> case Vector a
c Vector a -> Int -> Maybe a
forall a. Vector a -> Int -> Maybe a
V.!? Int
i of
        Maybe a
Nothing -> String -> Text
forall a. HasCallStack => String -> a
error (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ String
"Column index out of bounds at row " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i
        Just a
e
            | Just a :~: Text
Refl <- TypeRep a -> TypeRep Text -> Maybe (a :~: Text)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @T.Text) -> a
Text
e
            | App TypeRep a
t1 TypeRep b
t2 <- forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a
            , Just a :~~: Maybe
HRefl <- TypeRep a -> TypeRep Maybe -> Maybe (a :~~: Maybe)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
eqTypeRep TypeRep a
t1 (forall {k} (a :: k). Typeable a => TypeRep a
forall (a :: * -> *). Typeable a => TypeRep a
typeRep @Maybe) ->
                case TypeRep b -> TypeRep Text -> Maybe (b :~: Text)
forall (a :: k1) (b :: k1).
TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality TypeRep b
t2 (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @T.Text) of
                    Just b :~: Text
Refl -> Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
"null" a
Maybe Text
e
                    Maybe (b :~: Text)
Nothing -> Text -> Text
stripJust (String -> Text
T.pack (a -> String
forall a. Show a => a -> String
show a
e))
            | Bool
otherwise -> String -> Text
T.pack (a -> String
forall a. Show a => a -> String
show a
e)
showElement (UnboxedColumn Maybe Bitmap
bm Vector a
c) Int
i = case Maybe Bitmap
bm of
    Just Bitmap
b | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
b Int
i) -> Text
"null"
    Maybe Bitmap
_ -> case Vector a
c Vector a -> Int -> Maybe a
forall a. Unbox a => Vector a -> Int -> Maybe a
VU.!? Int
i of
        Maybe a
Nothing -> String -> Text
forall a. HasCallStack => String -> a
error (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ String
"Column index out of bounds at row " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i
        Just a
e -> String -> Text
T.pack (a -> String
forall a. Show a => a -> String
show a
e)
showElement (PackedText Maybe Bitmap
bm PackedTextData
p) Int
i = case Maybe Bitmap
bm of
    Just Bitmap
b | Bool -> Bool
not (Bitmap -> Int -> Bool
bitmapTestBit Bitmap
b Int
i) -> Text
"null"
    Maybe Bitmap
_ -> PackedTextData -> Int -> Text
packedIndexText PackedTextData
p Int
i

stripJust :: T.Text -> T.Text
stripJust :: Text -> Text
stripJust = Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
"null" (Maybe Text -> Text) -> (Text -> Maybe Text) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text -> Maybe Text
T.stripPrefix Text
"Just "