{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{- | Persist dataframe expressions and pipelines, and load them back for
inference in another process (or from Python — see "DataFrame.IR.ExprJson" for
the shared wire format).

  * A single 'Expr' round-trips via 'encodeExpr' / 'decodeExprAny' (or the
    type-asserting 'decodeExprAt') and the file helpers below.
  * A /pipeline/ is an ordered @['NamedExpr']@ — the same shape a fitted
    @Transform@ carries. 'encodeNamedExprs' / 'decodeNamedExprs' wrap it in a
    versioned envelope; replay a decoded pipeline for inference with
    @deriveMany@ / @applyTransform@.

All functions are pure 'Either'; the @save@ helpers report an encode failure
before touching the file (no partial write), and the @load@ helpers turn a
missing/unreadable file into a 'Left' rather than throwing.

Limitations inherited from the wire format: @CastWith@ / @CastExprWith@ (they
hold closures) and @collect@ (its list-typed output has no type tag) cannot be
serialized.
-}
module DataFrame.Expr.Serialize (
    -- * Single expressions (re-exported wire codec)
    SomeExpr (..),
    encodeExpr,
    encodeExprToBytes,
    decodeExprAny,
    decodeExprAt,
    decodeExprFromBytes,

    -- * Pipelines
    encodeNamedExprs,
    decodeNamedExprs,

    -- * File IO
    saveExprToFile,
    loadExprFromFile,
    loadExprAtFromFile,
    savePipelineToFile,
    loadPipelineFromFile,
) where

import Control.Exception (IOException, try)
import Control.Monad (when)
import Data.Aeson (object, (.:), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Types as Aeson
import Data.Bifunctor (first)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import qualified Data.Text as T

import DataFrame.IR.ExprJson (
    SomeExpr (..),
    decodeExprAny,
    decodeExprAt,
    encodeExpr,
    encodeExprToBytes,
    parseSomeExpr,
 )
import DataFrame.Internal.Column (Columnable)
import DataFrame.Internal.Expression (Expr, NamedExpr, UExpr (..))

-- | Decode a single expression from a strict JSON 'BS.ByteString'.
decodeExprFromBytes :: BS.ByteString -> Either String SomeExpr
decodeExprFromBytes :: ByteString -> Either String SomeExpr
decodeExprFromBytes ByteString
bs = ByteString -> Either String Value
forall a. FromJSON a => ByteString -> Either String a
Aeson.eitherDecodeStrict ByteString
bs Either String Value
-> (Value -> Either String SomeExpr) -> Either String SomeExpr
forall a b.
Either String a -> (a -> Either String b) -> Either String b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Value -> Either String SomeExpr
decodeExprAny

{- | Encode an ordered pipeline (named output expressions) into the versioned
envelope @{ "version": 1, "outputs": [ { "name", "expr" }, … ] }@.
-}
encodeNamedExprs :: [NamedExpr] -> Either String Aeson.Value
encodeNamedExprs :: [NamedExpr] -> Either String Value
encodeNamedExprs [NamedExpr]
nes = do
    [Value]
outs <- (NamedExpr -> Either String Value)
-> [NamedExpr] -> Either String [Value]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse NamedExpr -> Either String Value
encodeOne [NamedExpr]
nes
    Value -> Either String Value
forall a b. b -> Either a b
Right (Value -> Either String Value) -> Value -> Either String Value
forall a b. (a -> b) -> a -> b
$ [Pair] -> Value
object [Key
"version" Key -> Int -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= (Int
1 :: Int), Key
"outputs" Key -> [Value] -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= [Value]
outs]
  where
    encodeOne :: NamedExpr -> Either String Aeson.Value
    encodeOne :: NamedExpr -> Either String Value
encodeOne (Text
name, UExpr Expr a
e) = do
        Value
ev <- Expr a -> Either String Value
forall a. Columnable a => Expr a -> Either String Value
encodeExpr Expr a
e
        Value -> Either String Value
forall a b. b -> Either a b
Right (Value -> Either String Value) -> Value -> Either String Value
forall a b. (a -> b) -> a -> b
$ [Pair] -> Value
object [Key
"name" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Text
name, Key
"expr" Key -> Value -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Value
ev]

-- | Decode a pipeline produced by 'encodeNamedExprs'.
decodeNamedExprs :: Aeson.Value -> Either String [NamedExpr]
decodeNamedExprs :: Value -> Either String [NamedExpr]
decodeNamedExprs = (Value -> Parser [NamedExpr]) -> Value -> Either String [NamedExpr]
forall a b. (a -> Parser b) -> a -> Either String b
Aeson.parseEither Value -> Parser [NamedExpr]
parsePipeline
  where
    parsePipeline :: Value -> Parser [NamedExpr]
parsePipeline = String
-> (Object -> Parser [NamedExpr]) -> Value -> Parser [NamedExpr]
forall a. String -> (Object -> Parser a) -> Value -> Parser a
Aeson.withObject String
"Pipeline" ((Object -> Parser [NamedExpr]) -> Value -> Parser [NamedExpr])
-> (Object -> Parser [NamedExpr]) -> Value -> Parser [NamedExpr]
forall a b. (a -> b) -> a -> b
$ \Object
o -> do
        Int
ver <- Object
o Object -> Key -> Parser Int
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"version" :: Aeson.Parser Int
        Bool -> Parser () -> Parser ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ver Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
1) (Parser () -> Parser ()) -> Parser () -> Parser ()
forall a b. (a -> b) -> a -> b
$
            String -> Parser ()
forall a. String -> Parser a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Parser ()) -> String -> Parser ()
forall a b. (a -> b) -> a -> b
$
                String
