module Spark.Core.Internal.TypesFunctions(
isNullable,
iInnerStrictType,
columnType,
unsafeCastType,
intType,
arrayType,
compatibleTypes,
arrayType',
frameTypeFromCol,
colTypeFromFrame,
canNull,
structField,
structType,
structTypeFromFields,
tupleType,
structName,
iSingleField,
) where
import qualified Data.Text as T
import Data.List(sort, nub)
import qualified Data.Vector as V
import Data.Text(Text, intercalate)
import Formatting
import Spark.Core.Internal.TypesStructures
import Spark.Core.StructuresInternal
import Spark.Core.Internal.Utilities
import Spark.Core.Try
unsafeCastType :: SQLType a -> SQLType b
unsafeCastType (SQLType dt) = SQLType dt
columnType :: SQLType a -> DataType
columnType (SQLType dt) = dt
isNullable :: DataType -> Bool
isNullable (StrictType _) = False
isNullable (NullableType _) = True
frameTypeFromCol :: DataType -> StructType
frameTypeFromCol (StrictType (Struct struct)) = struct
frameTypeFromCol dt = _structFromUnfields [("value", dt)]
colTypeFromFrame :: StructType -> DataType
colTypeFromFrame st @ (StructType fs) = case V.toList fs of
[StructField {
structFieldName = fname,
structFieldType = (StrictType dt)}] | fname == "value" ->
StrictType dt
_ -> StrictType (Struct st)
compatibleTypes :: DataType -> DataType -> Bool
compatibleTypes (StrictType sdt) (StrictType sdt') = _compatibleTypesStrict sdt sdt'
compatibleTypes (NullableType sdt) (NullableType sdt') = _compatibleTypesStrict sdt sdt'
compatibleTypes _ _ = False
_compatibleTypesStrict :: StrictDataType -> StrictDataType -> Bool
_compatibleTypesStrict IntType IntType = True
_compatibleTypesStrict StringType StringType = True
_compatibleTypesStrict (ArrayType et) (ArrayType et') = compatibleTypes et et'
_compatibleTypesStrict (Struct (StructType v)) (Struct (StructType v')) =
(length v == length v') &&
and (V.zipWith compatibleTypes (structFieldType <$> v) (structFieldType <$> v'))
_compatibleTypesStrict _ _ = False
tupleType :: SQLType a -> SQLType b -> SQLType (a, b)
tupleType (SQLType dt1) (SQLType dt2) =
SQLType $ structType [structField "_1" dt1, structField "_2" dt2]
intType :: DataType
intType = StrictType IntType
structField :: T.Text -> DataType -> StructField
structField txt = StructField (FieldName txt)
structType :: [StructField] -> DataType
structType = StrictType . Struct . StructType . V.fromList
arrayType' :: DataType -> DataType
arrayType' = StrictType . ArrayType
canNull :: DataType -> DataType
canNull = NullableType . iInnerStrictType
arrayType :: SQLType a -> SQLType [a]
arrayType (SQLType dt) = SQLType (arrayType' dt)
iInnerStrictType :: DataType -> StrictDataType
iInnerStrictType (StrictType st) = st
iInnerStrictType (NullableType st) = st
iSingleField :: DataType -> Maybe DataType
iSingleField (StrictType (Struct (StructType fields))) = case V.toList fields of
[StructField _ dt] -> Just dt
_ -> Nothing
iSingleField _ = Nothing
structName :: StructType -> Text
structName (StructType fields) =
"struct(" <> intercalate "," (unFieldName . structFieldName <$> V.toList fields) <> ")"
structTypeFromFields :: [(FieldName, DataType)] -> Try StructType
structTypeFromFields [] = tryError "You cannot build an empty structure"
structTypeFromFields ((hfn, hdt):t) =
let fs = (hfn, hdt) : t
ct = StructType $ uncurry StructField <$> V.fromList fs
names = fst <$> fs
numNames = length names
numDistincts = length . nub $ names
in if numNames == numDistincts
then return ct
else tryError $ sformat ("Duplicate field names when building the struct: "%sh) (sort names)
_structFromUnfields :: [(T.Text, DataType)] -> StructType
_structFromUnfields l = StructType . V.fromList $ x where
x = [StructField (FieldName name) dt | (name, dt) <- l]