Conversion.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Conversions between 'Column' and vectors or lists, plus typed extraction.
module DataFrame.Internal.Column.Conversion where

import qualified Data.Text as T
import qualified Data.Vector as VB
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM

import Control.Exception (throw)
import Control.Monad.ST (ST)
import Data.Int (Int16, Int32, Int64, Int8)
import Data.Kind (Type)
import Data.Type.Equality (TestEquality (..))
import Data.Word (Word16, Word32, Word64, Word8)
import DataFrame.Errors (
    DataFrameException (ExpectedNonNullableException, TypeMismatchException),
    TypeErrorContext (
        MkTypeErrorContext,
        callingFunctionName,
        errorColumnName,
        expectedType,
        userType
    ),
 )
import DataFrame.Internal.Column.Base
import DataFrame.Internal.Column.Bitmap
import DataFrame.Internal.Column.Properties
import DataFrame.Internal.Column.Types
import DataFrame.Internal.Data.PackedText (packedIndexText, packedLength)
import System.Random (RandomGen, UniformRange, uniformR)
import Type.Reflection (TypeRep, Typeable, typeRep, type (:~:) (Refl))

-- | O(n) Convert a vector to a column, picking the best representation.
fromVector ::
    forall a.
    (Columnable a, ColumnifyRep (KindOf a) a) =>
    VB.Vector a -> Column
fromVector = toColumnRep @(KindOf a)

-- | O(n) Convert an unboxed vector to a column.
fromUnboxedVector ::
    forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Column
fromUnboxedVector = UnboxedColumn Nothing

-- | O(n) Convert a list to a column, picking the best representation.
fromList ::
    forall a.
    (Columnable a, ColumnifyRep (KindOf a) a) =>
    [a] -> Column
fromList = toColumnRep @(KindOf a) . VB.fromList

-- | O(n) Create a column of random elements within a range.
mkRandom ::
    (RandomGen g, Columnable a, ColumnifyRep (KindOf a) a, UniformRange a) =>
    g -> Int -> a -> a -> Column
