Inference.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}

module DataFrame.Operations.Inference (
    DateFormat,
    ParsingAssumption (..),
    makeParsingAssumption,
    makeParsingAssumptionBytes,
    byteStringDateParser,
    parseTimeOpt,
    promoteIntColumn,
    promoteIntColumnIndexed,
    readIntStrict,
) where

import qualified Data.ByteString as BS
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed.Mutable as VUM

import Control.Monad.ST (runST)
import Data.Maybe (isJust)
import Data.Time (Day, defaultTimeLocale, parseTimeM)
import DataFrame.Internal.Column (Column (..), finalizeParseResult)
import DataFrame.Internal.Parsing (readByteStringDate, readInt)
import DataFrame.Internal.Parsing.Fast (
    parseBoolField,
    parseDateField,
    parseDoubleField,
    parseIntField,
 )

type DateFormat = String

data ParsingAssumption
    = BoolAssumption
    | IntAssumption
    | DoubleAssumption
    | DateAssumption
    | NoAssumption
    | TextAssumption
    deriving (Eq, Show)

parseTimeOpt :: DateFormat -> T.Text -> Maybe Day
parseTimeOpt dateFormat s =
    parseTimeM {- Accept leading/trailing whitespace -}
        True
        defaultTimeLocale
        dateFormat
        (T.unpack s)

byteStringDateParser :: DateFormat -> BS.ByteString -> Maybe Day
byteStringDateParser "%Y-%m-%d" = parseDateField
byteStringDateParser fmt = readByteStringDate fmt
{-# INLINE byteStringDateParser #-}

readIntStrict :: T.Text -> Maybe Int
readIntStrict t
    | T.length t <= 18 = readInt t
    | otherwise = parseIntField (TE.encodeUtf8 t)
{-# INLINE readIntStrict #-}

pickAssumption ::
    Bool -> Bool -> Bool -> Bool -> Bool -> ParsingAssumption
pickAssumption seen b i d dt
    | not seen = NoAssumption
    | b = BoolAssumption
    | i && d = IntAssumption
    | d = DoubleAssumption
    | dt = DateAssumption
    | otherwise = TextAssumption

makeParsingAssumption ::
    DateFormat -> V.Vector (Maybe T.Text) -> ParsingAssumption
makeParsingAssumption dfmt cells = go 0 False True True True True
  where
    n = V.length cells
    go !i !seen !b !int !d !dt
        | i >= n || (seen && not (b || int || d || dt)) =
            pickAssumption seen b int d dt
        | otherwise = case V.unsafeIndex cells i of
            Nothing -> go (i + 1) seen b int d dt
            Just t ->
                let bs = TE.encodeUtf8 t
                 in go
                        (i + 1)
                        True
                        (b && isJust (parseBoolField bs))
                        (int && isJust (parseIntField bs))
                        (d && isJust (parseDoubleField bs))
                        (dt && isJust (parseTimeOpt dfmt t))

-- | Byte-level twin of 'makeParsingAssumption' for raw field bytes.
makeParsingAssumptionBytes ::
    DateFormat -> V.Vector (Maybe BS.ByteString) -> ParsingAssumption
makeParsingAssumptionBytes dfmt cells = go 0 False True True True True
  where
    n = V.length cells
    dateP = byteStringDateParser dfmt
    go !i !seen !b !int !d !dt
        | i >= n || (seen && not (b || int || d || dt)) =
            pickAssumption seen b int d dt
        | otherwise = case V.unsafeIndex cells i of
            Nothing -> go (i + 1) seen b int d dt
            Just bs ->
                go
                    (i + 1)
                    True
                    (b && isJust (parseBoolField bs))
                    (int && isJust (parseIntField bs))
                    (d && isJust (parseDoubleField bs))
                    (dt && isJust (dateP bs))

promoteIntColumn ::
    forall src.
    (Int -> src -> Bool) ->
    (src -> Maybe Int) ->
    (src -> Maybe Double) ->
    V.Vector src ->
    Maybe Column
promoteIntColumn isNullAt parseI parseD vec =
    promoteIntColumnIndexed
        (V.length vec)
        (\i -> isNullAt i (V.unsafeIndex vec i))
        (parseI . V.unsafeIndex vec)
        (parseD . V.unsafeIndex vec)
{-# INLINE promoteIntColumn #-}

{- | 'promoteIntColumn' over row indices instead of a materialized cell
vector (used by readers that keep cells as buffer offsets).
-}
promoteIntColumnIndexed ::
    Int ->
    (Int -> Bool) ->
    (Int -> Maybe Int) ->
    (Int -> Maybe Double) ->
    Maybe Column
promoteIntColumnIndexed n isNullAt parseI parseD = runST $ do
    values <- VUM.unsafeNew n
    vmask <- VUM.unsafeNew n
    let intLoop !i !anyNull
            | i >= n =
                fmap (uncurry UnboxedColumn)
                    <$> finalizeParseResult values vmask anyNull
            | isNullAt i = do
                VUM.unsafeWrite vmask i 0
                VUM.unsafeWrite values i 0
                intLoop (i + 1) True
            | otherwise = case parseI i of
                Just v -> do
                    VUM.unsafeWrite vmask i 1
                    VUM.unsafeWrite values i v
                    intLoop (i + 1) anyNull
                Nothing -> case parseD i of
                    Just dv -> promote i anyNull dv
                    Nothing -> return Nothing

        -- Switch to Double at row i: convert the Int prefix, then
        -- continue parsing the remaining cells as Double.
        promote !i !anyNull !dv = do
            dvalues <- VUM.unsafeNew n
            let copy !j
                    | j >= i = return ()
                    | otherwise = do
                        x <- VUM.unsafeRead values j
                        VUM.unsafeWrite dvalues j (fromIntegral x)
                        copy (j + 1)
            copy 0
            VUM.unsafeWrite vmask i 1
            VUM.unsafeWrite dvalues i dv
            dblLoop dvalues (i + 1) anyNull

        dblLoop dvalues !i !anyNull
            | i >= n =
                fmap (uncurry UnboxedColumn)
                    <$> finalizeParseResult dvalues vmask anyNull
            | isNullAt i = do
                VUM.unsafeWrite vmask i 0
                VUM.unsafeWrite dvalues i 0
                dblLoop dvalues (i + 1) True
            | otherwise = case parseD i of
                Just dv -> do
                    VUM.unsafeWrite vmask i 1
                    VUM.unsafeWrite dvalues i dv
                    dblLoop dvalues (i + 1) anyNull
                Nothing -> return Nothing
    intLoop 0 False
{-# INLINE promoteIntColumnIndexed #-}