{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module DataFrame.Typed.Freeze (
    -- * Safe boundary
    freeze,
    freezeWithError,
    freezeOrThrow,

    -- * Escape hatches
    thaw,
    unsafeFreeze,

    -- * Frame coercion
    ToDataFrame (..),
) where

import Control.Exception (throwIO)
import qualified Data.Text as T
import Type.Reflection (SomeTypeRep)

import Data.List (stripPrefix)
import DataFrame.Errors (DataFrameException (InternalException))
import qualified DataFrame.Internal.Column as C
import DataFrame.Internal.DataFrame (columnNames)
import qualified DataFrame.Internal.DataFrame as D
import DataFrame.Typed.Schema (KnownSchema (..))
import DataFrame.Typed.Types (TypedDataFrame (..))

{- | Validate that an untyped 'DataFrame' matches the expected schema @cols@,
then wrap it. Returns 'Nothing' on mismatch.
-}
freeze ::
    forall cols. (KnownSchema cols) => D.DataFrame -> Maybe (TypedDataFrame cols)
freeze :: forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> Maybe (TypedDataFrame cols)
freeze DataFrame
df = case forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> Either Text ()
validateSchema @cols DataFrame
df of
    Left Text
_ -> Maybe (TypedDataFrame cols)
forall a. Maybe a
Nothing
    Right ()
_ -> TypedDataFrame cols -> Maybe (TypedDataFrame cols)
forall a. a -> Maybe a
Just (DataFrame -> TypedDataFrame cols
forall (cols :: [(Symbol, *)]). DataFrame -> TypedDataFrame cols
TDF DataFrame
df)

-- | Like 'freeze' but returns a descriptive error message on failure.
freezeWithError ::
    forall cols.
    (KnownSchema cols) =>
    D.DataFrame -> Either T.Text (TypedDataFrame cols)
freezeWithError :: forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> Either Text (TypedDataFrame cols)
freezeWithError DataFrame
df = case forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> Either Text ()
validateSchema @cols DataFrame
df of
    Left Text
err -> Text -> Either Text (TypedDataFrame cols)
forall a b. a -> Either a b
Left Text
err
    Right ()
_ -> TypedDataFrame cols -> Either Text (TypedDataFrame cols)
forall a b. b -> Either a b
Right (DataFrame -> TypedDataFrame cols
forall (cols :: [(Symbol, *)]). DataFrame -> TypedDataFrame cols
TDF DataFrame
df)

{- | Validate and wrap like 'freezeWithError', but throw a 'DataFrameException'
in 'IO' on mismatch. The throwing boundary used by the typed readers
(@readCsv@ \/ @readParquet@).
-}
freezeOrThrow ::
    forall cols. (KnownSchema cols) => D.DataFrame -> IO (TypedDataFrame cols)
freezeOrThrow :: forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> IO (TypedDataFrame cols)
freezeOrThrow = (Text -> IO (TypedDataFrame cols))
-> (TypedDataFrame cols -> IO (TypedDataFrame cols))
-> Either Text (TypedDataFrame cols)
-> IO (TypedDataFrame cols)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (DataFrameException -> IO (TypedDataFrame cols)
forall e a. Exception e => e -> IO a
throwIO (DataFrameException -> IO (TypedDataFrame cols))
-> (Text -> DataFrameException) -> Text -> IO (TypedDataFrame cols)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> DataFrameException
InternalException) TypedDataFrame cols -> IO (TypedDataFrame cols)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Text (TypedDataFrame cols) -> IO (TypedDataFrame cols))
-> (DataFrame -> Either Text (TypedDataFrame cols))
-> DataFrame
-> IO (TypedDataFrame cols)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> Either Text (TypedDataFrame cols)
freezeWithError @cols

{- | Unwrap a typed DataFrame back to the untyped representation.
Always safe; discards type information.
-}
thaw :: TypedDataFrame cols -> D.DataFrame
thaw :: forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> DataFrame
thaw (TDF DataFrame
df) = DataFrame
df

class ToDataFrame f where
    toDataFrame :: f -> D.DataFrame