"DataFrame.Expr.Serialize: unsupported pipeline version " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
ver
        [Value]
outs <- Object
o Object -> Key -> Parser [Value]
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"outputs"
        (Value -> Parser NamedExpr) -> [Value] -> Parser [NamedExpr]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse Value -> Parser NamedExpr
parseOne [Value]
outs
    parseOne :: Aeson.Value -> Aeson.Parser NamedExpr
    parseOne :: Value -> Parser NamedExpr
parseOne = String -> (Object -> Parser NamedExpr) -> Value -> Parser NamedExpr
forall a. String -> (Object -> Parser a) -> Value -> Parser a
Aeson.withObject String
"PipelineOutput" ((Object -> Parser NamedExpr) -> Value -> Parser NamedExpr)
-> (Object -> Parser NamedExpr) -> Value -> Parser NamedExpr
forall a b. (a -> b) -> a -> b
$ \Object
o -> do
        Text
name <- Object
o Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"name" :: Aeson.Parser T.Text
        Value
rawExpr <- Object
o Object -> Key -> Parser Value
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"expr"
        SomeExpr TypeRep a
_ Expr a
e <- Value -> Parser SomeExpr
parseSomeExpr Value
rawExpr
        NamedExpr -> Parser NamedExpr
forall a. a -> Parser a
forall (m :: * -> *) a. Monad m => a -> m a
return (Text
name, Expr a -> UExpr
forall a. Columnable a => Expr a -> UExpr
UExpr Expr a
e)

-- | Encode an expression and write it to a file. No file is written on failure.
saveExprToFile :: (Columnable a) => FilePath -> Expr a -> IO (Either String ())
saveExprToFile :: forall a. Columnable a => String -> Expr a -> IO (Either String ())
saveExprToFile String
fp Expr a
e = case Expr a -> Either String ByteString
forall a. Columnable a => Expr a -> Either String ByteString
encodeExprToBytes Expr a
e of
    Left String
err -> Either String () -> IO (Either String ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ()
forall a b. a -> Either a b
Left String
err)
    Right ByteString
bs -> String -> ByteString -> IO (Either String ())
writeBytes String
fp ByteString
bs

-- | Load an expression of unknown output type from a file.
loadExprFromFile :: FilePath -> IO (Either String SomeExpr)
loadExprFromFile :: String -> IO (Either String SomeExpr)
loadExprFromFile String
fp = (Either String ByteString
-> (ByteString -> Either String SomeExpr) -> Either String SomeExpr
forall a b.
Either String a -> (a -> Either String b) -> Either String b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ByteString -> Either String SomeExpr
decodeExprFromBytes) (Either String ByteString -> Either String SomeExpr)
-> IO (Either String ByteString) -> IO (Either String SomeExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Either String ByteString)
readBytes String
fp

