{-# LANGUAGE FlexibleContexts #-}

{- |
Module      : DataFrame.TH.Parquet
License     : MIT

Parquet-file-based 'DataFrame.TH' splices. Splits out the Parquet ingest
path so @dataframe-th@ stays IO-agnostic.
-}
module DataFrame.TH.Parquet (
    declareColumnsFromParquetFile,
) where

import Control.Monad (filterM)
import Control.Monad.IO.Class (liftIO)
import Data.Int (Int64)
import qualified Data.Maybe as Maybe
import qualified Data.Set as S
import qualified Data.Text as T

import Language.Haskell.TH
import System.Directory (doesDirectoryExist)
import System.FilePath ((</>))
import System.FilePath.Glob (glob)

import qualified DataFrame.IO.Parquet as Parquet
import DataFrame.IO.Parquet.Schema (schemaToEmptyDataFrame)
import DataFrame.IO.Parquet.Thrift (
    cc_meta_data,
    cmd_path_in_schema,
    cmd_statistics,
    rg_columns,
    row_groups,
    schema,
    stats_null_count,
    unField,
 )
import qualified DataFrame.Internal.DataFrame as DI
import DataFrame.TH.Records (declareColumns)
import Prelude as P

{- | Splice a binding for every column of a parquet file (or directory of
parquet files). The schema is read from each file's metadata and merged.
-}
declareColumnsFromParquetFile :: String -> DecsQ
declareColumnsFromParquetFile :: String -> DecsQ
declareColumnsFromParquetFile String
path = do
    Bool
isDir <- IO Bool -> Q Bool
forall a. IO a -> Q a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bool -> Q Bool) -> IO Bool -> Q Bool
forall a b. (a -> b) -> a -> b
$ String -> IO Bool
doesDirectoryExist String
path
    let pat :: String
pat = if Bool
isDir then String
path String -> String -> String
</> String
"*.parquet" else String
path
    [String]
matches <- IO [String] -> Q [String]
forall a. IO a -> Q a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [String] -> Q [String]) -> IO [String] -> Q [String]
forall a b. (a -> b) -> a -> b
$ String -> IO [String]
glob String
pat
    [String]
files <- IO [String] -> Q [String]
forall a. IO a -> Q a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [String] -> Q [String]) -> IO [String] -> Q [String]
forall a b. (a -> b) -> a -> b
$ (String -> IO Bool) -> [String] -> IO [String]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM ((Bool -> Bool) -> IO Bool -> IO Bool
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Bool -> Bool
P.not (IO Bool -> IO Bool) -> (String -> IO Bool) -> String -> IO Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO Bool
doesDirectoryExist) [String]
matches
    [FileMetadata]
metas <- IO [FileMetadata] -> Q [FileMetadata]
forall a. IO a -> Q a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [FileMetadata] -> Q [FileMetadata])
-> IO [FileMetadata] -> Q [FileMetadata]
forall a b. (a -> b) -> a -> b
$ (String -> IO FileMetadata) -> [String] -> IO [FileMetadata]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM String -> IO FileMetadata
Parquet.readMetadataFromPath [String]
files
    let nullableCols :: S.Set T.Text
        nullableCols :: Set Text
nullableCols =
            [Text] -> Set Text
forall a. Ord a => [a] -> Set a
S.fromList
                [ String -> Text
T.pack ([String] -> String
forall a. HasCallStack => [a] -> a
last [String]
colPath)
                | FileMetadata
meta <- [FileMetadata]
metas
                , RowGroup
rg <- Field 4 [RowGroup] -> [RowGroup]
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (FileMetadata -> Field 4 [RowGroup]
row_groups FileMetadata
meta)
                , ColumnChunk
cc <- Field 1 [ColumnChunk] -> [ColumnChunk]
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (RowGroup -> Field 1 [ColumnChunk]
rg_columns RowGroup
rg)
                , Just ColumnMetaData
cm <- [Field 3 (Maybe ColumnMetaData) -> Maybe ColumnMetaData
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (ColumnChunk -> Field 3 (Maybe ColumnMetaData)
cc_meta_data ColumnChunk
cc)]
                , let colPath :: [String]
colPath = (Text -> String) -> [Text] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Text -> String
T.unpack (Field 3 [Text] -> [Text]
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (ColumnMetaData -> Field 3 [Text]
cmd_path_in_schema ColumnMetaData
cm))
                , Bool -> Bool
P.not ([String] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [String]
colPath)
                , let nc :: Int64
                      nc :: Int64
nc = case Field 12 (Maybe Statistics) -> Maybe Statistics
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (ColumnMetaData -> Field 12 (Maybe Statistics)
cmd_statistics ColumnMetaData
cm) of
                        Maybe Statistics
Nothing -> Int64
0
                        Just Statistics
stats ->
                            Int64 -> Maybe Int64 -> Int64
forall a. a -> Maybe a -> a
Maybe.fromMaybe Int64
0 (Field 3 (Maybe Int64) -> Maybe Int64
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (Field 3 (Maybe Int64) -> Maybe Int64)
-> Field 3 (Maybe Int64) -> Maybe Int64
forall a b. (a -> b) -> a -> b
$ Statistics -> Field 3 (Maybe Int64)
stats_null_count Statistics
stats)
                , Int64
nc Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
> Int64
0
                ]
    let df :: DataFrame
df =
            (DataFrame -> FileMetadata -> DataFrame)
-> DataFrame -> [FileMetadata] -> DataFrame
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
                ( \DataFrame
acc FileMetadata
meta ->
                    DataFrame
acc
                        DataFrame -> DataFrame -> DataFrame
forall a. Semigroup a => a -> a -> a
<> Set Text -> [SchemaElement] -> DataFrame
schemaToEmptyDataFrame
                            Set Text
nullableCols
                            (Field 2 [SchemaElement] -> [SchemaElement]
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField (FileMetadata -> Field 2 [SchemaElement]
schema FileMetadata
meta))
                )
                DataFrame
DI.empty
                [FileMetadata]
metas

    DataFrame -> DecsQ
declareColumns DataFrame
df