{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{- | Typed Parquet reading.

The reader validates the file against a type-level schema as it loads, supplied
by type application:

@
type Trips = '[ '(\"id\", Int), '(\"fare\", Double)]

trips <- readParquet \@Trips \"trips.parquet\"   -- IO (TypedDataFrame Trips)
@

'readParquet' (and 'readParquetWithOpts'\/'readParquetFiles') throw a
'DataFrameException' on schema mismatch; 'readParquetWithError' returns the
mismatch as an 'Either' instead.
-}
module DataFrame.Typed.IO.Parquet (
    readParquet,
    readParquetWithError,
    readParquetWithOpts,
    readParquetFiles,
) where

import Control.Applicative ((<|>))
import Control.Exception (SomeException, try)
import qualified Data.Text as T

import DataFrame.IO.Parquet (ParquetReadOptions (..), defaultParquetReadOptions)
import qualified DataFrame.IO.Parquet as Parquet
import DataFrame.Typed.Freeze (freezeOrThrow, freezeWithError)
import DataFrame.Typed.Schema (KnownSchema, schemaColumnNames)
import DataFrame.Typed.Types (TypedDataFrame)

{- | Read only the columns @cols@ names, unless the caller selected columns
themselves. Parquet supplies the element types, so the schema contributes the
projection and the freeze checks the types.
-}
schemaOptions ::
    forall cols. (KnownSchema cols) => ParquetReadOptions -> ParquetReadOptions
schemaOptions :: forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
ParquetReadOptions -> ParquetReadOptions
schemaOptions ParquetReadOptions
opts =
    ParquetReadOptions
opts
        { selectedColumns = selectedColumns opts <|> Just (schemaColumnNames @cols)
        }

{- | Read a Parquet file into a typed DataFrame, throwing on schema mismatch.
Reads only the columns @cols@ names.

==== __Example__
@
ghci> trips <- readParquet \@Trips \"trips.parquet\"
@
-}
readParquet ::
    forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols)
readParquet :: forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
FilePath -> IO (TypedDataFrame cols)
readParquet = forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
ParquetReadOptions -> FilePath -> IO (TypedDataFrame cols)
readParquetWithOpts @cols ParquetReadOptions
defaultParquetReadOptions

{- | Read a Parquet file, returning a descriptive error on schema mismatch or
a missing column instead of throwing.

==== __Example__
@
ghci> readParquetWithError \@Trips \"trips.parquet\"
Right (TDF ...)
@
-}
readParquetWithError ::
    forall cols.
    (KnownSchema cols) =>
    FilePath -> IO (Either T.Text (TypedDataFrame cols))
readParquetWithError :: forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
FilePath -> IO (Either Text (TypedDataFrame cols))
readParquetWithError FilePath
path = do
    Either SomeException DataFrame
r <-
        IO DataFrame -> IO (Either SomeException DataFrame)
forall e a. Exception e => IO a -> IO (Either e a)
try
            (ParquetReadOptions -> FilePath -> IO DataFrame
Parquet.readParquetWithOpts (forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
ParquetReadOptions -> ParquetReadOptions
schemaOptions @cols ParquetReadOptions
defaultParquetReadOptions) FilePath
path)
    Either Text (TypedDataFrame cols)
-> IO (Either Text (TypedDataFrame cols))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Text (TypedDataFrame cols)
 -> IO (Either Text (TypedDataFrame cols)))
-> Either Text (TypedDataFrame cols)
-> IO (Either Text (TypedDataFrame cols))
forall a b. (a -> b) -> a -> b
$ case Either SomeException DataFrame
r of
        Left (SomeException
e :: SomeException) -> Text -> Either Text (TypedDataFrame cols)
forall a b. a -> Either a b
Left (FilePath -> Text
T.pack (SomeException -> FilePath
forall a. Show a => a -> FilePath
show SomeException
e))
        Right DataFrame
df -> DataFrame -> Either Text (TypedDataFrame cols)
forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> Either Text (TypedDataFrame cols)
freezeWithError DataFrame
df

{- | Read a Parquet file with custom options, throwing on schema mismatch. The
schema still supplies the column selection; an explicit 'selectedColumns'
takes precedence.

==== __Example__
@
ghci> trips <- readParquetWithOpts \@Trips defaultParquetReadOptions{rowRange = Just (0, 10)} \"trips.parquet\"
@
-}
readParquetWithOpts ::
    forall cols.
    (KnownSchema cols) =>
    ParquetReadOptions -> FilePath -> IO (TypedDataFrame cols)
readParquetWithOpts :: forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
ParquetReadOptions -> FilePath -> IO (TypedDataFrame cols)
readParquetWithOpts ParquetReadOptions
opts FilePath
path =
    ParquetReadOptions -> FilePath -> IO DataFrame
Parquet.readParquetWithOpts (forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
ParquetReadOptions -> ParquetReadOptions
schemaOptions @cols ParquetReadOptions
opts) FilePath
path
        IO DataFrame
-> (DataFrame -> IO (TypedDataFrame cols))
-> IO (TypedDataFrame cols)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> IO (TypedDataFrame cols)
freezeOrThrow @cols

{- | Read a directory\/glob of Parquet files into a typed DataFrame.

==== __Example__
@
ghci> trips <- readParquetFiles \@Trips \".\/data\/trips\/*.parquet\"
@
-}
readParquetFiles ::
    forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols)
readParquetFiles :: forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
FilePath -> IO (TypedDataFrame cols)
readParquetFiles FilePath
path =
    ParquetReadOptions -> FilePath -> IO DataFrame
Parquet.readParquetFilesWithOpts
        (forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
ParquetReadOptions -> ParquetReadOptions
schemaOptions @cols ParquetReadOptions
defaultParquetReadOptions)
        FilePath
path
        IO DataFrame
-> (DataFrame -> IO (TypedDataFrame cols))
-> IO (TypedDataFrame cols)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> IO (TypedDataFrame cols)
freezeOrThrow @cols