{- | Load an expression from a file, asserting its output type. Fails with
'Left' if the stored expression has a different output type.
-}
loadExprAtFromFile ::
    forall a. (Columnable a) => FilePath -> IO (Either String (Expr a))
loadExprAtFromFile :: forall a. Columnable a => String -> IO (Either String (Expr a))
loadExprAtFromFile String
fp = (Either String ByteString
-> (ByteString -> Either String (Expr a)) -> Either String (Expr a)
forall a b.
Either String a -> (a -> Either String b) -> Either String b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ByteString -> Either String (Expr a)
decode) (Either String ByteString -> Either String (Expr a))
-> IO (Either String ByteString) -> IO (Either String (Expr a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Either String ByteString)
readBytes String
fp
  where
    decode :: ByteString -> Either String (Expr a)
decode ByteString
bs = ByteString -> Either String Value
forall a. FromJSON a => ByteString -> Either String a
Aeson.eitherDecodeStrict ByteString
bs Either String Value
-> (Value -> Either String (Expr a)) -> Either String (Expr a)
forall a b.
Either String a -> (a -> Either String b) -> Either String b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a. Columnable a => Value -> Either String (Expr a)
decodeExprAt @a

-- | Encode a pipeline and write it to a file. No file is written on failure.
savePipelineToFile :: FilePath -> [NamedExpr] -> IO (Either String ())
savePipelineToFile :: String -> [NamedExpr] -> IO (Either String ())
savePipelineToFile String
fp [NamedExpr]
nes = case [NamedExpr] -> Either String Value
encodeNamedExprs [NamedExpr]
nes of
    Left String
err -> Either String () -> IO (Either String ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ()
forall a b. a -> Either a b
Left String
err)
    Right Value
v -> String -> ByteString -> IO (Either String ())
writeBytes String
fp (LazyByteString -> ByteString
BL.toStrict (Value -> LazyByteString
forall a. ToJSON a => a -> LazyByteString
Aeson.encode Value
v))

-- | Load a pipeline produced by 'savePipelineToFile'.
loadPipelineFromFile :: FilePath -> IO (Either String [NamedExpr])
loadPipelineFromFile :: String -> IO (Either String [NamedExpr])
loadPipelineFromFile String
fp = (Either String ByteString
-> (ByteString -> Either String [NamedExpr])
-> Either String [NamedExpr]
forall a b.
Either String a -> (a -> Either String b) -> Either String b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ByteString -> Either String [NamedExpr]
decode) (Either String ByteString -> Either String [NamedExpr])
-> IO (Either String ByteString) -> IO (Either String [NamedExpr])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Either String ByteString)
readBytes String
fp
  where
    decode :: ByteString -> Either String [NamedExpr]
decode ByteString
bs = ByteString -> Either String Value
forall a. FromJSON a => ByteString -> Either String a
Aeson.eitherDecodeStrict ByteString
bs Either String Value
-> (Value -> Either String [NamedExpr])
-> Either String [NamedExpr]
forall a b.
Either String a -> (a -> Either String b) -> Either String b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Value -> Either String [NamedExpr]
decodeNamedExprs

writeBytes :: FilePath -> BS.ByteString -> IO (Either String ())
writeBytes :: String -> ByteString -> IO (Either String ())
writeBytes String
fp ByteString
bs = (IOException -> String)
-> Either IOException () -> Either String ()
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first IOException -> String
showIO (Either IOException () -> Either String ())
-> IO (Either IOException ()) -> IO (Either String ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO () -> IO (Either IOException ())
forall e a. Exception e => IO a -> IO (Either e a)
try (String -> ByteString -> IO ()
BS.writeFile String
fp ByteString
bs)

readBytes :: FilePath -> IO (Either String BS.ByteString)
readBytes :: String -> IO (Either String ByteString)
readBytes String
fp = (IOException -> String)
-> Either IOException ByteString -> Either String ByteString
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first IOException -> String
showIO (Either IOException ByteString -> Either String ByteString)
-> IO (Either IOException ByteString)
-> IO (Either String ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ByteString -> IO (Either IOException ByteString)
forall e a. Exception e => IO a -> IO (Either e a)
try (String -> IO ByteString
BS.readFile String
fp)

showIO :: IOException -> String
showIO :: IOException -> String
showIO = IOException -> String
forall a. Show a => a -> String
show