{-# LANGUAGE FlexibleContexts #-}
-- | Splices that derive column bindings from Parquet files.
module DataFrame.TH.Parquet (
declareColumnsFromParquetFile,
deriveSchemaValuesFromParquetFile,
) 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, declareSchemaValues)
import Prelude as P
-- | Splice a binding for every column of a parquet file or directory.
declareColumnsFromParquetFile :: String -> DecsQ
declareColumnsFromParquetFile path =
liftIO (emptyFrameFromParquet path) >>= declareColumns
deriveSchemaValuesFromParquetFile :: String -> String -> DecsQ
deriveSchemaValuesFromParquetFile prefix path =
liftIO (emptyFrameFromParquet path) >>= declareSchemaValues prefix
emptyFrameFromParquet :: String -> IO DI.DataFrame
emptyFrameFromParquet path = do
isDir <- doesDirectoryExist path
let pat = if isDir then path </> "*.parquet" else path
matches <- glob pat
files <- filterM (fmap P.not . doesDirectoryExist) matches
metas <- mapM Parquet.readMetadataFromPath files
let nullableCols :: S.Set T.Text
nullableCols =
S.fromList
[ T.pack (last colPath)
| meta <- metas
, rg <- unField (row_groups meta)
, cc <- unField (rg_columns rg)
, Just cm <- [unField (cc_meta_data cc)]
, let colPath = map T.unpack (unField (cmd_path_in_schema cm))
, P.not (null colPath)
, let nc :: Int64
nc = case unField (cmd_statistics cm) of
Nothing -> 0
Just stats ->
Maybe.fromMaybe 0 (unField $ stats_null_count stats)
, nc > 0
]
pure $
foldl
( \acc meta ->
acc
<> schemaToEmptyDataFrame
nullableCols
(unField (schema meta))
)
DI.empty
metas