module Spark.Core.Internal.TypesStructures where
import Data.Aeson
import Data.Vector(Vector)
import Control.Monad(guard)
import qualified Data.Vector as V
import qualified Data.Aeson as A
import qualified Data.Text as T
import GHC.Generics(Generic)
import Test.QuickCheck
import Spark.Core.StructuresInternal(FieldName(..))
import Spark.Core.Internal.Utilities
data StrictDataType =
IntType
| DoubleType
| StringType
| BoolType
| Struct !StructType
| ArrayType !DataType
deriving (Eq)
data DataType =
StrictType !StrictDataType
| NullableType !StrictDataType deriving (Eq)
data StructField = StructField {
structFieldName :: !FieldName,
structFieldType :: !DataType
} deriving (Eq)
data StructType = StructType {
structFields :: !(Vector StructField)
} deriving (Eq)
data Nullable = CanNull | NoNull deriving (Show, Eq)
data NullableDataType = NullableDataType !StrictDataType deriving (Eq)
data SQLType a = SQLType {
unSQLType :: !DataType
} deriving (Eq, Generic)
instance Show DataType where
show (StrictType x) = show x
show (NullableType x) = show x ++ "?"
instance Show StrictDataType where
show StringType = "string"
show DoubleType = "double"
show IntType = "int"
show BoolType = "bool"
show (Struct struct) = show struct
show (ArrayType at) = "[" ++ show at ++ "]"
instance Show StructField where
show field = (T.unpack . unFieldName . structFieldName) field ++ ":" ++ s where
s = show $ structFieldType field
instance Show StructType where
show struct = "{" ++ unwords (map show (V.toList . structFields $ struct)) ++ "}"
instance Show (SQLType a) where
show (SQLType dt) = show dt
instance Arbitrary StructField where
arbitrary = do
name <- elements ["_1", "a", "b", "abc"]
dt <- arbitrary :: Gen DataType
return $ StructField (FieldName $ T.pack name) dt
instance Arbitrary StructType where
arbitrary = do
fields <- listOf arbitrary
return . StructType . V.fromList $ fields
instance Arbitrary StrictDataType where
arbitrary = do
idx <- elements [1,2] :: Gen Int
return $ case idx of
1 -> StringType
2 -> IntType
_ -> failure "Arbitrary StrictDataType"
instance Arbitrary DataType where
arbitrary = do
x <- arbitrary
u <- arbitrary
return $ if x then
StrictType u
else
NullableType u
instance ToJSON StrictDataType where
toJSON IntType = "integer"
toJSON DoubleType = "double"
toJSON StringType = "string"
toJSON BoolType = "bool"
toJSON (Struct struct) = toJSON struct
toJSON (ArrayType (StrictType dt)) =
object [ "type" .= A.String "array"
, "elementType" .= toJSON dt
, "containsNull" .= A.Bool False ]
toJSON (ArrayType (NullableType dt)) =
object [ "type" .= A.String "array"
, "elementType" .= toJSON dt
, "containsNull" .= A.Bool True ]
instance ToJSON StructType where
toJSON (StructType fields) =
let
fs = (snd . _fieldToJson) <$> V.toList fields
in object [ "type" .= A.String "struct"
, "fields" .= fs ]
instance ToJSON DataType where
toJSON (StrictType dt) = object [
"nullable" .= A.Bool False,
"dt" .= toJSON dt]
toJSON (NullableType dt) = object [
"nullable" .= A.Bool True,
"dt" .= toJSON dt]
instance FromJSON DataType where
parseJSON = withObject "DataType" $ \o -> do
nullable <- o .: "nullable"
dt <- o .: "dt"
let c = if nullable then NullableType else StrictType
return (c dt)
instance FromJSON StructField where
parseJSON = withObject "StructField" $ \o -> do
n <- o .: "name"
dt <- o .: "type"
nullable <- o .: "nullable"
let c = if nullable then NullableType else StrictType
return $ StructField (FieldName n) (c dt)
instance FromJSON StructType where
parseJSON = withObject "StructType" $ \o -> do
tp <- o .: "type"
guard (tp == T.pack "struct")
fs <- o .: "fields"
return (StructType fs)
instance FromJSON StrictDataType where
parseJSON (A.String s) = case s of
"integer" -> return IntType
"double" -> return DoubleType
"string" -> return StringType
"bool" -> return BoolType
"boolean" -> return BoolType
_ -> fail ("StrictDataType: unknown type " ++ T.unpack s)
parseJSON (Object o) = do
tp <- o .: "type"
case T.pack tp of
"struct" -> Struct <$> parseJSON (Object o)
"array" -> do
dt <- o .: "elementType"
containsNull <- o .: "containsNull"
let c = if containsNull then NullableType else StrictType
return $ ArrayType (c dt)
s -> fail ("StrictDataType: unknown type " ++ T.unpack s)
parseJSON x = fail ("StrictDataType: cannot parse " ++ show x)
_fieldToJson :: StructField -> (T.Text, A.Value)
_fieldToJson (StructField (FieldName n) (StrictType dt)) =
(n, object [ "name" .= A.String n
, "type" .= toJSON dt
, "nullable" .= A.Bool False])
_fieldToJson (StructField (FieldName n) (NullableType dt)) =
(n, object [ "name" .= A.String n
, "type" .= toJSON dt
, "nullable" .= A.Bool True])