diff --git a/krapsh.cabal b/krapsh.cabal --- a/krapsh.cabal +++ b/krapsh.cabal @@ -1,5 +1,5 @@ name: krapsh -version: 0.1.6.2 +version: 0.1.9.0 cabal-version: >=1.10 build-type: Simple license: Apache-2.0 @@ -37,12 +37,15 @@ Spark.Core.Internal.DAGStructures Spark.Core.Internal.DatasetFunctions Spark.Core.Internal.DatasetStructures + Spark.Core.Internal.Groups + Spark.Core.Internal.Joins Spark.Core.Internal.LocalDataFunctions Spark.Core.Internal.OpFunctions Spark.Core.Internal.OpStructures Spark.Core.Internal.Paths Spark.Core.Internal.PathsUntyped Spark.Core.Internal.Utilities + Spark.Core.Internal.TypesGenerics Spark.Core.Internal.TypesStructures Spark.Core.Internal.TypesFunctions Spark.Core.Row @@ -54,7 +57,7 @@ aeson-pretty >=0.8.2 && <0.9, base >=4.8.1 && <5, base16-bytestring >=0.1.1.6 && <0.2, - binary ==0.8.3.0, + binary >=0.7 && <0.9, bytestring >=0.10.8.1 && <0.11, containers >=0.5.7.1 && <0.6, cryptohash-sha256 >=0.11.100.1 && <0.12, @@ -92,7 +95,6 @@ Spark.Core.Internal.RowGenericsFrom Spark.Core.Internal.RowStructures Spark.Core.Internal.RowUtils - Spark.Core.Internal.TypesGenerics ghc-options: -Wall test-suite krapsh-test @@ -104,7 +106,7 @@ bytestring >=0.10.8.1 && <0.11, containers >=0.5.7.1 && <0.6, formatting >=6.2.4 && <6.3, - krapsh >=0.1.6.2 && <0.2, + krapsh >=0.1.9.0 && <0.2, hspec ==2.*, text >=1.2.2.1 && <1.3, raw-strings-qq ==1.1.*, @@ -121,6 +123,7 @@ Spark.Core.Internal.RowUtilsSpec Spark.Core.Internal.DAGFunctionsSpec Spark.Core.Internal.PathsSpec + Spark.Core.Internal.GroupsSpec Spark.Core.PathSpec Spark.Core.ProjectionsSpec Spark.Core.RowToSQLSpec @@ -137,7 +140,7 @@ bytestring >=0.10.8.1 && <0.11, containers >=0.5.7.1 && <0.6, formatting >=6.2.4 && <6.3, - krapsh >=0.1.6.2 && <0.2, + krapsh >=0.1.9.0 && <0.2, hspec ==2.*, text >=1.2.2.1 && <1.3, raw-strings-qq ==1.1.*, @@ -148,6 +151,8 @@ other-modules: Spark.Core.CachingSpec Spark.Core.CollectSpec + Spark.Core.GroupsSpec Spark.Core.IntegrationUtilities + Spark.Core.JoinsSpec Spark.Core.SimpleAddSpec ghc-options: -fhpc -O0 -Wall diff --git a/src/Spark/Core/Column.hs b/src/Spark/Core/Column.hs --- a/src/Spark/Core/Column.hs +++ b/src/Spark/Core/Column.hs @@ -12,6 +12,7 @@ DynColumn, -- * Extractions and collations asCol, + asCol', pack1, pack, pack', diff --git a/src/Spark/Core/ColumnFunctions.hs b/src/Spark/Core/ColumnFunctions.hs --- a/src/Spark/Core/ColumnFunctions.hs +++ b/src/Spark/Core/ColumnFunctions.hs @@ -10,7 +10,8 @@ -- * Arithmetic operations (.+), -- * Reductions - colSum, + sumCol, + sumCol' ) where import Spark.Core.Internal.AlgebraStructures diff --git a/src/Spark/Core/Context.hs b/src/Spark/Core/Context.hs --- a/src/Spark/Core/Context.hs +++ b/src/Spark/Core/Context.hs @@ -18,9 +18,12 @@ FromSQL, defaultConf, executeCommand1, + executeCommand1', createSparkSessionDef, closeSparkSessionDef, - exec1Def) where + exec1Def, + exec1Def' + ) where import Data.Text(pack) diff --git a/src/Spark/Core/Dataset.hs b/src/Spark/Core/Dataset.hs --- a/src/Spark/Core/Dataset.hs +++ b/src/Spark/Core/Dataset.hs @@ -24,6 +24,8 @@ asDF, asDS, asLocalObservable, + castType, + castType', -- * Relations parents, untyped, diff --git a/src/Spark/Core/Functions.hs b/src/Spark/Core/Functions.hs --- a/src/Spark/Core/Functions.hs +++ b/src/Spark/Core/Functions.hs @@ -5,12 +5,15 @@ dataframe, constant, collect, + collect', count, identity, autocache, cache, uncache, - (@@) + (@@), + joinInner, + joinInner' ) where @@ -21,6 +24,7 @@ import Spark.Core.Types import Spark.Core.Row import Spark.Core.Internal.DatasetFunctions +import Spark.Core.Internal.Joins import Spark.Core.Internal.Utilities import Spark.Core.Internal.LocalDataFunctions import Spark.Core.Internal.FunctionsInternals() @@ -28,7 +32,7 @@ import Spark.Core.Internal.AggregationFunctions import Spark.Core.Internal.TypesStructures(SQLType(..)) -dataset :: (ToSQL a, SQLTypeable a) => [a] -> Dataset a +dataset :: (ToSQL a, SQLTypeable a, HasCallStack) => [a] -> Dataset a dataset l = emptyDataset op tp where tp = buildType op = NodeDistributedLit (unSQLType tp) (V.fromList ((toJSON . valueToCell) <$> l)) diff --git a/src/Spark/Core/Internal/AggregationFunctions.hs b/src/Spark/Core/Internal/AggregationFunctions.hs --- a/src/Spark/Core/Internal/AggregationFunctions.hs +++ b/src/Spark/Core/Internal/AggregationFunctions.hs @@ -5,39 +5,62 @@ -- A number of standard aggregation functions. -module Spark.Core.Internal.AggregationFunctions where +module Spark.Core.Internal.AggregationFunctions( + -- Standard library + collect, + collect', + count, + count', + sumCol, + sumCol', + -- Developer functions + AggTry, + UniversalAggregator(..), + applyUAOUnsafe, + applyUntypedUniAgg3 +) where import Data.Aeson(Value(Null)) import qualified Data.Text as T -import Formatting +import qualified Data.Vector as V import Spark.Core.Internal.DatasetStructures import Spark.Core.Internal.ColumnStructures +import Spark.Core.Internal.ColumnFunctions(colType, untypedCol) import Spark.Core.Internal.DatasetFunctions import Spark.Core.Internal.RowGenerics(ToSQL) import Spark.Core.Internal.LocalDataFunctions() import Spark.Core.Internal.FunctionsInternals import Spark.Core.Internal.OpStructures -import Spark.Core.Internal.Utilities(failure, HasCallStack) import Spark.Core.Internal.TypesStructures +import Spark.Core.Internal.Utilities +import Spark.Core.Internal.TypesFunctions(arrayType') +import Spark.Core.StructuresInternal(emptyFieldPath) import Spark.Core.Types +import Spark.Core.Try {-| The sum of all the elements in a column. If the data type is too small to represent the sum, the value being returned is undefined. -} -colSum :: forall ref a. (Num a, SQLTypeable a, ToSQL a) => +sumCol :: forall ref a. (Num a, SQLTypeable a, ToSQL a) => Column ref a -> LocalData a -colSum = applyUniAgg (_sumAgg :: UniversalAggregator a a) +sumCol = applyUAOUnsafe _sumAgg' +sumCol' :: DynColumn -> LocalFrame +sumCol' = applyUntypedUniAgg3 _sumAgg' + {-| The number of elements in a column. -} -- TODO use Long for the return data type. -count :: forall a. (SQLTypeable a) => Dataset a -> LocalData Int -count ds = applyUniAgg (_countAgg2 :: UniversalAggregator a Int) (asCol ds) +count :: forall a. Dataset a -> LocalData Int +count ds = applyUAOUnsafe _countAgg' (asCol ds) +count' :: DataFrame -> LocalFrame +count' df = applyUntypedUniAgg3 _countAgg' (asCol' df) + {-| Collects all the elements of a column into a list. NOTE: @@ -47,8 +70,14 @@ the returned data. -} collect :: forall ref a. (SQLTypeable a) => Column ref a -> LocalData [a] -collect = applyUniAgg (_collectAgg :: UniversalAggregator a [a]) +collect = applyUAOUnsafe _collectAgg' +{-| See the documentation of collect. -} +collect' :: DynColumn -> LocalFrame +collect' = applyUntypedUniAgg3 _collectAgg' + +type AggTry a = Either T.Text a + {-| This is the universal aggregator: the invariant aggregator and some extra laws to combine multiple outputs. @@ -56,6 +85,7 @@ A real implementation in Spark has also an inner pass. -} data UniversalAggregator a buff = UniversalAggregator { + uaMergeType :: SQLType buff, -- The result is partioning invariant uaInitialOuter :: Dataset a -> LocalData buff, -- This operation is associative and commutative @@ -63,101 +93,63 @@ uaMergeBuffer :: LocalData buff -> LocalData buff -> LocalData buff } --- | (internal) -univAggToOp :: forall a buff. (SQLTypeable a, SQLTypeable buff) => - UniversalAggregator a buff -> UniversalAggregatorOp -univAggToOp = univAggToOpTyped (buildType :: SQLType a) (buildType :: SQLType buff) - --- | (internal) -univAggToOpTyped :: forall a buff. - SQLType a -> - SQLType buff -> - UniversalAggregator a buff -> - UniversalAggregatorOp -univAggToOpTyped sqlta sqltm ua = - let - mt = unSQLType sqltm - outer = _unsafeExtractOp $ fun1ToOpTyped sqlta (uaInitialOuter ua) - merge = _unsafeExtractOp $ fun2ToOpTyped sqltm sqltm (uaMergeBuffer ua) - in UniversalAggregatorOp { - uaoMergeType = mt, - uaoInitialOuter = outer, - uaoMergeBuffer = merge +-- TODO(kps) check the coming type for non-summable types +_sumAgg' :: DataType -> AggTry UniversalAggregatorOp +_sumAgg' _ = pure UniversalAggregatorOp { + -- TODO(kps) switch to BigInt + uaoMergeType = StrictType IntType, + uaoInitialOuter = InnerAggOp $ AggFunction "SUM" (V.singleton emptyFieldPath), + uaoMergeBuffer = ColumnSemiGroupLaw "SUM" } --- | (internal) -applyUniAgg :: UniversalAggregator a b -> Column ref a -> LocalData b -applyUniAgg ua c = - let - ds = pack1 c - ld1 = uaInitialOuter ua ds - -- TODO understand how to pass this info - -- aggop = univAggToOpTyped (nodeType ds) (nodeType ld1) ua - -- ld = emptyLocalData (NodeUniversalAggregator aggop) (nodeType ld1) - in ld1 - - --- (internal) -simpleOp1Typed :: (IsLocality loca, IsLocality locb) => - SQLType b -> - T.Text -> - ComputeNode loca a -> ComputeNode locb b -simpleOp1Typed sqltb name = - let so = StandardOperator { - soName = name, - soOutputType = unSQLType sqltb, - soExtra = Null - } - no = NodeLocalOp so - in nodeOpToFun1Typed sqltb no - --- (internal) -simpleOp1 :: forall a b loca locb. (IsLocality loca, IsLocality locb, SQLTypeable a, SQLTypeable b) => - T.Text -> - ComputeNode loca a -> ComputeNode locb b -simpleOp1 = simpleOp1Typed (buildType :: SQLType b) - --- (internal) -simpleOp2 :: forall a1 a2 b loc1 loc2 locb. (SQLTypeable b, IsLocality loc1, IsLocality loc2, IsLocality locb) => - T.Text -> - ComputeNode loc1 a1 -> ComputeNode loc2 a2 -> ComputeNode locb b -simpleOp2 = simpleOp2Typed (buildType :: SQLType b) +_countAgg' :: DataType -> AggTry UniversalAggregatorOp +-- Counting will always succeed. +_countAgg' _ = pure UniversalAggregatorOp { + -- TODO(kps) switch to BigInt + uaoMergeType = StrictType IntType, + uaoInitialOuter = InnerAggOp $ AggFunction "COUNT" (V.singleton emptyFieldPath), + uaoMergeBuffer = ColumnSemiGroupLaw "SUM" + } --- (internal) -simpleOp2Typed :: (IsLocality loc1, IsLocality loc2, IsLocality locb) => - SQLType b -> - T.Text -> - ComputeNode loc1 a1 -> ComputeNode loc2 a2 -> ComputeNode locb b -simpleOp2Typed sqltb name = - let so = StandardOperator { - soName = name, - soOutputType = unSQLType sqltb, - soExtra = Null +_collectAgg' :: DataType -> AggTry UniversalAggregatorOp +-- Counting will always succeed. +_collectAgg' dt = + let ldt = arrayType' dt + soMerge = StandardOperator { + soName = "org.spark.Collect", + soOutputType = ldt, + soExtra = Null } - no = NodeLocalOp so - in nodeOpToFun2Typed sqltb no - -_unsafeExtractOp :: (HasCallStack) => NodeOp -> StandardOperator -_unsafeExtractOp (NodeLocalOp so) = so -_unsafeExtractOp (NodeOpaqueAggregator so) = so -_unsafeExtractOp (NodeDistributedOp so) = so -_unsafeExtractOp x = failure $ sformat ("Expected standard op, found "%shown) x - -_countAgg2 :: (SQLTypeable a) => UniversalAggregator a Int -_countAgg2 = UniversalAggregator { - uaInitialOuter = simpleOp1 "org.spark.Count", - uaMergeBuffer = (+) + soMono = StandardOperator { + soName = "org.spark.CatSorted", + soOutputType = ldt, + soExtra = Null + } + in pure UniversalAggregatorOp { + -- TODO(kps) switch to BigInt + uaoMergeType = ldt, + uaoInitialOuter = OpaqueAggTransform soMerge, + uaoMergeBuffer = OpaqueSemiGroupLaw soMono } -_sumAgg :: forall a. (SQLTypeable a, Num a, ToSQL a) => UniversalAggregator a a -_sumAgg = UniversalAggregator { - uaInitialOuter = simpleOp1 "org.spark.Sum", - uaMergeBuffer = (+) - } +applyUntypedUniAgg3 :: (DataType -> AggTry UniversalAggregatorOp) -> DynColumn -> LocalFrame +applyUntypedUniAgg3 f dc = do + c <- dc + let uaot = f . unSQLType . colType $ c + uao <- tryEither uaot + let no = NodeAggregatorReduction uao + let ds = pack1 c + return $ emptyLocalData no (SQLType (uaoMergeType uao)) `parents` [untyped ds] -_collectAgg :: forall a. SQLTypeable a => UniversalAggregator a [a] -_collectAgg = - UniversalAggregator { - uaInitialOuter = simpleOp1 "org.spark.Collect", - uaMergeBuffer = simpleOp2 "org.spark.CatSorted" - } +applyUAOUnsafe :: forall a b ref. (SQLTypeable b, HasCallStack) => (DataType -> AggTry UniversalAggregatorOp) -> Column ref a -> LocalData b +applyUAOUnsafe f c = + let lf = applyUntypedUniAgg3 f (untypedCol c) + in forceRight (asObservable lf) + +-- _guardType :: DataType -> (UntypedDataset -> UntypedLocalData) -> (UntypedDataset -> LocalFrame) +-- _guardType dt f ds = +-- if unSQLType (nodeType ds) == dt +-- then +-- pure $ f ds +-- else +-- tryError $ sformat ("Expected type "%sh%" but got type "%sh) dt (nodeType ds) diff --git a/src/Spark/Core/Internal/CachingUntyped.hs b/src/Spark/Core/Internal/CachingUntyped.hs --- a/src/Spark/Core/Internal/CachingUntyped.hs +++ b/src/Spark/Core/Internal/CachingUntyped.hs @@ -24,7 +24,8 @@ cachingType :: UntypedNode -> CacheTry NodeCachingType cachingType n = traceHint ("cachingType: n="<>show' (nodeOp n)<>" res=") $ case nodeOp n of NodeLocalOp _ -> pure Stop - NodeUniversalAggregator _ -> pure Stop + NodeAggregatorReduction _ -> pure Stop + NodeAggregatorLocalReduction _ -> pure Stop NodeOpaqueAggregator _ -> pure Stop NodeLocalLit _ _ -> pure Stop NodeStructuredTransform _ -> pure Through @@ -38,7 +39,9 @@ NodeDistributedOp so | soName so == opnameAutocache -> pure $ AutocacheOp (vertexToId n) NodeDistributedOp _ -> pure Through -- Nothing special for the other operations - + NodeBroadcastJoin -> pure Through + NodeGroupedReduction _ -> pure Stop + NodeReduction _ -> pure Stop autocacheGen :: AutocacheGen UntypedNode autocacheGen = AutocacheGen { diff --git a/src/Spark/Core/Internal/ColumnFunctions.hs b/src/Spark/Core/Internal/ColumnFunctions.hs --- a/src/Spark/Core/Internal/ColumnFunctions.hs +++ b/src/Spark/Core/Internal/ColumnFunctions.hs @@ -21,10 +21,12 @@ -- Developer API (projections) unsafeStaticProjection, dynamicProjection, + dropColReference, -- Public functions untypedCol, colFromObs, colFromObs', + castTypeCol, castCol, castCol', colRef @@ -39,7 +41,6 @@ import Data.List(find) import Formatting - import Spark.Core.Internal.ColumnStructures import Spark.Core.Internal.DatasetFunctions import Spark.Core.Internal.DatasetStructures @@ -89,7 +90,7 @@ {-| Converts a type column to an antyped column. -} untypedCol :: Column ref a -> DynColumn -untypedCol = pure . _unsafeCastColData . _dropReference +untypedCol = pure . _unsafeCastColData . dropColReference {-| Casts a dynamic column to a statically typed column. @@ -99,7 +100,7 @@ -} castCol :: ColumnReference ref -> SQLType a -> DynColumn -> Try (Column ref a) castCol r sqlt dc = - dc >>= _checkedCastColData sqlt >>= _checkedCastRefColData r + dc >>= castTypeCol sqlt >>= _checkedCastRefColData r {-| Casts a dynamic column to a statically typed column, but does not attempt to enforce a single origin at the type level. @@ -110,6 +111,16 @@ castCol' :: SQLType a -> DynColumn -> Try (Column UnknownReference a) castCol' = castCol ColumnReference + +-- | (internal) +castTypeCol :: SQLType b -> ColumnData ref a -> Try (ColumnData ref b) +castTypeCol sqlt cd = + if unSQLType sqlt == unSQLType (colType cd) + then pure (_unsafeCastColData cd) + else tryError $ sformat ("Cannot cast column "%sh%" to type "%sh) cd sqlt + + + -- (internal) colOrigin :: Column ref a -> UntypedDataset colOrigin = _cOrigin @@ -179,12 +190,6 @@ instance forall a to. Projection DynColumn (StaticColProjection a to) DynColumn where _performProjection dc proj = _projectDynCol dc (_colStaticProjToDynProj proj) -class StringStuff a where - stuffAsString :: a -> String - -instance StringStuff String where - stuffAsString = undefined - -- dyncolumn -> string -> dyncolumn instance Projection DynColumn String DynColumn where _performProjection dc s = _performProjection dc (_strToDynProj s) @@ -208,7 +213,7 @@ else pure (fp, dtTo) iUntypedColData :: Column ref a -> UntypedColumnData -iUntypedColData = _unsafeCastColData . _dropReference +iUntypedColData = _unsafeCastColData . dropColReference -- Recasts the column, trusting the user knows that the type is going to be compatible. _unsafeCastColData :: Column ref a -> Column ref b @@ -216,7 +221,7 @@ _checkedCastColData :: SQLType b -> ColumnData ref a -> Try (ColumnData ref b) _checkedCastColData sqlt cd = - if (unSQLType sqlt) == (unSQLType (colType cd)) + if unSQLType sqlt == unSQLType (colType cd) then pure (_unsafeCastColData cd) else tryError $ sformat ("Cannot cast column "%sh%" to type "%sh) cd sqlt @@ -270,7 +275,7 @@ _projectDynColData :: ColumnData ref a -> DynamicColProjection -> DynColumn _projectDynColData cd proj = - _dynProjTry proj (_cType cd) <&> uncurry (_projectColData0 . _dropReference $ cd) + _dynProjTry proj (_cType cd) <&> uncurry (_projectColData0 . dropColReference $ cd) _projectDynCol :: DynColumn -> DynamicColProjection -> DynColumn _projectDynCol c proj = do @@ -297,8 +302,8 @@ SQLType . structFieldType <$> z _extractField _ _ = Nothing -_dropReference :: ColumnData ref a -> ColumnData UnknownReference a -_dropReference c = c {_cOp = _cOp c} +dropColReference :: ColumnData ref a -> ColumnData UnknownReference a +dropColReference c = c {_cOp = _cOp c} -- | (internal) creates a new column with some empty data iEmptyCol :: Dataset a -> SQLType b -> FieldPath -> Column a b @@ -306,7 +311,7 @@ -- | (internal) Creates a new column with a dynamic type. _emptyDynCol :: Dataset a -> DataType -> FieldPath -> DynColumn -_emptyDynCol ds dt fp = Right $ _dropReference $ _emptyColData ds (SQLType dt) fp +_emptyDynCol ds dt fp = Right $ dropColReference $ _emptyColData ds (SQLType dt) fp -- A new column data structure. _emptyColData :: Dataset a -> SQLType b -> FieldPath -> ColumnData a b @@ -383,13 +388,13 @@ instance forall a. HomoBinaryOp2 a a a where - _liftFun f = BinaryOpFun id id f + _liftFun = BinaryOpFun id id instance forall ref a. HomoBinaryOp2 (Column ref a) DynColumn DynColumn where - _liftFun f = BinaryOpFun untypedCol id f + _liftFun = BinaryOpFun untypedCol id instance forall ref a. HomoBinaryOp2 DynColumn (Column ref a) DynColumn where - _liftFun f = BinaryOpFun id untypedCol f + _liftFun = BinaryOpFun id untypedCol instance (Num x) => Num (Column ref x) where diff --git a/src/Spark/Core/Internal/ContextIOInternal.hs b/src/Spark/Core/Internal/ContextIOInternal.hs --- a/src/Spark/Core/Internal/ContextIOInternal.hs +++ b/src/Spark/Core/Internal/ContextIOInternal.hs @@ -6,7 +6,8 @@ returnPure, createSparkSession, createSparkSession', - executeCommand1 + executeCommand1, + executeCommand1' ) where import Control.Concurrent(threadDelay) @@ -32,6 +33,8 @@ import Spark.Core.Internal.Client import Spark.Core.Internal.ContextInternal import Spark.Core.Internal.ContextStructures +import Spark.Core.Internal.DatasetFunctions(untypedLocalData) +import Spark.Core.Internal.DatasetStructures(UntypedLocalData) import Spark.Core.Row import Spark.Core.StructuresInternal import Spark.Core.Try @@ -78,6 +81,11 @@ executeCommand1 :: forall a. (FromSQL a, HasCallStack) => LocalData a -> SparkState (Try a) executeCommand1 ld = do + tcell <- executeCommand1' (untypedLocalData ld) + return $ tcell >>= (tryEither . cellToValue) + +executeCommand1' :: (HasCallStack) => UntypedLocalData -> SparkState (Try Cell) +executeCommand1' ld = do session <- get tcomp <- returnPure $ prepareExecution1 ld case tcomp of @@ -94,8 +102,7 @@ in do _ <- _sendComputation session comp nrs <- nodeResults - tcell <- returnPure $ storeResults comp nrs - return $ tcell >>= (tryEither . cellToValue) + returnPure $ storeResults comp nrs _randomSessionName :: IO Text _randomSessionName = do diff --git a/src/Spark/Core/Internal/ContextInteractive.hs b/src/Spark/Core/Internal/ContextInteractive.hs --- a/src/Spark/Core/Internal/ContextInteractive.hs +++ b/src/Spark/Core/Internal/ContextInteractive.hs @@ -14,6 +14,7 @@ SparkInteractiveException, createSparkSessionDef, exec1Def, + exec1Def', closeSparkSessionDef ) where @@ -30,8 +31,10 @@ import Spark.Core.Internal.ContextStructures import Spark.Core.Internal.DatasetStructures +import Spark.Core.Internal.DatasetFunctions(untypedLocalData) import Spark.Core.Internal.ContextIOInternal -import Spark.Core.Internal.RowGenericsFrom(FromSQL) +import Spark.Core.Internal.RowGenericsFrom(FromSQL, cellToValue) +import Spark.Core.Internal.RowStructures(Cell) import Spark.Core.Internal.Utilities import Spark.Core.StructuresInternal import Spark.Core.Try @@ -58,7 +61,7 @@ If a session already exists, an exception will be thrown. -} -createSparkSessionDef :: (HasCallStack) => SparkSessionConf -> IO () +createSparkSessionDef :: SparkSessionConf -> IO () createSparkSessionDef conf = do current <- _currentSession case current of @@ -79,18 +82,29 @@ -} exec1Def :: (FromSQL a, HasCallStack) => LocalData a -> IO a exec1Def ld = do + c <- exec1Def' (pure (untypedLocalData ld)) + case cellToValue c of + Right x -> return x + Left txt -> _throw txt + +exec1Def' :: (HasCallStack) => LocalFrame -> IO Cell +exec1Def' lf = do mCtx <- _currentSession case mCtx of Nothing -> _throw "No default context found. You must first create a default spark context with createSparkSessionDef" - Just ctx -> do - (res, newSt) <- (runStateT . runStdoutLoggingT . executeCommand1) ld ctx - _setSession newSt - case res of - Right x -> - return x + Just ctx -> + case lf of Left err -> throwM (SparkInteractiveException err) + Right ld -> do + (res, newSt) <- (runStateT . runStdoutLoggingT . executeCommand1') ld ctx + _setSession newSt + case res of + Right x -> + return x + Left err -> + throwM (SparkInteractiveException err) {-| Closes the default session. The default session is empty after this call completes. @@ -98,24 +112,24 @@ NOTE: This does not currently clear up the resources! It is a stub implementation used in testing. -} -closeSparkSessionDef :: (HasCallStack) => IO () +closeSparkSessionDef :: IO () closeSparkSessionDef = do _ <- _removeSession return () -_currentSession :: (HasCallStack) => IO (Maybe SparkSession) +_currentSession :: IO (Maybe SparkSession) _currentSession = readIORef _globalSessionRef -_setSession :: (HasCallStack) => SparkSession -> IO () +_setSession :: SparkSession -> IO () _setSession st = writeIORef _globalSessionRef (Just st) -_removeSession :: (HasCallStack) => IO (Maybe SparkSession) +_removeSession :: IO (Maybe SparkSession) _removeSession = do current <- _currentSession _ <- writeIORef _globalSessionRef Nothing return current -_throw :: (HasCallStack) => Text -> IO a +_throw :: Text -> IO a _throw txt = throwM $ SparkInteractiveException Error { ePath = NodePath V.empty, diff --git a/src/Spark/Core/Internal/ContextInternal.hs b/src/Spark/Core/Internal/ContextInternal.hs --- a/src/Spark/Core/Internal/ContextInternal.hs +++ b/src/Spark/Core/Internal/ContextInternal.hs @@ -137,7 +137,7 @@ -- Like the type, remove the row wrapper in the case of basic elements -- TODO(kps) figure out what the exact semantics are. -- It seems collect is behaving differently than the other nodes. -_postprocessBasic :: (HasCallStack) => Cell -> Cell +_postprocessBasic :: Cell -> Cell _postprocessBasic (RowArray rows) = RowArray (process <$> rows) where process (RowArray arr) = case V.toList arr of diff --git a/src/Spark/Core/Internal/DAGFunctions.hs b/src/Spark/Core/Internal/DAGFunctions.hs --- a/src/Spark/Core/Internal/DAGFunctions.hs +++ b/src/Spark/Core/Internal/DAGFunctions.hs @@ -97,7 +97,7 @@ All the vertices referred by edges must be present in the list of vertices. -} -buildGraphFromList :: forall v e. (Show v, Show e) => +buildGraphFromList :: forall v e. (Show v) => [Vertex v] -> [Edge e] -> DagTry (Graph v e) buildGraphFromList vxs eds = do -- 1. Group the edges by start point @@ -217,7 +217,7 @@ in filter f (toList (gVertices g)) -- | Flips the edges of this graph (it is also a DAG) -reverseGraph :: forall v e. (HasCallStack, Show v, Show e) => Graph v e -> Graph v e +reverseGraph :: forall v e. Graph v e -> Graph v e reverseGraph g = let vxMap = M.fromList ((vertexId &&& id) <$> toList (gVertices g)) @@ -239,7 +239,7 @@ -- | A generic transform over the graph that may account for potential failures -- in the process. -graphMapVertices :: forall m v e v2. (HasCallStack, Show v2, Show v, Show e, Monad m) => +graphMapVertices :: forall m v e v2. (HasCallStack, Show v2, Monad m) => Graph v e -> -- The start graph (v -> [(v2,e)] -> m v2) -> -- The transform m (Graph v2 e) diff --git a/src/Spark/Core/Internal/DatasetFunctions.hs b/src/Spark/Core/Internal/DatasetFunctions.hs --- a/src/Spark/Core/Internal/DatasetFunctions.hs +++ b/src/Spark/Core/Internal/DatasetFunctions.hs @@ -14,6 +14,7 @@ asDF, asDS, asLocalObservable, + asObservable, -- Standard functions identity, autocache, @@ -35,6 +36,7 @@ nodeParents, nodeType, untypedDataset, + untypedLocalData, updateNode, -- Developer conversions fun1ToOpTyped, @@ -44,6 +46,11 @@ nodeOpToFun1Untyped, nodeOpToFun2, nodeOpToFun2Typed, + nodeOpToFun2Untyped, + unsafeCastDataset, + placeholder, + castType, + castType', -- Internal opnameCache, opnameUnpersist, @@ -74,6 +81,7 @@ import Spark.Core.Internal.Utilities import Spark.Core.Internal.RowUtils import Spark.Core.Internal.TypesGenerics +import Spark.Core.Internal.TypesFunctions -- | (developer) The operation performed by this node. nodeOp :: ComputeNode loc a -> NodeOp @@ -207,13 +215,7 @@ -- operation is not correct. -- This operation assumes that both field names and types are correct. asDS :: forall a. (SQLTypeable a) => DataFrame -> Try (Dataset a) -asDS df = do - n <- df - let dt = unSQLType (buildType :: SQLType a) - let dt' = unSQLType (nodeType n) - if dt == dt' - then pure (_unsafeCastNode n) - else tryError $ sformat ("Casting error: dataframe has type "%sh%" incompatible with type "%sh) dt' dt +asDS = _asTyped -- | Converts a local node to a local frame. @@ -221,6 +223,9 @@ asLocalObservable :: ComputeNode LocLocal a -> LocalFrame asLocalObservable = pure . _unsafeCastNode +asObservable :: forall a. (SQLTypeable a) => LocalFrame -> Try (LocalData a) +asObservable = _asTyped + -- | Converts any node to an untyped node untyped :: ComputeNode loc a -> UntypedNode untyped = _unsafeCastNode @@ -228,6 +233,10 @@ untypedDataset :: ComputeNode LocDistributed a -> UntypedDataset untypedDataset = _unsafeCastNode +{-| Removes type informatino from an observable. -} +untypedLocalData :: ComputeNode LocLocal a -> UntypedLocalData +untypedLocalData = _unsafeCastNode + {-| Adds parents to the node. It is assumed the parents are the unique set of nodes required by the operation defined in this node. @@ -273,8 +282,7 @@ -- (internal) -- Tries to update the locality of a node. This is a checked cast. -- TODO: remove, it is only used to cast to local frame -castLocality :: forall a loc loc'. ( - CheckedLocalityCast loc, CheckedLocalityCast loc') => +castLocality :: forall a loc loc'. (CheckedLocalityCast loc') => ComputeNode loc a -> Try (ComputeNode loc' a) castLocality node = let @@ -332,14 +340,19 @@ -- | (internal) placeholderTyped :: forall a loc. (IsLocality loc) => SQLType a -> ComputeNode loc a -placeholderTyped tp = +placeholderTyped tp = _unsafeCastNode n where + n = placeholder (unSQLType tp) :: ComputeNode loc Cell + +placeholder :: forall loc. (IsLocality loc) => DataType -> ComputeNode loc Cell +placeholder tp = let - so = makeOperator "org.spark.Placeholder" tp + t = SQLType tp + so = makeOperator "org.spark.Placeholder" t (TypedLocality l) = _getTypedLocality :: TypedLocality loc op = case l of Local -> NodeLocalOp so Distributed -> NodeDistributedOp so - in _emptyNode op tp + in _emptyNode op t -- | (internal) conversion fun1ToOpTyped :: forall a loc a' loc'. (IsLocality loc) => @@ -352,36 +365,44 @@ fun2ToOpTyped sqlt1 sqlt2 f = nodeOp $ f (placeholderTyped sqlt1) (placeholderTyped sqlt2) -- | (internal) conversion -nodeOpToFun1 :: forall a1 a2 loc1 loc2. (IsLocality loc1, SQLTypeable a2, IsLocality loc2) => +nodeOpToFun1 :: forall a1 a2 loc1 loc2. (SQLTypeable a2, IsLocality loc2) => NodeOp -> ComputeNode loc1 a1 -> ComputeNode loc2 a2 nodeOpToFun1 = nodeOpToFun1Typed (buildType :: SQLType a2) -- | (internal) conversion -nodeOpToFun1Typed :: forall a1 a2 loc1 loc2. (HasCallStack, IsLocality loc1, IsLocality loc2) => +nodeOpToFun1Typed :: forall a1 a2 loc1 loc2. (IsLocality loc2) => SQLType a2 -> NodeOp -> ComputeNode loc1 a1 -> ComputeNode loc2 a2 nodeOpToFun1Typed sqlt no node = let n2 = _emptyNode no sqlt :: ComputeNode loc2 a2 in n2 `parents` [untyped node] -- | (internal) conversion -nodeOpToFun1Untyped :: forall loc1 loc2. (HasCallStack, IsLocality loc1, IsLocality loc2) => +nodeOpToFun1Untyped :: forall loc1 loc2. (IsLocality loc2) => DataType -> NodeOp -> ComputeNode loc1 Cell -> ComputeNode loc2 Cell nodeOpToFun1Untyped dt no node = let n2 = _emptyNode no (SQLType dt) :: ComputeNode loc2 Cell in n2 `parents` [untyped node] -- | (internal) conversion -nodeOpToFun2 :: forall a a1 a2 loc loc1 loc2. (SQLTypeable a, IsLocality loc, IsLocality loc1, IsLocality loc2) => +nodeOpToFun2 :: forall a a1 a2 loc loc1 loc2. (SQLTypeable a, IsLocality loc) => NodeOp -> ComputeNode loc1 a1 -> ComputeNode loc2 a2 -> ComputeNode loc a nodeOpToFun2 = nodeOpToFun2Typed (buildType :: SQLType a) -- | (internal) conversion -nodeOpToFun2Typed :: forall a a1 a2 loc loc1 loc2. (IsLocality loc, IsLocality loc1, IsLocality loc2) => +nodeOpToFun2Typed :: forall a a1 a2 loc loc1 loc2. (IsLocality loc) => SQLType a -> NodeOp -> ComputeNode loc1 a1 -> ComputeNode loc2 a2 -> ComputeNode loc a nodeOpToFun2Typed sqlt no node1 node2 = let n2 = _emptyNode no sqlt :: ComputeNode loc a in n2 `parents` [untyped node1, untyped node2] +-- | (internal) conversion +nodeOpToFun2Untyped :: forall loc1 loc2 loc3. (IsLocality loc3) => + DataType -> NodeOp -> ComputeNode loc1 Cell -> ComputeNode loc2 Cell -> ComputeNode loc3 Cell +nodeOpToFun2Untyped dt no node1 node2 = + let n2 = _emptyNode no (SQLType dt) :: ComputeNode loc3 Cell + in n2 `parents` [untyped node1, untyped node2] + + -- ******* INSTANCES ********* -- Put here because it depends on some private functions. @@ -411,11 +432,34 @@ toJSON (TypedLocality Local) = A.String "local" toJSON (TypedLocality Distributed) = A.String "distributed" +unsafeCastDataset :: ComputeNode LocDistributed a -> ComputeNode LocDistributed b +unsafeCastDataset ds = ds { _cnType = _cnType ds } + +-- TODO: figure out the story around haskell types vs datatypes +-- Should we have equivalence classes for haskell, so that a tuple has the +-- same type as a structure? +-- Probably not, it breaks the correspondence. +-- Probably, it makes the metadata story easier. +castType :: SQLType a -> ComputeNode loc b -> Try (ComputeNode loc a) +castType sqlt n = do + let dt = unSQLType sqlt + let dt' = unSQLType (nodeType n) + if dt `compatibleTypes` dt' + then let n' = updateNode n (\node -> node { _cnType = dt }) in + pure (_unsafeCastNode n') + else tryError $ sformat ("Casting error: dataframe has type "%sh%" incompatible with type "%sh) dt' dt + +castType' :: SQLType a -> Try (ComputeNode loc Cell) -> Try (ComputeNode loc a) +castType' sqlt df = df >>= castType sqlt + +_asTyped :: forall loc a. (SQLTypeable a) => Try (ComputeNode loc Cell) -> Try (ComputeNode loc a) +_asTyped = castType' (buildType :: SQLType a) + -- Performs an unsafe type recast. -- This is useful for internal code that knows whether -- this operation is legal or not through some other means. -- This may still throw an error if the cast is illegal. -_unsafeCastNode :: CheckedLocalityCast loc2 => ComputeNode loc1 a -> ComputeNode loc2 b +_unsafeCastNode :: ComputeNode loc1 a -> ComputeNode loc2 b _unsafeCastNode x = x { _cnType = _cnType x, _cnLocality = _cnLocality x diff --git a/src/Spark/Core/Internal/FunctionsInternals.hs b/src/Spark/Core/Internal/FunctionsInternals.hs --- a/src/Spark/Core/Internal/FunctionsInternals.hs +++ b/src/Spark/Core/Internal/FunctionsInternals.hs @@ -14,17 +14,19 @@ NameTuple(..), TupleEquivalence(..), asCol, + asCol', pack1, pack, pack', struct', struct, + -- Developer tools + checkOrigin ) where import Control.Arrow import qualified Data.Vector as V import qualified Data.Text as T -import Data.List(sort, nub) import Formatting import Spark.Core.Internal.ColumnStructures @@ -32,7 +34,6 @@ import Spark.Core.Internal.DatasetFunctions import Spark.Core.Internal.DatasetStructures import Spark.Core.Internal.Utilities -import Spark.Core.Internal.TypesGenerics import Spark.Core.Internal.TypesFunctions import Spark.Core.Internal.TypesStructures import Spark.Core.Internal.OpStructures @@ -94,14 +95,17 @@ -- fun' = undefined -- | Represents a dataframe as a single column. -asCol :: (HasCallStack) => Dataset a -> Column a a +asCol :: Dataset a -> Column a a asCol ds = -- Simply recast the dataset as a column. -- The empty path indicates that we are wrapping the whole thing. iEmptyCol ds (unsafeCastType $ nodeType ds) (FieldPath V.empty) +asCol' :: DataFrame -> DynColumn +asCol' = ((iUntypedColData . asCol) <$>) + -- | Packs a single column into a dataframe. -pack1 :: (HasCallStack) => Column ref a -> Dataset a +pack1 :: Column ref a -> Dataset a pack1 c = emptyDataset (NodeStructuredTransform (colOp c)) (colType c) `parents` [untyped (colOrigin c)] @@ -123,7 +127,7 @@ TODO: example. -} -pack :: forall ref a b. (StaticColPackable2 ref a b, HasCallStack) => a -> Dataset b +pack :: forall ref a b. (StaticColPackable2 ref a b) => a -> Dataset b pack z = let c = _staticPackAsColumn2 z :: ColumnData ref b in pack1 c @@ -132,7 +136,7 @@ Columns must have different names, or an error is returned. -} -struct' :: (HasCallStack) => [DynColumn] -> DynColumn +struct' :: [DynColumn] -> DynColumn struct' cols = do l <- sequence cols let fields = (colFieldName &&& id) <$> l @@ -143,10 +147,21 @@ The field names of the columns are discarded, and replaced by the field names of the structure. -} -struct :: forall ref a b. (StaticColPackable2 ref a b, HasCallStack) => a -> Column ref b +struct :: forall ref a b. (StaticColPackable2 ref a b) => a -> Column ref b struct = _staticPackAsColumn2 +checkOrigin :: [DynColumn] -> Try [UntypedColumnData] +checkOrigin x = _checkOrigin =<< sequence x + +_checkOrigin :: [UntypedColumnData] -> Try [UntypedColumnData] +_checkOrigin [] = pure [] +_checkOrigin l = + case _columnOrigin l of + [_] -> pure l + l' -> tryError $ sformat ("Too many distinct origins: "%sh) l' + + instance forall x. (DynColPackable x) => DynColPackable [x] where _packAsColumn = struct' . (_packAsColumn <$>) @@ -170,7 +185,6 @@ -- The equations that bind column packable stuff through their tuple equivalents instance forall ref b a1 a2 z1 z2. ( - SQLTypeable b, TupleEquivalence b (a1, a2), StaticColPackable2 ref z1 a1, StaticColPackable2 ref z2 a2) => @@ -183,7 +197,6 @@ in _unsafeBuildStruct [x1, x2] names instance forall ref b a1 a2 a3 z1 z2 z3. ( - SQLTypeable b, TupleEquivalence b (a1, a2, a3), StaticColPackable2 ref z1 a1, StaticColPackable2 ref z2 a2, @@ -197,8 +210,7 @@ names = tupleFieldNames :: NameTuple b in _unsafeBuildStruct [x1, x2, x3] names -_unsafeBuildStruct :: (HasCallStack, SQLTypeable x) => - [UntypedColumnData] -> NameTuple x -> Column ref x +_unsafeBuildStruct :: [UntypedColumnData] -> NameTuple x -> Column ref x _unsafeBuildStruct cols (NameTuple names) = if length cols /= length names then failure $ sformat ("The number of columns and names differs:"%sh%" and "%sh) cols names @@ -210,27 +222,19 @@ _buildStruct :: [(FieldName, UntypedColumnData)] -> Try UntypedColumnData -_buildStruct [] = tryError "You cannot build an empty structure" -_buildStruct ((hfn, hcol):t) = - let cols = ((hfn, hcol):t) - cols' = V.fromList cols - fields = ColStruct $ (uncurry TransformField .(fst &&& colOp . snd)) <$> cols' - ct = StructType $ (uncurry StructField . (fst &&& unSQLType . colType . snd)) <$> cols' - name = "struct(" <> T.intercalate "," (unFieldName . fst <$> cols) <> ")" - names = fst <$> cols - numNames = length names - numDistincts = length . nub $ names - origins = _columnOrigin (snd <$> cols) - in case (origins, numNames == numDistincts) of - ([_], True) -> - pure ColumnData { - _cOrigin = _cOrigin hcol, - _cType = StrictType $ Struct ct, - _cOp = fields, - _cReferingPath = Just $ unsafeFieldName name - } - (l, True) -> tryError $ sformat ("Too many distinct origins: "%sh) l - (_, False) -> tryError $ sformat ("Duplicate field names when building the struct: "%sh) (sort names) +_buildStruct cols = do + let fields = ColStruct $ (uncurry TransformField . (fst &&& colOp . snd)) <$> V.fromList cols + st <- structTypeFromFields $ (fst &&& unSQLType . colType . snd) <$> cols + let name = structName st + case _columnOrigin (snd <$> cols) of + [ds] -> + pure ColumnData { + _cOrigin = ds, + _cType = StrictType (Struct st), + _cOp = fields, + _cReferingPath = Just $ unsafeFieldName name + } + l -> tryError $ sformat ("Too many distinct origins: "%sh) l _columnOrigin :: [UntypedColumnData] -> [UntypedDataset] _columnOrigin l = diff --git a/src/Spark/Core/Internal/Groups.hs b/src/Spark/Core/Internal/Groups.hs new file mode 100644 --- /dev/null +++ b/src/Spark/Core/Internal/Groups.hs @@ -0,0 +1,300 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE FlexibleContexts #-} + +-- A number of standard aggregation functions. + +module Spark.Core.Internal.Groups( + GroupData, + LogicalGroupData, + -- Typed functions + groupByKey, + mapGroup, + aggKey, + groupAsDS + -- Developer + +) where + +import qualified Data.Text as T +import qualified Data.Vector as V +import Formatting +import Debug.Trace(trace) + +import Spark.Core.Internal.DatasetStructures +import Spark.Core.Internal.ColumnStructures +import Spark.Core.Internal.ColumnFunctions(untypedCol, colType, colOp, iUntypedColData, colOrigin, castTypeCol, dropColReference) +import Spark.Core.Internal.DatasetFunctions +import Spark.Core.Internal.LocalDataFunctions() +import Spark.Core.Internal.FunctionsInternals +import Spark.Core.Internal.TypesFunctions(tupleType, structTypeFromFields) +import Spark.Core.Internal.OpStructures +import Spark.Core.Internal.TypesStructures +import Spark.Core.Internal.Utilities +import Spark.Core.Internal.RowStructures(Cell) +import Spark.Core.Try +import Spark.Core.StructuresInternal + +{-| A dataset that has been partitioned according to some given field. +-} +data GroupData key val = GroupData { + -- The dataset of reference for this group + _gdRef :: !UntypedDataset, + -- The columns used to partition the data by keys. + _gdKey :: !UntypedColumnData, + -- The columns that contain the values. + _gdValue :: !UntypedColumnData +} + +type LogicalGroupData = Try UntypedGroupData + +{-| (developper) + +A group data type with no typing information. +-} +type UntypedGroupData = GroupData Cell Cell + +type GroupTry a = Either T.Text a + +-- A useful type when chaining operations withing groups. +data PipedTrans = + PipedError !T.Text + | PipedDataset !UntypedDataset + | PipedGroup !UntypedGroupData + deriving (Show) + + +{-| Performs a logical group of data based on a key. +-} +groupByKey :: (HasCallStack) => Column ref key -> Column ref val -> GroupData key val +groupByKey keys vals = forceRight $ _castGroup (colType keys) (colType vals) =<< _groupByKey (iUntypedColData keys) (iUntypedColData vals) + +{-| Transforms the values in a group. +-} +-- This only allows direct transforms, so it is probably valid in all cases. +mapGroup :: GroupData key val -> (forall ref. Column ref val -> Column ref val') -> GroupData key val' +mapGroup g f = + let c = _unsafeCastColData (_gdValue g) + -- TODO: this is wrong, an aggregation may have been forced in between. + in g { _gdValue = iUntypedColData (f c) } + +{-| The generalized value transform. + +This generalizes mapGroup to allow more complex transforms involving joins, +groups, etc. +-} +-- TODO: this can fail +-- magGroupGen :: (forall ref. Column ref val -> Dataset val') -> GroupData key val -> GroupData key val' +-- magGroupGen _ _ = undefined + +{-| Given a group and an aggregation function, aggregates the data. + +Note: not all the reduction functions may be used in this case. The analyzer +will fail if the function is not universal. +-} +-- TODO: it should be a try, this can fail +aggKey :: (HasCallStack) => GroupData key val -> (forall ref. Column ref val -> LocalData val') -> Dataset (key, val') +aggKey gd f = trace "aggKey" $ + let ugd = _untypedGroup gd + keyt = traceHint "aggKey: keyt: " $ mapGroupKeys gd colType + valt = traceHint "aggKey: valt: " $ mapGroupValues gd colType + -- We call the function twice: the first one to recover the type info, + -- and the second time to perform the unrolling. + -- TODO we should be able to do it in one pass instead. + fOut = traceHint "aggKey: fOut: " $ f (mapGroupValues gd dropColReference) + valt' = traceHint "aggKey: valt': " $ nodeType fOut + t = traceHint "aggKey: t: " $ tupleType keyt valt' + f' c = untypedLocalData . f <$> castTypeCol valt c + tud = traceHint "aggKey: tud: " $ _aggKey ugd f' + res = castType' t tud + in forceRight res + +{-| Creates a group by 'expanding' a value into a potentially large collection. + +Note on performance: this function is optimized to work at any scale and may not +be the most efficient when the generated collections are small (a few elements). +-} +-- TODO: it should be a try, this can fail +-- expand :: Column ref key -> Column ref val -> (LocalData val -> Dataset val') -> GroupData key val' +-- expand = undefined + +{-| Builds groups within groups. + +This function allows groups to be constructed from each collections inside a +group. + +This function is usually not used directly by the user, but rather as part of +more complex pipelines that may involve multiple levels of nesting. +-} +-- groupInGroup :: GroupData key val -> (forall ref. Column ref val -> GroupData key' val') -> GroupData (key', key) val' +-- groupInGroup _ _ = undefined + +{-| Reduces a group in group into a single group. +-} +-- aggGroup :: GroupData (key, key') val -> (forall ref. LocalData key -> Column ref val -> LocalData val') -> GroupData key val +-- aggGroup _ _ = undefined + +{-| Returns the collapsed representation of a grouped dataset, discarding group +information. +-} +groupAsDS :: forall key val. GroupData key val -> Dataset (key, val) +groupAsDS g = pack s where + c1 = _unsafeCastColData (_gdKey g) :: Column UnknownReference key + c2 = _unsafeCastColData (_gdValue g) :: Column UnknownReference val + s = struct (c1, c2) :: Column UnknownReference (key, val) + +mapGroupKeys :: GroupData key val -> (forall ref. Column ref key -> a) -> a +mapGroupKeys gd f = + f (_unsafeCastColData (_gdKey gd)) + +mapGroupValues :: GroupData key val -> (forall ref. Column ref val -> a) -> a +mapGroupValues gd f = + f (_unsafeCastColData (_gdValue gd)) + +-- ******** INSTANCES *********** + + +instance Show (GroupData key val) where + show gd = T.unpack s where + s = sformat ("GroupData[key="%sh%", val="%sh%"]") (_gdKey gd) (_gdValue gd) + +-- ******** PRIVATE METHODS ******** + +_mapStructuredTransform :: ColOp -> LogicalGroupData -> GroupTry LogicalGroupData +_mapStructuredTransform = undefined + +_mapAggTransform :: AggTransform -> LogicalGroupData -> GroupTry LogicalGroupData +_mapAggTransform = undefined + +_pError :: T.Text -> PipedTrans +_pError = PipedError + +_unrollTransform :: PipedTrans -> NodeId -> UntypedNode -> PipedTrans +_unrollTransform start nid un | nodeId un == nid = start +_unrollTransform start nid un = case nodeParents un of + [p] -> + let pt' = _unrollTransform start nid p in _unrollStep pt' un + _ -> + _pError $ sformat (sh%": operations with multiple parents cannot be used in groups yet.") un + +_unrollStep :: PipedTrans -> UntypedNode -> PipedTrans +_unrollStep pt un = traceHint ("_unrollStep: pt=" <> show' pt <> " un=" <> show' un <> " res=") $ + let op = nodeOp un + dt = unSQLType (nodeType un) in case nodeParents un of + [p] -> + case (pt, op) of + (PipedError e, _) -> PipedError e + (PipedDataset ds, NodeStructuredTransform _) -> + -- This is simply dointg a DS -> DS transform. + -- TODO: this breaks the encapsulation of ComputeNode + let ds' = updateNode un (\un' -> un' { _cnParents = V.singleton (untyped ds)}) + in PipedDataset ds' + (PipedGroup g, NodeStructuredTransform co) -> + _unrollGroupTrans g co + (PipedGroup g, NodeAggregatorReduction uao) -> + case uaoInitialOuter uao of + OpaqueAggTransform x -> _pError $ sformat ("Cannot apply opaque transform in the context of an aggregation: "%sh) x + InnerAggOp ao -> + PipedDataset $ _applyAggOp dt ao g + _ -> _pError $ sformat (sh%": Operation not supported with trans="%sh%" and parents="%sh) op pt p + l -> _pError $ sformat (sh%": expected one parent but got "%sh) un l + +-- dt: output type of the aggregation op +_applyAggOp :: (HasCallStack) => DataType -> AggOp -> UntypedGroupData -> UntypedDataset +_applyAggOp dt ao ugd = traceHint ("_applyAggOp dt=" <> show' dt <> " ao=" <> show' ao <> " ugd=" <> show' ugd <> " res=") $ + -- Reset the names to make sure there are no collision. + let c1 = untypedCol (_gdKey ugd) @@ T.unpack "_1" + c2 = untypedCol (_gdValue ugd) @@ T.unpack "_2" + s = struct' [c1, c2] + p = pack1 <$> s + ds = forceRight p + -- The structure of the result dataframe + keyDt = unSQLType (colType (_gdKey ugd)) + st' = structTypeFromFields [(unsafeFieldName "key", keyDt), (unsafeFieldName "agg", dt)] + -- The keys are different, so we know we this operation is legit: + st = forceRight st' + resDt = SQLType . StrictType . Struct $ st + ds2 = emptyDataset (NodeGroupedReduction ao) resDt `parents` [untyped ds] + in ds2 + +_unrollGroupTrans :: UntypedGroupData -> ColOp -> PipedTrans +_unrollGroupTrans ugd co = case _combineColOp (colOp (_gdValue ugd)) co of + -- TODO: this is ugly, we are loosing the error structure. + Left x -> _pError $ "_unrollGroupTrans: failure with " <> show' x + Right co' -> PipedGroup $ ugd { _gdValue = _transformCol co' (_gdValue ugd) } + + +-- TODO: this should be moved to ColumnFunctions +_transformCol :: ColOp -> UntypedColumnData -> UntypedColumnData +-- TODO: at this point, it should be checked for correctness (the fields +-- being extracted should exist) +_transformCol co ucd = ucd { _cOp = co } + +-- Takes a column operation and chain it with another column operation. +_combineColOp :: ColOp -> ColOp -> Try ColOp +_combineColOp _ (x @ (ColLit _ _)) = pure x +_combineColOp x (ColFunction fn v) = + ColFunction fn <$> sequence (_combineColOp x <$> v) +_combineColOp x (ColExtraction fp) = _extractColOp x (V.toList (unFieldPath fp)) +_combineColOp x (ColStruct v) = + ColStruct <$> sequence (f <$> v) where + f (TransformField n val) = TransformField n <$> _combineColOp x val + +_extractColOp :: ColOp -> [FieldName] -> Try ColOp +_extractColOp x [] = pure x +_extractColOp (ColStruct s) (fn : t) = + case V.find (\x -> tfName x == fn) s of + Just (TransformField _ co) -> + _extractColOp co t + Nothing -> + tryError $ sformat ("Expected to find field "%sh%" in structure "%sh) fn s +_extractColOp x y = + tryError $ sformat ("Cannot perform extraction "%sh%" on column operation "%sh) y x + +_aggKey :: UntypedGroupData -> (UntypedColumnData -> Try UntypedLocalData) -> Try UntypedDataset +_aggKey ugd f = + let inputDt = unSQLType . colType . _gdValue $ ugd + p = placeholder inputDt :: UntypedDataset + startNid = nodeId p in do + uld <- f (_unsafeCastColData (asCol p)) + case _unrollTransform (PipedGroup ugd) startNid (untyped uld) of + PipedError t -> tryError t + PipedGroup g -> + -- This is a programming error + tryError $ sformat ("Expected a dataframe at the output but got a group: "%sh) g + PipedDataset ds -> pure ds + +_unsafeCastColData :: Column ref a -> Column ref' a' +_unsafeCastColData c = c { _cType = _cType c } + +{-| Checks that the group can be cast. +-} +_castGroup :: + SQLType key -> SQLType val -> UntypedGroupData -> Try (GroupData key val) +_castGroup (SQLType keyType) (SQLType valType) ugd = + let keyType' = unSQLType . colType . _gdKey $ ugd + valType' = unSQLType . colType . _gdValue $ ugd in + if keyType == keyType' + then if valType == valType' + then + pure ugd { _gdRef = _gdRef ugd } + else + tryError $ sformat ("The value column (of type "%sh%") cannot be cast to type "%sh) valType' valType + else + tryError $ sformat ("The value column (of type "%sh%") cannot be cast to type "%sh) keyType' keyType + +_untypedGroup :: GroupData key val -> UntypedGroupData +_untypedGroup gd = gd { _gdRef = _gdRef gd } + +_groupByKey :: UntypedColumnData -> UntypedColumnData -> LogicalGroupData +_groupByKey keys vals = + if nodeId (colOrigin keys) == nodeId (colOrigin vals) + then + pure GroupData { + _gdRef = colOrigin keys, + _gdKey = keys, + _gdValue = vals + } + else + tryError $ sformat ("The columns have different origin: "%sh%" and "%sh) keys vals diff --git a/src/Spark/Core/Internal/Joins.hs b/src/Spark/Core/Internal/Joins.hs new file mode 100644 --- /dev/null +++ b/src/Spark/Core/Internal/Joins.hs @@ -0,0 +1,96 @@ +{-# LANGUAGE OverloadedStrings #-} +{-| Exposes some of Spark's joining algorithms. +-} +module Spark.Core.Internal.Joins( + join, + join', + joinInner, + joinInner', + joinObs, + joinObs' +) where + +import qualified Data.Aeson as A +import qualified Data.Vector as V +import Control.Arrow + +import Spark.Core.Internal.ColumnStructures +import Spark.Core.Internal.ColumnFunctions +import Spark.Core.Internal.DatasetStructures +import Spark.Core.Internal.DatasetFunctions +import Spark.Core.Internal.FunctionsInternals +import Spark.Core.Internal.OpStructures +import Spark.Core.Internal.TypesStructures +import Spark.Core.Internal.Utilities +import Spark.Core.Internal.TypesFunctions(structTypeFromFields) +import Spark.Core.Try +import Spark.Core.StructuresInternal(unsafeFieldName) + +{-| Standard (inner) join on two sets of data. +-} +join :: Column ref1 key -> Column ref1 value1 -> Column ref2 key -> Column ref2 value2 -> Dataset (key, value1, value2) +join = joinInner + +{-| Untyped version of the standard join. +-} +join' :: DynColumn -> DynColumn -> DynColumn -> DynColumn -> DataFrame +join' = joinInner' + +{-| Explicit inner join. +-} +joinInner :: Column ref1 key -> Column ref1 value1 -> Column ref2 key -> Column ref2 value2 -> Dataset (key, value1, value2) +joinInner key1 val1 key2 val2 = unsafeCastDataset (forceRight df) where + df = joinInner' (untypedCol key1) (untypedCol val1) (untypedCol key2) (untypedCol val2) + +{-| Untyped version of the inner join. +-} +joinInner' :: DynColumn -> DynColumn -> DynColumn -> DynColumn -> DataFrame +joinInner' key1 val1 key2 val2 = do + df1 <- pack' (struct' [key1, val1]) + df2 <- pack' (struct' [key2, val2]) + dt <- _joinTypeInner key1 val1 val2 + let so = StandardOperator { soName = "org.spark.Join", soOutputType = dt, soExtra = A.String "inner" } + let ds = emptyDataset (NodeDistributedOp so) (SQLType dt) + let f ds' = ds' { _cnParents = V.fromList [untyped df1, untyped df2] } + return $ updateNode ds f + +{-| Broadcasts an observable alongside a dataset to make it available as an +extra column. +-} +-- This is the low-level operation that is used to implement the other +-- broadcast operations. +joinObs :: (HasCallStack) => Column ref val -> LocalData val' -> Dataset (val, val') +joinObs c ld = + -- TODO: has a forcing at the last moment so that we can at least + -- have stronger guarantees in the type coercion. + unsafeCastDataset $ forceRight $ joinObs' (untypedCol c) (pure (untypedLocalData ld)) + +{-| Broadcasts an observable along side a dataset to make it available as +an extra column. + +The resulting dataframe has 2 columns: + - one column called 'values' + - one column called 'broadcast' + + Note: this is a low-level operation. Users may want to use broadcastObs instead. +-} +joinObs' :: DynColumn -> LocalFrame -> DataFrame +joinObs' dc lf = do + let df = pack' dc + dc' <- df + c <- asCol' df + o <- lf + st <- structTypeFromFields [(unsafeFieldName "values", unSQLType (colType c)), (unsafeFieldName "broadcast", unSQLType (nodeType o))] + let sqlt = SQLType (StrictType (Struct st)) + return $ emptyDataset NodeBroadcastJoin sqlt `parents` [untyped dc', untyped o] + +-- {-| Broadcasts an observable along the axis of a dataset. +-- -} +-- broadcastObs :: ColumnReference ref -> LocalData val -> Column ref val +-- broadcastObs = missing "broadcastObs" + +_joinTypeInner :: DynColumn -> DynColumn -> DynColumn -> Try DataType +_joinTypeInner kcol col1 col2 = do + cs <- sequence [kcol, col1, col2] + st <- structTypeFromFields $ (colFieldName &&& unSQLType . colType) <$> cs + return $ StrictType (Struct st) diff --git a/src/Spark/Core/Internal/OpFunctions.hs b/src/Spark/Core/Internal/OpFunctions.hs --- a/src/Spark/Core/Internal/OpFunctions.hs +++ b/src/Spark/Core/Internal/OpFunctions.hs @@ -15,6 +15,7 @@ import qualified Data.Vector as V import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as LBS +import Data.Text(Text) import Data.Aeson((.=), toJSON) import Data.Char(isSymbol) import qualified Crypto.Hash.SHA256 as SHA @@ -28,21 +29,44 @@ simpleShowOp :: NodeOp -> T.Text simpleShowOp (NodeLocalOp op) = soName op simpleShowOp (NodeDistributedOp op) = soName op -simpleShowOp (NodeLocalLit _ _) = T.pack "org.spark.LocalConstant" +simpleShowOp (NodeLocalLit _ _) = "org.spark.LocalConstant" simpleShowOp (NodeOpaqueAggregator op) = soName op -simpleShowOp (NodeUniversalAggregator ua) = soName . uaoMergeBuffer $ ua -simpleShowOp (NodeStructuredTransform _) = T.pack "org.spark.Select" -simpleShowOp (NodeDistributedLit _ _) = T.pack "org.spark.Constant" +simpleShowOp (NodeAggregatorReduction uao) = + case uaoInitialOuter uao of + OpaqueAggTransform so -> soName so + _ -> "org.spark.StructuredReduction" +simpleShowOp (NodeAggregatorLocalReduction ua) = _prettyShowSGO . uaoMergeBuffer $ ua +simpleShowOp (NodeStructuredTransform _) = "org.spark.Select" +simpleShowOp (NodeDistributedLit _ _) = "org.spark.Constant" +simpleShowOp (NodeGroupedReduction _) = "org.spark.GroupedReduction" +simpleShowOp (NodeReduction _) = "org.spark.Reduction" +simpleShowOp NodeBroadcastJoin = "org.spark.BroadcastJoin" -- A human-readable string that represents column operations. prettyShowColOp :: ColOp -> T.Text prettyShowColOp (ColExtraction fpath) = T.pack (show fpath) prettyShowColOp (ColFunction txt cols) = - _prettyShowColFun txt (V.toList cols) + _prettyShowColFun txt (V.toList (prettyShowColOp <$> cols)) prettyShowColOp (ColLit _ cell) = T.pack (show cell) prettyShowColOp (ColStruct s) = "struct(" <> T.intercalate "," (prettyShowColOp . tfValue <$> V.toList s) <> ")" +_prettyShowAggOp :: AggOp -> T.Text +_prettyShowAggOp (AggUdaf _ ucn fp) = ucn <> "(" <> show' fp <> ")" +_prettyShowAggOp (AggFunction sfn v) = _prettyShowColFun sfn r where + r = V.toList (show' <$> v) +_prettyShowAggOp (AggStruct v) = + "struct(" <> T.intercalate "," (_prettyShowAggOp . afValue <$> V.toList v) <> ")" + +_prettyShowAggTrans :: AggTransform -> Text +_prettyShowAggTrans (OpaqueAggTransform op) = soName op +_prettyShowAggTrans (InnerAggOp ao) = _prettyShowAggOp ao + +_prettyShowSGO :: SemiGroupOperator -> Text +_prettyShowSGO (OpaqueSemiGroupLaw so) = soName so +_prettyShowSGO (UdafSemiGroupOperator ucn) = ucn +_prettyShowSGO (ColumnSemiGroupLaw sfn) = sfn + -- (internal) -- The extra data associated with the operation, and that is required -- by the backend to successfully perform the operation. @@ -59,6 +83,12 @@ -- as a SQL datatype. A.object [ "cellType" .= toJSON dt, "content" .= toJSON lst] +extraNodeOpData (NodeDistributedOp so) = soExtra so +extraNodeOpData (NodeGroupedReduction ao) = toJSON ao +extraNodeOpData (NodeAggregatorReduction ua) = + case uaoInitialOuter ua of + OpaqueAggTransform so -> toJSON (soExtra so) + InnerAggOp ao -> toJSON ao extraNodeOpData _ = A.Null -- Adds the content of a node op to a hash. @@ -73,16 +103,16 @@ "extra" .= extraNodeOpData op] -_prettyShowColFun :: T.Text -> [ColOp] -> T.Text +_prettyShowColFun :: T.Text -> [Text] -> T.Text _prettyShowColFun txt [col] | _isSym txt = - T.concat [txt, prettyShowColOp col] + T.concat [txt, col] _prettyShowColFun txt [col1, col2] | _isSym txt = -- This is not perfect for complex operations, but it should get the job done -- for now. -- TODO eventually use operator priority here - T.concat [prettyShowColOp col1, txt, prettyShowColOp col2] + T.concat [col1, txt, col2] _prettyShowColFun txt cols = - let vals = T.intercalate ", " (prettyShowColOp <$> cols) in + let vals = T.intercalate ", " cols in T.concat [txt, "(", vals, ")"] _isSym :: T.Text -> Bool @@ -108,6 +138,30 @@ A.object ["name" .= T.pack (show fn), "op" .= toJSON colOp] in A.Array $ fun <$> v +-- instance A.ToJSON AggTransform where +-- toJSON (OpaqueAggTransform so) = A.object [ +-- "aggOpaqueTrans" .= toJSON so +-- ] + +instance A.ToJSON UdafApplication where + toJSON Algebraic = toJSON (T.pack "algebraic") + toJSON Complete = toJSON (T.pack "complete") + +instance A.ToJSON AggField where + toJSON (AggField fn aggOp) = + A.object ["name" .= show' fn, "op" .= toJSON aggOp] + +instance A.ToJSON AggOp where + toJSON (AggUdaf ua ucn fp) = A.object [ + "aggOp" .= T.pack "udaf", + "udafApplication" .= toJSON ua, + "className" .= ucn, + "field" .= toJSON fp] + toJSON (AggFunction sfn v) = A.object [ + "aggOp" .= toJSON (T.pack "function"), + "functionName" .= toJSON sfn, + "fields" .= toJSON (V.toList v)] + toJSON (AggStruct v) = toJSON (V.toList v) _hashUpdateJson :: SHA.Ctx -> A.Value -> SHA.Ctx _hashUpdateJson ctx val = SHA.update ctx bs where diff --git a/src/Spark/Core/Internal/OpStructures.hs b/src/Spark/Core/Internal/OpStructures.hs --- a/src/Spark/Core/Internal/OpStructures.hs +++ b/src/Spark/Core/Internal/OpStructures.hs @@ -11,6 +11,20 @@ import Spark.Core.StructuresInternal import Spark.Core.Internal.TypesStructures(DataType, SQLType, SQLType(unSQLType)) +{-| The name of a SQL function. + +It is one of the predefined SQL functions available in Spark. +-} +type SqlFunctionName = T.Text + +{-| The classpath of a UDAF. +-} +type UdafClassName = T.Text + +{-| The name of an operator defined in Kraps. +-} +type OperatorName = T.Text + {-| The invariant respected by a transform. Depending on the value of the invariant, different optimizations @@ -51,7 +65,7 @@ -- | An operator defined by default in the release of Krapsh. -- All other physical operators can be converted to a standard operators. data StandardOperator = StandardOperator { - soName :: !T.Text, + soName :: !OperatorName, soOutputType :: !DataType, soExtra :: !Value } deriving (Eq, Show) @@ -76,7 +90,7 @@ -- In this case, the other columns may matter -- TODO(kps) add if this function is partition invariant. -- It should be the case most of the time. - | ColFunction !T.Text !(Vector ColOp) + | ColFunction !SqlFunctionName !(Vector ColOp) -- | A constant defined for each element. -- The type should be the same as for the column -- A literal is always direct @@ -97,7 +111,54 @@ | InnerStruct !(Vector TransformField) deriving (Eq, Show) +{-| When applying a UDAF, determines if it should only perform the algebraic +portion of the UDAF (initialize+update+merge), or if it also performs the final, +non-algebraic step. +-} +data UdafApplication = Algebraic | Complete deriving (Eq, Show) +data AggOp = + -- The name of the UDAF and the field path to apply it onto. + AggUdaf !UdafApplication !UdafClassName !FieldPath + -- A column function that can be applied (sum, max, etc.) + | AggFunction !SqlFunctionName !(Vector FieldPath) + | AggStruct !(Vector AggField) + deriving (Eq, Show) + +{-| A field in the resulting aggregation transform. +-} +data AggField = AggField { + afName :: !FieldName, + afValue :: !AggOp +} deriving (Eq, Show) + +{-| +-} +data AggTransform = + OpaqueAggTransform !StandardOperator + | InnerAggOp !AggOp deriving (Eq, Show) + +{-| The representation of a semi-group law in Spark. + +This is the basic law used in universal aggregators. It is a function on +observables that must respect the following laws: + +f :: X -> X -> X +commutative +associative + +A neutral element is not required for the semi-group laws. However, if used in +the context of a universal aggregator, such an element implicitly exists and +corresponds to the empty dataset. +-} +data SemiGroupOperator = + -- | A standard operator that happens to respect the semi-group laws. + OpaqueSemiGroupLaw !StandardOperator + -- | The merging portion of a UDAF + | UdafSemiGroupOperator !UdafClassName + -- | A SQL operator that happens to respect the semi-group laws. + | ColumnSemiGroupLaw !SqlFunctionName deriving (Eq, Show) + -- ********* DATASET OPERATORS ************ -- These describe Dataset -> Dataset transforms. @@ -119,10 +180,25 @@ -- Dataset -> Local data transform data UniversalAggregatorOp = UniversalAggregatorOp { uaoMergeType :: !DataType, - uaoInitialOuter :: !StandardOperator, - uaoMergeBuffer :: !StandardOperator + uaoInitialOuter :: !AggTransform, + uaoMergeBuffer :: !SemiGroupOperator } deriving (Eq, Show) + +data NodeOp2 = + -- empty -> local + NodeLocalLiteral !DataType !Value + -- empty -> distributed + | NodeDistributedLiteral !DataType !(Vector Value) + -- distributed -> local + | NodeStructuredAggregation !AggOp !(Maybe UniversalAggregatorOp) + -- distributed -> distributed or local -> local + | NodeStructuredTransform2 !Locality !ColOp + -- [distributed, local] -> [local, distributed] opaque + | NodeOpaqueTransform !Locality StandardOperator + deriving (Eq, Show) + + {- A node operation. A description of all the operations between nodes. @@ -137,15 +213,26 @@ Additionally, some operations are associated with algebraic invariants to enable programmatic transformations. -} +-- TODO: way too many different ops. Restructure into a few fundamental ops with +-- options. data NodeOp = -- | An operation between local nodes: [Observable] -> Observable NodeLocalOp StandardOperator -- | An observable literal | NodeLocalLit !DataType !Value + -- | A special join that broadcasts a value along a dataset. + | NodeBroadcastJoin -- | Some aggregator that does not respect any particular invariant. | NodeOpaqueAggregator StandardOperator + -- It implicicty expects a dataframe with 2 fields: + -- - the first field is used as a key + -- - the second field is passed to the reducer + | NodeGroupedReduction !AggOp + | NodeReduction !AggTransform + -- TODO: remove these -- | A universal aggregator. - | NodeUniversalAggregator UniversalAggregatorOp + | NodeAggregatorReduction UniversalAggregatorOp + | NodeAggregatorLocalReduction UniversalAggregatorOp -- | A structured transform, performed either on a local node or a -- distributed node. | NodeStructuredTransform !ColOp diff --git a/src/Spark/Core/Internal/Paths.hs b/src/Spark/Core/Internal/Paths.hs --- a/src/Spark/Core/Internal/Paths.hs +++ b/src/Spark/Core/Internal/Paths.hs @@ -29,7 +29,6 @@ import Spark.Core.Try import Spark.Core.Internal.DAGStructures -import Spark.Core.Internal.Utilities import Spark.Core.Internal.ComputeDag import Spark.Core.StructuresInternal @@ -47,13 +46,13 @@ -- Assigns paths in a graph. -- -computePaths :: (HasCallStack, HasNodeName v) => +computePaths :: (HasNodeName v) => ComputeDag v PathEdge -> Try (M.Map VertexId NodePath) computePaths cd = let nodecg = mapVertexData getNodeName cd in _computePaths nodecg -assignPaths' :: (HasCallStack, HasNodeName v) => +assignPaths' :: (HasNodeName v) => M.Map VertexId NodePath -> ComputeDag v e -> ComputeDag v e assignPaths' m cd = let f vx = @@ -65,8 +64,7 @@ -- The main function to perform the pass assignments. -- It starts from the graph of dependencies and from the local name info, -- and computes the complete paths (if possible), starting from the fringe. -_computePaths :: (HasCallStack) => - ComputeDag NodeName PathEdge -> Try (M.Map VertexId NodePath) +_computePaths :: ComputeDag NodeName PathEdge -> Try (M.Map VertexId NodePath) _computePaths cg = let scopes = iGetScopes0 (toList . cdOutputs $ cg) (_splitParents' (cdEdges cg)) @@ -168,17 +166,17 @@ innerParents = psInner split -- A fold on the parents parF :: Vertex a -> Scopes -> Scopes - parF v s = + parF = -- Same boundary and parent, but update the scopes - _getScopes' splitter mScopeId boundary v s + _getScopes' splitter mScopeId boundary scopesPar = foldr' parF scopes logParents -- Now work on the inner nodes: vid = vertexId un boundary' = S.fromList (vertexId <$> logParents) inF :: Vertex a -> Scopes -> Scopes - inF v s = + inF = -- parent is current, boundary is current logical - _getScopes' splitter (Just vid) boundary' v s + _getScopes' splitter (Just vid) boundary' scopesIn = foldr' inF scopesPar innerParents scopesFinal = scopesIn `mergeScopes` _singleScope mScopeId vid diff --git a/src/Spark/Core/Internal/RowGenerics.hs b/src/Spark/Core/Internal/RowGenerics.hs --- a/src/Spark/Core/Internal/RowGenerics.hs +++ b/src/Spark/Core/Internal/RowGenerics.hs @@ -20,7 +20,7 @@ import GHC.Generics import qualified Data.Vector as V -import Data.Text(pack) +import Data.Text(pack, Text) import Spark.Core.Internal.RowStructures import Spark.Core.Internal.Utilities @@ -53,9 +53,16 @@ _valueToCell (Just x) = _valueToCell x _valueToCell Nothing = Empty +instance (ToSQL a, ToSQL b) => ToSQL (a, b) where + _valueToCell (x, y) = RowArray (V.fromList [valueToCell x, valueToCell y]) + instance ToSQL Int where _valueToCell = IntElement +instance ToSQL Text where + _valueToCell = StringElement + + class GToSQL r where _g2buffer :: r a -> CurrentBuffer _g2cell :: r a -> Cell @@ -79,13 +86,13 @@ -- _g2cell !(M1 x) = let !y = _g2cell x in -- trace ("GToSQL M1: y = " ++ show y) y -instance (GToSQL a, Constructor c) => GToSQL (M1 C c a) where +instance (GToSQL a) => GToSQL (M1 C c a) where _g2buffer (M1 x) = let !y = _g2buffer x in y -instance (GToSQL a, Selector c) => GToSQL (M1 S c a) where +instance (GToSQL a) => GToSQL (M1 S c a) where _g2buffer (M1 x) = let !y = ConsData [_g2cell x] in y -instance (GToSQL a, Datatype c) => GToSQL (M1 D c a) where +instance (GToSQL a) => GToSQL (M1 D c a) where _g2buffer (M1 x) = case _g2buffer x of ConsData cs -> BuiltCell $ RowArray (V.fromList cs) diff --git a/src/Spark/Core/Internal/RowGenericsFrom.hs b/src/Spark/Core/Internal/RowGenericsFrom.hs --- a/src/Spark/Core/Internal/RowGenericsFrom.hs +++ b/src/Spark/Core/Internal/RowGenericsFrom.hs @@ -20,6 +20,7 @@ import GHC.Generics import Data.Text(Text, pack) +import Control.Applicative(liftA2) import Control.Monad.Except import Formatting import qualified Data.Vector as V @@ -58,12 +59,24 @@ _cellToValue (IntElement x) = pure x _cellToValue x = throwError $ sformat ("FromSQL: Decoding an int from "%shown) x +instance FromSQL Text where + _cellToValue (StringElement txt) = pure txt + _cellToValue x = throwError $ sformat ("FromSQL: Decoding a unicode text from "%shown) x + instance FromSQL Cell where _cellToValue = pure instance FromSQL a => FromSQL [a] where _cellToValue (RowArray xs) = sequence (_cellToValue <$> V.toList xs) _cellToValue x = throwError $ sformat ("FromSQL: Decoding array from "%shown) x + +instance (FromSQL a1, FromSQL a2) => FromSQL (a1, a2) where + _cellToValue (RowArray xs) = case V.toList xs of + [x1, x2] -> + liftA2 (,) (_cellToValue x1) (_cellToValue x2) + l -> throwError $ sformat ("FromSQL: Expected 2 elements but got "%sh) l + _cellToValue x = throwError $ sformat ("FromSQL: Decoding array from "%shown) x + -- ******* GENERIC ******** class GFromSQL r where @@ -72,24 +85,18 @@ instance GFromSQL U1 where _gFcell x = failure $ pack $ "GFromSQL UI called" ++ show x +_f :: Monad m => m (f p) -> m (g p) -> m ((:*:) f g p) +_f x1t x2t = do + x1 <- x1t + x2 <- x2t + return (x1 :*: x2) + instance (GFromSQL a, GFromSQL b) => GFromSQL (a :*: b) where _gFcell (D2Normal (RowArray arr)) | not (V.null arr) = let (cell : l) = V.toList arr - x1t = _gFcell (D2Normal cell) - x2t = _gFcell (D2Cons l) - x = do - x1 <- x1t - x2 <- x2t - return (x1 :*: x2) - in x + in _f (_gFcell (D2Normal cell)) (_gFcell (D2Cons l)) _gFcell (D2Cons (cell : l)) = - let x1t = _gFcell (D2Cons [cell]) - x2t = _gFcell (D2Cons l) - x = do - x1 <- x1t - x2 <- x2t - return (x1 :*: x2) - in x + _f (_gFcell (D2Cons [cell])) (_gFcell (D2Cons l)) _gFcell x = failure $ pack ("GFromSQL (a :*: b) " ++ show x) @@ -102,10 +109,10 @@ let xt = _gFcell (D2Normal cell) in M1 <$> xt -instance (GFromSQL a, Constructor c) => GFromSQL (M1 C c a) where +instance (GFromSQL a) => GFromSQL (M1 C c a) where _gFcell = _m1 "GFromSQL (M1 C c a)" -instance (GFromSQL a, Selector c) => GFromSQL (M1 S c a) where +instance (GFromSQL a) => GFromSQL (M1 S c a) where _gFcell (D2Normal (RowArray arr)) | V.length arr == 1 = M1 <$> _gFcell (D2Cons [cell]) where cell = V.head arr diff --git a/src/Spark/Core/Internal/RowUtils.hs b/src/Spark/Core/Internal/RowUtils.hs --- a/src/Spark/Core/Internal/RowUtils.hs +++ b/src/Spark/Core/Internal/RowUtils.hs @@ -3,6 +3,7 @@ module Spark.Core.Internal.RowUtils( jsonToCell, checkCell, + rowArray ) where import Data.Aeson @@ -20,6 +21,7 @@ import Spark.Core.StructuresInternal(FieldName(..)) import Spark.Core.Internal.Utilities +type TryCell = Either Text Cell -- | Decodes a JSON into a row. -- This operation requires a SQL type that describes @@ -34,8 +36,11 @@ Nothing -> pure c Just txt -> throwError txt +{-| Convenience constructor for an array of cells. +-} +rowArray :: [Cell] -> Cell +rowArray = RowArray . V.fromList -type TryCell = Either Text Cell -- Returns an error message if something wrong is found _checkCell :: DataType -> Cell -> Maybe Text diff --git a/src/Spark/Core/Internal/TypesFunctions.hs b/src/Spark/Core/Internal/TypesFunctions.hs --- a/src/Spark/Core/Internal/TypesFunctions.hs +++ b/src/Spark/Core/Internal/TypesFunctions.hs @@ -7,21 +7,31 @@ unsafeCastType, intType, arrayType, + compatibleTypes, arrayType', frameTypeFromCol, colTypeFromFrame, canNull, structField, structType, + structTypeFromFields, + tupleType, + structName, iSingleField, -- cellType, ) where -import Data.Text as T +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 -- Performs a cast of the type. -- This may throw an error if the required type b is not @@ -66,6 +76,25 @@ -- The strict int type + +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 @@ -100,6 +129,25 @@ [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 diff --git a/src/Spark/Core/Internal/TypesGenerics.hs b/src/Spark/Core/Internal/TypesGenerics.hs --- a/src/Spark/Core/Internal/TypesGenerics.hs +++ b/src/Spark/Core/Internal/TypesGenerics.hs @@ -12,11 +12,9 @@ module Spark.Core.Internal.TypesGenerics where import qualified Data.Vector as V -import Data.Text(Text, pack) -import Data.Proxy +import qualified Data.Text as T import GHC.Generics import Formatting -import Debug.Trace import Spark.Core.Internal.TypesStructures import Spark.Core.Internal.TypesFunctions @@ -28,7 +26,7 @@ -- Given a tag on a type, returns the equivalent SQL type. -- This is the type for a cell, not for a row. -- TODO(kps) more documentation -buildType :: (SQLTypeable a) => SQLType a +buildType :: (HasCallStack, SQLTypeable a) => SQLType a buildType = _buildType @@ -38,44 +36,41 @@ -- used by Spark. -- See also buildType on how to use it. class SQLTypeable a where - _genericTypeFromValue :: a -> GenericType - default _genericTypeFromValue :: (Generic a, GenSQLTypeable (Rep a)) => a -> GenericType - _genericTypeFromValue _ = genBuildType (Proxy :: Proxy a) - - -- | The only function that should matter for users in this file. - -- Given a type, returns the SQL representation of this type. - _buildType :: SQLType a - _buildType = - let !dt = _genericTypeFromValue (undefined :: a) - SQLType u = dt in SQLType u + _genericTypeFromValue :: (HasCallStack) => a -> GenericType + default _genericTypeFromValue :: (HasCallStack, Generic a, GenSQLTypeable (Rep a)) => a -> GenericType + _genericTypeFromValue x = genTypeFromProxy (from x) --- These a private types that should not be used elsewhere. -data GenericRow -type GenericType = SQLType GenericRow +-- Generic SQLTypeable +class GenSQLTypeable f where + genTypeFromProxy :: (HasCallStack) => f a -> GenericType --- Generic building type. -genBuildType :: forall a. (Generic a, GenSQLTypeable (Rep a)) => Proxy a -> GenericType -genBuildType _ = genTypeFromProxy (Proxy :: Proxy (Rep a)) +-- | The only function that should matter for users in this file. +-- Given a type, returns the SQL representation of this type. +_buildType :: forall a. (HasCallStack, SQLTypeable a) => SQLType a +_buildType = + let dt = _genericTypeFromValue (undefined :: a) + in SQLType dt +type GenericType = DataType instance SQLTypeable Int where - _genericTypeFromValue _ = SQLType (StrictType IntType) + _genericTypeFromValue _ = StrictType IntType -instance SQLTypeable Text where - _genericTypeFromValue _ = SQLType (StrictType StringType) +instance SQLTypeable T.Text where + _genericTypeFromValue _ = StrictType StringType instance {-# INCOHERENT #-} SQLTypeable String where - _genericTypeFromValue _ = SQLType (StrictType StringType) + _genericTypeFromValue _ = StrictType StringType instance SQLTypeable a => SQLTypeable (Maybe a) where _genericTypeFromValue _ = let SQLType dt = buildType :: (SQLType a) in - (SQLType . NullableType . iInnerStrictType) dt + (NullableType . iInnerStrictType) dt instance {-# OVERLAPPABLE #-} SQLTypeable a => SQLTypeable [a] where _genericTypeFromValue _ = let SQLType dt = buildType :: (SQLType a) in - (SQLType . StrictType . ArrayType) dt + (StrictType . ArrayType) dt instance forall a1 a2. ( SQLTypeable a2, @@ -86,69 +81,62 @@ SQLType t2 = buildType :: SQLType a2 in _buildTupleStruct [t1, t2] -_buildTupleStruct :: [DataType] -> SQLType x +_buildTupleStruct :: [GenericType] -> GenericType _buildTupleStruct dts = - let fnames = unsafeFieldName . pack. ("_" ++) . show <$> ([1..] :: [Int]) + let fnames = unsafeFieldName . T.pack. ("_" ++) . show <$> ([1..] :: [Int]) fs = uncurry StructField <$> zip fnames dts - in SQLType . StrictType . Struct . StructType $ V.fromList fs + in StrictType . Struct . StructType $ V.fromList fs -- instance (SQLTypeable a, SQLTypeable b) => SQLTypeable (a,b) where -- _genericTypeFromValue _ = _genericTypeFromValue (undefined :: a) ++ _genericTypeFromValue (undefined :: b) --- Generic SQLTypeable -class GenSQLTypeable a where - genTypeFromProxy :: Proxy a -> GenericType - --- Datatype -instance GenSQLTypeable f => GenSQLTypeable (M1 D x f) where - genTypeFromProxy _ = genTypeFromProxy (Proxy :: Proxy f) +instance (GenSQLTypeable f) => GenSQLTypeable (M1 D c f) where + genTypeFromProxy m = genTypeFromProxy (unM1 m) --- Constructor Metadata instance (GenSQLTypeable f, Constructor c) => GenSQLTypeable (M1 C c f) where - genTypeFromProxy _ - | conIsRecord (undefined :: t c f a) = - let !dt = genTypeFromProxy (Proxy :: Proxy f) in + genTypeFromProxy m + | conIsRecord m = + let x = unM1 m + dt = genTypeFromProxy x in dt | otherwise = -- It is assumed to be a newtype and we are going to unwrap it - let !dt1 = genTypeFromProxy (Proxy :: Proxy f) - in case iSingleField (unSQLType dt1) of - Just dt -> SQLType dt + let !dt1 = genTypeFromProxy (unM1 m) + in case iSingleField dt1 of + Just dt -> dt Nothing -> failure $ sformat ("M1 C "%sh%" dt1="%sh) n dt1 - where m = undefined :: t c f a - n = conName m + where n = conName m -- Selector Metadata instance (GenSQLTypeable f, Selector c) => GenSQLTypeable (M1 S c f) where - genTypeFromProxy _ = - let !st = genTypeFromProxy (Proxy :: Proxy f) - m = undefined :: t c f a + genTypeFromProxy m = + let st = genTypeFromProxy (unM1 m) n = selName m - SQLType innerdt = st - field = StructField { structFieldName = FieldName $ pack n, structFieldType = innerdt } + field = StructField { structFieldName = FieldName $ T.pack n, structFieldType = st } st2 = StructType (V.singleton field) in - SQLType (StrictType $ Struct st2) + StrictType $ Struct st2 --- Constructor Paramater -instance (GenSQLTypeable (Rep f), SQLTypeable f) => GenSQLTypeable (K1 R f) where - genTypeFromProxy _ = _genericTypeFromValue (undefined :: f) +instance (SQLTypeable a) => GenSQLTypeable (K1 R a) where + genTypeFromProxy m = _genericTypeFromValue (unK1 m) -- Sum branch instance (GenSQLTypeable a, GenSQLTypeable b) => GenSQLTypeable (a :+: b) where - genTypeFromProxy _ = - let !y1 = genTypeFromProxy (Proxy :: Proxy a) - !y2 = genTypeFromProxy (Proxy :: Proxy b) in - -- TODO: need to prune the branch and throw an error here - trace ("SUM: y1=" ++ show y1 ++ " y2=" ++ show y2) y1 + genTypeFromProxy (L1 x) = genTypeFromProxy x + genTypeFromProxy (R1 x) = genTypeFromProxy x -- Product branch instance (GenSQLTypeable a, GenSQLTypeable b) => GenSQLTypeable (a :*: b) where - genTypeFromProxy _ = - let y1 = genTypeFromProxy (Proxy :: Proxy a) - y2 = genTypeFromProxy (Proxy :: Proxy b) in case (y1, y2) of - (SQLType (StrictType (Struct s1)), SQLType (StrictType (Struct s2))) -> - (SQLType . StrictType . Struct) s where + genTypeFromProxy z = + -- Due to optimizations that I do not understand, the decomposition has to + -- be done inside the function. + -- Otherwise, the value (which is undefined) gets to be evaluated, and breaks + -- the code. + let (x1 :*: x2) = z + y1 = genTypeFromProxy x1 + y2 = genTypeFromProxy x2 in case (y1, y2) of + (StrictType (Struct s1), StrictType (Struct s2)) -> + (StrictType . Struct) s where fs = structFields s1 V.++ structFields s2 s = StructType fs _ -> failure $ sformat ("should not happen: left="%sh%" right="%sh) y1 y2 diff --git a/src/Spark/Core/Row.hs b/src/Spark/Core/Row.hs --- a/src/Spark/Core/Row.hs +++ b/src/Spark/Core/Row.hs @@ -4,7 +4,8 @@ FromSQL, valueToCell, cellToValue, - jsonToCell + jsonToCell, + rowArray ) where import Spark.Core.Internal.RowStructures diff --git a/src/Spark/Core/StructuresInternal.hs b/src/Spark/Core/StructuresInternal.hs --- a/src/Spark/Core/StructuresInternal.hs +++ b/src/Spark/Core/StructuresInternal.hs @@ -13,6 +13,9 @@ catNodePath, fieldName, unsafeFieldName, + emptyFieldPath, + nullFieldPath, + headFieldPath, fieldPath, ) where @@ -61,6 +64,16 @@ -- TODO: proper implementation fieldPath :: T.Text -> Either String FieldPath fieldPath x = Right . FieldPath . V.singleton $ FieldName x + +emptyFieldPath :: FieldPath +emptyFieldPath = FieldPath V.empty + +nullFieldPath :: FieldPath -> Bool +nullFieldPath = V.null . unFieldPath + +headFieldPath :: FieldPath -> Maybe FieldName +headFieldPath (FieldPath v) | V.null v = Nothing +headFieldPath (FieldPath v) = Just $ V.head v -- | The concatenated path. This is the inverse function of fieldPath. catNodePath :: NodePath -> T.Text diff --git a/src/Spark/Core/Types.hs b/src/Spark/Core/Types.hs --- a/src/Spark/Core/Types.hs +++ b/src/Spark/Core/Types.hs @@ -21,33 +21,17 @@ buildType, StructField, StructType, - castType, + -- castType, catNodePath ) where -import Formatting - import Spark.Core.Internal.TypesStructures import Spark.Core.Internal.TypesGenerics import Spark.Core.Internal.TypesFunctions import Spark.Core.StructuresInternal import Spark.Core.Internal.FunctionsInternals(TupleEquivalence(..), NameTuple(..)) -import Spark.Core.Try -- | Description of types supported in DataSets -- Krapsh supports a restrictive subset of Algebraic Datatypes that is amenable to SQL -- transformations. This file contains the description of all the supported types, and some -- conversion tools. - - --- -- Converts an (untyped) datatype to a generic tagged SQLType. --- cellType :: DataType -> CellType --- cellType = SQLType - --- Takes a given type and attempts to cast it to another type, --- which is known by the type system. -castType :: forall a b. (SQLTypeable b) => SQLType a -> Try (SQLType b) -castType sqlt = - let sqlt' = buildType :: SQLType b in - if unSQLType sqlt == unSQLType sqlt' then Right sqlt' - else tryError $ sformat ("castType: tried to cast "%shown%" into incompatible type "%shown) sqlt sqlt' diff --git a/test-integration/Spark/Core/CachingSpec.hs b/test-integration/Spark/Core/CachingSpec.hs --- a/test-integration/Spark/Core/CachingSpec.hs +++ b/test-integration/Spark/Core/CachingSpec.hs @@ -19,7 +19,7 @@ let ds = dataset l let ds' = autocache ds let c1 = asCol ds' - let s1 = colSum c1 + let s1 = sumCol c1 let s2 = count ds' let x = s1 + s2 l2 <- exec1Def x diff --git a/test-integration/Spark/Core/CollectSpec.hs b/test-integration/Spark/Core/CollectSpec.hs --- a/test-integration/Spark/Core/CollectSpec.hs +++ b/test-integration/Spark/Core/CollectSpec.hs @@ -12,6 +12,7 @@ import Spark.Core.Row import Spark.Core.Functions import Spark.Core.Column +import Spark.Core.IntegrationUtilities -- Collecting a dataset made from a list should yield the same list (modulo @@ -39,18 +40,16 @@ collectIdempotent ([] :: [Int]) run "ints1" $ collectIdempotent ([4,5,1,2,3] :: [Int]) - -- TODO(kps) in Spark 2.0.2, this fails!!! - -- Works with Spark 2.0.1 -> report - -- run "ints1_opt" $ - -- collectIdempotent ([Just 1, Nothing] :: [Maybe Int]) - -- run "nothing_ints_opt" $ - -- collectIdempotent ([Nothing] :: [Maybe Int]) run "ints1_opt" $ + collectIdempotent ([Just 1, Nothing] :: [Maybe Int]) + run "nothing_ints_opt" $ + collectIdempotent ([Nothing] :: [Maybe Int]) + run "ints1_opt" $ collectIdempotent ([Just 1, Just 2] :: [Maybe Int]) run "empty_ints_opt" $ collectIdempotent ([] :: [Maybe Int]) - -- describe "Integration test - collect on TestStruct5" $ do - -- run "empty_TestStruct5" $ - -- collectIdempotent ([] :: [TestStruct5]) - -- run "empty_TestStruct5" $ - -- collectIdempotent ([TestStruct5 1 2] :: [TestStruct5]) + describe "Integration test - collect on TestStruct5" $ do + run "empty_TestStruct5" $ + collectIdempotent ([] :: [TestStruct5]) + run "single_TestStruct5" $ + collectIdempotent ([TestStruct5 1 2] :: [TestStruct5]) diff --git a/test-integration/Spark/Core/GroupsSpec.hs b/test-integration/Spark/Core/GroupsSpec.hs new file mode 100644 --- /dev/null +++ b/test-integration/Spark/Core/GroupsSpec.hs @@ -0,0 +1,35 @@ +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} + +module Spark.Core.GroupsSpec where + +import Test.Hspec +import Data.Text(Text) + +import Spark.Core.Context +import Spark.Core.Functions +import Spark.Core.ColumnFunctions +import Spark.Core.Column +import Spark.Core.IntegrationUtilities +import Spark.Core.CollectSpec(run) +import Spark.Core.Internal.Groups + +sumGroup :: [MyPair] -> [(Text, Int)] -> IO () +sumGroup l lexp = do + let ds = dataset l + let keys = ds // myKey' + let values = ds // myVal' + let g = groupByKey keys values + let ds2 = g `aggKey` sumCol + l2 <- exec1Def $ collect (asCol ds2) + l2 `shouldBe` lexp + +spec :: Spec +spec = do + describe "Integration test - groups on (text, int)" $ do + run "empty" $ + sumGroup [] [] + run "one" $ + sumGroup [MyPair "x" 1] [("x", 1)] + run "two" $ + sumGroup [MyPair "x" 1, MyPair "x" 2, MyPair "y" 1] [("x", 3), ("y", 1)] diff --git a/test-integration/Spark/Core/IntegrationUtilities.hs b/test-integration/Spark/Core/IntegrationUtilities.hs --- a/test-integration/Spark/Core/IntegrationUtilities.hs +++ b/test-integration/Spark/Core/IntegrationUtilities.hs @@ -6,10 +6,12 @@ module Spark.Core.IntegrationUtilities where import GHC.Generics (Generic) +import Data.Text(Text) import Spark.Core.Context import Spark.Core.Types import Spark.Core.Row +import Spark.Core.Column data TestStruct1 = TestStruct1 { ts1f1 :: Int, @@ -36,7 +38,7 @@ data TestStruct5 = TestStruct5 { ts5f1 :: Int, ts5f2 :: Int -} deriving (Show, Eq, Generic) +} deriving (Show, Eq, Generic, Ord) -- instance ToJSON TestStruct5 instance SQLTypeable TestStruct5 instance FromSQL TestStruct5 @@ -51,3 +53,14 @@ newtype TestT1 = TestT1 { unTestT1 :: Int } deriving (Eq, Show, Generic, Num) + +data MyPair = MyPair { + myKey :: Text, + myVal :: Int } deriving (Generic, Show) + +myKey' :: StaticColProjection MyPair Text +myKey' = unsafeStaticProjection buildType "myKey" +myVal' :: StaticColProjection MyPair Int +myVal' = unsafeStaticProjection buildType "myVal" +instance SQLTypeable MyPair +instance ToSQL MyPair diff --git a/test-integration/Spark/Core/JoinsSpec.hs b/test-integration/Spark/Core/JoinsSpec.hs new file mode 100644 --- /dev/null +++ b/test-integration/Spark/Core/JoinsSpec.hs @@ -0,0 +1,24 @@ +{-# LANGUAGE MultiParamTypeClasses #-} + +module Spark.Core.JoinsSpec where + +import Test.Hspec + +import Spark.Core.Context +import Spark.Core.Dataset +import Spark.Core.Column +import Spark.Core.Row +import Spark.Core.Functions +import Spark.Core.SimpleAddSpec(run) + +spec :: Spec +spec = do + describe "Join test - join on ints" $ do + run "empty_ints1" $ do + let ds1 = dataset [(1,2)] :: Dataset (Int, Int) + let ds2 = dataset [(1,3)] :: Dataset (Int, Int) + let df1 = asDF ds1 + let df2 = asDF ds2 + let df = joinInner' (df1//"_1") (df1//"_2") (df2//"_1") (df2//"_2" @@ "_3") + res <- exec1Def' (collect' (asCol' df)) + res `shouldBe` rowArray [rowArray [IntElement 1, IntElement 2, IntElement 3]] diff --git a/test/Spark/Core/Internal/CachingSpec.hs b/test/Spark/Core/Internal/CachingSpec.hs --- a/test/Spark/Core/Internal/CachingSpec.hs +++ b/test/Spark/Core/Internal/CachingSpec.hs @@ -157,7 +157,7 @@ let ds = dataset l let ds' = autocache ds let c1 = asCol ds' - let s1 = colSum c1 + let s1 = sumCol c1 let s2 = count ds' let x = s1 + s2 let g = traceHint "g=" (intErrors x) diff --git a/test/Spark/Core/Internal/GroupsSpec.hs b/test/Spark/Core/Internal/GroupsSpec.hs new file mode 100644 --- /dev/null +++ b/test/Spark/Core/Internal/GroupsSpec.hs @@ -0,0 +1,56 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE MultiParamTypeClasses #-} + +module Spark.Core.Internal.GroupsSpec where + +import Data.Text(Text) +import Test.Hspec +import GHC.Generics +import Data.Either(isRight) + +import Spark.Core.Functions +import Spark.Core.ColumnFunctions +import Spark.Core.Dataset +import Spark.Core.Column +import Spark.Core.Row +import Spark.Core.Types +import Spark.Core.Internal.Groups + + +data MyPair = MyPair { + myKey :: Text, + myVal :: Int } deriving (Generic, Show) + +myKey' :: StaticColProjection MyPair Text +myKey' = unsafeStaticProjection buildType "myKey" +myVal' :: StaticColProjection MyPair Int +myVal' = unsafeStaticProjection buildType "myVal" +instance SQLTypeable MyPair +instance ToSQL MyPair + +-- The tests are really light for now, and just check that the code passes the +-- dynamic type checker. +spec :: Spec +spec = do + describe "typed grouping tests" $ do + let ds = dataset [MyPair "1" 1, MyPair "2" 2] + let keys = ds // myKey' + let values = ds // myVal' + let g = groupByKey keys values + let sqlt1 = buildType :: SQLType MyPair + it "group" $ do + let tds2 = castType sqlt1 (groupAsDS g) + tds2 `shouldSatisfy` isRight + it "map group" $ do + let g2 = g `mapGroup` \c -> c + c + let tds2 = castType sqlt1 (groupAsDS g2) + tds2 `shouldSatisfy` isRight + it "simple reduce" $ do + let ds2 = g `aggKey` sumCol + let tds3 = castType sqlt1 ds2 + tds3 `shouldSatisfy` isRight + it "complex reduce" $ do + let ds2 = g `aggKey` \c -> sumCol (c + c) + let tds3 = castType sqlt1 ds2 + tds3 `shouldSatisfy` isRight diff --git a/test/Spark/Core/SimpleExamplesSpec.hs b/test/Spark/Core/SimpleExamplesSpec.hs --- a/test/Spark/Core/SimpleExamplesSpec.hs +++ b/test/Spark/Core/SimpleExamplesSpec.hs @@ -30,13 +30,13 @@ describe "Simple examples" $ do it "Precdence of renaming" $ do let numbers = asCol ds1 - let s = colSum numbers + let s = sumCol numbers let numCount = count ds1 let avg = s `div` numCount @@ "myaverage" _cnName avg `shouldSatisfy` isJust it "name for simple integers" $ do let numbers = asCol ds1 - let s = colSum numbers + let s = sumCol numbers let numCount = count ds1 let avg = s `div` numCount @@ "myaverage" -- TODO: should it show "value: int" instead? diff --git a/test/Spark/Core/TypesSpec.hs b/test/Spark/Core/TypesSpec.hs --- a/test/Spark/Core/TypesSpec.hs +++ b/test/Spark/Core/TypesSpec.hs @@ -10,10 +10,13 @@ import Spark.Core.Types import Spark.Core.Internal.TypesFunctions +import Spark.Core.Internal.TypesGenerics() data TestStruct1 = TestStruct1 { ts1f1 :: Int, - ts1f2 :: Maybe Int } deriving (Show, Generic, SQLTypeable) + ts1f2 :: Maybe Int } deriving (Show, Generic) + +instance SQLTypeable TestStruct1 data TestStruct2 = TestStruct2 { ts2f1 :: [Int] } deriving (Show, Generic, SQLTypeable)