instance ToDataFrame D.DataFrame where
    toDataFrame :: DataFrame -> DataFrame
toDataFrame = DataFrame -> DataFrame
forall a. a -> a
id

instance ToDataFrame (TypedDataFrame cols) where
    toDataFrame :: TypedDataFrame cols -> DataFrame
toDataFrame = TypedDataFrame cols -> DataFrame
forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> DataFrame
thaw

{- | Wrap an untyped DataFrame without any validation.
Used internally after delegation where the library guarantees schema correctness.
-}
unsafeFreeze :: D.DataFrame -> TypedDataFrame cols
unsafeFreeze :: forall (cols :: [(Symbol, *)]). DataFrame -> TypedDataFrame cols
unsafeFreeze = DataFrame -> TypedDataFrame cols
forall (cols :: [(Symbol, *)]). DataFrame -> TypedDataFrame cols
TDF

validateSchema ::
    forall cols.
    (KnownSchema cols) =>
    D.DataFrame -> Either T.Text ()
validateSchema :: forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> Either Text ()
validateSchema DataFrame
df = ((Text, SomeTypeRep) -> Either Text ())
-> [(Text, SomeTypeRep)] -> Either Text ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Text, SomeTypeRep) -> Either Text ()
checkCol (forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
[(Text, SomeTypeRep)]
schemaEvidence @cols)
  where
    checkCol :: (T.Text, SomeTypeRep) -> Either T.Text ()
    checkCol :: (Text, SomeTypeRep) -> Either Text ()
checkCol (Text
name, SomeTypeRep
expectedRep) = case Text -> DataFrame -> Maybe Column
D.getColumn Text
name DataFrame
df of
        Maybe Column
Nothing ->
            Text -> Either Text ()
forall a b. a -> Either a b
Left (Text -> Either Text ()) -> Text -> Either Text ()
forall a b. (a -> b) -> a -> b
$
                Text
"Column '"
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"' not found in DataFrame. "
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"Available columns: "
                    Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack ([Text] -> String
forall a. Show a => a -> String
show (DataFrame -> [Text]
columnNames DataFrame
df))
        Just Column
col ->
            if SomeTypeRep -> Column -> Bool
matchesType SomeTypeRep
expectedRep Column
col
                then () -> Either Text ()
forall a b. b -> Either a b
Right ()
                else
                    Text -> Either Text ()
forall a b. a -> Either a b
Left (Text -> Either Text ()) -> Text -> Either Text ()
forall a b. (a -> b) -> a -> b
$
                        Text
"Type mismatch on column '"
                            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name
                            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"': expected "
                            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (SomeTypeRep -> String
forall a. Show a => a -> String
show SomeTypeRep
expectedRep)
                            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
", got "
                            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Column -> String
C.columnTypeString Column
col)

{- | Check if a Column's element type matches the expected SomeTypeRep.
For nullable columns (those with a bitmap), @Maybe a@ in the schema matches
a column whose inner type is @a@, since we store nullable data as
@BoxedColumn (Just bm) a@ or @UnboxedColumn (Just bm) a@ rather than
@Column (Maybe a)@.
-}
matchesType :: SomeTypeRep -> C.Column -> Bool
matchesType :: SomeTypeRep -> Column -> Bool
matchesType SomeTypeRep
expected Column
col =
    let expectedStr :: String
expectedStr = SomeTypeRep -> String
forall a. Show a => a -> String
show SomeTypeRep
expected
        colTypeStr :: String
colTypeStr = Column -> String
C.columnTypeString Column
col
     in String
expectedStr String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
colTypeStr
            Bool -> Bool -> Bool
|| ( Column -> Bool
C.hasMissing Column
col -- nullable column: schema says "Maybe X", column stores "X" with a bitmap
                    Bool -> Bool -> Bool
&& String -> Maybe String
forall a. a -> Maybe a
Just String
colTypeStr Maybe String -> Maybe String -> Bool
forall a. Eq a => a -> a -> Bool
== String -> String -> Maybe String
forall a. Eq a => [a] -> [a] -> Maybe [a]
stripPrefix String
"Maybe " String
expectedStr
               )