module Spark.Core.Internal.TypesStructures where
import Data.Aeson
import Data.Vector(Vector)
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
| StringType
| Struct !StructType
| ArrayType { elementType :: !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 IntType = "int"
show (Struct struct) = show struct
show array = "[" ++ show (elementType array) ++ "]"
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 StringType = "string"
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]
_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])