mkRandom pureGen k lo hi = fromList $ go pureGen k
  where
    go _g 0 = []
    go g n =
        let
            (!v, !g') = uniformR (lo, hi) g
         in
            v : go g' (n - 1)

-- | O(n) Convert a column to a list; throws if the type is wrong.
toList :: forall a. (Columnable a) => Column -> [a]
toList xs = case toVector @a xs of
    Left err -> throw err
    Right val -> VB.toList val

-- | Type-safe conversion of a column to a vector of element type @a@.
toVector ::
    forall a v.
    (VG.Vector v a, Columnable a) => Column -> Either DataFrameException (v a)
toVector col = case col of
    PackedText _ _ -> toVector (materializePacked col)
    MergedColumn _ _ -> toVector (materializeMerged col)
    BoxedColumn bm (inner :: VB.Vector c) ->
        -- Check if user wants Maybe c (nullable) or c directly
        case testEquality (typeRep @a) (typeRep @c) of
            Just Refl
                | hasNulls (VB.length inner) bm -> Left (nullsMismatch @c)
                | otherwise -> Right $ VG.convert inner
            Nothing ->
                -- Try: a = Maybe c
                case testEquality (typeRep @a) (typeRep @(Maybe c)) of
                    Just Refl ->
                        -- Use VB.generate to avoid fusion forcing null slots
                        let !n = VB.length inner
                            maybeVec = case bm of
                                Nothing -> VB.generate n (Just . VB.unsafeIndex inner)
                                Just bitmap -> VB.generate n $ \i ->
                                    if bitmapTestBit bitmap i then Just (VB.unsafeIndex inner i) else Nothing
                         in Right $ VG.convert maybeVec
                    Nothing ->
                        Left $
                            TypeMismatchException
                                ( MkTypeErrorContext
                                    { userType = Right (typeRep @a)
                                    , expectedType = Right (typeRep @c)
                                    , callingFunctionName = Just "toVector"
                                    , errorColumnName = Nothing
                                    }
                                )
    UnboxedColumn bm (inner :: VU.Vector c) ->
        case testEquality (typeRep @a) (typeRep @c) of
            Just Refl
                | hasNulls (VU.length inner) bm -> Left (nullsMismatch @c)
                | otherwise -> Right $ VG.convert inner
            Nothing ->
                case testEquality (typeRep @a) (typeRep @(Maybe c)) of
                    Just Refl ->
                        let maybeVec = case bm of
                                Nothing -> VB.generate (VU.length inner) (Just . VU.unsafeIndex inner)
                                Just bitmap -> VB.generate (VU.length inner) $ \i ->
                                    if bitmapTestBit bitmap i then Just (VU.unsafeIndex inner i) else Nothing
                         in Right $ VG.convert maybeVec
                    Nothing ->
                        Left $
                            TypeMismatchException
                                ( MkTypeErrorContext
                                    { userType = Right (typeRep @a)
                                    , expectedType = Right (typeRep @c)
                                    , callingFunctionName = Just "toVector"
                                    , errorColumnName = Nothing
                                    }
                                )
{-# INLINEABLE toVector #-}

hasNulls :: Int -> Maybe Bitmap -> Bool
hasNulls n = maybe False (bitmapHasNulls n)

-- | The column has nulls, so it must be read as @Maybe c@.
nullsMismatch :: forall c. (Columnable c) => DataFrameException
nullsMismatch =
    TypeMismatchException
        ( MkTypeErrorContext
            { userType = Right (typeRep @c)
            , expectedType = Right (typeRep @(Maybe c))
            , callingFunctionName = Just "toVector"
            , errorColumnName = Nothing
            }
        )

-- Some common types we will use for numerical computing.

-- | Convert a column to an unboxed 'Double' vector, coercing numeric types.
toDoubleVector :: Column -> Either DataFrameException (VU.Vector Double)
toDoubleVector column =
    case column of
        PackedText _ _ -> toDoubleVector (materializePacked column)
        MergedColumn _ _ -> toDoubleVector (materializeMerged column)
        UnboxedColumn (Just _) _ -> Left ExpectedNonNullableException
        UnboxedColumn Nothing (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Double) of
            Just Refl -> Right f
            Nothing -> case sFloating @a of
                STrue -> Right (VU.map realToFrac f)
                SFalse -> case sIntegral @a of
                    STrue -> Right (VU.map fromIntegral f)
                    SFalse ->
                        Left $
                            TypeMismatchException
                                ( MkTypeErrorContext
                                    { userType = Right (typeRep @Double)
                                    , expectedType = Right (typeRep @a)
                                    , callingFunctionName = Just "toDoubleVector"
                                    , errorColumnName = Nothing
                                    }
                                )
        BoxedColumn (Just _) (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of
            Just Refl -> Left ExpectedNonNullableException
            Nothing ->
                Left $
                    TypeMismatchException
                        ( MkTypeErrorContext
                            { userType = Right (typeRep @Double)
                            , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
                            , callingFunctionName = Just "toDoubleVector"
                            , errorColumnName = Nothing
                            }
                        )
        BoxedColumn Nothing (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of
            Just Refl -> Right (VB.convert $ VB.map fromIntegral f)
            Nothing ->
                Left $
                    TypeMismatchException
                        ( MkTypeErrorContext
                            { userType = Right (typeRep @Double)
                            , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
                            , callingFunctionName = Just "toDoubleVector"
                            , errorColumnName = Nothing
                            }
                        )

-- | Convert a column to an unboxed 'Float' vector, coercing numeric types.
toFloatVector :: Column -> Either DataFrameException (VU.Vector Float)
toFloatVector column =
    case column of
        PackedText _ _ -> toFloatVector (materializePacked column)
        MergedColumn _ _ -> toFloatVector (materializeMerged column)
        UnboxedColumn bm (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Float) of
            Just Refl -> case bm of
                Nothing -> Right f
                Just bitmap -> Right $ VU.imap (\i x -> if bitmapTestBit bitmap i then x else read "NaN") f
            Nothing -> case sFloating @a of
                STrue ->
                    Right
                        ( VU.imap
                            ( \i x -> case bm of
                                Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
                                _ -> realToFrac x
                            )
                            f
                        )
                SFalse -> case sIntegral @a of
                    STrue ->
                        Right
                            ( VU.imap
                                ( \i x -> case bm of
                                    Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
                                    _ -> fromIntegral x
                                )
                                f
                            )
                    SFalse ->
                        Left $
                            TypeMismatchException
                                ( MkTypeErrorContext
                                    { userType = Right (typeRep @Float)
                                    , expectedType = Right (typeRep @a)
                                    , callingFunctionName = Just "toFloatVector"
                                    , errorColumnName = Nothing
                                    }
                                )
        BoxedColumn bm (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of
            Just Refl ->
                Right
                    ( VB.convert $
                        VB.imap
                            ( \i x -> case bm of
                                Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
                                _ -> fromIntegral x
                            )
                            f
                    )
            Nothing ->
                Left $
                    TypeMismatchException
                        ( MkTypeErrorContext
                            { userType = Right (typeRep @Float)
                            , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
                            , callingFunctionName = Just "toFloatVector"
                            , errorColumnName = Nothing
                            }
                        )

-- | Convert a column to an unboxed 'Int' vector, coercing numeric types.
toIntVector :: Column -> Either DataFrameException (VU.Vector Int)
toIntVector column =
    case column of
        PackedText _ _ -> toIntVector (materializePacked column)
        MergedColumn _ _ -> toIntVector (materializeMerged column)
        UnboxedColumn _ (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Int) of
            Just Refl -> Right f
            Nothing -> case sFloating @a of
                STrue -> Right (VU.map (round . (realToFrac :: a -> Double)) f)
                SFalse -> case widenToInt f of
                    Just v -> Right v
                    Nothing ->
                        Left $
                            TypeMismatchException
                                ( MkTypeErrorContext
                                    { userType = Right (typeRep @Int)
                                    , expectedType = Right (typeRep @a)
                                    , callingFunctionName = Just "toIntVector"
                                    , errorColumnName = Nothing
                                    }
                                )
        BoxedColumn _ (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of
            Just Refl -> Right (VB.convert $ VB.map fromIntegral f)
            Nothing ->
                Left $
                    TypeMismatchException
                        ( MkTypeErrorContext
                            { userType = Right (typeRep @Int)
                            , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
                            , callingFunctionName = Just "toIntVector"
                            , errorColumnName = Nothing
                            }
                        )

-- | Widen an integral vector to 'Int', or 'Nothing' if not integral.
widenToInt ::
    forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Maybe (VU.Vector Int)
widenToInt v
    | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = Just v
    | Just Refl <- testEquality (typeRep @a) (typeRep @Int64) =
        Just (VU.map fromIntegral v)
    | Just Refl <- testEquality (typeRep @a) (typeRep @Int32) =
        Just (VU.map fromIntegral v)
    | Just Refl <- testEquality (typeRep @a) (typeRep @Int16) =
        Just (VU.map fromIntegral v)
    | Just Refl <- testEquality (typeRep @a) (typeRep @Int8) =
        Just (VU.map fromIntegral v)
    | Just Refl <- testEquality (typeRep @a) (typeRep @Word) =
        Just (VU.map fromIntegral v)
    | Just Refl <- testEquality (typeRep @a) (typeRep @Word64) =
        Just (VU.map fromIntegral v)
    | Just Refl <- testEquality (typeRep @a) (typeRep @Word32) =
        Just (VU.map fromIntegral v)
    | Just Refl <- testEquality (typeRep @a) (typeRep @Word16) =
        Just (VU.map fromIntegral v)
    | Just Refl <- testEquality (typeRep @a) (typeRep @Word8) =
        Just (VU.map fromIntegral v)
    | otherwise = case sIntegral @a of
        STrue -> Just (VU.map fromIntegral v)
        SFalse -> Nothing

toUnboxedVector ::
    forall a.
    (Columnable a, VU.Unbox a) => Column -> Either DataFrameException (VU.Vector a)
toUnboxedVector column =
    case column of
        UnboxedColumn _ (f :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of
            Just Refl -> Right f
            Nothing ->
                Left $
                    TypeMismatchException
                        ( MkTypeErrorContext
                            { userType = Right (typeRep @a)
                            , expectedType = Right (typeRep @b)
                            , callingFunctionName = Just "toUnboxedVector"
                            , errorColumnName = Nothing
                            }
                        )
        _ ->
            Left $
                TypeMismatchException
                    ( MkTypeErrorContext
                        { userType = Right (typeRep @a)
                        , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
                        , callingFunctionName = Just "toUnboxedVector"
                        , errorColumnName = Nothing
                        }
                    )
{-# INLINE toUnboxedVector #-}

-- Freezes the data vector; builds the bitmap only if the column had nulls.
{-# INLINE finalizeParseResult #-}
finalizeParseResult ::
    (VU.Unbox a) =>
    VUM.STVector s a ->
    VUM.STVector s Word8 ->
    Bool ->
    ST s (Maybe (Maybe Bitmap, VU.Vector a))
finalizeParseResult values vmask anyNull
    | anyNull = do
        vs <- VU.unsafeFreeze values
        vm <- VU.unsafeFreeze vmask
        return (Just (Just (buildBitmapFromValid vm), vs))
    | otherwise = do
        vs <- VU.unsafeFreeze values
        return (Just (Nothing, vs))

-- | Convert any Column to a vector of Text labels (one per row).
columnToTextVec :: Column -> VB.Vector T.Text
columnToTextVec c@(MergedColumn _ _) = columnToTextVec (materializeMerged c)
columnToTextVec (BoxedColumn bm (col' :: VB.Vector a)) =
    case bm of
        Nothing -> case testEquality (typeRep @a) (typeRep @T.Text) of
            Just Refl -> col'
            Nothing -> VB.map (T.pack . show) col'
        Just bitmap -> case testEquality (typeRep @a) (typeRep @T.Text) of
            Just Refl ->
                VB.imap
                    (\i x -> if bitmapTestBit bitmap i then x else "null")
                    col'
            Nothing ->
                VB.imap
                    (\i x -> if bitmapTestBit bitmap i then T.pack (show x) else "null")
                    col'
columnToTextVec (UnboxedColumn bm col') =
    case bm of
        Nothing -> VB.map (T.pack . show) (VB.convert col')
        Just bitmap ->
            VB.generate (VU.length col') $ \i ->
                if bitmapTestBit bitmap i then T.pack (show (col' VU.! i)) else "null"
columnToTextVec (PackedText bm p) =
    VB.generate (packedLength p) $ \i -> case bm of
        Just bitmap | not (bitmapTestBit bitmap i) -> "null"
        _ -> packedIndexText p i

-- An internal helper for type errors
throwTypeMismatch ::
    forall (a :: Type) (b :: Type).
    (Typeable a, Typeable b) => Either DataFrameException Column
throwTypeMismatch =
    Left $
        TypeMismatchException
            MkTypeErrorContext
                { userType = Right (typeRep @b)
                , expectedType = Right (typeRep @a)
                , callingFunctionName = Nothing
                , errorColumnName = Nothing
                }