diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Revision history for dataframe +## 0.4.0.5 +* Faster groupby: does less allocations by keeping everything in a mutable vector. +* declareColumnsFromCsvFile now infers types from a sample rather than reading the whole dataframe. +* Decision trees API is now more configurable. +* Add annotation to show what expressions were used to derive a column. + ## 0.4.0.4 * More robust synthesis based decision tree * Improved performance on sum and mean. diff --git a/dataframe.cabal b/dataframe.cabal --- a/dataframe.cabal +++ b/dataframe.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: dataframe -version: 0.4.0.4 +version: 0.4.0.5 synopsis: A fast, safe, and intuitive DataFrame library. @@ -48,6 +48,7 @@ DataFrame.Display.Web.Plot, DataFrame.Internal.Types, DataFrame.Internal.Expression, + DataFrame.Internal.Interpreter, DataFrame.Internal.Parsing, DataFrame.Internal.Column, DataFrame.Internal.Statistics, diff --git a/src/DataFrame/DecisionTree.hs b/src/DataFrame/DecisionTree.hs --- a/src/DataFrame/DecisionTree.hs +++ b/src/DataFrame/DecisionTree.hs @@ -14,8 +14,9 @@ import qualified DataFrame.Functions as F import DataFrame.Internal.Column import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn) -import DataFrame.Internal.Expression (Expr (..), eSize, interpret) -import DataFrame.Internal.Statistics (percentileOrd') +import DataFrame.Internal.Expression (Expr (..), eSize) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.Internal.Statistics (percentile', percentileOrd') import DataFrame.Internal.Types import DataFrame.Operations.Core (columnNames, nRows) import DataFrame.Operations.Statistics (percentile) @@ -27,6 +28,7 @@ import Data.Function (on) import Data.List (foldl', maximumBy, sortBy) import qualified Data.Map.Strict as M +import Data.Maybe import qualified Data.Text as T import Data.Type.Equality import qualified Data.Vector as V @@ -35,25 +37,33 @@ import DataFrame.Functions ((.<), (.<=), (.==), (.>), (.>=)) -data TreeConfig = TreeConfig +data TreeConfig + = TreeConfig { maxTreeDepth :: Int , minSamplesSplit :: Int + , minLeafSize :: Int , synthConfig :: SynthConfig } + deriving (Eq, Show) data SynthConfig = SynthConfig { maxExprDepth :: Int + , boolExpansion :: Int , percentiles :: [Int] + , complexityPenalty :: Double , enableStringOps :: Bool , enableCrossCols :: Bool , enableArithOps :: Bool } + deriving (Eq, Show) defaultSynthConfig :: SynthConfig defaultSynthConfig = SynthConfig { maxExprDepth = 2 + , boolExpansion = 2 , percentiles = [0, 10 .. 100] + , complexityPenalty = 0.05 , enableStringOps = True , enableCrossCols = True , enableArithOps = True @@ -62,8 +72,9 @@ defaultTreeConfig :: TreeConfig defaultTreeConfig = TreeConfig - { maxTreeDepth = 10 + { maxTreeDepth = 4 , minSamplesSplit = 5 + , minLeafSize = 1 , synthConfig = defaultSynthConfig } @@ -98,7 +109,7 @@ | depth <= 0 || nRows df <= minSamplesSplit cfg = Lit (majorityValue @a target df) | otherwise = - case findBestSplit @a target conds df of + case findBestSplit @a cfg target conds df of Nothing -> Lit (majorityValue @a target df) Just bestCond -> let (dfTrue, dfFalse) = partitionDataFrame bestCond df @@ -208,11 +219,54 @@ percentiles = map (Lit . (`percentileOrd'` col)) [1, 25, 75, 99] in map (Col @a colName .==) percentiles - (OptionalColumn (col :: V.Vector a)) -> - let - percentiles = map (Lit . (`percentileOrd'` col)) [1, 25, 75, 99] - in - map (Col @a colName .==) percentiles + (OptionalColumn (col :: V.Vector (Maybe a))) -> case sFloating @a of + STrue -> + let + doubleCol = + VU.convert + (V.map fromJust (V.filter isJust (V.map (fmap (realToFrac @a @Double)) col))) + in + zipWith + ($) + [ (Col @(Maybe a) colName .==) + , (Col @(Maybe a) colName .<=) + , (Col @(Maybe a) colName .>=) + ] + ( Lit Nothing + : map + ( Lit + . Just + . realToFrac + . (`percentile'` doubleCol) + ) + (percentiles cfg) + ) + SFalse -> case sIntegral @a of + STrue -> + let + doubleCol = + VU.convert + (V.map fromJust (V.filter isJust (V.map (fmap (fromIntegral @a @Double)) col))) + in + zipWith + ($) + [ (Col @(Maybe a) colName .==) + , (Col @(Maybe a) colName .<=) + , (Col @(Maybe a) colName .>=) + ] + ( Lit Nothing + : map + ( Lit + . Just + . round + . (`percentile'` doubleCol) + ) + (percentiles cfg) + ) + SFalse -> + map + ((Col @(Maybe a) colName .==) . Lit . (`percentileOrd'` col)) + [1, 25, 75, 99] (UnboxedColumn (col :: VU.Vector a)) -> [] columnConds = concatMap colConds [(l, r) | l <- columnNames df, r <- columnNames df] where @@ -238,11 +292,9 @@ findBestSplit :: forall a. (Columnable a) => - T.Text -> [Expr Bool] -> DataFrame -> Maybe (Expr Bool) -findBestSplit target conds df = + TreeConfig -> T.Text -> [Expr Bool] -> DataFrame -> Maybe (Expr Bool) +findBestSplit cfg target conds df = let - minLeafSize = 1 - lambda = 0.05 initialImpurity = calculateGini @a target df evalGain cond = let (t, f) = partitionDataFrame cond df @@ -252,7 +304,8 @@ newImpurity = (weightT * calculateGini @a target t) + (weightF * calculateGini @a target f) - in ( (initialImpurity - newImpurity) - lambda * fromIntegral (eSize cond) + in ( (initialImpurity - newImpurity) + - complexityPenalty (synthConfig cfg) * fromIntegral (eSize cond) , negate (eSize cond) ) @@ -262,7 +315,7 @@ let (t, f) = partitionDataFrame c df in - nRows t >= minLeafSize && nRows f >= minLeafSize + nRows t >= minLeafSize cfg && nRows f >= minLeafSize cfg ) (nubOrd conds) sortedConditions = take 10 (sortBy (flip compare `on` evalGain) validConds) @@ -273,7 +326,13 @@ Just $ maximumBy (compare `on` evalGain) - (boolExprs df sortedConditions sortedConditions 0 3) + ( boolExprs + df + sortedConditions + sortedConditions + 0 + (boolExpansion (synthConfig cfg)) + ) calculateGini :: forall a. diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs --- a/src/DataFrame/Functions.hs +++ b/src/DataFrame/Functions.hs @@ -398,7 +398,9 @@ declareColumnsFromCsvFile :: String -> DecsQ declareColumnsFromCsvFile path = do - df <- liftIO (CSV.readCsv path) + df <- + liftIO + (CSV.readSeparated (CSV.defaultReadOptions{CSV.numColumns = Just 100}) path) declareColumns df -- TODO: We don't have to read the whole file, we can just read the schema. diff --git a/src/DataFrame/IO/CSV.hs b/src/DataFrame/IO/CSV.hs --- a/src/DataFrame/IO/CSV.hs +++ b/src/DataFrame/IO/CSV.hs @@ -87,7 +87,7 @@ VM.unsafeWrite active count val writeIORef countRef $! count + 1 else do - frozen <- V.freeze active + frozen <- V.unsafeFreeze active modifyIORef' chunksRef (frozen :) newActive <- VM.unsafeNew chunkSize @@ -107,7 +107,7 @@ VUM.unsafeWrite active count val writeIORef countRef $! count + 1 else do - frozen <- VU.freeze active + frozen <- VU.unsafeFreeze active modifyIORef' chunksRef (frozen :) newActive <- VUM.unsafeNew chunkSize @@ -123,7 +123,7 @@ active <- readIORef activeRef chunks <- readIORef chunksRef - lastChunk <- V.freeze (VM.slice 0 count active) + lastChunk <- V.unsafeFreeze (VM.slice 0 count active) return $! V.concat (reverse (lastChunk : chunks)) @@ -134,7 +134,7 @@ active <- readIORef activeRef chunks <- readIORef chunksRef - lastChunk <- VU.freeze (VUM.slice 0 count active) + lastChunk <- VU.unsafeFreeze (VUM.slice 0 count active) return $! VU.concat (reverse (lastChunk : chunks)) -- | STANDARD CONFIG TYPES @@ -165,6 +165,8 @@ -} , columnSeparator :: Char -- ^ Character that separates column values. + , numColumns :: Maybe Int + -- ^ Number of columns to read. } shouldInferFromSample :: TypeSpec -> Bool @@ -187,6 +189,7 @@ , safeRead = True , dateFormat = "%Y-%m-%d" , columnSeparator = ',' + , numColumns = Nothing } {- | Read CSV file from path and load it into a dataframe. @@ -260,7 +263,7 @@ (sampleRow, _) <- peekStream rowsToProcess builderCols <- initializeColumns (V.toList sampleRow) opts - processStream rowsToProcess builderCols + processStream rowsToProcess builderCols (numColumns opts) frozenCols <- V.fromList <$> mapM freezeBuilderColumn builderCols let numRows = maybe 0 columnLength (frozenCols V.!? 0) @@ -270,7 +273,7 @@ frozenCols (M.fromList (zip columnNames [0 ..])) (numRows, V.length frozenCols) - + M.empty -- TODO give typed column references return $ if shouldInferFromSample (typeSpec opts) then @@ -302,10 +305,14 @@ Nothing -> BuilderText <$> newPagedVector <*> pure validityRef processStream :: - CsvStream.Records (V.Vector BL.ByteString) -> [BuilderColumn] -> IO () -processStream (Cons (Right row) rest) cols = processRow row cols >> processStream rest cols -processStream (Cons (Left err) _) _ = error ("CSV Parse Error: " ++ err) -processStream (Nil _ _) _ = return () + CsvStream.Records (V.Vector BL.ByteString) -> + [BuilderColumn] -> + Maybe Int -> + IO () +processStream _ _ (Just 0) = return () +processStream (Cons (Right row) rest) cols n = processRow row cols >> processStream rest cols (fmap (flip (-) 1) n) +processStream (Cons (Left err) _) _ _ = error ("CSV Parse Error: " ++ err) +processStream (Nil _ _) _ _ = return () processRow :: V.Vector BL.ByteString -> [BuilderColumn] -> IO () processRow !vals !cols = V.zipWithM_ processValue vals (V.fromList cols) diff --git a/src/DataFrame/IO/Unstable/CSV.hs b/src/DataFrame/IO/Unstable/CSV.hs --- a/src/DataFrame/IO/Unstable/CSV.hs +++ b/src/DataFrame/IO/Unstable/CSV.hs @@ -133,7 +133,7 @@ zip (Vector.toList columnNames) [0 ..] dataframeDimensions = (numRow, numCol) return $ - DataFrame columns columnIndices dataframeDimensions + DataFrame columns columnIndices dataframeDimensions M.empty {-# INLINE extractField #-} extractField :: diff --git a/src/DataFrame/Internal/Column.hs b/src/DataFrame/Internal/Column.hs --- a/src/DataFrame/Internal/Column.hs +++ b/src/DataFrame/Internal/Column.hs @@ -333,13 +333,13 @@ atIndices :: S.Set Int -> Column -> Column atIndices indexes (BoxedColumn column) = BoxedColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column atIndices indexes (OptionalColumn column) = OptionalColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column -atIndices indexes (UnboxedColumn column) = UnboxedColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column +atIndices indexes (UnboxedColumn column) = UnboxedColumn $ VU.ifilter (\i _ -> i `S.member` indexes) column {-# INLINE atIndices #-} -- | O(n) Selects the elements at a given set of indices. Does not change the order. atIndicesStable :: VU.Vector Int -> Column -> Column atIndicesStable indexes (BoxedColumn column) = BoxedColumn $ VG.unsafeBackpermute column (VG.convert indexes) -atIndicesStable indexes (UnboxedColumn column) = UnboxedColumn $ VG.unsafeBackpermute column indexes +atIndicesStable indexes (UnboxedColumn column) = UnboxedColumn $ VU.unsafeBackpermute column indexes atIndicesStable indexes (OptionalColumn column) = OptionalColumn $ VG.unsafeBackpermute column (VG.convert indexes) {-# INLINE atIndicesStable #-} diff --git a/src/DataFrame/Internal/DataFrame.hs b/src/DataFrame/Internal/DataFrame.hs --- a/src/DataFrame/Internal/DataFrame.hs +++ b/src/DataFrame/Internal/DataFrame.hs @@ -20,6 +20,7 @@ import DataFrame.Display.Terminal.PrettyPrint import DataFrame.Errors import DataFrame.Internal.Column +import DataFrame.Internal.Expression import Text.Printf import Type.Reflection (typeRep) @@ -32,6 +33,7 @@ -- ^ Keeps the column names in the order they were inserted in. , dataframeDimensions :: (Int, Int) -- ^ (rows, columns) + , derivingExpressions :: M.Map T.Text UExpr } {- | A record that contains information about how and what @@ -121,6 +123,7 @@ { columns = V.empty , columnIndices = M.empty , dataframeDimensions = (0, 0) + , derivingExpressions = M.empty } {- | Safely retrieves a column by name from the dataframe. diff --git a/src/DataFrame/Internal/Expression.hs b/src/DataFrame/Internal/Expression.hs --- a/src/DataFrame/Internal/Expression.hs +++ b/src/DataFrame/Internal/Expression.hs @@ -1,5 +1,4 @@ {-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} @@ -7,7 +6,6 @@ {-# LANGUAGE InstanceSigs #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} @@ -15,23 +13,13 @@ module DataFrame.Internal.Expression where -import Control.Monad.ST (runST) -import Data.Bifunctor -import qualified Data.Map as M -import Data.Maybe (fromMaybe, isJust) import Data.String import qualified Data.Text as T import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) -import qualified Data.Vector as V import qualified Data.Vector.Generic as VG -import qualified Data.Vector.Mutable as VM import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM -import DataFrame.Errors import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame -import DataFrame.Internal.Types -import Type.Reflection (TypeRep, Typeable, typeOf, typeRep, pattern App) +import Type.Reflection (Typeable, typeOf, typeRep) data Expr a where Col :: (Columnable a) => T.Text -> Expr a @@ -57,774 +45,11 @@ data UExpr where Wrap :: (Columnable a) => Expr a -> UExpr -type NamedExpr = (T.Text, UExpr) - -interpret :: - forall a. - (Columnable a) => - DataFrame -> Expr a -> Either DataFrameException (TypedColumn a) -interpret df (Lit value) = case sUnbox @a of - -- Specialize the creation of unboxed columns to avoid an extra allocation. - STrue -> pure $ TColumn $ fromUnboxedVector $ VU.replicate (numRows df) value - SFalse -> pure $ TColumn $ fromVector $ V.replicate (numRows df) value -interpret df (Col name) = maybe columnNotFound (pure . TColumn) (getColumn name df) - where - columnNotFound = Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) --- Unary operations. -interpret df expr@(UnaryOp _ (f :: c -> d) value) = first (handleInterpretException (show expr)) $ do - (TColumn value') <- interpret @c df value - fmap TColumn (mapColumn f value') --- Variations of binary operations. -interpret df expr@(BinaryOp _ (f :: c -> d -> e) left right) = first (handleInterpretException (show expr)) $ case (left, right) of - (Lit left, Lit right) -> interpret df (Lit (f left right)) - (Lit left, right) -> do - -- If we have a literal then we don't have to materialise - -- the column. - (TColumn value') <- interpret @d df right - fmap TColumn (mapColumn (f left) value') - (left, Lit right) -> do - -- Same as the above except the right side is the - -- literl. - (TColumn value') <- interpret @c df left - fmap TColumn (mapColumn (`f` right) value') - (_, _) -> do - -- In the general case we interpret and zip. - (TColumn left') <- interpret @c df left - (TColumn right') <- interpret @d df right - fmap TColumn (zipWithColumns f left' right') --- Conditionals -interpret df expr@(If cond l r) = first (handleInterpretException (show expr)) $ do - (TColumn conditions) <- interpret @Bool df cond - (TColumn left) <- interpret @a df l - (TColumn right) <- interpret @a df r - let branch (c :: Bool) (l' :: a, r' :: a) = if c then l' else r' - fmap TColumn (zipWithColumns branch conditions (zipColumns left right)) -interpret df expression@(AggVector expr op (f :: v b -> c)) = do - (TColumn column) <- interpret @b df expr - -- Helper for errors. Should probably find a way of throwing this - -- without leaking the fact that we use `Vector` to users. - let aggTypeError expected = - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @(v b)) - , expectedType = Left expected :: Either String (TypeRep ()) - , callingFunctionName = Just "interpret" - , errorColumnName = Nothing - } - ) - let processColumn :: - (Columnable d) => d -> Either DataFrameException (TypedColumn a) - processColumn col = case testEquality (typeRep @(v b)) (typeOf col) of - Just Refl -> interpret @c df (Lit (f col)) - Nothing -> Left $ aggTypeError (show (typeOf col)) - case column of - (BoxedColumn col) -> processColumn col - (OptionalColumn col) -> processColumn col - (UnboxedColumn col) -> processColumn col -interpret df expression@(AggReduce expr op (f :: a -> a -> a)) = first (handleInterpretException (show expr)) $ do - (TColumn column) <- interpret @a df expr - value <- foldl1Column f column - pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value -interpret df expression@(AggNumericVector expr op (f :: VU.Vector b -> c)) = first (handleInterpretException (show expression)) $ do - (TColumn column) <- interpret @b df expr - case column of - (UnboxedColumn (v :: VU.Vector d)) -> case testEquality (typeRep @d) (typeRep @b) of - Just Refl -> - pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) (f v) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - (Right (typeRep @b)) - (Right (typeRep @d)) - (Just "interpret") - (Just (show expression)) - ) - _ -> error "Trying to apply numeric computation to non-numeric column" -interpret df expression@(AggFold expr op start (f :: (a -> b -> a))) = first (handleInterpretException (show expression)) $ do - (TColumn column) <- interpret @b df expr - value <- foldlColumn f start column - pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value - -data AggregationResult a - = UnAggregated Column - | Aggregated (TypedColumn a) +instance Show UExpr where + show :: UExpr -> String + show (Wrap expr) = show expr -interpretAggregation :: - forall a. - (Columnable a) => - GroupedDataFrame -> Expr a -> Either DataFrameException (AggregationResult a) -interpretAggregation gdf (Lit value) = - Right $ - Aggregated $ - TColumn $ - fromVector $ - V.replicate (VG.length (offsets gdf) - 1) value -interpretAggregation gdf@(Grouped df names indices os) (Col name) = case getColumn name df of - Nothing -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) - Just (BoxedColumn col) -> Right $ UnAggregated $ fromVector $ mkUnaggregatedColumnBoxed col os indices - Just (OptionalColumn col) -> Right $ UnAggregated $ fromVector $ mkUnaggregatedColumnBoxed col os indices - Just (UnboxedColumn col) -> - Right $ UnAggregated $ fromVector $ mkUnaggregatedColumnUnboxed col os indices -interpretAggregation gdf expression@(UnaryOp _ (f :: c -> d) expr) = - case interpretAggregation @c gdf expr of - Left (TypeMismatchException context) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expr) - } - ) - Left e -> Left e - Right (UnAggregated unaggregated) -> case unaggregated of - BoxedColumn (col :: V.Vector b) -> case testEquality (typeRep @b) (typeRep @(V.Vector c)) of - Just Refl -> case sUnbox @d of - SFalse -> Right $ UnAggregated $ fromVector $ V.map (V.map f) col - STrue -> - Right $ - UnAggregated $ - fromVector $ - V.map (V.convert @V.Vector @d @VU.Vector . V.map f) col - Nothing -> case testEquality (typeRep @b) (typeRep @(VU.Vector c)) of - Nothing -> Left $ nestedTypeException @b @c (show expression) - Just Refl -> case (sUnbox @c, sUnbox @a) of - (SFalse, _) -> Left $ InternalException "Boxed type inside an unboxed column" - (STrue, STrue) -> Right $ UnAggregated $ fromVector $ V.map (VU.map f) col - (STrue, _) -> Right $ UnAggregated $ fromVector $ V.map (V.map f . VU.convert) col - _ -> Left $ InternalException "Aggregated into a non-boxed column" - Right (Aggregated (TColumn aggregated)) -> case mapColumn f aggregated of - Left e -> Left e - Right col -> Right $ Aggregated $ TColumn col -interpretAggregation gdf expression@(BinaryOp name (f :: c -> d -> e) left (Lit (right :: g))) = case testEquality (typeRep @g) (typeRep @d) of - Just Refl -> interpretAggregation gdf (UnaryOp name (`f` right) left) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @g) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) -interpretAggregation gdf expression@(BinaryOp name (f :: c -> d -> e) (Lit (left :: g)) right) = case testEquality (typeRep @g) (typeRep @c) of - Just Refl -> interpretAggregation gdf (UnaryOp name (f left) right) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @g) - , expectedType = Right (typeRep @c) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) -interpretAggregation gdf expression@(BinaryOp _ (f :: c -> d -> e) left right) = - case (interpretAggregation @c gdf left, interpretAggregation @d gdf right) of - (Right (Aggregated (TColumn left')), Right (Aggregated (TColumn right'))) -> case zipWithColumns f left' right' of - Left e -> Left e - Right col -> Right $ Aggregated $ TColumn col - (Right (UnAggregated left'), Right (UnAggregated right')) -> case (left', right') of - (BoxedColumn (l :: V.Vector m), BoxedColumn (r :: V.Vector n)) -> case testEquality (typeRep @m) (typeRep @(VU.Vector c)) of - Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector d)) of - Just Refl -> case (sUnbox @c, sUnbox @d, sUnbox @e) of - (STrue, STrue, STrue) -> - Right $ UnAggregated $ fromVector $ V.zipWith (VU.zipWith f) l r - (STrue, STrue, SFalse) -> - Right $ - UnAggregated $ - fromVector $ - V.zipWith (\l' r' -> V.zipWith f (V.convert l') (V.convert r')) l r - (_, _, _) -> Left $ InternalException "Boxed vectors contain unboxed types" - Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector d)) of - Just Refl -> case sUnbox @c of - STrue -> - Right $ - UnAggregated $ - fromVector $ - V.zipWith (V.zipWith f . V.convert) l r - SFalse -> Left $ InternalException "Unboxed vectors contain boxed types" - Nothing -> Left $ nestedTypeException @n @d (show right) - Nothing -> case testEquality (typeRep @m) (typeRep @(V.Vector c)) of - Nothing -> Left $ nestedTypeException @m @c (show left) - Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector d)) of - Just Refl -> case (sUnbox @d, sUnbox @e) of - (STrue, STrue) -> - Right $ - UnAggregated $ - fromVector $ - V.zipWith - (\l' r' -> V.convert @V.Vector @e @VU.Vector $ V.zipWith f l' (V.convert r')) - l - r - (STrue, SFalse) -> - Right $ - UnAggregated $ - fromVector $ - V.zipWith (\l' r' -> V.zipWith f l' (V.convert r')) l r - (_, _) -> Left $ InternalException "Unboxed vectors contain boxed types" - Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector d)) of - Just Refl -> case sUnbox @e of - SFalse -> - Right $ - UnAggregated $ - fromVector $ - V.zipWith (V.zipWith f . V.convert) l r - STrue -> - Right $ - UnAggregated $ - fromVector $ - V.zipWith (\l' r' -> V.convert @V.Vector @e @VU.Vector $ V.zipWith f l' r') l r - Nothing -> Left $ nestedTypeException @n @d (show right) - _ -> Left $ InternalException "Aggregated into a non-boxed column" - (Right _, Right _) -> - Left $ - AggregatedAndNonAggregatedException (T.pack $ show left) (T.pack $ show right) - (Left (TypeMismatchException context), _) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show left) - } - ) - (Left e, _) -> Left e - (_, Left (TypeMismatchException context)) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show right) - } - ) - (_, Left e) -> Left e -interpretAggregation gdf expression@(If cond (Lit l) (Lit r)) = - case interpretAggregation @Bool gdf cond of - Right (Aggregated (TColumn conditions)) -> case mapColumn - (\(c :: Bool) -> if c then l else r) - conditions of - Left e -> Left e - Right v -> Right $ Aggregated (TColumn v) - Right (UnAggregated conditions) -> case sUnbox @a of - STrue -> case mapColumn - (\(c :: VU.Vector Bool) -> VU.map (\c' -> if c' then l else r) c) - conditions of - Left (TypeMismatchException context) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - Left e -> Left e - Right v -> Right $ UnAggregated v - SFalse -> case mapColumn - (\(c :: VU.Vector Bool) -> V.map (\c' -> if c' then l else r) (VU.convert c)) - conditions of - Left (TypeMismatchException context) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - Left e -> Left e - Right v -> Right $ UnAggregated v - Left (TypeMismatchException context) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show cond) - } - ) - Left e -> Left e -interpretAggregation gdf expression@(If cond (Lit l) r) = - case ( interpretAggregation @Bool gdf cond - , interpretAggregation @a gdf r - ) of - ( Right (Aggregated (TColumn conditions)) - , Right (Aggregated (TColumn right)) - ) -> case zipWithColumns - (\(c :: Bool) (r' :: a) -> if c then l else r') - conditions - right of - Left e -> Left e - Right v -> Right $ Aggregated (TColumn v) - ( Right (UnAggregated conditions) - , Right (UnAggregated right@(BoxedColumn (right' :: V.Vector c))) - ) -> case testEquality (typeRep @(V.Vector a)) (typeRep @c) of - Just Refl -> case zipWithColumns - ( \(c :: VU.Vector Bool) (r' :: V.Vector a) -> - V.zipWith - (\c' r'' -> if c' then l else r'') - (V.convert c) - r' - ) - conditions - right of - Left (TypeMismatchException context) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - Left e -> Left e - Right v -> Right $ UnAggregated v - Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @c) of - Nothing -> Left $ nestedTypeException @c @a (show expression) - Just Refl -> case sUnbox @a of - SFalse -> Left $ InternalException "Boxed type in unboxed column" - STrue -> case zipWithColumns - ( \(c :: VU.Vector Bool) (r' :: VU.Vector a) -> - VU.zipWith - (\c' r'' -> if c' then l else r'') - c - r' - ) - conditions - right of - Left (TypeMismatchException context) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - Left e -> Left e - Right v -> Right $ UnAggregated v - (Right _, Right _) -> - Left $ - AggregatedAndNonAggregatedException (T.pack $ show l) (T.pack $ show r) - (Left (TypeMismatchException context), _) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show cond) - } - ) - (Left e, _) -> Left e - (_, Left (TypeMismatchException context)) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show r) - } - ) - (_, Left e) -> Left e -interpretAggregation gdf expression@(If cond l (Lit r)) = - case ( interpretAggregation @Bool gdf cond - , interpretAggregation @a gdf l - ) of - ( Right (Aggregated (TColumn conditions)) - , Right (Aggregated (TColumn left)) - ) -> case zipWithColumns - (\(c :: Bool) (l' :: a) -> if c then l' else r) - conditions - left of - Left e -> Left e - Right v -> Right $ Aggregated (TColumn v) - ( Right (UnAggregated conditions) - , Right (UnAggregated left@(BoxedColumn (left' :: V.Vector c))) - ) -> case testEquality (typeRep @(V.Vector a)) (typeRep @c) of - Just Refl -> case zipWithColumns - ( \(c :: VU.Vector Bool) (l' :: V.Vector a) -> - V.zipWith - (\c' l'' -> if c' then l'' else r) - (V.convert c) - l' - ) - conditions - left of - Left (TypeMismatchException context) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - Left e -> Left e - Right v -> Right $ UnAggregated v - Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @c) of - Nothing -> Left $ nestedTypeException @c @a (show expression) - Just Refl -> case sUnbox @a of - SFalse -> Left $ InternalException "Boxed type in unboxed column" - STrue -> case zipWithColumns - ( \(c :: VU.Vector Bool) (l' :: VU.Vector a) -> - VU.zipWith - (\c' l'' -> if c' then l'' else r) - c - l' - ) - conditions - left of - Left (TypeMismatchException context) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - Left e -> Left e - Right v -> Right $ UnAggregated v - (Right _, Right _) -> - Left $ - AggregatedAndNonAggregatedException (T.pack $ show l) (T.pack $ show r) - (Left (TypeMismatchException context), _) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show cond) - } - ) - (Left e, _) -> Left e - (_, Left (TypeMismatchException context)) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show r) - } - ) - (_, Left e) -> Left e -interpretAggregation gdf expression@(If cond l r) = - case ( interpretAggregation @Bool gdf cond - , interpretAggregation @a gdf l - , interpretAggregation @a gdf r - ) of - ( Right (Aggregated (TColumn conditions)) - , Right (Aggregated (TColumn left)) - , Right (Aggregated (TColumn right)) - ) -> case zipWithColumns - (\(c :: Bool) (l' :: a, r' :: a) -> if c then l' else r') - conditions - (zipColumns left right) of - Left e -> Left e - Right v -> Right $ Aggregated (TColumn v) - ( Right (UnAggregated conditions) - , Right (UnAggregated left@(BoxedColumn (left' :: V.Vector b))) - , Right (UnAggregated right@(BoxedColumn (right' :: V.Vector c))) - ) -> case testEquality (typeRep @b) (typeRep @c) of - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @c) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - Just Refl -> case testEquality (typeRep @(V.Vector a)) (typeRep @b) of - Just Refl -> case zipWithColumns - ( \(c :: VU.Vector Bool) (l' :: V.Vector a, r' :: V.Vector a) -> - V.zipWith - (\c' (l'', r'') -> if c' then l'' else r'') - (V.convert c) - (V.zip l' r') - ) - conditions - (zipColumns left right) of - Left (TypeMismatchException context) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - Left e -> Left e - Right v -> Right $ UnAggregated v - Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @b) of - Nothing -> Left $ nestedTypeException @b @a (show expression) - Just Refl -> case sUnbox @a of - SFalse -> Left $ InternalException "Boxed type in unboxed column" - STrue -> case zipWithColumns - ( \(c :: VU.Vector Bool) (l' :: VU.Vector a, r' :: VU.Vector a) -> - VU.zipWith - (\c' (l'', r'') -> if c' then l'' else r'') - c - (VU.zip l' r') - ) - conditions - (zipColumns left right) of - Left (TypeMismatchException context) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - Left e -> Left e - Right v -> Right $ UnAggregated v - (Right _, Right _, Right _) -> - Left $ - AggregatedAndNonAggregatedException (T.pack $ show l) (T.pack $ show r) - (Left (TypeMismatchException context), _, _) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show cond) - } - ) - (Left e, _, _) -> Left e - (_, Left (TypeMismatchException context), _) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show l) - } - ) - (_, Left e, _) -> Left e - (_, _, Left (TypeMismatchException context)) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show r) - } - ) - (_, _, Left e) -> Left e -interpretAggregation gdf@(Grouped df names indices os) expression@(AggVector expr op (f :: v b -> c)) = - case interpretAggregation @b gdf expr of - Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(v b)) (typeRep @d) of - Nothing -> Left $ nestedTypeException @d @b (show expr) - Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of - Nothing -> Right $ Aggregated $ TColumn $ fromVector $ V.map (f . V.convert) col - Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map f col - Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column" - Right (Aggregated (TColumn (BoxedColumn (col :: V.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of - Just Refl -> interpretAggregation @c gdf (Lit (f col)) - Nothing -> interpretAggregation @c gdf (Lit ((f . V.convert) col)) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expr) - } - ) - Right (Aggregated (TColumn (UnboxedColumn (col :: VU.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> case testEquality (typeRep @v) (typeRep @VU.Vector) of - Just Refl -> interpretAggregation @c gdf (Lit (f col)) - Nothing -> interpretAggregation @c gdf (Lit ((f . VU.convert) col)) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expr) - } - ) - Right (Aggregated (TColumn (OptionalColumn (col :: V.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of - Just Refl -> interpretAggregation @c gdf (Lit (f col)) - Nothing -> interpretAggregation @c gdf (Lit ((f . V.convert) col)) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expr) - } - ) - (Left (TypeMismatchException context)) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - (Left e) -> Left e -interpretAggregation gdf@(Grouped df names indices os) expression@(AggNumericVector (Col name) op (f :: VU.Vector b -> c)) = - case getColumn name df of - -- TODO(mchavinda): Fix the compedium of type errors here - -- This is mostly done help with the benchmarking. - Nothing -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) - Just (BoxedColumn col) -> error "Type mismatch." - Just (OptionalColumn col) -> error "Type mismatch." - Just (UnboxedColumn (col :: VU.Vector d)) -> case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> case testEquality (typeRep @c) (typeRep @a) of - Just Refl -> - Right $ - Aggregated $ - TColumn $ - fromUnboxedVector $ - mkAggregatedColumnUnboxed col os indices f - Nothing -> error "Type mismatch" - Nothing -> error "Type mismatch" -interpretAggregation gdf@(Grouped df names indices os) expression@(AggNumericVector expr op (f :: VU.Vector b -> c)) = - case interpretAggregation @b gdf expr of - (Left (TypeMismatchException context)) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - (Left e) -> Left e - Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(VU.Vector b)) (typeRep @d) of - Nothing -> case testEquality (typeRep @(VU.Vector Int)) (typeRep @d) of - Nothing -> case testEquality (typeRep @(V.Vector Integer)) (typeRep @d) of - Nothing -> Left $ nestedTypeException @d @b (show expr) - Just Refl -> - Right $ - Aggregated $ - TColumn $ - fromVector $ - V.map (f . VU.convert . V.map fromIntegral) col - Just Refl -> - Right $ - Aggregated $ - TColumn $ - fromVector $ - V.map f (VG.map (VG.map fromIntegral) col) - Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map f col - Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column" - Right (Aggregated (TColumn (BoxedColumn (col :: V.Vector d)))) -> case testEquality (typeRep @Integer) (typeRep @d) of - Just Refl -> interpretAggregation @c gdf (Lit ((f . V.convert . V.map fromIntegral) col)) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expr) - } - ) - Right (Aggregated (TColumn (UnboxedColumn (col :: VU.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> interpretAggregation @c gdf (Lit (f col)) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expr) - } - ) - Right (Aggregated (TColumn (OptionalColumn (col :: V.Vector (Maybe d))))) -> case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> - interpretAggregation @c - gdf - (Lit ((f . V.convert . V.map (fromMaybe 0) . V.filter isJust) col)) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expr) - } - ) -interpretAggregation gdf@(Grouped df names indices os) expression@(AggReduce (Col name) op (f :: a -> a -> a)) = - case getColumn name df of - Nothing -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) - Just (BoxedColumn (col :: V.Vector d)) -> case testEquality (typeRep @a) (typeRep @d) of - Nothing -> error "Type mismatch" - Just Refl -> - Right $ - Aggregated $ - TColumn $ - fromVector $ - mkReducedColumnBoxed col os indices f - Just (OptionalColumn (col :: V.Vector d)) -> case testEquality (typeRep @a) (typeRep @d) of - Nothing -> error "Type mismatch" - Just Refl -> - Right $ - Aggregated $ - TColumn $ - fromVector $ - mkReducedColumnBoxed col os indices f - Just (UnboxedColumn (col :: VU.Vector d)) -> case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> - Right $ - Aggregated $ - TColumn $ - fromUnboxedVector $ - mkReducedColumnUnboxed col os indices f - Nothing -> error "Type mismatch" -interpretAggregation gdf@(Grouped df names indices os) expression@(AggReduce expr op (f :: a -> a -> a)) = - case interpretAggregation @a gdf expr of - (Left (TypeMismatchException context)) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - (Left e) -> Left e - Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(V.Vector a)) (typeRep @d) of - Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @d) of - Nothing -> Left $ nestedTypeException @d @a (show expr) - Just Refl -> case sUnbox @a of - STrue -> - Right $ - Aggregated $ - TColumn $ - fromVector $ - V.map (VU.foldl1' f) col - SFalse -> Left $ InternalException "Boxed type inside an unboxed column" - Just Refl -> - Right $ - Aggregated $ - TColumn $ - fromVector $ - V.map (VG.foldl1' f) col - Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column" - Right (Aggregated (TColumn column)) -> case foldl1Column f column of - Left e -> Left e - Right value -> interpretAggregation @a gdf (Lit value) -interpretAggregation gdf@(Grouped df names indices os) expression@(AggFold expr op s (f :: (a -> b -> a))) = - case interpretAggregation @b gdf expr of - (Left (TypeMismatchException context)) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - (Left e) -> Left e - Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(V.Vector b)) (typeRep @d) of - Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map (V.foldl' f s) col - Nothing -> case testEquality (typeRep @(VU.Vector b)) (typeRep @d) of - Just Refl -> case sUnbox @b of - STrue -> - Right $ Aggregated $ TColumn $ fromVector $ V.map (VU.foldl' f s) col - SFalse -> Left $ InternalException "Boxed type inside an unboxed column" - Nothing -> Left $ nestedTypeException @d @b (show expr) - Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column" - Right (Aggregated (TColumn column)) -> case foldlColumn f s column of - Left e -> Left e - Right value -> interpretAggregation @a gdf (Lit value) +type NamedExpr = (T.Text, UExpr) instance (Num a, Columnable a) => Num (Expr a) where (+) :: Expr a -> Expr a -> Expr a @@ -1075,158 +300,3 @@ eSize (AggVector expr op _) = eSize expr + 1 eSize (AggReduce expr op _) = eSize expr + 1 eSize (AggFold expr op _ _) = eSize expr + 1 - --- Helpers -mkTypeMismatchException :: - (Typeable a, Typeable b) => - Maybe String -> Maybe String -> TypeErrorContext a b -> DataFrameException -mkTypeMismatchException callPoint errorLocation context = - TypeMismatchException - ( context - { callingFunctionName = callPoint - , errorColumnName = errorLocation - } - ) - -handleInterpretException :: String -> DataFrameException -> DataFrameException -handleInterpretException errorLocation (TypeMismatchException context) = mkTypeMismatchException (Just "interpret") (Just errorLocation) context -handleInterpretException _ e = e - -numRows :: DataFrame -> Int -numRows df = fst (dataframeDimensions df) - -mkUnaggregatedColumnBoxed :: - forall a. - (Columnable a) => - V.Vector a -> VU.Vector Int -> VU.Vector Int -> V.Vector (V.Vector a) -mkUnaggregatedColumnBoxed col os indices = - let - sorted = V.unsafeBackpermute col (V.convert indices) - n i = os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i) - start i = os `VG.unsafeIndex` i - in - V.generate - (VU.length os - 1) - ( \i -> - V.unsafeSlice (start i) (n i) sorted - ) - -mkUnaggregatedColumnUnboxed :: - forall a. - (Columnable a, VU.Unbox a) => - VU.Vector a -> VU.Vector Int -> VU.Vector Int -> V.Vector (VU.Vector a) -mkUnaggregatedColumnUnboxed col os indices = - let - sorted = VU.unsafeBackpermute col indices - n i = os `VU.unsafeIndex` (i + 1) - (os `VU.unsafeIndex` i) - start i = os `VG.unsafeIndex` i - in - V.generate - (VU.length os - 1) - ( \i -> - VU.unsafeSlice (start i) (n i) sorted - ) - -mkAggregatedColumnUnboxed :: - forall a b. - (Columnable a, VU.Unbox a, Columnable b, VU.Unbox b) => - VU.Vector a -> - VU.Vector Int -> - VU.Vector Int -> - (VU.Vector a -> b) -> - VU.Vector b -mkAggregatedColumnUnboxed col os indices f = - let - sorted = VU.unsafeBackpermute col indices - n i = os `VU.unsafeIndex` (i + 1) - (os `VU.unsafeIndex` i) - start i = os `VG.unsafeIndex` i - in - VU.generate - (VU.length os - 1) - ( \i -> - f (VU.unsafeSlice (start i) (n i) sorted) - ) - -mkReducedColumnUnboxed :: - forall a. - (VU.Unbox a) => - VU.Vector a -> - VU.Vector Int -> - VU.Vector Int -> - (a -> a -> a) -> - VU.Vector a -mkReducedColumnUnboxed col os indices f = runST $ do - let len = VU.length os - 1 - mvec <- VUM.unsafeNew len - - let loopOut i - | i == len = return () - | otherwise = do - let start = os `VU.unsafeIndex` i - let end = os `VU.unsafeIndex` (i + 1) - let initVal = col `VU.unsafeIndex` (indices `VU.unsafeIndex` start) - - let loopIn !acc idx - | idx == end = acc - | otherwise = - let val = col `VU.unsafeIndex` (indices `VU.unsafeIndex` idx) - in loopIn (f acc val) (idx + 1) - let !finalVal = loopIn initVal (start + 1) - VUM.unsafeWrite mvec i finalVal - loopOut (i + 1) - - loopOut 0 - VU.unsafeFreeze mvec -{-# INLINE mkReducedColumnUnboxed #-} - -mkReducedColumnBoxed :: - V.Vector a -> - VU.Vector Int -> - VU.Vector Int -> - (a -> a -> a) -> - V.Vector a -mkReducedColumnBoxed col os indices f = runST $ do - let len = VU.length os - 1 - mvec <- VM.unsafeNew len - - let loopOut i - | i == len = return () - | otherwise = do - let start = os `VU.unsafeIndex` i - let end = os `VU.unsafeIndex` (i + 1) - let initVal = col `V.unsafeIndex` (indices `VU.unsafeIndex` start) - - let loopIn !acc idx - | idx == end = acc - | otherwise = - let val = col `V.unsafeIndex` (indices `VU.unsafeIndex` idx) - in loopIn (f acc val) (idx + 1) - let !finalVal = loopIn initVal (start + 1) - VM.unsafeWrite mvec i finalVal - loopOut (i + 1) - - loopOut 0 - V.unsafeFreeze mvec -{-# INLINE mkReducedColumnBoxed #-} - -nestedTypeException :: - forall a b. (Typeable a, Typeable b) => String -> DataFrameException -nestedTypeException expression = case typeRep @a of - App t1 t2 -> - TypeMismatchException - ( MkTypeErrorContext - { userType = Left (show (typeRep @b)) :: Either String (TypeRep ()) - , expectedType = Left (show (typeRep @a)) :: Either String (TypeRep ()) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just expression - } - ) - t -> - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @(VU.Vector b)) - , expectedType = Right (typeRep @b) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just expression - } - ) diff --git a/src/DataFrame/Internal/Interpreter.hs b/src/DataFrame/Internal/Interpreter.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Internal/Interpreter.hs @@ -0,0 +1,953 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UndecidableInstances #-} + +module DataFrame.Internal.Interpreter where + +import Control.Monad.ST (runST) +import Data.Bifunctor +import qualified Data.Map as M +import Data.Maybe (fromMaybe, isJust) +import qualified Data.Text as T +import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as VM +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import DataFrame.Errors +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame +import DataFrame.Internal.Expression +import DataFrame.Internal.Types +import Type.Reflection (TypeRep, Typeable, typeOf, typeRep, pattern App) + +interpret :: + forall a. + (Columnable a) => + DataFrame -> Expr a -> Either DataFrameException (TypedColumn a) +interpret df (Lit value) = case sUnbox @a of + -- Specialize the creation of unboxed columns to avoid an extra allocation. + STrue -> pure $ TColumn $ fromUnboxedVector $ VU.replicate (numRows df) value + SFalse -> pure $ TColumn $ fromVector $ V.replicate (numRows df) value +interpret df (Col name) = maybe columnNotFound (pure . TColumn) (getColumn name df) + where + columnNotFound = Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) +-- Unary operations. +interpret df expr@(UnaryOp _ (f :: c -> d) value) = first (handleInterpretException (show expr)) $ do + (TColumn value') <- interpret @c df value + fmap TColumn (mapColumn f value') +-- Variations of binary operations. +interpret df expr@(BinaryOp _ (f :: c -> d -> e) left right) = first (handleInterpretException (show expr)) $ case (left, right) of + (Lit left, Lit right) -> interpret df (Lit (f left right)) + (Lit left, right) -> do + -- If we have a literal then we don't have to materialise + -- the column. + (TColumn value') <- interpret @d df right + fmap TColumn (mapColumn (f left) value') + (left, Lit right) -> do + -- Same as the above except the right side is the + -- literl. + (TColumn value') <- interpret @c df left + fmap TColumn (mapColumn (`f` right) value') + (_, _) -> do + -- In the general case we interpret and zip. + (TColumn left') <- interpret @c df left + (TColumn right') <- interpret @d df right + fmap TColumn (zipWithColumns f left' right') +-- Conditionals +interpret df expr@(If cond l r) = first (handleInterpretException (show expr)) $ do + (TColumn conditions) <- interpret @Bool df cond + (TColumn left) <- interpret @a df l + (TColumn right) <- interpret @a df r + let branch (c :: Bool) (l' :: a, r' :: a) = if c then l' else r' + fmap TColumn (zipWithColumns branch conditions (zipColumns left right)) +interpret df expression@(AggVector expr op (f :: v b -> c)) = do + (TColumn column) <- interpret @b df expr + -- Helper for errors. Should probably find a way of throwing this + -- without leaking the fact that we use `Vector` to users. + let aggTypeError expected = + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @(v b)) + , expectedType = Left expected :: Either String (TypeRep ()) + , callingFunctionName = Just "interpret" + , errorColumnName = Nothing + } + ) + let processColumn :: + (Columnable d) => d -> Either DataFrameException (TypedColumn a) + processColumn col = case testEquality (typeRep @(v b)) (typeOf col) of + Just Refl -> interpret @c df (Lit (f col)) + Nothing -> Left $ aggTypeError (show (typeOf col)) + case column of + (BoxedColumn col) -> processColumn col + (OptionalColumn col) -> processColumn col + (UnboxedColumn col) -> processColumn col +interpret df expression@(AggReduce expr op (f :: a -> a -> a)) = first (handleInterpretException (show expr)) $ do + (TColumn column) <- interpret @a df expr + value <- foldl1Column f column + pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value +interpret df expression@(AggNumericVector expr op (f :: VU.Vector b -> c)) = first (handleInterpretException (show expression)) $ do + (TColumn column) <- interpret @b df expr + case column of + (UnboxedColumn (v :: VU.Vector d)) -> case testEquality (typeRep @d) (typeRep @b) of + Just Refl -> + pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) (f v) + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + (Right (typeRep @b)) + (Right (typeRep @d)) + (Just "interpret") + (Just (show expression)) + ) + _ -> error "Trying to apply numeric computation to non-numeric column" +interpret df expression@(AggFold expr op start (f :: (a -> b -> a))) = first (handleInterpretException (show expression)) $ do + (TColumn column) <- interpret @b df expr + value <- foldlColumn f start column + pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value + +data AggregationResult a + = UnAggregated Column + | Aggregated (TypedColumn a) + +interpretAggregation :: + forall a. + (Columnable a) => + GroupedDataFrame -> Expr a -> Either DataFrameException (AggregationResult a) +interpretAggregation gdf (Lit value) = + Right $ + Aggregated $ + TColumn $ + fromVector $ + V.replicate (VU.length (offsets gdf) - 1) value +interpretAggregation gdf@(Grouped df names indices os) (Col name) = case getColumn name df of + Nothing -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) + Just (BoxedColumn col) -> Right $ UnAggregated $ fromVector $ mkUnaggregatedColumnBoxed col os indices + Just (OptionalColumn col) -> Right $ UnAggregated $ fromVector $ mkUnaggregatedColumnBoxed col os indices + Just (UnboxedColumn col) -> + Right $ UnAggregated $ fromVector $ mkUnaggregatedColumnUnboxed col os indices +interpretAggregation gdf expression@(UnaryOp _ (f :: c -> d) expr) = + case interpretAggregation @c gdf expr of + Left (TypeMismatchException context) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expr) + } + ) + Left e -> Left e + Right (UnAggregated unaggregated) -> case unaggregated of + BoxedColumn (col :: V.Vector b) -> case testEquality (typeRep @b) (typeRep @(V.Vector c)) of + Just Refl -> case sUnbox @d of + SFalse -> Right $ UnAggregated $ fromVector $ V.map (V.map f) col + STrue -> + Right $ + UnAggregated $ + fromVector $ + V.map (V.convert @V.Vector @d @VU.Vector . V.map f) col + Nothing -> case testEquality (typeRep @b) (typeRep @(VU.Vector c)) of + Nothing -> Left $ nestedTypeException @b @c (show expression) + Just Refl -> case (sUnbox @c, sUnbox @a) of + (SFalse, _) -> Left $ InternalException "Boxed type inside an unboxed column" + (STrue, STrue) -> Right $ UnAggregated $ fromVector $ V.map (VU.map f) col + (STrue, _) -> Right $ UnAggregated $ fromVector $ V.map (V.map f . VU.convert) col + _ -> Left $ InternalException "Aggregated into a non-boxed column" + Right (Aggregated (TColumn aggregated)) -> case mapColumn f aggregated of + Left e -> Left e + Right col -> Right $ Aggregated $ TColumn col +interpretAggregation gdf expression@(BinaryOp name (f :: c -> d -> e) left (Lit (right :: g))) = case testEquality (typeRep @g) (typeRep @d) of + Just Refl -> interpretAggregation gdf (UnaryOp name (`f` right) left) + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @g) + , expectedType = Right (typeRep @d) + , callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) +interpretAggregation gdf expression@(BinaryOp name (f :: c -> d -> e) (Lit (left :: g)) right) = case testEquality (typeRep @g) (typeRep @c) of + Just Refl -> interpretAggregation gdf (UnaryOp name (f left) right) + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @g) + , expectedType = Right (typeRep @c) + , callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) +interpretAggregation gdf expression@(BinaryOp _ (f :: c -> d -> e) left right) = + case (interpretAggregation @c gdf left, interpretAggregation @d gdf right) of + (Right (Aggregated (TColumn left')), Right (Aggregated (TColumn right'))) -> case zipWithColumns f left' right' of + Left e -> Left e + Right col -> Right $ Aggregated $ TColumn col + (Right (UnAggregated left'), Right (UnAggregated right')) -> case (left', right') of + (BoxedColumn (l :: V.Vector m), BoxedColumn (r :: V.Vector n)) -> case testEquality (typeRep @m) (typeRep @(VU.Vector c)) of + Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector d)) of + Just Refl -> case (sUnbox @c, sUnbox @d, sUnbox @e) of + (STrue, STrue, STrue) -> + Right $ UnAggregated $ fromVector $ V.zipWith (VU.zipWith f) l r + (STrue, STrue, SFalse) -> + Right $ + UnAggregated $ + fromVector $ + V.zipWith (\l' r' -> V.zipWith f (V.convert l') (V.convert r')) l r + (_, _, _) -> Left $ InternalException "Boxed vectors contain unboxed types" + Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector d)) of + Just Refl -> case sUnbox @c of + STrue -> + Right $ + UnAggregated $ + fromVector $ + V.zipWith (V.zipWith f . V.convert) l r + SFalse -> Left $ InternalException "Unboxed vectors contain boxed types" + Nothing -> Left $ nestedTypeException @n @d (show right) + Nothing -> case testEquality (typeRep @m) (typeRep @(V.Vector c)) of + Nothing -> Left $ nestedTypeException @m @c (show left) + Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector d)) of + Just Refl -> case (sUnbox @d, sUnbox @e) of + (STrue, STrue) -> + Right $ + UnAggregated $ + fromVector $ + V.zipWith + (\l' r' -> V.convert @V.Vector @e @VU.Vector $ V.zipWith f l' (V.convert r')) + l + r + (STrue, SFalse) -> + Right $ + UnAggregated $ + fromVector $ + V.zipWith (\l' r' -> V.zipWith f l' (V.convert r')) l r + (_, _) -> Left $ InternalException "Unboxed vectors contain boxed types" + Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector d)) of + Just Refl -> case sUnbox @e of + SFalse -> + Right $ + UnAggregated $ + fromVector $ + V.zipWith (V.zipWith f . V.convert) l r + STrue -> + Right $ + UnAggregated $ + fromVector $ + V.zipWith (\l' r' -> V.convert @V.Vector @e @VU.Vector $ V.zipWith f l' r') l r + Nothing -> Left $ nestedTypeException @n @d (show right) + _ -> Left $ InternalException "Aggregated into a non-boxed column" + (Right _, Right _) -> + Left $ + AggregatedAndNonAggregatedException (T.pack $ show left) (T.pack $ show right) + (Left (TypeMismatchException context), _) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show left) + } + ) + (Left e, _) -> Left e + (_, Left (TypeMismatchException context)) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show right) + } + ) + (_, Left e) -> Left e +interpretAggregation gdf expression@(If cond (Lit l) (Lit r)) = + case interpretAggregation @Bool gdf cond of + Right (Aggregated (TColumn conditions)) -> case mapColumn + (\(c :: Bool) -> if c then l else r) + conditions of + Left e -> Left e + Right v -> Right $ Aggregated (TColumn v) + Right (UnAggregated conditions) -> case sUnbox @a of + STrue -> case mapColumn + (\(c :: VU.Vector Bool) -> VU.map (\c' -> if c' then l else r) c) + conditions of + Left (TypeMismatchException context) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + Left e -> Left e + Right v -> Right $ UnAggregated v + SFalse -> case mapColumn + (\(c :: VU.Vector Bool) -> V.map (\c' -> if c' then l else r) (VU.convert c)) + conditions of + Left (TypeMismatchException context) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + Left e -> Left e + Right v -> Right $ UnAggregated v + Left (TypeMismatchException context) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show cond) + } + ) + Left e -> Left e +interpretAggregation gdf expression@(If cond (Lit l) r) = + case ( interpretAggregation @Bool gdf cond + , interpretAggregation @a gdf r + ) of + ( Right (Aggregated (TColumn conditions)) + , Right (Aggregated (TColumn right)) + ) -> case zipWithColumns + (\(c :: Bool) (r' :: a) -> if c then l else r') + conditions + right of + Left e -> Left e + Right v -> Right $ Aggregated (TColumn v) + ( Right (UnAggregated conditions) + , Right (UnAggregated right@(BoxedColumn (right' :: V.Vector c))) + ) -> case testEquality (typeRep @(V.Vector a)) (typeRep @c) of + Just Refl -> case zipWithColumns + ( \(c :: VU.Vector Bool) (r' :: V.Vector a) -> + V.zipWith + (\c' r'' -> if c' then l else r'') + (V.convert c) + r' + ) + conditions + right of + Left (TypeMismatchException context) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + Left e -> Left e + Right v -> Right $ UnAggregated v + Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @c) of + Nothing -> Left $ nestedTypeException @c @a (show expression) + Just Refl -> case sUnbox @a of + SFalse -> Left $ InternalException "Boxed type in unboxed column" + STrue -> case zipWithColumns + ( \(c :: VU.Vector Bool) (r' :: VU.Vector a) -> + VU.zipWith + (\c' r'' -> if c' then l else r'') + c + r' + ) + conditions + right of + Left (TypeMismatchException context) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + Left e -> Left e + Right v -> Right $ UnAggregated v + (Right _, Right _) -> + Left $ + AggregatedAndNonAggregatedException (T.pack $ show l) (T.pack $ show r) + (Left (TypeMismatchException context), _) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show cond) + } + ) + (Left e, _) -> Left e + (_, Left (TypeMismatchException context)) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show r) + } + ) + (_, Left e) -> Left e +interpretAggregation gdf expression@(If cond l (Lit r)) = + case ( interpretAggregation @Bool gdf cond + , interpretAggregation @a gdf l + ) of + ( Right (Aggregated (TColumn conditions)) + , Right (Aggregated (TColumn left)) + ) -> case zipWithColumns + (\(c :: Bool) (l' :: a) -> if c then l' else r) + conditions + left of + Left e -> Left e + Right v -> Right $ Aggregated (TColumn v) + ( Right (UnAggregated conditions) + , Right (UnAggregated left@(BoxedColumn (left' :: V.Vector c))) + ) -> case testEquality (typeRep @(V.Vector a)) (typeRep @c) of + Just Refl -> case zipWithColumns + ( \(c :: VU.Vector Bool) (l' :: V.Vector a) -> + V.zipWith + (\c' l'' -> if c' then l'' else r) + (V.convert c) + l' + ) + conditions + left of + Left (TypeMismatchException context) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + Left e -> Left e + Right v -> Right $ UnAggregated v + Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @c) of + Nothing -> Left $ nestedTypeException @c @a (show expression) + Just Refl -> case sUnbox @a of + SFalse -> Left $ InternalException "Boxed type in unboxed column" + STrue -> case zipWithColumns + ( \(c :: VU.Vector Bool) (l' :: VU.Vector a) -> + VU.zipWith + (\c' l'' -> if c' then l'' else r) + c + l' + ) + conditions + left of + Left (TypeMismatchException context) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + Left e -> Left e + Right v -> Right $ UnAggregated v + (Right _, Right _) -> + Left $ + AggregatedAndNonAggregatedException (T.pack $ show l) (T.pack $ show r) + (Left (TypeMismatchException context), _) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show cond) + } + ) + (Left e, _) -> Left e + (_, Left (TypeMismatchException context)) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show r) + } + ) + (_, Left e) -> Left e +interpretAggregation gdf expression@(If cond l r) = + case ( interpretAggregation @Bool gdf cond + , interpretAggregation @a gdf l + , interpretAggregation @a gdf r + ) of + ( Right (Aggregated (TColumn conditions)) + , Right (Aggregated (TColumn left)) + , Right (Aggregated (TColumn right)) + ) -> case zipWithColumns + (\(c :: Bool) (l' :: a, r' :: a) -> if c then l' else r') + conditions + (zipColumns left right) of + Left e -> Left e + Right v -> Right $ Aggregated (TColumn v) + ( Right (UnAggregated conditions) + , Right (UnAggregated left@(BoxedColumn (left' :: V.Vector b))) + , Right (UnAggregated right@(BoxedColumn (right' :: V.Vector c))) + ) -> case testEquality (typeRep @b) (typeRep @c) of + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @b) + , expectedType = Right (typeRep @c) + , callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + Just Refl -> case testEquality (typeRep @(V.Vector a)) (typeRep @b) of + Just Refl -> case zipWithColumns + ( \(c :: VU.Vector Bool) (l' :: V.Vector a, r' :: V.Vector a) -> + V.zipWith + (\c' (l'', r'') -> if c' then l'' else r'') + (V.convert c) + (V.zip l' r') + ) + conditions + (zipColumns left right) of + Left (TypeMismatchException context) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + Left e -> Left e + Right v -> Right $ UnAggregated v + Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @b) of + Nothing -> Left $ nestedTypeException @b @a (show expression) + Just Refl -> case sUnbox @a of + SFalse -> Left $ InternalException "Boxed type in unboxed column" + STrue -> case zipWithColumns + ( \(c :: VU.Vector Bool) (l' :: VU.Vector a, r' :: VU.Vector a) -> + VU.zipWith + (\c' (l'', r'') -> if c' then l'' else r'') + c + (VU.zip l' r') + ) + conditions + (zipColumns left right) of + Left (TypeMismatchException context) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + Left e -> Left e + Right v -> Right $ UnAggregated v + (Right _, Right _, Right _) -> + Left $ + AggregatedAndNonAggregatedException (T.pack $ show l) (T.pack $ show r) + (Left (TypeMismatchException context), _, _) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show cond) + } + ) + (Left e, _, _) -> Left e + (_, Left (TypeMismatchException context), _) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show l) + } + ) + (_, Left e, _) -> Left e + (_, _, Left (TypeMismatchException context)) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show r) + } + ) + (_, _, Left e) -> Left e +interpretAggregation gdf@(Grouped df names indices os) expression@(AggVector expr op (f :: v b -> c)) = + case interpretAggregation @b gdf expr of + Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(v b)) (typeRep @d) of + Nothing -> Left $ nestedTypeException @d @b (show expr) + Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of + Nothing -> Right $ Aggregated $ TColumn $ fromVector $ V.map (f . V.convert) col + Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map f col + Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column" + Right (Aggregated (TColumn (BoxedColumn (col :: V.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of + Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of + Just Refl -> interpretAggregation @c gdf (Lit (f col)) + Nothing -> interpretAggregation @c gdf (Lit ((f . V.convert) col)) + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @b) + , expectedType = Right (typeRep @d) + , callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expr) + } + ) + Right (Aggregated (TColumn (UnboxedColumn (col :: VU.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of + Just Refl -> case testEquality (typeRep @v) (typeRep @VU.Vector) of + Just Refl -> interpretAggregation @c gdf (Lit (f col)) + Nothing -> interpretAggregation @c gdf (Lit ((f . VU.convert) col)) + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @b) + , expectedType = Right (typeRep @d) + , callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expr) + } + ) + Right (Aggregated (TColumn (OptionalColumn (col :: V.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of + Just Refl -> case testEquality (typeRep @v) (typeRep @V.Vector) of + Just Refl -> interpretAggregation @c gdf (Lit (f col)) + Nothing -> interpretAggregation @c gdf (Lit ((f . V.convert) col)) + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @b) + , expectedType = Right (typeRep @d) + , callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expr) + } + ) + (Left (TypeMismatchException context)) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + (Left e) -> Left e +interpretAggregation gdf@(Grouped df names indices os) expression@(AggNumericVector (Col name) op (f :: VU.Vector b -> c)) = + case getColumn name df of + -- TODO(mchavinda): Fix the compedium of type errors here + -- This is mostly done help with the benchmarking. + Nothing -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) + Just (BoxedColumn col) -> error "Type mismatch." + Just (OptionalColumn col) -> error "Type mismatch." + Just (UnboxedColumn (col :: VU.Vector d)) -> case testEquality (typeRep @b) (typeRep @d) of + Just Refl -> case testEquality (typeRep @c) (typeRep @a) of + Just Refl -> + Right $ + Aggregated $ + TColumn $ + fromUnboxedVector $ + mkAggregatedColumnUnboxed col os indices f + Nothing -> error "Type mismatch" + Nothing -> error "Type mismatch" +interpretAggregation gdf@(Grouped df names indices os) expression@(AggNumericVector expr op (f :: VU.Vector b -> c)) = + case interpretAggregation @b gdf expr of + (Left (TypeMismatchException context)) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + (Left e) -> Left e + Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(VU.Vector b)) (typeRep @d) of + Nothing -> case testEquality (typeRep @(VU.Vector Int)) (typeRep @d) of + Nothing -> case testEquality (typeRep @(V.Vector Integer)) (typeRep @d) of + Nothing -> Left $ nestedTypeException @d @b (show expr) + Just Refl -> + Right $ + Aggregated $ + TColumn $ + fromVector $ + V.map (f . VU.convert . V.map fromIntegral) col + Just Refl -> + Right $ + Aggregated $ + TColumn $ + fromVector $ + V.map f (V.map (VU.map fromIntegral) col) + Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map f col + Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column" + Right (Aggregated (TColumn (BoxedColumn (col :: V.Vector d)))) -> case testEquality (typeRep @Integer) (typeRep @d) of + Just Refl -> interpretAggregation @c gdf (Lit ((f . V.convert . V.map fromIntegral) col)) + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @b) + , expectedType = Right (typeRep @d) + , callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expr) + } + ) + Right (Aggregated (TColumn (UnboxedColumn (col :: VU.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of + Just Refl -> interpretAggregation @c gdf (Lit (f col)) + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @b) + , expectedType = Right (typeRep @d) + , callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expr) + } + ) + Right (Aggregated (TColumn (OptionalColumn (col :: V.Vector (Maybe d))))) -> case testEquality (typeRep @b) (typeRep @d) of + Just Refl -> + interpretAggregation @c + gdf + (Lit ((f . V.convert . V.map (fromMaybe 0) . V.filter isJust) col)) + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @b) + , expectedType = Right (typeRep @d) + , callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expr) + } + ) +interpretAggregation gdf@(Grouped df names indices os) expression@(AggReduce (Col name) op (f :: a -> a -> a)) = + case getColumn name df of + Nothing -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) + Just (BoxedColumn (col :: V.Vector d)) -> case testEquality (typeRep @a) (typeRep @d) of + Nothing -> error "Type mismatch" + Just Refl -> + Right $ + Aggregated $ + TColumn $ + fromVector $ + mkReducedColumnBoxed col os indices f + Just (OptionalColumn (col :: V.Vector d)) -> case testEquality (typeRep @a) (typeRep @d) of + Nothing -> error "Type mismatch" + Just Refl -> + Right $ + Aggregated $ + TColumn $ + fromVector $ + mkReducedColumnBoxed col os indices f + Just (UnboxedColumn (col :: VU.Vector d)) -> case testEquality (typeRep @a) (typeRep @d) of + Just Refl -> + Right $ + Aggregated $ + TColumn $ + fromUnboxedVector $ + mkReducedColumnUnboxed col os indices f + Nothing -> error "Type mismatch" +interpretAggregation gdf@(Grouped df names indices os) expression@(AggReduce expr op (f :: a -> a -> a)) = + case interpretAggregation @a gdf expr of + (Left (TypeMismatchException context)) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + (Left e) -> Left e + Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(V.Vector a)) (typeRep @d) of + Nothing -> case testEquality (typeRep @(VU.Vector a)) (typeRep @d) of + Nothing -> Left $ nestedTypeException @d @a (show expr) + Just Refl -> case sUnbox @a of + STrue -> + Right $ + Aggregated $ + TColumn $ + fromVector $ + V.map (VU.foldl1' f) col + SFalse -> Left $ InternalException "Boxed type inside an unboxed column" + Just Refl -> + Right $ + Aggregated $ + TColumn $ + fromVector $ + V.map (V.foldl1' f) col + Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column" + Right (Aggregated (TColumn column)) -> case foldl1Column f column of + Left e -> Left e + Right value -> interpretAggregation @a gdf (Lit value) +interpretAggregation gdf@(Grouped df names indices os) expression@(AggFold expr op s (f :: (a -> b -> a))) = + case interpretAggregation @b gdf expr of + (Left (TypeMismatchException context)) -> + Left $ + TypeMismatchException + ( context + { callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just (show expression) + } + ) + (Left e) -> Left e + Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(V.Vector b)) (typeRep @d) of + Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map (V.foldl' f s) col + Nothing -> case testEquality (typeRep @(VU.Vector b)) (typeRep @d) of + Just Refl -> case sUnbox @b of + STrue -> + Right $ Aggregated $ TColumn $ fromVector $ V.map (VU.foldl' f s) col + SFalse -> Left $ InternalException "Boxed type inside an unboxed column" + Nothing -> Left $ nestedTypeException @d @b (show expr) + Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column" + Right (Aggregated (TColumn column)) -> case foldlColumn f s column of + Left e -> Left e + Right value -> interpretAggregation @a gdf (Lit value) + +handleInterpretException :: String -> DataFrameException -> DataFrameException +handleInterpretException errorLocation (TypeMismatchException context) = mkTypeMismatchException (Just "interpret") (Just errorLocation) context +handleInterpretException _ e = e + +numRows :: DataFrame -> Int +numRows df = fst (dataframeDimensions df) + +mkUnaggregatedColumnBoxed :: + forall a. + (Columnable a) => + V.Vector a -> VU.Vector Int -> VU.Vector Int -> V.Vector (V.Vector a) +mkUnaggregatedColumnBoxed col os indices = + let + sorted = V.unsafeBackpermute col (V.convert indices) + n i = os `VU.unsafeIndex` (i + 1) - (os `VU.unsafeIndex` i) + start i = os `VU.unsafeIndex` i + in + V.generate + (VU.length os - 1) + ( \i -> + V.unsafeSlice (start i) (n i) sorted + ) + +mkUnaggregatedColumnUnboxed :: + forall a. + (Columnable a, VU.Unbox a) => + VU.Vector a -> VU.Vector Int -> VU.Vector Int -> V.Vector (VU.Vector a) +mkUnaggregatedColumnUnboxed col os indices = + let + sorted = VU.unsafeBackpermute col indices + n i = os `VU.unsafeIndex` (i + 1) - (os `VU.unsafeIndex` i) + start i = os `VU.unsafeIndex` i + in + V.generate + (VU.length os - 1) + ( \i -> + VU.unsafeSlice (start i) (n i) sorted + ) + +mkAggregatedColumnUnboxed :: + forall a b. + (Columnable a, VU.Unbox a, Columnable b, VU.Unbox b) => + VU.Vector a -> + VU.Vector Int -> + VU.Vector Int -> + (VU.Vector a -> b) -> + VU.Vector b +mkAggregatedColumnUnboxed col os indices f = + let + sorted = VU.unsafeBackpermute col indices + n i = os `VU.unsafeIndex` (i + 1) - (os `VU.unsafeIndex` i) + start i = os `VU.unsafeIndex` i + in + VU.generate + (VU.length os - 1) + ( \i -> + f (VU.unsafeSlice (start i) (n i) sorted) + ) + +mkReducedColumnUnboxed :: + forall a. + (VU.Unbox a) => + VU.Vector a -> + VU.Vector Int -> + VU.Vector Int -> + (a -> a -> a) -> + VU.Vector a +mkReducedColumnUnboxed col os indices f = runST $ do + let len = VU.length os - 1 + mvec <- VUM.unsafeNew len + + let loopOut i + | i == len = return () + | otherwise = do + let !start = os `VU.unsafeIndex` i + let !end = os `VU.unsafeIndex` (i + 1) + let !initVal = col `VU.unsafeIndex` (indices `VU.unsafeIndex` start) + + let loopIn !acc !idx + | idx == end = acc + | otherwise = + let val = col `VU.unsafeIndex` (indices `VU.unsafeIndex` idx) + in loopIn (f acc val) (idx + 1) + let !finalVal = loopIn initVal (start + 1) + VUM.unsafeWrite mvec i finalVal + loopOut (i + 1) + + loopOut 0 + VU.unsafeFreeze mvec +{-# INLINE mkReducedColumnUnboxed #-} + +mkReducedColumnBoxed :: + V.Vector a -> + VU.Vector Int -> + VU.Vector Int -> + (a -> a -> a) -> + V.Vector a +mkReducedColumnBoxed col os indices f = runST $ do + let len = VU.length os - 1 + mvec <- VM.unsafeNew len + + let loopOut i + | i == len = return () + | otherwise = do + let start = os `VU.unsafeIndex` i + let end = os `VU.unsafeIndex` (i + 1) + let initVal = col `V.unsafeIndex` (indices `VU.unsafeIndex` start) + + let loopIn !acc idx + | idx == end = acc + | otherwise = + let val = col `V.unsafeIndex` (indices `VU.unsafeIndex` idx) + in loopIn (f acc val) (idx + 1) + let !finalVal = loopIn initVal (start + 1) + VM.unsafeWrite mvec i finalVal + loopOut (i + 1) + + loopOut 0 + V.unsafeFreeze mvec +{-# INLINE mkReducedColumnBoxed #-} + +nestedTypeException :: + forall a b. (Typeable a, Typeable b) => String -> DataFrameException +nestedTypeException expression = case typeRep @a of + App t1 t2 -> + TypeMismatchException + ( MkTypeErrorContext + { userType = Left (show (typeRep @b)) :: Either String (TypeRep ()) + , expectedType = Left (show (typeRep @a)) :: Either String (TypeRep ()) + , callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just expression + } + ) + t -> + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @(VU.Vector b)) + , expectedType = Right (typeRep @b) + , callingFunctionName = Just "interpretAggregation" + , errorColumnName = Just expression + } + ) + +mkTypeMismatchException :: + (Typeable a, Typeable b) => + Maybe String -> Maybe String -> TypeErrorContext a b -> DataFrameException +mkTypeMismatchException callPoint errorLocation context = + TypeMismatchException + ( context + { callingFunctionName = callPoint + , errorColumnName = errorLocation + } + ) diff --git a/src/DataFrame/Lazy/IO/CSV.hs b/src/DataFrame/Lazy/IO/CSV.hs --- a/src/DataFrame/Lazy/IO/CSV.hs +++ b/src/DataFrame/Lazy/IO/CSV.hs @@ -136,6 +136,7 @@ { columns = cols , columnIndices = M.fromList (zip columnNames [0 ..]) , dataframeDimensions = (maybe 0 columnLength (cols V.!? 0), V.length cols) + , derivingExpressions = M.empty -- TODO create base expressions. } , (pos, unconsumed, r + 1) ) diff --git a/src/DataFrame/Operations/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs --- a/src/DataFrame/Operations/Aggregation.hs +++ b/src/DataFrame/Operations/Aggregation.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} @@ -5,6 +6,7 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE Strict #-} {-# LANGUAGE TypeApplications #-} module DataFrame.Operations.Aggregation where @@ -14,7 +16,6 @@ import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Algorithms.Merge as VA -import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM @@ -31,6 +32,7 @@ ) import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..)) import DataFrame.Internal.Expression +import DataFrame.Internal.Interpreter import DataFrame.Internal.Types import DataFrame.Operations.Core import DataFrame.Operations.Subset @@ -54,22 +56,103 @@ Grouped df names - (VG.map fst valueIndices) - (VU.fromList (reverse (changingPoints valueIndices))) + (VU.map fst valueIndices) + (changingPoints valueIndices) where indicesToGroup = M.elems $ M.filterWithKey (\k _ -> k `elem` names) (columnIndices df) - rowRepresentations = computeRowHashes indicesToGroup df - + doubleToInt :: Double -> Int + doubleToInt = floor . (* 1000) valueIndices = runST $ do - withIndexes <- VG.thaw $ VG.indexed rowRepresentations - VA.sortBy (\(a, b) (a', b') -> compare b' b) withIndexes - VG.unsafeFreeze withIndexes + let n = fst (dimensions df) + mv <- VUM.new n -changingPoints :: (Eq a, VU.Unbox a) => VU.Vector (Int, a) -> [Int] -changingPoints vs = VG.length vs : fst (VU.ifoldl findChangePoints initialState vs) + let selectedCols = map (columns df V.!) indicesToGroup + + forM_ selectedCols $ \case + UnboxedColumn (v :: VU.Vector a) -> + case testEquality (typeRep @a) (typeRep @Int) of + Just Refl -> + VU.imapM_ + ( \i x -> do + (_, !h) <- VUM.unsafeRead mv i + VUM.unsafeWrite mv i (i, hashWithSalt h x) + ) + v + Nothing -> + case testEquality (typeRep @a) (typeRep @Double) of + Just Refl -> + VU.imapM_ + ( \i d -> do + (_, !h) <- VUM.unsafeRead mv i + VUM.unsafeWrite mv i (i, hashWithSalt h (doubleToInt d)) + ) + v + Nothing -> + case sIntegral @a of + STrue -> + VU.imapM_ + ( \i d -> do + let x :: Int + x = fromIntegral @a @Int d + (_, !h) <- VUM.unsafeRead mv i + VUM.unsafeWrite mv i (i, hashWithSalt h x) + ) + v + SFalse -> + case sFloating @a of + STrue -> + VU.imapM_ + ( \i d -> do + let x :: Int + x = doubleToInt (realToFrac d :: Double) + (_, !h) <- VUM.unsafeRead mv i + VUM.unsafeWrite mv i (i, hashWithSalt h x) + ) + v + SFalse -> + VU.imapM_ + ( \i d -> do + let x = hash (show d) + (_, !h) <- VUM.unsafeRead mv i + VUM.unsafeWrite mv i (i, hashWithSalt h x) + ) + v + BoxedColumn (v :: V.Vector a) -> + case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> + V.imapM_ + ( \i t -> do + (_, !h) <- VUM.unsafeRead mv i + VUM.unsafeWrite mv i (i, hashWithSalt h t) + ) + v + Nothing -> + V.imapM_ + ( \i d -> do + let x = hash (show d) + (_, !h) <- VUM.unsafeRead mv i + VUM.unsafeWrite mv i (i, hashWithSalt h x) + ) + v + OptionalColumn v -> + V.imapM_ + ( \i d -> do + let x = hash (show d) + (_, !h) <- VUM.unsafeRead mv i + VUM.unsafeWrite mv i (i, hashWithSalt h x) + ) + v + + VA.sortBy (\(!a, !b) (!a', !b') -> compare b' b) mv + VU.unsafeFreeze mv + +changingPoints :: VU.Vector (Int, Int) -> VU.Vector Int +changingPoints vs = + VU.reverse + (VU.fromList (VU.length vs : fst (VU.ifoldl' findChangePoints initialState vs))) where - initialState = ([0], snd (VG.head vs)) - findChangePoints (offsets, currentVal) index (_, newVal) + initialState = ([0], snd (VU.head vs)) + findChangePoints (!offsets, !currentVal) index (_, !newVal) | currentVal == newVal = (offsets, currentVal) | otherwise = (index : offsets, newVal) @@ -168,7 +251,7 @@ let df' = selectIndices - (VG.map (valueIndices VG.!) (VG.init offsets)) + (VU.map (valueIndices VU.!) (VU.init offsets)) (select groupingColumns df) f (name, Wrap (expr :: Expr a)) d = @@ -185,12 +268,12 @@ selectIndices :: VU.Vector Int -> DataFrame -> DataFrame selectIndices xs df = df - { columns = VG.map (atIndicesStable xs) (columns df) - , dataframeDimensions = (VG.length xs, VG.length (columns df)) + { columns = V.map (atIndicesStable xs) (columns df) + , dataframeDimensions = (VU.length xs, V.length (columns df)) } -- | Filter out all non-unique values in a dataframe. distinct :: DataFrame -> DataFrame -distinct df = selectIndices (VG.map (indices VG.!) (VG.init os)) df +distinct df = selectIndices (VU.map (indices VU.!) (VU.init os)) df where (Grouped _ _ indices os) = groupBy (columnNames df) df diff --git a/src/DataFrame/Operations/Core.hs b/src/DataFrame/Operations/Core.hs --- a/src/DataFrame/Operations/Core.hs +++ b/src/DataFrame/Operations/Core.hs @@ -45,6 +45,7 @@ getColumn, ) import DataFrame.Internal.Expression +import DataFrame.Internal.Interpreter import DataFrame.Internal.Parsing (isNullish) import DataFrame.Internal.Row (Any, mkColumnFromRow) import Type.Reflection @@ -350,11 +351,13 @@ (V.map (expandColumn n) (columns d V.// [(i, column)])) (columnIndices d) (n, c) + M.empty Nothing -> DataFrame (V.map (expandColumn n) (columns d `V.snoc` column)) (M.insert name c (columnIndices d)) (n, c + 1) + M.empty {- | /O(n)/ Clones a column and places it under a new name in the dataframe. diff --git a/src/DataFrame/Operations/Statistics.hs b/src/DataFrame/Operations/Statistics.hs --- a/src/DataFrame/Operations/Statistics.hs +++ b/src/DataFrame/Operations/Statistics.hs @@ -29,6 +29,7 @@ getColumn, ) import DataFrame.Internal.Expression +import DataFrame.Internal.Interpreter import DataFrame.Internal.Row (showValue, toAny) import DataFrame.Internal.Statistics import DataFrame.Internal.Types diff --git a/src/DataFrame/Operations/Subset.hs b/src/DataFrame/Operations/Subset.hs --- a/src/DataFrame/Operations/Subset.hs +++ b/src/DataFrame/Operations/Subset.hs @@ -24,6 +24,7 @@ import DataFrame.Internal.Column import DataFrame.Internal.DataFrame (DataFrame (..), empty, getColumn) import DataFrame.Internal.Expression +import DataFrame.Internal.Interpreter import DataFrame.Operations.Core import DataFrame.Operations.Transformations (apply) import System.Random diff --git a/src/DataFrame/Operations/Transformations.hs b/src/DataFrame/Operations/Transformations.hs --- a/src/DataFrame/Operations/Transformations.hs +++ b/src/DataFrame/Operations/Transformations.hs @@ -26,6 +26,7 @@ ) import DataFrame.Internal.DataFrame (DataFrame (..), getColumn) import DataFrame.Internal.Expression +import DataFrame.Internal.Interpreter import DataFrame.Operations.Core -- | O(k) Apply a function to a given column in a dataframe. @@ -68,7 +69,10 @@ derive :: forall a. (Columnable a) => T.Text -> Expr a -> DataFrame -> DataFrame derive name expr df = case interpret @a df (normalize expr) of Left e -> throw e - Right (TColumn value) -> insertColumn name value df + Right (TColumn value) -> + (insertColumn name value df) + { derivingExpressions = M.insert name (Wrap expr) (derivingExpressions df) + } {- | O(k) Apply a function to an expression in a dataframe and add the result into `alias` column but diff --git a/src/DataFrame/Synthesis.hs b/src/DataFrame/Synthesis.hs --- a/src/DataFrame/Synthesis.hs +++ b/src/DataFrame/Synthesis.hs @@ -20,8 +20,8 @@ import DataFrame.Internal.Expression ( Expr (..), eSize, - interpret, ) +import DataFrame.Internal.Interpreter (interpret) import DataFrame.Internal.Statistics import DataFrame.Operations.Core (columnAsDoubleVector) import qualified DataFrame.Operations.Statistics as Stats