{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

{- |
Module      : DataFrame.IO.Parquet.Schema
License     : MIT

Helpers for converting a parquet schema (a list of 'SchemaElement') into an
empty 'DataFrame' whose columns have the right types but no rows. Used by
'DataFrame.TH.declareColumnsFromParquetFile' and by any tooling that wants
to inspect a parquet schema as a 'DataFrame'.
-}
module DataFrame.IO.Parquet.Schema (
    schemaToEmptyDataFrame,
    schemaElemToColumn,
    emptyColumnForType,
    emptyNullableColumnForType,
) where

import Data.Int (Int32, Int64)
import qualified Data.Maybe as Maybe
import qualified Data.Set as S
import qualified Data.Text as T

import DataFrame.IO.Parquet.Thrift (
    SchemaElement,
    ThriftType (..),
    name,
    num_children,
    schematype,
    unField,
 )
import DataFrame.Internal.Column (Column, fromList)
import DataFrame.Internal.DataFrame (DataFrame, fromNamedColumns)

{- | Build an empty 'DataFrame' from a flat list of parquet 'SchemaElement's.
Only leaf elements (those with no children) become columns. Columns whose
name is in @nullableCols@ are typed as @Maybe a@; the rest are typed as @a@.
-}
schemaToEmptyDataFrame :: S.Set T.Text -> [SchemaElement] -> DataFrame
schemaToEmptyDataFrame :: Set Text -> [SchemaElement] -> DataFrame
schemaToEmptyDataFrame Set Text
nullableCols [SchemaElement]
elems =
    let leafElems :: [SchemaElement]
leafElems =
            (SchemaElement -> Bool) -> [SchemaElement] -> [SchemaElement]
forall a. (a -> Bool) -> [a] -> [a]
filter (\SchemaElement
e -> Int32 -> Maybe Int32 -> Int32
forall a. a -> Maybe a -> a
Maybe.fromMaybe Int32
0 (Field 5 (Maybe Int32) -> Maybe Int32
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField SchemaElement
e.num_children) Int32 -> Int32 -> Bool
forall a. Eq a => a -> a -> Bool
== Int32
0) [SchemaElement]
elems
     in [(Text, Column)] -> DataFrame
fromNamedColumns ((SchemaElement -> (Text, Column))
-> [SchemaElement] -> [(Text, Column)]
forall a b. (a -> b) -> [a] -> [b]
map (Set Text -> SchemaElement -> (Text, Column)
schemaElemToColumn Set Text
nullableCols) [SchemaElement]
leafElems)

{- | Convert a single parquet 'SchemaElement' into a named empty 'Column',
picking a nullable or non-nullable representation based on @nullableCols@.
-}
schemaElemToColumn :: S.Set T.Text -> SchemaElement -> (T.Text, Column)
schemaElemToColumn :: Set Text -> SchemaElement -> (Text, Column)
schemaElemToColumn Set Text
nullableCols SchemaElement
element =
    let colName :: Text
colName = Field 4 Text -> Text
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField SchemaElement
element.name
        isNull :: Bool
isNull = Text
colName Text -> Set Text -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Text
nullableCols
        column :: Column
column =
            if Bool
isNull
                then Maybe ThriftType -> Column
emptyNullableColumnForType (Field 1 (Maybe ThriftType) -> Maybe ThriftType
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField SchemaElement
element.schematype)
                else Maybe ThriftType -> Column
emptyColumnForType (Field 1 (Maybe ThriftType) -> Maybe ThriftType
forall (n :: Nat) a. KnownNat n => Field n a -> a
unField SchemaElement
element.schematype)
     in (Text
colName, Column
column)

-- | An empty 'Column' of the given parquet physical type.
emptyColumnForType :: Maybe ThriftType -> Column
emptyColumnForType :: Maybe ThriftType -> Column
emptyColumnForType = \case
    Just (BOOLEAN Enumeration 0
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @Bool []
    Just (INT32 Enumeration 1
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @Int32 []
    Just (INT64 Enumeration 2
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @Int64 []
    Just (INT96 Enumeration 3
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @Int64 []
    Just (FLOAT Enumeration 4
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @Float []
    Just (DOUBLE Enumeration 5
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @Double []
    Just (BYTE_ARRAY Enumeration 6
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @T.Text []
    Just (FIXED_LEN_BYTE_ARRAY Enumeration 7
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @T.Text []
    Maybe ThriftType
other -> [Char] -> Column
forall a. HasCallStack => [Char] -> a
error ([Char] -> Column) -> [Char] -> Column
forall a b. (a -> b) -> a -> b
$ [Char]
"Unsupported parquet type for column: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Maybe ThriftType -> [Char]
forall a. Show a => a -> [Char]
show Maybe ThriftType
other

-- | Like 'emptyColumnForType' but produces a nullable @Maybe a@ column.
emptyNullableColumnForType :: Maybe ThriftType -> Column
emptyNullableColumnForType :: Maybe ThriftType -> Column
emptyNullableColumnForType = \case
    Just (BOOLEAN Enumeration 0
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @(Maybe Bool) []
    Just (INT32 Enumeration 1
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @(Maybe Int32) []
    Just (INT64 Enumeration 2
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @(Maybe Int64) []
    Just (INT96 Enumeration 3
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @(Maybe Int64) []
    Just (FLOAT Enumeration 4
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @(Maybe Float) []
    Just (DOUBLE Enumeration 5
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @(Maybe Double) []
    Just (BYTE_ARRAY Enumeration 6
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @(Maybe T.Text) []
    Just (FIXED_LEN_BYTE_ARRAY Enumeration 7
_) -> forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
fromList @(Maybe T.Text) []
    Maybe ThriftType
other -> [Char] -> Column
forall a. HasCallStack => [Char] -> a
error ([Char] -> Column) -> [Char] -> Column
forall a b. (a -> b) -> a -> b
$ [Char]
"Unsupported parquet type for column: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Maybe ThriftType -> [Char]
forall a. Show a => a -> [Char]
show Maybe ThriftType
other