diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2026 Michael Chavinda + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/dataframe-operations.cabal b/dataframe-operations.cabal new file mode 100644 --- /dev/null +++ b/dataframe-operations.cabal @@ -0,0 +1,65 @@ +cabal-version: 2.4 +name: dataframe-operations +version: 1.0.0.0 + +synopsis: Column operations, expression DSL, and statistics for the dataframe ecosystem. +description: + Untyped column operations (select, filter, sort, join, groupBy, + aggregate, etc.), the expression DSL ('DataFrame.Functions' and + 'DataFrame.Monad'), basic statistics, and the typed wrapper layer + (@DataFrame.Typed.{Access,Operations,Join,Aggregate,Expr}@). Built + on top of @dataframe-core@ and @dataframe-parsing@; pulled in by + every higher-level dataframe satellite. + +bug-reports: https://github.com/mchav/dataframe/issues +license: MIT +license-file: LICENSE +author: Michael Chavinda +maintainer: mschavinda@gmail.com +copyright: (c) 2024-2025 Michael Chavinda +category: Data +tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2 + +common warnings + ghc-options: + -Wincomplete-patterns + -Wincomplete-uni-patterns + -Wunused-imports + -Wunused-local-binds + +library + import: warnings + exposed-modules: + DataFrame.Functions + DataFrame.Monad + DataFrame.Internal.Statistics + DataFrame.Operations.Aggregation + DataFrame.Operations.Core + DataFrame.Operations.Join + DataFrame.Operations.Merge + DataFrame.Operations.Permutation + DataFrame.Operations.Statistics + DataFrame.Operations.Subset + DataFrame.Operations.Transformations + DataFrame.Operations.Typing + DataFrame.Typed.Access + DataFrame.Typed.Aggregate + DataFrame.Typed.Expr + DataFrame.Typed.Join + DataFrame.Typed.Operations + build-depends: base >= 4 && < 5, + bytestring >= 0.11 && < 0.13, + containers >= 0.6.7 && < 0.9, + dataframe-core ^>= 1.0, + dataframe-parsing ^>= 1.0, + deepseq >= 1 && < 2, + hashable >= 1.2 && < 2, + random >= 1.2 && < 2, + regex-tdfa >= 1.3.0 && < 2, + text >= 2.0 && < 3, + time >= 1.12 && < 2, + unordered-containers >= 0.1 && < 1, + vector ^>= 0.13, + vector-algorithms ^>= 0.9 + hs-source-dirs: src + default-language: Haskell2010 diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Functions.hs @@ -0,0 +1,676 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE IncoherentInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +module DataFrame.Functions (module DataFrame.Functions, module DataFrame.Operators) where + +import DataFrame.Internal.Column +import DataFrame.Internal.Expression +import DataFrame.Internal.Statistics + +import Control.Applicative +import qualified Data.Char as Char +import Data.Either +import Data.Function (on) +import Data.Int +import qualified Data.List as L +import qualified Data.Map as M +import qualified Data.Maybe as Maybe +import qualified Data.Text as T +import Data.Time +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU + +import DataFrame.Internal.Nullable ( + BaseType, + NullLift1Op (applyNull1), + NullLift1Result, + NullLift2Op (applyNull2), + NullLift2Result, + ) +import DataFrame.Operators +import Text.Regex.TDFA +import Prelude hiding (maximum, minimum) +import Prelude as P + +lift :: (Columnable a, Columnable b) => (a -> b) -> Expr a -> Expr b +lift f = + Unary (MkUnaryOp{unaryFn = f, unaryName = "unaryUdf", unarySymbol = Nothing}) + +lift2 :: + (Columnable c, Columnable b, Columnable a) => + (c -> b -> a) -> + Expr c -> + Expr b -> + Expr a +lift2 f = + Binary + ( MkBinaryOp + { binaryFn = f + , binaryName = "binaryUdf" + , binarySymbol = Nothing + , binaryCommutative = False + , binaryPrecedence = 0 + } + ) + +{- | Lift a unary function over a nullable or non-nullable column expression. +When the input is @Maybe a@, 'Nothing' short-circuits (like 'fmap'). +When the input is plain @a@, the function is applied directly. + +The return type is inferred via 'NullLift1Result': no annotation needed. +-} +nullLift :: + (NullLift1Op a r (NullLift1Result a r), Columnable (NullLift1Result a r)) => + (BaseType a -> r) -> + Expr a -> + Expr (NullLift1Result a r) +nullLift f = + Unary + (MkUnaryOp{unaryFn = applyNull1 f, unaryName = "nullLift", unarySymbol = Nothing}) + +{- | Lift a binary function over nullable or non-nullable column expressions. +Any 'Nothing' operand short-circuits to 'Nothing' in the result. + +The return type is inferred via 'NullLift2Result': no annotation needed. +-} +nullLift2 :: + (NullLift2Op a b r (NullLift2Result a b r), Columnable (NullLift2Result a b r)) => + (BaseType a -> BaseType b -> r) -> + Expr a -> + Expr b -> + Expr (NullLift2Result a b r) +nullLift2 f = + Binary + ( MkBinaryOp + { binaryFn = applyNull2 f + , binaryName = "nullLift2" + , binarySymbol = Nothing + , binaryCommutative = False + , binaryPrecedence = 0 + } + ) + +{- | Lenient numeric \/ text coercion returning @Maybe a@. Looks up column +@name@ and coerces its values to @a@. Values that cannot be converted +(parse failures, type mismatches) become 'Nothing'; successfully converted +values are wrapped in 'Just'. Existing 'Nothing' in optional source columns +stays as 'Nothing'. +-} +cast :: forall a. (Columnable a, Read a) => T.Text -> Expr (Maybe a) +cast colName = CastWith colName "cast" (either (const Nothing) Just) + +{- | Lenient coercion that substitutes a default for unconvertible values. +Looks up column @name@, coerces its values to @a@, and uses @def@ wherever +conversion fails or the source value is 'Nothing'. +-} +castWithDefault :: forall a. (Columnable a, Read a) => a -> T.Text -> Expr a +castWithDefault def colName = + CastWith colName ("castWithDefault:" <> T.pack (show def)) (fromRight def) + +{- | Lenient coercion returning @Either T.Text a@. Successfully converted +values are 'Right'; values that cannot be parsed are kept as 'Left' with +their original string representation, so the caller can inspect or handle +them downstream. Existing 'Nothing' in optional source columns becomes +@Left \"null\"@. +-} +castEither :: + forall a. (Columnable a, Read a) => T.Text -> Expr (Either T.Text a) +castEither colName = CastWith colName "castEither" (either (Left . T.pack) Right) + +{- | Lenient coercion for assertedly non-nullable columns. +Substitutes @error@ for @Nothing@, so it will crash at evaluation time if +any @Nothing@ is actually encountered. For non-nullable and +fully-populated nullable columns no cost is paid. +-} +unsafeCast :: forall a. (Columnable a, Read a) => T.Text -> Expr a +unsafeCast colName = + CastWith + colName + "unsafeCast" + (fromRight (error "unsafeCast: unexpected Nothing in column")) + +castExpr :: + forall b src. + (Columnable b, Columnable src, Read b) => + Expr src -> + Expr (Maybe b) +castExpr = CastExprWith @b @(Maybe b) @src "castExpr" (either (const Nothing) Just) + +castExprWithDefault :: + forall b src. (Columnable b, Columnable src, Read b) => b -> Expr src -> Expr b +castExprWithDefault def = + CastExprWith @b @b @src + ("castExprWithDefault:" <> T.pack (show def)) + (fromRight def) + +castExprEither :: + forall b src. + (Columnable b, Columnable src, Read b) => + Expr src -> + Expr (Either T.Text b) +castExprEither = + CastExprWith @b @(Either T.Text b) @src + "castExprEither" + (either (Left . T.pack) Right) + +unsafeCastExpr :: + forall b src. (Columnable b, Columnable src, Read b) => Expr src -> Expr b +unsafeCastExpr = + CastExprWith @b @b @src + "unsafeCastExpr" + (fromRight (error "unsafeCastExpr: unexpected Nothing in column")) + +toDouble :: (Columnable a, Real a) => Expr a -> Expr Double +toDouble = + Unary + ( MkUnaryOp + { unaryFn = realToFrac + , unaryName = "toDouble" + , unarySymbol = Nothing + } + ) + +infix 8 `div` +div :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a +div = lift2Decorated Prelude.div "div" (Just "//") False 7 + +mod :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a +mod = lift2Decorated Prelude.mod "mod" Nothing False 7 + +eq :: (Columnable a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool +eq = lift2Decorated (==) "eq" (Just "==") True 4 + +lt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool +lt = lift2Decorated (<) "lt" (Just "<") False 4 + +gt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool +gt = lift2Decorated (>) "gt" (Just ">") False 4 + +leq :: + (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool +leq = lift2Decorated (<=) "leq" (Just "<=") False 4 + +geq :: + (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool +geq = lift2Decorated (>=) "geq" (Just ">=") False 4 + +and :: Expr Bool -> Expr Bool -> Expr Bool +and = (.&&) + +or :: Expr Bool -> Expr Bool -> Expr Bool +or = (.||) + +not :: Expr Bool -> Expr Bool +not = + Unary + (MkUnaryOp{unaryFn = Prelude.not, unaryName = "not", unarySymbol = Just "~"}) + +count :: (Columnable a) => Expr a -> Expr Int +count = Agg (MergeAgg "count" (0 :: Int) (\c _ -> c + 1) (+) id) +{-# SPECIALIZE count :: Expr Double -> Expr Int #-} +{-# SPECIALIZE count :: Expr Float -> Expr Int #-} +{-# SPECIALIZE count :: Expr Int -> Expr Int #-} +{-# SPECIALIZE count :: Expr Int8 -> Expr Int #-} +{-# SPECIALIZE count :: Expr Int16 -> Expr Int #-} +{-# SPECIALIZE count :: Expr Int32 -> Expr Int #-} +{-# SPECIALIZE count :: Expr Int64 -> Expr Int #-} +{-# INLINEABLE count #-} + +-- | Row count, the equivalent of SQL's @COUNT(*)@. +countAll :: Expr Int +countAll = count (Lit (0 :: Int)) +{-# INLINE countAll #-} + +collect :: (Columnable a) => Expr a -> Expr [a] +collect = Agg (FoldAgg "collect" (Just []) (flip (:))) +{-# SPECIALIZE collect :: Expr Double -> Expr [Double] #-} +{-# SPECIALIZE collect :: Expr Float -> Expr [Float] #-} +{-# SPECIALIZE collect :: Expr Int -> Expr [Int] #-} +{-# INLINEABLE collect #-} + +mode :: (Ord a, Columnable a, Eq a) => Expr a -> Expr a +mode = + Agg + ( CollectAgg + "mode" + ( fst + . L.maximumBy (compare `on` snd) + . M.toList + . V.foldl' (\m e -> M.insertWith (+) e (1 :: Int) m) M.empty + ) + ) +{-# SPECIALIZE mode :: Expr Double -> Expr Double #-} +{-# SPECIALIZE mode :: Expr Float -> Expr Float #-} +{-# SPECIALIZE mode :: Expr Int -> Expr Int #-} +{-# SPECIALIZE mode :: Expr Int8 -> Expr Int8 #-} +{-# SPECIALIZE mode :: Expr Int16 -> Expr Int16 #-} +{-# SPECIALIZE mode :: Expr Int32 -> Expr Int32 #-} +{-# SPECIALIZE mode :: Expr Int64 -> Expr Int64 #-} +{-# INLINEABLE mode #-} + +minimum :: (Columnable a, Ord a) => Expr a -> Expr a +minimum = Agg (FoldAgg "minimum" Nothing Prelude.min) +{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Double -> Expr Double #-} +{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Float -> Expr Float #-} +{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int -> Expr Int #-} +{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int8 -> Expr Int8 #-} +{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int16 -> Expr Int16 #-} +{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int32 -> Expr Int32 #-} +{-# SPECIALIZE DataFrame.Functions.minimum :: Expr Int64 -> Expr Int64 #-} +{-# INLINEABLE DataFrame.Functions.minimum #-} + +maximum :: (Columnable a, Ord a) => Expr a -> Expr a +maximum = Agg (FoldAgg "maximum" Nothing Prelude.max) +{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Double -> Expr Double #-} +{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Float -> Expr Float #-} +{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int -> Expr Int #-} +{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int8 -> Expr Int8 #-} +{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int16 -> Expr Int16 #-} +{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int32 -> Expr Int32 #-} +{-# SPECIALIZE DataFrame.Functions.maximum :: Expr Int64 -> Expr Int64 #-} +{-# INLINEABLE DataFrame.Functions.maximum #-} + +sum :: forall a. (Columnable a, Num a) => Expr a -> Expr a +sum = Agg (FoldAgg "sum" Nothing (+)) +{-# SPECIALIZE DataFrame.Functions.sum :: Expr Double -> Expr Double #-} +{-# SPECIALIZE DataFrame.Functions.sum :: Expr Float -> Expr Float #-} +{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int -> Expr Int #-} +{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int8 -> Expr Int8 #-} +{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int16 -> Expr Int16 #-} +{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int32 -> Expr Int32 #-} +{-# SPECIALIZE DataFrame.Functions.sum :: Expr Int64 -> Expr Int64 #-} +{-# INLINEABLE DataFrame.Functions.sum #-} + +sumMaybe :: forall a. (Columnable a, Num a) => Expr (Maybe a) -> Expr a +sumMaybe = Agg (CollectAgg "sumMaybe" (P.sum . Maybe.catMaybes . V.toList)) +{-# SPECIALIZE sumMaybe :: Expr (Maybe Double) -> Expr Double #-} +{-# SPECIALIZE sumMaybe :: Expr (Maybe Float) -> Expr Float #-} +{-# SPECIALIZE sumMaybe :: Expr (Maybe Int) -> Expr Int #-} +{-# SPECIALIZE sumMaybe :: Expr (Maybe Int8) -> Expr Int8 #-} +{-# SPECIALIZE sumMaybe :: Expr (Maybe Int16) -> Expr Int16 #-} +{-# SPECIALIZE sumMaybe :: Expr (Maybe Int32) -> Expr Int32 #-} +{-# SPECIALIZE sumMaybe :: Expr (Maybe Int64) -> Expr Int64 #-} +{-# INLINEABLE sumMaybe #-} + +mean :: (Columnable a, Real a) => Expr a -> Expr Double +mean = + Agg + ( MergeAgg + "mean" + (MeanAcc 0.0 0) + (\(MeanAcc s c) x -> MeanAcc (s + realToFrac x) (c + 1)) + (\(MeanAcc s1 c1) (MeanAcc s2 c2) -> MeanAcc (s1 + s2) (c1 + c2)) + (\(MeanAcc s c) -> if c == 0 then 0 / 0 else s / fromIntegral c) + ) +{-# SPECIALIZE mean :: Expr Double -> Expr Double #-} +{-# SPECIALIZE mean :: Expr Float -> Expr Double #-} +{-# SPECIALIZE mean :: Expr Int -> Expr Double #-} +{-# SPECIALIZE mean :: Expr Int8 -> Expr Double #-} +{-# SPECIALIZE mean :: Expr Int16 -> Expr Double #-} +{-# SPECIALIZE mean :: Expr Int32 -> Expr Double #-} +{-# SPECIALIZE mean :: Expr Int64 -> Expr Double #-} +{-# INLINEABLE mean #-} + +meanMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double +meanMaybe = Agg (CollectAgg "meanMaybe" (mean' . optionalToDoubleVector)) +{-# SPECIALIZE meanMaybe :: Expr (Maybe Double) -> Expr Double #-} +{-# SPECIALIZE meanMaybe :: Expr (Maybe Float) -> Expr Double #-} +{-# SPECIALIZE meanMaybe :: Expr (Maybe Int) -> Expr Double #-} +{-# SPECIALIZE meanMaybe :: Expr (Maybe Int8) -> Expr Double #-} +{-# SPECIALIZE meanMaybe :: Expr (Maybe Int16) -> Expr Double #-} +{-# SPECIALIZE meanMaybe :: Expr (Maybe Int32) -> Expr Double #-} +{-# SPECIALIZE meanMaybe :: Expr (Maybe Int64) -> Expr Double #-} +{-# INLINEABLE meanMaybe #-} + +variance :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double +variance = Agg (CollectAgg "variance" variance') +{-# SPECIALIZE variance :: Expr Double -> Expr Double #-} +{-# SPECIALIZE variance :: Expr Float -> Expr Double #-} +{-# SPECIALIZE variance :: Expr Int -> Expr Double #-} +{-# SPECIALIZE variance :: Expr Int8 -> Expr Double #-} +{-# SPECIALIZE variance :: Expr Int16 -> Expr Double #-} +{-# SPECIALIZE variance :: Expr Int32 -> Expr Double #-} +{-# SPECIALIZE variance :: Expr Int64 -> Expr Double #-} +{-# INLINEABLE variance #-} + +median :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double +median = Agg (CollectAgg "median" median') +{-# SPECIALIZE median :: Expr Double -> Expr Double #-} +{-# SPECIALIZE median :: Expr Float -> Expr Double #-} +{-# SPECIALIZE median :: Expr Int -> Expr Double #-} +{-# SPECIALIZE median :: Expr Int8 -> Expr Double #-} +{-# SPECIALIZE median :: Expr Int16 -> Expr Double #-} +{-# SPECIALIZE median :: Expr Int32 -> Expr Double #-} +{-# SPECIALIZE median :: Expr Int64 -> Expr Double #-} +{-# INLINEABLE median #-} + +medianMaybe :: (Columnable a, Real a) => Expr (Maybe a) -> Expr Double +medianMaybe = Agg (CollectAgg "meanMaybe" (median' . optionalToDoubleVector)) +{-# SPECIALIZE medianMaybe :: Expr (Maybe Double) -> Expr Double #-} +{-# SPECIALIZE medianMaybe :: Expr (Maybe Float) -> Expr Double #-} +{-# SPECIALIZE medianMaybe :: Expr (Maybe Int) -> Expr Double #-} +{-# SPECIALIZE medianMaybe :: Expr (Maybe Int8) -> Expr Double #-} +{-# SPECIALIZE medianMaybe :: Expr (Maybe Int16) -> Expr Double #-} +{-# SPECIALIZE medianMaybe :: Expr (Maybe Int32) -> Expr Double #-} +{-# SPECIALIZE medianMaybe :: Expr (Maybe Int64) -> Expr Double #-} +{-# INLINEABLE medianMaybe #-} + +optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double +optionalToDoubleVector = + VU.fromList + . V.foldl' + (\acc e -> if Maybe.isJust e then realToFrac (Maybe.fromMaybe 0 e) : acc else acc) + [] + +percentile :: Int -> Expr Double -> Expr Double +percentile n = + Agg + ( CollectAgg + (T.pack $ "percentile " ++ show n) + (percentile' n) + ) + +stddev :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double +stddev = Agg (CollectAgg "stddev" (sqrt . variance')) +{-# SPECIALIZE stddev :: Expr Double -> Expr Double #-} +{-# SPECIALIZE stddev :: Expr Float -> Expr Double #-} +{-# SPECIALIZE stddev :: Expr Int -> Expr Double #-} +{-# SPECIALIZE stddev :: Expr Int8 -> Expr Double #-} +{-# SPECIALIZE stddev :: Expr Int16 -> Expr Double #-} +{-# SPECIALIZE stddev :: Expr Int32 -> Expr Double #-} +{-# SPECIALIZE stddev :: Expr Int64 -> Expr Double #-} +{-# INLINEABLE stddev #-} + +stddevMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double +stddevMaybe = Agg (CollectAgg "stddevMaybe" (sqrt . variance' . optionalToDoubleVector)) +{-# SPECIALIZE stddevMaybe :: Expr (Maybe Double) -> Expr Double #-} +{-# SPECIALIZE stddevMaybe :: Expr (Maybe Float) -> Expr Double #-} +{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int) -> Expr Double #-} +{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int8) -> Expr Double #-} +{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int16) -> Expr Double #-} +{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int32) -> Expr Double #-} +{-# SPECIALIZE stddevMaybe :: Expr (Maybe Int64) -> Expr Double #-} +{-# INLINEABLE stddevMaybe #-} + +zScore :: Expr Double -> Expr Double +zScore c = (c - mean c) / stddev c + +pow :: (Columnable a, Num a) => Expr a -> Int -> Expr a +pow expr i = lift2Decorated (^) "pow" (Just "^") True 8 expr (Lit i) +{-# SPECIALIZE pow :: Expr Double -> Int -> Expr Double #-} +{-# SPECIALIZE pow :: Expr Float -> Int -> Expr Float #-} +{-# SPECIALIZE pow :: Expr Int -> Int -> Expr Int #-} +{-# INLINEABLE pow #-} + +relu :: (Columnable a, Num a, Ord a) => Expr a -> Expr a +relu = liftDecorated (Prelude.max 0) "relu" Nothing +{-# SPECIALIZE relu :: Expr Double -> Expr Double #-} +{-# SPECIALIZE relu :: Expr Float -> Expr Float #-} +{-# SPECIALIZE relu :: Expr Int -> Expr Int #-} +{-# INLINEABLE relu #-} + +min :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr a +min = lift2Decorated Prelude.min "min" Nothing True 1 +{-# SPECIALIZE DataFrame.Functions.min :: + Expr Double -> Expr Double -> Expr Double + #-} +{-# SPECIALIZE DataFrame.Functions.min :: + Expr Float -> Expr Float -> Expr Float + #-} +{-# SPECIALIZE DataFrame.Functions.min :: Expr Int -> Expr Int -> Expr Int #-} +{-# INLINEABLE DataFrame.Functions.min #-} + +max :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr a +max = lift2Decorated Prelude.max "max" Nothing True 1 +{-# SPECIALIZE DataFrame.Functions.max :: + Expr Double -> Expr Double -> Expr Double + #-} +{-# SPECIALIZE DataFrame.Functions.max :: + Expr Float -> Expr Float -> Expr Float + #-} +{-# SPECIALIZE DataFrame.Functions.max :: Expr Int -> Expr Int -> Expr Int #-} +{-# INLINEABLE DataFrame.Functions.max #-} + +reduce :: + forall a b. + (Columnable a, Columnable b) => + Expr b -> + a -> + (a -> b -> a) -> + Expr a +reduce expr start f = Agg (FoldAgg "foldUdf" (Just start) f) expr +{-# INLINEABLE reduce #-} + +toMaybe :: (Columnable a) => Expr a -> Expr (Maybe a) +toMaybe = liftDecorated Just "toMaybe" Nothing +{-# SPECIALIZE toMaybe :: Expr Double -> Expr (Maybe Double) #-} +{-# SPECIALIZE toMaybe :: Expr Float -> Expr (Maybe Float) #-} +{-# SPECIALIZE toMaybe :: Expr Int -> Expr (Maybe Int) #-} +{-# INLINEABLE toMaybe #-} + +fromMaybe :: (Columnable a) => a -> Expr (Maybe a) -> Expr a +fromMaybe d = liftDecorated (Maybe.fromMaybe d) "fromMaybe" Nothing +{-# SPECIALIZE fromMaybe :: Double -> Expr (Maybe Double) -> Expr Double #-} +{-# SPECIALIZE fromMaybe :: Float -> Expr (Maybe Float) -> Expr Float #-} +{-# SPECIALIZE fromMaybe :: Int -> Expr (Maybe Int) -> Expr Int #-} +{-# INLINEABLE fromMaybe #-} + +isJust :: (Columnable a) => Expr (Maybe a) -> Expr Bool +isJust = liftDecorated Maybe.isJust "isJust" Nothing +{-# SPECIALIZE isJust :: Expr (Maybe Double) -> Expr Bool #-} +{-# SPECIALIZE isJust :: Expr (Maybe Int) -> Expr Bool #-} +{-# INLINEABLE isJust #-} + +isNothing :: (Columnable a) => Expr (Maybe a) -> Expr Bool +isNothing = liftDecorated Maybe.isNothing "isNothing" Nothing +{-# SPECIALIZE isNothing :: Expr (Maybe Double) -> Expr Bool #-} +{-# SPECIALIZE isNothing :: Expr (Maybe Int) -> Expr Bool #-} +{-# INLINEABLE isNothing #-} + +fromJust :: (Columnable a) => Expr (Maybe a) -> Expr a +fromJust = liftDecorated Maybe.fromJust "fromJust" Nothing +{-# SPECIALIZE fromJust :: Expr (Maybe Double) -> Expr Double #-} +{-# SPECIALIZE fromJust :: Expr (Maybe Int) -> Expr Int #-} +{-# INLINEABLE fromJust #-} + +whenPresent :: + forall a b. + (Columnable a, Columnable b) => + (a -> b) -> + Expr (Maybe a) -> + Expr (Maybe b) +whenPresent f = liftDecorated (fmap f) "whenPresent" Nothing +{-# INLINEABLE whenPresent #-} + +whenBothPresent :: + forall a b c. + (Columnable a, Columnable b, Columnable c) => + (a -> b -> c) -> + Expr (Maybe a) -> + Expr (Maybe b) -> + Expr (Maybe c) +whenBothPresent f = lift2Decorated (\l r -> f <$> l <*> r) "whenBothPresent" Nothing False 0 +{-# INLINEABLE whenBothPresent #-} + +recode :: + forall a b. + (Columnable a, Columnable b, Show (a, b)) => + [(a, b)] -> + Expr a -> + Expr (Maybe b) +recode mapping = + Unary + ( MkUnaryOp + { unaryFn = (`lookup` mapping) + , unaryName = "recode " <> T.pack (show mapping) + , unarySymbol = Nothing + } + ) + +recodeWithCondition :: + forall a b. + (Columnable a, Columnable b) => + Expr b -> + [(Expr a -> Expr Bool, b)] -> + Expr a -> + Expr b +recodeWithCondition fallback [] _val = fallback +recodeWithCondition fallback ((cond, val) : rest) expr = ifThenElse (cond expr) (lit val) (recodeWithCondition fallback rest expr) + +recodeWithDefault :: + forall a b. + (Columnable a, Columnable b, Show (a, b)) => + b -> + [(a, b)] -> + Expr a -> + Expr b +recodeWithDefault d mapping = + Unary + ( MkUnaryOp + { unaryFn = Maybe.fromMaybe d . (`lookup` mapping) + , unaryName = + "recodeWithDefault " <> T.pack (show d) <> " " <> T.pack (show mapping) + , unarySymbol = Nothing + } + ) + +firstOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a) +firstOrNothing = liftDecorated Maybe.listToMaybe "firstOrNothing" Nothing + +lastOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a) +lastOrNothing = liftDecorated (Maybe.listToMaybe . reverse) "lastOrNothing" Nothing + +splitOn :: T.Text -> Expr T.Text -> Expr [T.Text] +splitOn delim = liftDecorated (T.splitOn delim) "splitOn" Nothing + +match :: T.Text -> Expr T.Text -> Expr (Maybe T.Text) +match regex = + liftDecorated + ((\r -> if T.null r then Nothing else Just r) . (=~ regex)) + ("match " <> T.pack (show regex)) + Nothing + +matchAll :: T.Text -> Expr T.Text -> Expr [T.Text] +matchAll regex = + liftDecorated + (getAllTextMatches . (=~ regex)) + ("matchAll " <> T.pack (show regex)) + Nothing + +parseDate :: + (ParseTime t, Columnable t) => T.Text -> Expr T.Text -> Expr (Maybe t) +parseDate format = + liftDecorated + (parseTimeM True defaultTimeLocale (T.unpack format) . T.unpack) + ("parseDate " <> format) + Nothing + +daysBetween :: Expr Day -> Expr Day -> Expr Int +daysBetween = + lift2Decorated + (\d1 d2 -> fromIntegral (diffDays d1 d2)) + "daysBetween" + Nothing + True + 2 + +bind :: + forall a b m. + (Columnable a, Columnable (m a), Monad m, Columnable b, Columnable (m b)) => + (a -> m b) -> + Expr (m a) -> + Expr (m b) +bind f = liftDecorated (>>= f) "bind" Nothing + +{- | Window function: evaluate an expression partitioned by the given columns. + +Each partition computes the inner expression independently, and the result +is broadcast back to every row in that partition. This is analogous to +Polars' @.over()@ or SQL @OVER (PARTITION BY ...)@. + +@ +-- Per-country median, broadcast to every row: +F.over [\"country\"] (F.median (F.col \@Double \"amount\")) + +-- Deviation from group mean: +F.col \@Double \"amount\" - F.over [\"group\"] (F.mean (F.col \@Double \"amount\")) +@ +-} +over :: (Columnable a) => [T.Text] -> Expr a -> Expr a +over = Over + +-- See Section 2.4 of the Haskell Report https://www.haskell.org/definition/haskell2010.pdf +isReservedId :: T.Text -> Bool +isReservedId t = case t of + "case" -> True + "class" -> True + "data" -> True + "default" -> True + "deriving" -> True + "do" -> True + "else" -> True + "foreign" -> True + "if" -> True + "import" -> True + "in" -> True + "infix" -> True + "infixl" -> True + "infixr" -> True + "instance" -> True + "let" -> True + "module" -> True + "newtype" -> True + "of" -> True + "then" -> True + "type" -> True + "where" -> True + _ -> False + +isVarId :: T.Text -> Bool +isVarId t = case T.uncons t of + -- We might want to check c == '_' || Char.isLower c + -- since the haskell report considers '_' a lowercase character + -- However, to prevent an edge case where a user may have a + -- "Name" and an "_Name_" in the same scope, wherein we'd end up + -- with duplicate "_Name_"s, we eschew the check for '_' here. + Just (c, _) -> Char.isLower c && Char.isAlpha c + Nothing -> False + +isHaskellIdentifier :: T.Text -> Bool +isHaskellIdentifier t = Prelude.not (isVarId t) || isReservedId t + +sanitize :: T.Text -> T.Text +sanitize t + | isValid = t + | isHaskellIdentifier t' = "_" <> t' <> "_" + | otherwise = t' + where + isValid = + Prelude.not (isHaskellIdentifier t) + && isVarId t + && T.all Char.isAlphaNum t + t' = T.map replaceInvalidCharacters . T.filter (Prelude.not . parentheses) $ t + replaceInvalidCharacters c + | Char.isUpper c = Char.toLower c + | Char.isSpace c = '_' + | Char.isPunctuation c = '_' -- '-' will also become a '_' + | Char.isSymbol c = '_' + | Char.isAlphaNum c = c -- Blanket condition + | otherwise = '_' -- If we're unsure we'll default to an underscore + parentheses c = case c of + '(' -> True + ')' -> True + '{' -> True + '}' -> True + '[' -> True + ']' -> True + _ -> False diff --git a/src/DataFrame/Internal/Statistics.hs b/src/DataFrame/Internal/Statistics.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Internal/Statistics.hs @@ -0,0 +1,285 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module DataFrame.Internal.Statistics where + +import qualified Data.Vector as V +import qualified Data.Vector.Algorithms.Intro as VA +import qualified Data.Vector.Mutable as VM +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM + +import Control.Exception (throw) +import Control.Monad.ST (runST) +import DataFrame.Errors (DataFrameException (..)) + +mean' :: (Real a, VU.Unbox a) => VU.Vector a -> Double +mean' samp + | VU.null samp = throw $ EmptyDataSetException "mean" + | otherwise = rtf (VU.sum samp) / fromIntegral (VU.length samp) +{-# INLINE [0] mean' #-} + +meanDouble' :: VU.Vector Double -> Double +meanDouble' samp + | VU.null samp = throw $ EmptyDataSetException "mean" + | otherwise = VU.sum samp / fromIntegral (VU.length samp) +{-# INLINE meanDouble' #-} + +meanInt' :: VU.Vector Int -> Double +meanInt' samp + | VU.null samp = throw $ EmptyDataSetException "mean" + | otherwise = fromIntegral (VU.sum samp) / fromIntegral (VU.length samp) +{-# INLINE meanInt' #-} + +{-# RULES +"mean'/Double" [1] forall (xs :: VU.Vector Double). + mean' xs = + meanDouble' xs +"mean'/Int" [1] forall (xs :: VU.Vector Int). + mean' xs = + meanInt' xs + #-} + +median' :: (Real a, VU.Unbox a) => VU.Vector a -> Double +median' samp + | VU.null samp = throw $ EmptyDataSetException "median" + | otherwise = runST $ do + mutableSamp <- VU.thaw samp + VA.sort mutableSamp + let len = VU.length samp + middleIndex = len `div` 2 + middleElement <- VUM.read mutableSamp middleIndex + if odd len + then pure (rtf middleElement) + else do + prev <- VUM.read mutableSamp (middleIndex - 1) + pure (rtf (middleElement + prev) / 2) +{-# INLINE median' #-} + +-- accumulator: count, mean, m2 +data VarAcc + = VarAcc {-# UNPACK #-} !Int {-# UNPACK #-} !Double {-# UNPACK #-} !Double + deriving (Show) + +varianceStep :: VarAcc -> Double -> VarAcc +varianceStep (VarAcc !n !meanVal !m2) !x = + let !n' = n + 1 + !delta = x - meanVal + !meanVal' = meanVal + delta / fromIntegral n' + !m2' = m2 + delta * (x - meanVal') + in VarAcc n' meanVal' m2' +{-# INLINE varianceStep #-} + +computeVariance :: VarAcc -> Double +computeVariance (VarAcc !n _ !m2) + | n < 2 = 0 -- or error "variance of <2 samples" + | otherwise = m2 / fromIntegral (n - 1) +{-# INLINE computeVariance #-} + +variance' :: (Real a, VU.Unbox a) => VU.Vector a -> Double +variance' = computeVariance . VU.foldl' varianceStep (VarAcc 0 0 0) . VU.map rtf +{-# INLINE variance' #-} + +varianceDouble' :: VU.Vector Double -> Double +varianceDouble' = computeVariance . VU.foldl' varianceStep (VarAcc 0 0 0) +{-# INLINE varianceDouble' #-} + +-- accumulator: count, mean, m2, m3 +data SkewAcc = SkewAcc !Int !Double !Double !Double deriving (Show) + +skewnessStep :: (VU.Unbox a, Num a, Real a) => SkewAcc -> a -> SkewAcc +skewnessStep (SkewAcc !n !meanVal !m2 !m3) !x' = + let !n' = n + 1 + x = rtf x' + !k = fromIntegral n' + !delta = x - meanVal + !meanVal' = meanVal + delta / k + !m2' = m2 + (delta ^ (2 :: Int) * (k - 1)) / k + !m3' = + m3 + + (delta ^ (3 :: Int) * (k - 1) * (k - 2)) / k ^ (2 :: Int) + - (3 * delta * m2) / k + in SkewAcc n' meanVal' m2' m3' +{-# INLINE skewnessStep #-} + +computeSkewness :: SkewAcc -> Double +computeSkewness (SkewAcc n _ m2 m3) + | n < 3 = 0 -- or error "skewness of <3 samples" + | otherwise = (sqrt (fromIntegral n - 1) * m3) / sqrt (m2 ^ (3 :: Int)) +{-# INLINE computeSkewness #-} + +skewness' :: (VU.Unbox a, Real a, Num a) => VU.Vector a -> Double +skewness' = computeSkewness . VU.foldl' skewnessStep (SkewAcc 0 0 0 0) +{-# INLINE skewness' #-} + +data CorrelationStats + = CorrelationStats + {-# UNPACK #-} !Double + {-# UNPACK #-} !Double + {-# UNPACK #-} !Double + {-# UNPACK #-} !Double + {-# UNPACK #-} !Double + +correlation' :: VU.Vector Double -> VU.Vector Double -> Maybe Double +correlation' xs ys + | n < 2 = Nothing + | VU.length xs /= VU.length ys = Nothing + | otherwise = + let nf = fromIntegral n + initial = CorrelationStats 0 0 0 0 0 + (CorrelationStats sumX sumY sumXX sumYY sumXY) = VU.ifoldl' step initial xs + + !num = nf * sumXY - sumX * sumY + !den = sqrt ((nf * sumXX - sumX * sumX) * (nf * sumYY - sumY * sumY)) + in Just (num / den) + where + n = VU.length xs + step (CorrelationStats sx sy sxx syy sxy) i x = + let !y = VU.unsafeIndex ys i + in CorrelationStats (sx + x) (sy + y) (sxx + x * x) (syy + y * y) (sxy + x * y) +{-# INLINE correlation' #-} + +quantiles' :: + (VU.Unbox a, Num a, Real a) => + VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double +quantiles' qs q samp + | VU.null samp = throw $ EmptyDataSetException "quantiles" + | q < 2 = throw $ WrongQuantileNumberException q + | VU.any (\i -> i < 0 || i > q) qs = throw $ WrongQuantileIndexException qs q + | otherwise = runST $ do + let !n = VU.length samp + mutableSamp <- VU.thaw samp + VA.sort mutableSamp + VU.mapM + ( \i -> do + let !p = fromIntegral i / fromIntegral q + !position = p * fromIntegral (n - 1) :: Double + !index = floor position :: Int + !f = position - fromIntegral index + x <- fmap rtf (VUM.read mutableSamp index) + if f == 0 + then return x + else do + y <- fmap rtf (VUM.read mutableSamp (index + 1)) + return $ (1 - f) * x + f * y + ) + qs +{-# INLINE quantiles' #-} + +percentile' :: (VU.Unbox a, Num a, Real a) => Int -> VU.Vector a -> Double +percentile' n = VU.head . quantiles' (VU.fromList [n]) 100 + +quantilesOrd' :: + (Ord a, Eq a) => + VU.Vector Int -> Int -> V.Vector a -> V.Vector a +quantilesOrd' qs q samp + | V.null samp = throw $ EmptyDataSetException "quantiles" + | q < 2 = throw $ WrongQuantileNumberException q + | VU.any (\i -> i < 0 || i > q) qs = throw $ WrongQuantileIndexException qs q + | otherwise = runST $ do + let !n = V.length samp + mutableSamp <- V.thaw samp + VA.sort mutableSamp + V.mapM + ( \i -> do + let !p = fromIntegral i / fromIntegral q :: Double + !position = p * fromIntegral (n - 1) + !index = floor position :: Int + -- This is not exact for Ord instances. + -- Figure out how to make it so. + VM.read mutableSamp index + ) + (V.convert qs) + +percentileOrd' :: (Ord a, Eq a) => Int -> V.Vector a -> a +percentileOrd' n = V.head . quantilesOrd' (VU.fromList [n]) 100 + +interQuartileRange' :: (VU.Unbox a, Num a, Real a) => VU.Vector a -> Double +interQuartileRange' samp = + let quartiles = quantiles' (VU.fromList [1, 3]) 4 samp + in quartiles VU.! 1 - quartiles VU.! 0 +{-# INLINE interQuartileRange' #-} + +meanSquaredError :: VU.Vector Double -> VU.Vector Double -> Maybe Double +meanSquaredError target prediction = + let + squareDiff = VU.ifoldl' (\sq i e -> (e - target VU.! i) ^ (2 :: Int) + sq) 0 prediction + in + Just $ squareDiff / fromIntegral (max (VU.length target) (VU.length prediction)) +{-# INLINE meanSquaredError #-} + +mutualInformationBinned :: + Int -> VU.Vector Double -> VU.Vector Double -> Maybe Double +mutualInformationBinned k xs ys + | VU.length xs /= VU.length ys = Nothing + | VU.null xs = Nothing + | k < 2 = Nothing + | rx <= 0 || ry <= 0 = Just 0 + | otherwise = + let bx = VU.map (binIndex xmin xmax k) xs + by = VU.map (binIndex ymin ymax k) ys + n = fromIntegral (VU.length xs) :: Double + mx = bincount k bx + my = bincount k by + mxy = jointBincount k bx by + in Just $ + sum + [ let !cxy = fromIntegral c + !pxy = cxy / n + !px = fromIntegral (mx VU.! i) / n + !py = fromIntegral (my VU.! j) / n + in if c == 0 then 0 else pxy * logBase 2 (pxy / (px * py)) + | i <- [0 .. k - 1] + , j <- [0 .. k - 1] + , let !c = mxy VU.! (i * k + j) + ] + where + (xmin, xmax) = (VU.minimum xs, VU.maximum xs) + (ymin, ymax) = (VU.minimum ys, VU.maximum ys) + rx = xmax - xmin + ry = ymax - ymin + +binIndex :: Double -> Double -> Int -> Double -> Int +binIndex lo hi k x + | hi == lo = 0 + | otherwise = + let !t = (x - lo) / (hi - lo) + !ix = floor (fromIntegral k * t) :: Int + in max 0 (min (k - 1) ix) +{-# INLINE binIndex #-} + +bincount :: Int -> VU.Vector Int -> VU.Vector Int +bincount k bs = VU.create $ do + mv <- VU.thaw (VU.replicate k 0) + VU.forM_ bs $ \b -> do + let i + | b < 0 = 0 + | b >= k = k - 1 + | otherwise = b + x <- VUM.read mv i + VUM.write mv i (x + 1) + pure mv +{-# INLINE bincount #-} + +jointBincount :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int +jointBincount k bx by = VU.create $ do + mv <- VU.thaw (VU.replicate (k * k) 0) + VU.forM_ (VU.zip bx by) $ \(i, j) -> do + let ii = clamp i 0 (k - 1) + jj = clamp j 0 (k - 1) + ix = ii * k + jj + x <- VUM.read mv ix + VUM.write mv ix (x + 1) + pure mv + where + clamp z a b = max a (min b z) +{-# INLINE jointBincount #-} + +rtf :: (Real a) => a -> Double +rtf = realToFrac +{-# NOINLINE [1] rtf #-} + +{-# RULES +"rtf/Double" [2] forall (x :: Double). rtf x = x + #-} diff --git a/src/DataFrame/Monad.hs b/src/DataFrame/Monad.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Monad.hs @@ -0,0 +1,101 @@ +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} + +module DataFrame.Monad where + +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.DataFrame (DataFrame) +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Nullable (BaseType) +import qualified DataFrame.Operations.Core as D +import qualified DataFrame.Operations.Subset as D +import DataFrame.Operations.Transformations (ImputeOp) +import qualified DataFrame.Operations.Transformations as D + +import qualified Data.Text as T +import System.Random + +-- A re-implementation of the state monad. +-- `mtl` might be too heavy a dependency just to get +-- a single monad instance. +newtype FrameM a = FrameM {runFrameM_ :: DataFrame -> (DataFrame, a)} + +instance Functor FrameM where + fmap :: (a -> b) -> FrameM a -> FrameM b + fmap f (FrameM g) = FrameM $ \df -> + let (df', x) = g df + in (df', f x) + +instance Applicative FrameM where + pure x = FrameM (,x) + (<*>) :: FrameM (a -> b) -> FrameM a -> FrameM b + FrameM ff <*> FrameM fx = FrameM $ \df -> + let (df1, f) = ff df + (df2, x) = fx df1 + in (df2, f x) + +instance Monad FrameM where + (>>=) :: FrameM a -> (a -> FrameM b) -> FrameM b + FrameM g >>= f = FrameM $ \df -> + let (df1, x) = g df + FrameM h = f x + in h df1 + +modifyM :: (DataFrame -> DataFrame) -> FrameM () +modifyM f = FrameM $ \df -> (f df, ()) + +inspectM :: (DataFrame -> b) -> FrameM b +inspectM f = FrameM $ \df -> (df, f df) + +deriveM :: (Columnable a) => T.Text -> Expr a -> FrameM (Expr a) +deriveM name expr = FrameM $ \df -> + let df' = D.derive name expr df + in (df', Col name) + +renameM :: (Columnable a) => Expr a -> T.Text -> FrameM (Expr a) +renameM (Col oldName) newName = FrameM $ \df -> + let df' = D.rename oldName newName df + in (df', Col newName) +renameM expr newName = deriveM newName expr + +filterWhereM :: Expr Bool -> FrameM () +filterWhereM p = modifyM (D.filterWhere p) + +sampleM :: (RandomGen g) => g -> Double -> FrameM () +sampleM pureGen p = modifyM (D.sample pureGen p) + +takeM :: Int -> FrameM () +takeM n = modifyM (D.take n) + +filterJustM :: (Columnable a) => Expr (Maybe a) -> FrameM (Expr a) +filterJustM (Col name) = FrameM $ \df -> + let df' = D.filterJust name df + in (df', Col name) +filterJustM expr = + error $ "Cannot filter on compound expression: " ++ show expr + +imputeM :: + (ImputeOp a, Columnable (BaseType a)) => + Expr a -> + BaseType a -> + FrameM (Expr (BaseType a)) +imputeM expr@(Col name) value = FrameM $ \df -> + let df' = D.impute expr value df + in (df', Col name) +imputeM expr _ = error $ "Cannot impute on compound expression: " ++ show expr + +runFrameM :: DataFrame -> FrameM a -> (a, DataFrame) +runFrameM df (FrameM action) = + let (df', a) = action df + in (a, df') + +evalFrameM :: DataFrame -> FrameM a -> a +evalFrameM df m = fst (runFrameM df m) + +execFrameM :: DataFrame -> FrameM a -> DataFrame +execFrameM df m = snd (runFrameM df m) diff --git a/src/DataFrame/Operations/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Aggregation.hs @@ -0,0 +1,165 @@ +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE Strict #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.Operations.Aggregation ( + module DataFrame.Operations.Aggregation, + groupBy, + buildRowToGroup, + changingPoints, +) where + +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM + +import Control.Exception (throw) +import Control.Monad +import Control.Monad.ST (runST) +import Data.Hashable +import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) +import DataFrame.Errors +import DataFrame.Internal.Column ( + Column (..), + TypedColumn (..), + atIndicesStable, + bitmapTestBit, + ) +import DataFrame.Internal.DataFrame ( + DataFrame (..), + GroupedDataFrame (..), + columnNames, + insertColumn, + ) +import DataFrame.Internal.Expression +import DataFrame.Internal.Grouping (buildRowToGroup, changingPoints, groupBy) +import DataFrame.Internal.Interpreter +import DataFrame.Internal.Types +import DataFrame.Operations.Core +import DataFrame.Operations.Subset +import Type.Reflection (typeRep) + +computeRowHashes :: [Int] -> DataFrame -> VU.Vector Int +computeRowHashes indices df = runST $ do + let n = fst (dimensions df) + mv <- VUM.new n + + let selectedCols = map (columns df V.!) indices + + forM_ selectedCols $ \case + UnboxedColumn _ (v :: VU.Vector a) -> + case testEquality (typeRep @a) (typeRep @Int) of + Just Refl -> + VU.imapM_ + ( \i (x :: Int) -> do + h <- VUM.unsafeRead mv i + VUM.unsafeWrite mv i (hashWithSalt h x) + ) + v + Nothing -> + case testEquality (typeRep @a) (typeRep @Double) of + Just Refl -> + VU.imapM_ + ( \i (d :: Double) -> do + h <- VUM.unsafeRead mv i + VUM.unsafeWrite mv 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 (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 (hashWithSalt h x) + ) + v + SFalse -> + VU.imapM_ + ( \i d -> do + let x = hash (show d) + h <- VUM.unsafeRead mv i + VUM.unsafeWrite mv i (hashWithSalt h x) + ) + v + BoxedColumn bm (v :: V.Vector a) -> + case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> + V.imapM_ + ( \i (t :: T.Text) -> do + h <- VUM.unsafeRead mv i + let h' = case bm of + Just bm' | not (bitmapTestBit bm' i) -> hashWithSalt h (0 :: Int) + _ -> hashWithSalt h t + VUM.unsafeWrite mv i h' + ) + v + Nothing -> + V.imapM_ + ( \i d -> do + let x = case bm of + Just bm' | not (bitmapTestBit bm' i) -> 0 :: Int + _ -> hash (show d) + h <- VUM.unsafeRead mv i + VUM.unsafeWrite mv i (hashWithSalt h x) + ) + v + + VU.unsafeFreeze mv + where + doubleToInt :: Double -> Int + doubleToInt = floor . (* 1000) + +{- | Aggregate a grouped dataframe using the expressions given. +All ungrouped columns will be dropped. +-} +aggregate :: [NamedExpr] -> GroupedDataFrame -> DataFrame +aggregate aggs gdf@(Grouped df groupingColumns valIndices offs _rowToGroup) = + let + df' = + selectIndices + (VU.map (valIndices VU.!) (VU.init offs)) + (select groupingColumns df) + + f (name, UExpr (expr :: Expr a)) d = + let + value = case interpretAggregation @a gdf expr of + Left e -> throw e + Right (UnAggregated _) -> throw $ UnaggregatedException (T.pack $ show expr) + Right (Aggregated (TColumn col)) -> col + in + insertColumn name value d + in + fold f aggs df' + +selectIndices :: VU.Vector Int -> DataFrame -> DataFrame +selectIndices xs df = + 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 (VU.map (indices VU.!) (VU.init os)) df + where + (Grouped _ _ indices os _rtg) = groupBy (columnNames df) df diff --git a/src/DataFrame/Operations/Core.hs b/src/DataFrame/Operations/Core.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Core.hs @@ -0,0 +1,955 @@ +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.Operations.Core where + +import qualified Data.List as L +import qualified Data.Map as M +import qualified Data.Map.Strict as MS +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Unboxed as VU + +import Control.Exception (throw) +import Data.Bits (popCount) +import Data.Either +import qualified Data.Foldable as Fold +import Data.Function (on, (&)) +import Data.Maybe +import Data.Type.Equality (TestEquality (..)) +import DataFrame.Errors +import DataFrame.Internal.Column ( + Column (..), + Columnable, + TypedColumn (..), + columnLength, + columnTypeString, + expandColumn, + fromList, + fromVector, + toDoubleVector, + toFloatVector, + toIntVector, + toUnboxedVector, + toVector, + ) +import DataFrame.Internal.DataFrame ( + DataFrame (..), + columnIndices, + columnNames, + derivingExpressions, + empty, + fromNamedColumns, + getColumn, + insertColumn, + null, + ) +import DataFrame.Internal.Expression +import DataFrame.Internal.Interpreter +import DataFrame.Internal.Parsing (isNullish) +import DataFrame.Internal.Row (Any, mkColumnFromRow) +import Type.Reflection +import Prelude hiding (null) + +{- | O(1) Get DataFrame dimensions i.e. (rows, columns) + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified DataFrame as D +>>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])] +>>> D.dimensions df + +(100, 3) +@ +-} +dimensions :: DataFrame -> (Int, Int) +dimensions = dataframeDimensions +{-# INLINE dimensions #-} + +{- | O(1) Get number of rows in a dataframe. + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified DataFrame as D +>>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])] +>>> D.nRows df +100 +@ +-} +nRows :: DataFrame -> Int +nRows = fst . dataframeDimensions + +{- | O(1) Get number of columns in a dataframe. + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified DataFrame as D +>>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])] +>>> D.nColumns df +3 +@ +-} +nColumns :: DataFrame -> Int +nColumns = snd . dataframeDimensions + +{- | O(k) Get column names of the DataFrame in order of insertion. + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified DataFrame as D +>>> df = D.fromNamedColumns [("a", D.fromList [1..100]), ("b", D.fromList [1..100]), ("c", D.fromList [1..100])] +>>> D.columnNames df + +["a", "b", "c"] +@ +-} + +-- 'columnNames' is now defined in "DataFrame.Internal.DataFrame" and +-- re-exported from "DataFrame" at the top level. + +{- | Adds a vector to the dataframe. If the vector has less elements than the dataframe and the dataframe is not empty +the vector is converted to type `Maybe a` filled with `Nothing` to match the size of the dataframe. Similarly, +if the vector has more elements than what's currently in the dataframe, the other columns in the dataframe are +change to `Maybe ` and filled with `Nothing`. + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified DataFrame as D +>>> import qualified Data.Vector as V +>>> D.insertVector "numbers" (V.fromList [(1 :: Int)..10]) D.empty + +-------- + numbers +-------- + Int +-------- + 1 + 2 + 3 + 4 + 5 + 6 + 7 + 8 + 9 + 10 + +@ +-} +insertVector :: + forall a. + (Columnable a) => + -- | Column Name + T.Text -> + -- | Vector to add to column + V.Vector a -> + -- | DataFrame to add column to + DataFrame -> + DataFrame +insertVector name xs = insertColumn name (fromVector xs) +{-# INLINE insertVector #-} + +{- | Adds a foldable collection to the dataframe. If the collection has less elements than the +dataframe and the dataframe is not empty +the collection is converted to type `Maybe a` filled with `Nothing` to match the size of the dataframe. Similarly, +if the collection has more elements than what's currently in the dataframe, the other columns in the dataframe are +change to `Maybe ` and filled with `Nothing`. + +Be careful not to insert infinite collections with this function as that will crash the program. + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified DataFrame as D +>>> D.insert "numbers" [(1 :: Int)..10] D.empty + +-------- + numbers +-------- + Int +-------- + 1 + 2 + 3 + 4 + 5 + 6 + 7 + 8 + 9 + 10 + +@ +-} +insert :: + forall a t. + (Columnable a, Foldable t) => + -- | Column Name + T.Text -> + -- | Sequence to add to dataframe + t a -> + -- | DataFrame to add column to + DataFrame -> + DataFrame +insert name xs = insertColumn name (fromList (Fold.foldr' (:) [] xs)) -- TODO: Do reflection on container type so we can sometimes avoid the list construction. +{-# INLINE insert #-} + +{- | Adds a vector to the dataframe and pads it with a default value if it has less elements than the number of rows. + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified Data.Vector as V +>>> import qualified DataFrame as D +>>> df = D.fromNamedColumns [("x", D.fromList [(1 :: Int)..10])] +>>> D.insertVectorWithDefault 0 "numbers" (V.fromList [(1 :: Int),2,3]) df + +------------- + x | numbers +----|-------- +Int | Int +----|-------- +1 | 1 +2 | 2 +3 | 3 +4 | 0 +5 | 0 +6 | 0 +7 | 0 +8 | 0 +9 | 0 +10 | 0 + +@ +-} +insertVectorWithDefault :: + forall a. + (Columnable a) => + -- | Default Value + a -> + -- | Column name + T.Text -> + -- | Data to add to column + V.Vector a -> + -- | DataFrame to add the column to + DataFrame -> + DataFrame +insertVectorWithDefault defaultValue name xs d = + let (rows, _) = dataframeDimensions d + values = xs V.++ V.replicate (rows - V.length xs) defaultValue + in insertColumn name (fromVector values) d + +{- | Adds a list to the dataframe and pads it with a default value if it has less elements than the number of rows. + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified DataFrame as D +>>> df = D.fromNamedColumns [("x", D.fromList [(1 :: Int)..10])] +>>> D.insertWithDefault 0 "numbers" [(1 :: Int),2,3] df + +------------- + x | numbers +----|-------- +Int | Int +----|-------- +1 | 1 +2 | 2 +3 | 3 +4 | 0 +5 | 0 +6 | 0 +7 | 0 +8 | 0 +9 | 0 +10 | 0 + +@ +-} +insertWithDefault :: + forall a t. + (Columnable a, Foldable t) => + -- | Default Value + a -> + -- | Column name + T.Text -> + -- | Data to add to column + t a -> + -- | DataFrame to add the column to + DataFrame -> + DataFrame +insertWithDefault defaultValue name xs d = + let (rows, _) = dataframeDimensions d + xs' = Fold.foldr' (:) [] xs + values = xs' ++ replicate (rows - length xs') defaultValue + in insertColumn name (fromList values) d + +{- | /O(n)/ Adds an unboxed vector to the dataframe. + +Same as insertVector but takes an unboxed vector. If you insert a vector of numbers through insertVector it will either way be converted +into an unboxed vector so this function saves that extra work/conversion. +-} +insertUnboxedVector :: + forall a. + (Columnable a, VU.Unbox a) => + -- | Column Name + T.Text -> + -- | Unboxed vector to add to column + VU.Vector a -> + -- | DataFrame to add the column to + DataFrame -> + DataFrame +insertUnboxedVector name xs = insertColumn name (UnboxedColumn Nothing xs) + +{- | /O(n)/ Add a column to the dataframe. + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified DataFrame as D +>>> D.insertColumn "numbers" (D.fromList [(1 :: Int)..10]) D.empty + +-------- + numbers +-------- + Int +-------- + 1 + 2 + 3 + 4 + 5 + 6 + 7 + 8 + 9 + 10 + +@ +-} + +-- 'insertColumn' is now defined in "DataFrame.Internal.DataFrame". + +{- | /O(n)/ Clones a column and places it under a new name in the dataframe. + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified Data.Vector as V +>>> df = insertVector "numbers" (V.fromList [1..10]) D.empty +>>> D.cloneColumn "numbers" "others" df + +----------------- + numbers | others +---------|------- + Int | Int +---------|------- + 1 | 1 + 2 | 2 + 3 | 3 + 4 | 4 + 5 | 5 + 6 | 6 + 7 | 7 + 8 | 8 + 9 | 9 + 10 | 10 + +@ +-} +cloneColumn :: T.Text -> T.Text -> DataFrame -> DataFrame +cloneColumn original new df + | null df = throw (EmptyDataSetException "cloneColumn") + | otherwise = fromMaybe + ( throw $ + ColumnsNotFoundException [original] "cloneColumn" (M.keys $ columnIndices df) + ) + $ do + column <- getColumn original df + return $ insertColumn new column df + +{- | /O(n)/ Renames a single column. + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified DataFrame as D +>>> import qualified Data.Vector as V +>>> df = insertVector "numbers" (V.fromList [1..10]) D.empty +>>> D.rename "numbers" "others" df + +------- + others +------- + Int +------- + 1 + 2 + 3 + 4 + 5 + 6 + 7 + 8 + 9 + 10 + +@ +-} +rename :: T.Text -> T.Text -> DataFrame -> DataFrame +rename orig new df = either throw id (renameSafe orig new df) + +{- | /O(n)/ Renames many columns. + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified DataFrame as D +>>> import qualified Data.Vector as V +>>> df = D.insertVector "others" (V.fromList [11..20]) (D.insertVector "numbers" (V.fromList [1..10]) D.empty) +>>> df + +----------------- + numbers | others +---------|------- + Int | Int +---------|------- + 1 | 11 + 2 | 12 + 3 | 13 + 4 | 14 + 5 | 15 + 6 | 16 + 7 | 17 + 8 | 18 + 9 | 19 + 10 | 20 + +>>> D.renameMany [("numbers", "first_10"), ("others", "next_10")] df + +------------------- + first_10 | next_10 +----------|-------- + Int | Int +----------|-------- + 1 | 11 + 2 | 12 + 3 | 13 + 4 | 14 + 5 | 15 + 6 | 16 + 7 | 17 + 8 | 18 + 9 | 19 + 10 | 20 + +@ +-} +renameMany :: [(T.Text, T.Text)] -> DataFrame -> DataFrame +renameMany = fold (uncurry rename) + +renameSafe :: + T.Text -> T.Text -> DataFrame -> Either DataFrameException DataFrame +renameSafe orig new df + | null df = throw (EmptyDataSetException "rename") + | otherwise = fromMaybe + (Left $ ColumnsNotFoundException [orig] "rename" (M.keys $ columnIndices df)) + $ do + columnIndex <- M.lookup orig (columnIndices df) + let origRemoved = M.delete orig (columnIndices df) + let newAdded = M.insert new columnIndex origRemoved + return (Right df{columnIndices = newAdded}) + +data ColumnInfo = ColumnInfo + { nameOfColumn :: !T.Text + , nonNullValues :: !Int + , nullValues :: !Int + , typeOfColumn :: !T.Text + } + +{- | O(n * k ^ 2) Returns the number of non-null columns in the dataframe and the type associated with each column. + +==== __Example__ +@ +>>> import qualified Data.Vector as V +>>> df = D.insertVector "others" (V.fromList [11..20]) (D.insertVector "numbers" (V.fromList [1..10]) D.empty) +>>> D.describeColumns df + +-------------------------------------------------------- + Column Name | # Non-null Values | # Null Values | Type +-------------|-------------------|---------------|----- + Text | Int | Int | Text +-------------|-------------------|---------------|----- + others | 10 | 0 | Int + numbers | 10 | 0 | Int + +@ +-} +describeColumns :: DataFrame -> DataFrame +describeColumns df = + empty + & insertColumn "Column Name" (fromList (map nameOfColumn infos)) + & insertColumn "# Non-null Values" (fromList (map nonNullValues infos)) + & insertColumn "# Null Values" (fromList (map nullValues infos)) + & insertColumn "Type" (fromList (map typeOfColumn infos)) + where + infos = + L.sortBy (compare `on` nonNullValues) (V.ifoldl' go [] (columns df)) :: + [ColumnInfo] + indexMap = M.fromList (map (\(a, b) -> (b, a)) $ M.toList (columnIndices df)) + columnName i = M.lookup i indexMap + go acc i col@(BoxedColumn _bm (_c :: V.Vector a)) = + let + cname = columnName i + countNulls = nulls col + columnType = T.pack $ columnTypeString col + in + if isNothing cname + then acc + else + ColumnInfo + (fromMaybe "" cname) + (columnLength col - countNulls) + countNulls + columnType + : acc + go acc i col@(UnboxedColumn _bm _c) = + let + cname = columnName i + countNulls = nulls col + columnType = T.pack $ columnTypeString col + in + if isNothing cname + then acc + else + ColumnInfo + (fromMaybe "" cname) + (columnLength col - countNulls) + countNulls + columnType + : acc + +nulls :: Column -> Int +nulls (BoxedColumn (Just bm) xs) = + -- count null bits in bitmap + let n = VG.length xs + in n - VU.foldl' (\acc b -> acc + popCount b) 0 bm +nulls (BoxedColumn Nothing (xs :: V.Vector a)) = case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> VG.length $ VG.filter isNullish xs + Nothing -> case testEquality (typeRep @a) (typeRep @String) of + Just Refl -> VG.length $ VG.filter (isNullish . T.pack) xs + Nothing -> 0 +nulls (UnboxedColumn (Just bm) xs) = + let n = VG.length xs + in n - VU.foldl' (\acc b -> acc + popCount b) 0 bm +nulls _ = 0 + +{- | Creates a dataframe from a list of tuples with name and column. + +==== __Example__ +@ +>>> df = D.fromNamedColumns [("numbers", D.fromList [1..10]), ("others", D.fromList [11..20])] +>>> df +----------------- + numbers | others +---------|------- + Int | Int +---------|------- + 1 | 11 + 2 | 12 + 3 | 13 + 4 | 14 + 5 | 15 + 6 | 16 + 7 | 17 + 8 | 18 + 9 | 19 + 10 | 20 + +@ +-} + +-- 'fromNamedColumns' is now defined in "DataFrame.Internal.DataFrame". + +{- | Create a dataframe from a list of columns. The column names are "0", "1"... etc. +Useful for quick exploration but you should probably always rename the columns after +or drop the ones you don't want. + +==== __Example__ +@ +>>> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]] +>>> df +----------------- + 0 | 1 +-----|---- + Int | Int +-----|---- + 1 | 11 + 2 | 12 + 3 | 13 + 4 | 14 + 5 | 15 + 6 | 16 + 7 | 17 + 8 | 18 + 9 | 19 + 10 | 20 + +@ +-} +fromUnnamedColumns :: [Column] -> DataFrame +fromUnnamedColumns = fromNamedColumns . zip (map (T.pack . show) [(0 :: Int) ..]) + +{- | Create a dataframe from a list of column names and rows. + +==== __Example__ +@ +>>> df = D.fromRows ["A", "B"] [[D.toAny 1, D.toAny 11], [D.toAny 2, D.toAny 12], [D.toAny 3, D.toAny 13]] + +>>> df + +---------- + A | B +-----|---- + Int | Int +-----|---- + 1 | 11 + 2 | 12 + 3 | 13 + +@ +-} +fromRows :: [T.Text] -> [[Any]] -> DataFrame +fromRows names rows = + L.foldl' + (\df i -> insertColumn (names !! i) (mkColumnFromRow i rows) df) + empty + [0 .. length names - 1] + +{- | O (k * n) Counts the occurences of each value in a given column. + +==== __Example__ +@ +>>> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]] + +>>> D.valueCounts @Int "0" df + +[(1,1),(2,1),(3,1),(4,1),(5,1),(6,1),(7,1),(8,1),(9,1),(10,1)] + +@ +-} +valueCounts :: + forall a. (Ord a, Columnable a) => Expr a -> DataFrame -> [(a, Int)] +valueCounts expr df + | null df = throw (EmptyDataSetException "valueCounts") + | otherwise = case columnAsVector expr df of + Left e -> throw e + Right column' -> + let + column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column' + in + M.toAscList column + +{- | O (k * n) Shows the proportions of each value in a given column. + +==== __Example__ +@ +>>> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]] + +>>> D.valueCounts @Int "0" df + +[(1,0.1),(2,0.1),(3,0.1),(4,0.1),(5,0.1),(6,0.1),(7,0.1),(8,0.1),(9,0.1),(10,0.1)] + +@ +-} +valueProportions :: + forall a. (Ord a, Columnable a) => Expr a -> DataFrame -> [(a, Double)] +valueProportions expr df + | null df = throw (EmptyDataSetException "valueCounts") + | otherwise = case columnAsVector expr df of + Left e -> throw e + Right column' -> + let + counts = + M.toAscList + (V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column') + total = fromIntegral (sum (map snd counts)) + in + map (fmap ((/ total) . fromIntegral)) counts + +{- | A left fold for dataframes that takes the dataframe as the last object. +This makes it easier to chain operations. + +==== __Example__ +@ +>>> df = D.fromNamedColumns [("x", D.fromList [1..100]), ("y", D.fromList [11..110])] +>>> D.fold D.dropLast [1..5] df + +--------- + x | y +----|---- +Int | Int +----|---- +1 | 11 +2 | 12 +3 | 13 +4 | 14 +5 | 15 +6 | 16 +7 | 17 +8 | 18 +9 | 19 +10 | 20 +11 | 21 +12 | 22 +13 | 23 +14 | 24 +15 | 25 +16 | 26 +17 | 27 +18 | 28 +19 | 29 +20 | 30 + +Showing 20 rows out of 85 + +@ +-} +fold :: (a -> DataFrame -> DataFrame) -> [a] -> DataFrame -> DataFrame +fold f xs acc = L.foldl' (flip f) acc xs + +{- | Returns a dataframe as a two dimensional vector of floats. + +Converts all columns in the dataframe to float vectors and transposes them +into a row-major matrix representation. + +This is useful for handing data over into ML systems. + +Returns 'Left' with an error if any column cannot be converted to floats. +-} +toFloatMatrix :: + DataFrame -> Either DataFrameException (V.Vector (VU.Vector Float)) +toFloatMatrix df = case V.foldl' + (\acc c -> V.snoc <$> acc <*> toFloatVector c) + (Right V.empty :: Either DataFrameException (V.Vector (VU.Vector Float))) + (columns df) of + Left e -> Left e + Right m -> + pure $ + V.generate + (fst (dataframeDimensions df)) + ( \i -> + foldl + (\acc j -> acc `VU.snoc` ((m VG.! j) VG.! i)) + VU.empty + [0 .. (V.length m - 1)] + ) + +{- | Returns a dataframe as a two dimensional vector of doubles. + +Converts all columns in the dataframe to double vectors and transposes them +into a row-major matrix representation. + +This is useful for handing data over into ML systems. + +Returns 'Left' with an error if any column cannot be converted to doubles. +-} +toDoubleMatrix :: + DataFrame -> Either DataFrameException (V.Vector (VU.Vector Double)) +toDoubleMatrix df = case V.foldl' + (\acc c -> V.snoc <$> acc <*> toDoubleVector c) + (Right V.empty :: Either DataFrameException (V.Vector (VU.Vector Double))) + (columns df) of + Left e -> Left e + Right m -> + pure $ + V.generate + (fst (dataframeDimensions df)) + ( \i -> + foldl + (\acc j -> acc `VU.snoc` ((m VG.! j) VG.! i)) + VU.empty + [0 .. (V.length m - 1)] + ) + +{- | Returns a dataframe as a two dimensional vector of ints. + +Converts all columns in the dataframe to int vectors and transposes them +into a row-major matrix representation. + +This is useful for handing data over into ML systems. + +Returns 'Left' with an error if any column cannot be converted to ints. +-} +toIntMatrix :: DataFrame -> Either DataFrameException (V.Vector (VU.Vector Int)) +toIntMatrix df = case V.foldl' + (\acc c -> V.snoc <$> acc <*> toIntVector c) + (Right V.empty :: Either DataFrameException (V.Vector (VU.Vector Int))) + (columns df) of + Left e -> Left e + Right m -> + pure $ + V.generate + (fst (dataframeDimensions df)) + ( \i -> + foldl + (\acc j -> acc `VU.snoc` ((m VG.! j) VG.! i)) + VU.empty + [0 .. (V.length m - 1)] + ) + +{- | Get a specific column as a vector. + +You must specify the type via type applications. + +==== __Examples__ + +>>> columnAsVector (F.col @Int "age") df +Right [25, 30, 35, ...] + +>>> columnAsVector (F.col @Text "name") df +Right ["Alice", "Bob", "Charlie", ...] +-} +columnAsVector :: + forall a. + (Columnable a) => Expr a -> DataFrame -> Either DataFrameException (V.Vector a) +columnAsVector expr df + | null df = throw (EmptyDataSetException "columnAsVector") + | otherwise = case expr of + (Col name) -> case getColumn name df of + Just col -> toVector col + Nothing -> + Left $ + ColumnsNotFoundException [name] "columnAsVector" (M.keys $ columnIndices df) + _ -> case interpret df expr of + Left e -> throw e + Right (TColumn col) -> toVector col + +{- | Retrieves a column as an unboxed vector of 'Int' values. + +Returns 'Left' with a 'DataFrameException' if the column cannot be converted to ints. +This may occur if the column contains non-numeric data or values outside the 'Int' range. +-} +columnAsIntVector :: + (Columnable a, Num a) => + Expr a -> DataFrame -> Either DataFrameException (VU.Vector Int) +columnAsIntVector (Col name) df = case getColumn name df of + Just col -> toIntVector col + Nothing -> + Left $ + ColumnsNotFoundException [name] "columnAsIntVector" (M.keys $ columnIndices df) +columnAsIntVector expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> toIntVector col + +{- | Retrieves a column as an unboxed vector of 'Double' values. + +Returns 'Left' with a 'DataFrameException' if the column cannot be converted to doubles. +This may occur if the column contains non-numeric data. +-} +columnAsDoubleVector :: + (Columnable a, Num a) => + Expr a -> DataFrame -> Either DataFrameException (VU.Vector Double) +columnAsDoubleVector (Col name) df = case getColumn name df of + Just col -> toDoubleVector col + Nothing -> + Left $ + ColumnsNotFoundException + [name] + "columnAsDoubleVector" + (M.keys $ columnIndices df) +columnAsDoubleVector expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> toDoubleVector col + +{- | Retrieves a column as an unboxed vector of 'Float' values. + +Returns 'Left' with a 'DataFrameException' if the column cannot be converted to floats. +This may occur if the column contains non-numeric data. +-} +columnAsFloatVector :: + (Columnable a, Num a) => + Expr a -> DataFrame -> Either DataFrameException (VU.Vector Float) +columnAsFloatVector (Col name) df = case getColumn name df of + Just col -> toFloatVector col + Nothing -> + Left $ + ColumnsNotFoundException + [name] + "columnAsFloatVector" + (M.keys $ columnIndices df) +columnAsFloatVector expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> toFloatVector col + +columnAsUnboxedVector :: + forall a. + (Columnable a, VU.Unbox a) => + Expr a -> DataFrame -> Either DataFrameException (VU.Vector a) +columnAsUnboxedVector (Col name) df = case getColumn name df of + Just col -> toUnboxedVector col + Nothing -> + Left $ + ColumnsNotFoundException + [name] + "columnAsFloatVector" + (M.keys $ columnIndices df) +columnAsUnboxedVector expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> toUnboxedVector col +{-# SPECIALIZE columnAsUnboxedVector :: + Expr Double -> DataFrame -> Either DataFrameException (VU.Vector Double) + #-} +{-# INLINE columnAsUnboxedVector #-} + +{- | Get a specific column as a list. + +You must specify the type via type applications. + +==== __Examples__ + +>>> columnAsList @Int "age" df +[25, 30, 35, ...] + +>>> columnAsList @Text "name" df +["Alice", "Bob", "Charlie", ...] + +==== __Throws__ + +* 'error' - if the column type doesn't match the requested type +-} +columnAsList :: forall a. (Columnable a) => Expr a -> DataFrame -> [a] +columnAsList expr df = either throw V.toList (columnAsVector expr df) + +{- | Returns the provenance of all columns in the DataFrame as a list of +@(name, expression)@ pairs. Derived columns show their expression; +raw columns show an identity @col \@type name@ expression. +-} + +-- TODO: mchavinda - Expand out these expressions if possible. +showDerivedExpressions :: DataFrame -> [NamedExpr] +showDerivedExpressions df = + let exprs = derivingExpressions df + names = columnNames df + toNamedExpr name = case M.lookup name exprs of + Just uexpr -> (name, uexpr) + Nothing -> (name, identityUExpr name) + in map toNamedExpr names + where + identityUExpr name = case getColumn name df of + Just (BoxedColumn (Just _) (_ :: V.Vector a)) -> UExpr (Col @(Maybe a) name) + Just (BoxedColumn Nothing (_ :: V.Vector a)) -> UExpr (Col @a name) + Just (UnboxedColumn (Just _) (_ :: VU.Vector a)) -> UExpr (Col @(Maybe a) name) + Just (UnboxedColumn Nothing (_ :: VU.Vector a)) -> UExpr (Col @a name) + Nothing -> error $ "showDerivedExpressions: column not found: " ++ T.unpack name diff --git a/src/DataFrame/Operations/Join.hs b/src/DataFrame/Operations/Join.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Join.hs @@ -0,0 +1,1102 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE NumericUnderscores #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.Operations.Join where + +import Control.Applicative ((<|>)) +import Control.Exception (throw) +import Control.Monad (forM_, when) +import Control.Monad.ST (ST, runST) +import qualified Data.HashMap.Strict as HM +#if !MIN_VERSION_base(4,20,0) +import Data.List (foldl') +#endif +import qualified Data.Map.Strict as M +import Data.Maybe (fromMaybe) +import Data.STRef (newSTRef, readSTRef, writeSTRef) +import qualified Data.Set as S +import qualified Data.Text as T +import Data.Type.Equality (TestEquality (..)) +import qualified Data.Vector as VB +import qualified Data.Vector.Algorithms.Merge as VA +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import DataFrame.Errors ( + DataFrameException (ColumnsNotFoundException), + ) +import DataFrame.Internal.Column as D +import DataFrame.Internal.DataFrame as D +import DataFrame.Operations.Aggregation as D +import DataFrame.Operations.Core as D +import Type.Reflection + +-- | Equivalent to SQL join types. +data JoinType + = INNER + | LEFT + | RIGHT + | FULL_OUTER + deriving (Show) + +-- | Join two dataframes using SQL join semantics. +join :: + JoinType -> + [T.Text] -> + DataFrame -> -- Right hand side + DataFrame -> -- Left hand side + DataFrame +join INNER xs right = innerJoin xs right +join LEFT xs right = leftJoin xs right +join RIGHT xs right = rightJoin xs right +join FULL_OUTER xs right = fullOuterJoin xs right + +{- | Row-count threshold for the build side. +When the build side exceeds this, sort-merge join is used +instead of hash join to avoid L3 cache thrashing. +-} +joinStrategyThreshold :: Int +joinStrategyThreshold = 500_000 + +{- | A compact index mapping hash values to contiguous slices of +original row indices. All indices live in a single unboxed vector; +the HashMap stores @(offset, length)@ into that vector. +-} +data CompactIndex = CompactIndex + { ciSortedIndices :: {-# UNPACK #-} !(VU.Vector Int) + , ciOffsets :: !(HM.HashMap Int (Int, Int)) + } + +{- | Build a compact index from a vector of row hashes. +Sorts @(hash, originalIndex)@ pairs by hash, then scans for +contiguous runs to populate the offset map. +-} +buildCompactIndex :: VU.Vector Int -> CompactIndex +buildCompactIndex hashes = + let n = VU.length hashes + (sortedHashes, sortedIndices) = sortWithIndices hashes + !offs = buildOffsets sortedHashes n 0 HM.empty + in CompactIndex sortedIndices offs + where + buildOffsets :: + VU.Vector Int -> + Int -> + Int -> + HM.HashMap Int (Int, Int) -> + HM.HashMap Int (Int, Int) + buildOffsets !sh !n !i !acc + | i >= n = acc + | otherwise = + let !h = sh `VU.unsafeIndex` i + !end = findGroupEnd sh h (i + 1) n + in buildOffsets sh n end (HM.insert h (i, end - i) acc) + +-- | Find the end of a contiguous run of equal values starting at @j@. +findGroupEnd :: VU.Vector Int -> Int -> Int -> Int -> Int +findGroupEnd !v !h !j !n + | j >= n = j + | v `VU.unsafeIndex` j == h = findGroupEnd v h (j + 1) n + | otherwise = j +{-# INLINE findGroupEnd #-} + +{- | Sort a hash vector, returning sorted hashes and corresponding original indices. +Sorts an index array using hash values as the comparison key, avoiding the +intermediate pair vector used by the naive zip-then-sort approach. +-} +sortWithIndices :: VU.Vector Int -> (VU.Vector Int, VU.Vector Int) +sortWithIndices hashes = runST $ do + let n = VU.length hashes + mv <- VU.thaw (VU.enumFromN 0 n) + VA.sortBy + (\i j -> compare (hashes `VU.unsafeIndex` i) (hashes `VU.unsafeIndex` j)) + mv + sortedIdxs <- VU.unsafeFreeze mv + return (VU.unsafeBackpermute hashes sortedIdxs, sortedIdxs) + +-- | Write the cross product of two index ranges into mutable vectors. +fillCrossProduct :: + VU.Vector Int -> + VU.Vector Int -> + Int -> + Int -> + Int -> + Int -> + VUM.MVector s Int -> + VUM.MVector s Int -> + Int -> + ST s () +fillCrossProduct !leftSI !rightSI !lStart !lEnd !rStart !rEnd !lv !rv !pos = goL lStart pos + where + !rLen = rEnd - rStart + goL !li !p + | li >= lEnd = return () + | otherwise = do + let !lOrigIdx = leftSI `VU.unsafeIndex` li + goR lOrigIdx rStart p + goL (li + 1) (p + rLen) + goR !lOrigIdx !ri !q + | ri >= rEnd = return () + | otherwise = do + VUM.unsafeWrite lv q lOrigIdx + VUM.unsafeWrite rv q (rightSI `VU.unsafeIndex` ri) + goR lOrigIdx (ri + 1) (q + 1) +{-# INLINE fillCrossProduct #-} + +-- | Compute key-column indices from the column index map. +keyColIndices :: S.Set T.Text -> DataFrame -> [Int] +keyColIndices csSet df = M.elems $ M.restrictKeys (D.columnIndices df) csSet + +-- | Validate that all requested join keys exist, then return their indices. +validatedKeyColIndices :: T.Text -> S.Set T.Text -> DataFrame -> [Int] +validatedKeyColIndices callPoint csSet df = + let columnIdxs = D.columnIndices df + missingKeys = S.toAscList (csSet `S.difference` M.keysSet columnIdxs) + in case missingKeys of + [] -> M.elems $ M.restrictKeys columnIdxs csSet + _ -> throw (ColumnsNotFoundException missingKeys callPoint (M.keys columnIdxs)) + +-- ============================================================ +-- Inner Join +-- ============================================================ + +{- | Performs an inner join on two dataframes using the specified key columns. +Returns only rows where the key values exist in both dataframes. + +==== __Example__ +@ +ghci> df = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2", "K3"]), ("A", D.fromList ["A0", "A1", "A2", "A3"])] +ghci> other = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2"]), ("B", D.fromList ["B0", "B1", "B2"])] +ghci> D.innerJoin ["key"] df other + +----------------- + key | A | B +------|-----|---- + Text | Text| Text +------|-----|---- + K0 | A0 | B0 + K1 | A1 | B1 + K2 | A2 | B2 + +@ +-} +innerJoin :: [T.Text] -> DataFrame -> DataFrame -> DataFrame +innerJoin cs left right + | D.null right || D.null left = D.empty + | otherwise = innerJoinNonEmpty cs left right + +innerJoinNonEmpty :: [T.Text] -> DataFrame -> DataFrame -> DataFrame +innerJoinNonEmpty cs left right = + let + csSet = S.fromList cs + leftRows = fst (D.dimensions left) + rightRows = fst (D.dimensions right) + + leftKeyIdxs = validatedKeyColIndices "innerJoin" csSet left + rightKeyIdxs = validatedKeyColIndices "innerJoin" csSet right + leftHashes = D.computeRowHashes leftKeyIdxs left + rightHashes = D.computeRowHashes rightKeyIdxs right + + buildRows = min leftRows rightRows + (leftIxs, rightIxs) + | buildRows > joinStrategyThreshold = + sortMergeInnerKernel leftHashes rightHashes + | rightRows <= leftRows = + -- Build on right (smaller or equal), probe with left + hashInnerKernel leftHashes rightHashes + | otherwise = + -- Build on left (smaller), probe with right, swap result + let (!rIxs, !lIxs) = hashInnerKernel rightHashes leftHashes + in (lIxs, rIxs) + in + assembleInner csSet left right leftIxs rightIxs + +-- | Compute hashes for the given key column names in a DataFrame. +buildHashColumn :: [T.Text] -> DataFrame -> VU.Vector Int +buildHashColumn keys df = + let csSet = S.fromList keys + keyIdxs = validatedKeyColIndices "buildHashColumn" csSet df + in D.computeRowHashes keyIdxs df + +{- | Probe one batch of rows against a pre-built 'CompactIndex'. +Returns @(probeExpandedIxs, buildExpandedIxs)@. +Unlike 'hashInnerKernel', does not build the index (it is pre-built once) +and has no cross-product row guard — the caller controls probe batch size. +-} +hashProbeKernel :: + -- | Built once from the full right\/build side. + CompactIndex -> + -- | Probe hashes (one batch). + VU.Vector Int -> + (VU.Vector Int, VU.Vector Int) +hashProbeKernel ci probeHashes = + let ciIxs = ciSortedIndices ci + ciOff = ciOffsets ci + (pFrozen, bFrozen) = runST $ do + let !probeN = VU.length probeHashes + initCap = max 1 (min probeN 1_000_000) + + initPv <- VUM.unsafeNew initCap + initBv <- VUM.unsafeNew initCap + pvRef <- newSTRef initPv + bvRef <- newSTRef initBv + capRef <- newSTRef initCap + posRef <- newSTRef (0 :: Int) + + let ensureCapacity needed = do + cap <- readSTRef capRef + when (needed > cap) $ do + let newCap = max needed (cap * 2) + delta = newCap - cap + pv <- readSTRef pvRef + bv <- readSTRef bvRef + newPv <- VUM.unsafeGrow pv delta + newBv <- VUM.unsafeGrow bv delta + writeSTRef pvRef newPv + writeSTRef bvRef newBv + writeSTRef capRef newCap + + go !i + | i >= probeN = return () + | otherwise = do + let !h = probeHashes `VU.unsafeIndex` i + case HM.lookup h ciOff of + Nothing -> go (i + 1) + Just (!start, !len) -> do + !p <- readSTRef posRef + ensureCapacity (p + len) + pv <- readSTRef pvRef + bv <- readSTRef bvRef + fillBuild i start len p 0 pv bv + writeSTRef posRef (p + len) + go (i + 1) + fillBuild !probeIdx !start !len !p !j !pv !bv + | j >= len = return () + | otherwise = do + VUM.unsafeWrite pv (p + j) probeIdx + VUM.unsafeWrite bv (p + j) (ciIxs `VU.unsafeIndex` (start + j)) + fillBuild probeIdx start len p (j + 1) pv bv + go 0 + + !total <- readSTRef posRef + pv <- readSTRef pvRef + bv <- readSTRef bvRef + (,) + <$> VU.unsafeFreeze (VUM.slice 0 total pv) + <*> VU.unsafeFreeze (VUM.slice 0 total bv) + in (VU.force pFrozen, VU.force bFrozen) + +{- | Hash-based inner join kernel. +Builds compact index on @buildHashes@ (second arg), probes with +@probeHashes@ (first arg). +Returns @(probeExpandedIndices, buildExpandedIndices)@. +Uses a dynamically growing output buffer to avoid pre-allocating the full +cross-product size (which can be astronomically large for low-cardinality keys). +-} + +{- | Maximum number of output rows allowed from a join kernel. +Exceeding this limit indicates a cross-product explosion (e.g. low-cardinality keys). +-} +maxJoinOutputRows :: Int +maxJoinOutputRows = 500_000_000 + +hashInnerKernel :: + VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) +hashInnerKernel probeHashes buildHashes = + let (pFrozen, bFrozen) = runST $ do + let ci = buildCompactIndex buildHashes + ciIxs = ciSortedIndices ci + ciOff = ciOffsets ci + !probeN = VU.length probeHashes + !buildN = VU.length buildHashes + initCap = max 1 (min (probeN + buildN) 1_000_000) + + initPv <- VUM.unsafeNew initCap + initBv <- VUM.unsafeNew initCap + pvRef <- newSTRef initPv + bvRef <- newSTRef initBv + capRef <- newSTRef initCap + posRef <- newSTRef (0 :: Int) + + let ensureCapacity needed = do + cap <- readSTRef capRef + when (needed > cap) $ do + let newCap = max needed (cap * 2) + delta = newCap - cap + pv <- readSTRef pvRef + bv <- readSTRef bvRef + newPv <- VUM.unsafeGrow pv delta + newBv <- VUM.unsafeGrow bv delta + writeSTRef pvRef newPv + writeSTRef bvRef newBv + writeSTRef capRef newCap + + go !i + | i >= probeN = return () + | otherwise = do + let !h = probeHashes `VU.unsafeIndex` i + case HM.lookup h ciOff of + Nothing -> go (i + 1) + Just (!start, !len) -> do + !p <- readSTRef posRef + when (p + len > maxJoinOutputRows) $ + error $ + "Join output would exceed " + ++ show maxJoinOutputRows + ++ " rows (cross-product explosion). " + ++ "Consider filtering or using higher-cardinality join keys or using the lazy API." + ensureCapacity (p + len) + pv <- readSTRef pvRef + bv <- readSTRef bvRef + fillBuild i start len p 0 pv bv + writeSTRef posRef (p + len) + go (i + 1) + fillBuild !probeIdx !start !len !p !j !pv !bv + | j >= len = return () + | otherwise = do + VUM.unsafeWrite pv (p + j) probeIdx + VUM.unsafeWrite bv (p + j) (ciIxs `VU.unsafeIndex` (start + j)) + fillBuild probeIdx start len p (j + 1) pv bv + go 0 + + !total <- readSTRef posRef + pv <- readSTRef pvRef + bv <- readSTRef bvRef + (,) + <$> VU.unsafeFreeze (VUM.slice 0 total pv) + <*> VU.unsafeFreeze (VUM.slice 0 total bv) + in -- VU.force copies the slice into a compact array, releasing the oversized + -- backing buffer allocated by the doubling strategy. + (VU.force pFrozen, VU.force bFrozen) + +{- | Sort-merge inner join kernel. +Sorts both sides by hash, walks in lockstep. +Returns @(leftExpandedIndices, rightExpandedIndices)@. +Uses a dynamically growing output buffer instead of a two-pass count-then-allocate +strategy, which OOMs when low-cardinality keys produce large cross products. +-} +sortMergeInnerKernel :: + VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) +sortMergeInnerKernel leftHashes rightHashes = + let (lFrozen, rFrozen) = runST $ do + let (leftSH, leftSI) = sortWithIndices leftHashes + (rightSH, rightSI) = sortWithIndices rightHashes + !leftN = VU.length leftHashes + !rightN = VU.length rightHashes + initCap = max 1 (min (leftN + rightN) 1_000_000) + + initLv <- VUM.unsafeNew initCap + initRv <- VUM.unsafeNew initCap + lvRef <- newSTRef initLv + rvRef <- newSTRef initRv + capRef <- newSTRef initCap + posRef <- newSTRef (0 :: Int) + + let ensureCapacity needed = do + cap <- readSTRef capRef + when (needed > cap) $ do + let newCap = max needed (cap * 2) + delta = newCap - cap + lv <- readSTRef lvRef + rv <- readSTRef rvRef + newLv <- VUM.unsafeGrow lv delta + newRv <- VUM.unsafeGrow rv delta + writeSTRef lvRef newLv + writeSTRef rvRef newRv + writeSTRef capRef newCap + + fillGroup !li !lEnd !ri !rEnd = do + let !lLen = lEnd - li + !rLen = rEnd - ri + !groupSize = lLen * rLen + !p <- readSTRef posRef + when (p + groupSize > maxJoinOutputRows) $ + error $ + "Join output would exceed " + ++ show maxJoinOutputRows + ++ " rows (cross-product explosion with group sizes " + ++ show lLen + ++ " × " + ++ show rLen + ++ "). Consider filtering or using higher-cardinality join keys." + ensureCapacity (p + groupSize) + lv <- readSTRef lvRef + rv <- readSTRef rvRef + let goL !lIdx !pos + | lIdx >= lEnd = return () + | otherwise = do + let !lOrig = leftSI `VU.unsafeIndex` lIdx + goR lOrig ri pos + goL (lIdx + 1) (pos + rLen) + goR !lOrig !rIdx !pos + | rIdx >= rEnd = return () + | otherwise = do + VUM.unsafeWrite lv pos lOrig + VUM.unsafeWrite rv pos (rightSI `VU.unsafeIndex` rIdx) + goR lOrig (rIdx + 1) (pos + 1) + goL li p + writeSTRef posRef (p + groupSize) + + fill !li !ri + | li >= leftN || ri >= rightN = return () + | lh < rh = fill (li + 1) ri + | lh > rh = fill li (ri + 1) + | otherwise = do + let !lEnd = findGroupEnd leftSH lh (li + 1) leftN + !rEnd = findGroupEnd rightSH rh (ri + 1) rightN + fillGroup li lEnd ri rEnd + fill lEnd rEnd + where + !lh = leftSH `VU.unsafeIndex` li + !rh = rightSH `VU.unsafeIndex` ri + + fill 0 0 + + !total <- readSTRef posRef + lv <- readSTRef lvRef + rv <- readSTRef rvRef + (,) + <$> VU.unsafeFreeze (VUM.slice 0 total lv) + <*> VU.unsafeFreeze (VUM.slice 0 total rv) + in -- VU.force copies the slice into a compact array, releasing the oversized + -- backing buffer allocated by the doubling strategy. + (VU.force lFrozen, VU.force rFrozen) + +-- | Assemble the result DataFrame for an inner join from expanded index vectors. +assembleInner :: + S.Set T.Text -> + DataFrame -> + DataFrame -> + VU.Vector Int -> + VU.Vector Int -> + DataFrame +assembleInner csSet left right leftIxs rightIxs = + let !resultLen = VU.length leftIxs + leftColSet = S.fromList (D.columnNames left) + rightColNames = D.columnNames right + + -- Pre-expand every column once + expandedLeftCols = VB.map (D.atIndicesStable leftIxs) (D.columns left) + expandedRightCols = VB.map (D.atIndicesStable rightIxs) (D.columns right) + + getExpandedLeft name = do + idx <- M.lookup name (D.columnIndices left) + return (expandedLeftCols `VB.unsafeIndex` idx) + + getExpandedRight name = do + idx <- M.lookup name (D.columnIndices right) + return (expandedRightCols `VB.unsafeIndex` idx) + + -- Base DataFrame: all left columns, expanded + baseDf = + left + { columns = expandedLeftCols + , dataframeDimensions = (resultLen, snd (D.dataframeDimensions left)) + , derivingExpressions = M.empty + } + + insertIfPresent _ Nothing df = df + insertIfPresent name (Just c) df = D.insertColumn name c df + in D.fold + ( \name df -> + if S.member name csSet + then df -- Key column already present from left side + else + if S.member name leftColSet + then -- Overlapping non-key column: merge with These + insertIfPresent + name + (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name) + df + else -- Right-only column + insertIfPresent name (getExpandedRight name) df + ) + rightColNames + baseDf + +-- ============================================================ +-- Left Join +-- ============================================================ + +{- | Performs a left join on two dataframes using the specified key columns. +Returns all rows from the left dataframe, with matching rows from the right dataframe. +Non-matching rows will have Nothing/null values for columns from the right dataframe. + +==== __Example__ +@ +ghci> df = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2", "K3"]), ("A", D.fromList ["A0", "A1", "A2", "A3"])] +ghci> other = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2"]), ("B", D.fromList ["B0", "B1", "B2"])] +ghci> D.leftJoin ["key"] df other + +------------------------ + key | A | B +------|-----|---------- + Text | Text| Maybe Text +------|-----|---------- + K0 | A0 | Just "B0" + K1 | A1 | Just "B1" + K2 | A2 | Just "B2" + K3 | A3 | Nothing + +@ +-} +leftJoin :: [T.Text] -> DataFrame -> DataFrame -> DataFrame +leftJoin = leftJoinWithCallPoint "leftJoin" + +leftJoinWithCallPoint :: + T.Text -> [T.Text] -> DataFrame -> DataFrame -> DataFrame +leftJoinWithCallPoint callPoint cs left right + | D.null right || D.nRows right == 0 = left + | D.null left || D.nRows left == 0 = D.empty + | otherwise = leftJoinNonEmpty callPoint cs left right + +leftJoinNonEmpty :: T.Text -> [T.Text] -> DataFrame -> DataFrame -> DataFrame +leftJoinNonEmpty callPoint cs left right = + let + csSet = S.fromList cs + rightRows = fst (D.dimensions right) + + leftKeyIdxs = validatedKeyColIndices callPoint csSet left + rightKeyIdxs = validatedKeyColIndices callPoint csSet right + leftHashes = D.computeRowHashes leftKeyIdxs left + rightHashes = D.computeRowHashes rightKeyIdxs right + + -- Right is always the build side for left join + (leftIxs, rightIxs) + | rightRows > joinStrategyThreshold = + sortMergeLeftKernel leftHashes rightHashes + | otherwise = + hashLeftKernel leftHashes rightHashes + in + -- rightIxs uses -1 as sentinel for "no match" + assembleLeft csSet left right leftIxs rightIxs + +{- | Hash-based left join kernel. +Returns @(leftExpandedIndices, rightExpandedIndices)@ where +right indices use @-1@ as sentinel for unmatched rows. +Uses a dynamically growing output buffer to avoid pre-allocating the full +cross-product size (which can be astronomically large for low-cardinality keys). +-} +hashLeftKernel :: + VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) +hashLeftKernel leftHashes rightHashes = runST $ do + let ci = buildCompactIndex rightHashes + ciIxs = ciSortedIndices ci + ciOff = ciOffsets ci + !leftN = VU.length leftHashes + !rightN = VU.length rightHashes + initCap = max 1 (min (leftN + rightN) 1_000_000) + + initLv <- VUM.unsafeNew initCap + initRv <- VUM.unsafeNew initCap + lvRef <- newSTRef initLv + rvRef <- newSTRef initRv + capRef <- newSTRef initCap + posRef <- newSTRef (0 :: Int) + + let ensureCapacity needed = do + cap <- readSTRef capRef + when (needed > cap) $ do + let newCap = max needed (cap * 2) + delta = newCap - cap + lv <- readSTRef lvRef + rv <- readSTRef rvRef + newLv <- VUM.unsafeGrow lv delta + newRv <- VUM.unsafeGrow rv delta + writeSTRef lvRef newLv + writeSTRef rvRef newRv + writeSTRef capRef newCap + + go !i + | i >= leftN = return () + | otherwise = do + let !h = leftHashes `VU.unsafeIndex` i + !p <- readSTRef posRef + case HM.lookup h ciOff of + Nothing -> do + ensureCapacity (p + 1) + lv <- readSTRef lvRef + rv <- readSTRef rvRef + VUM.unsafeWrite lv p i + VUM.unsafeWrite rv p (-1) + writeSTRef posRef (p + 1) + Just (!start, !len) -> do + ensureCapacity (p + len) + lv <- readSTRef lvRef + rv <- readSTRef rvRef + fillBuild i start len p 0 lv rv + writeSTRef posRef (p + len) + go (i + 1) + fillBuild !leftIdx !start !len !p !j !lv !rv + | j >= len = return () + | otherwise = do + VUM.unsafeWrite lv (p + j) leftIdx + VUM.unsafeWrite rv (p + j) (ciIxs `VU.unsafeIndex` (start + j)) + fillBuild leftIdx start len p (j + 1) lv rv + go 0 + + !total <- readSTRef posRef + lv <- readSTRef lvRef + rv <- readSTRef rvRef + (,) + <$> VU.unsafeFreeze (VUM.slice 0 total lv) + <*> VU.unsafeFreeze (VUM.slice 0 total rv) + +{- | Sort-merge left join kernel. +Returns @(leftExpandedIndices, rightExpandedIndices)@ with @-1@ sentinel. +Uses a dynamically growing output buffer instead of a two-pass count-then-allocate +strategy, which OOMs when low-cardinality keys produce large cross products. +-} +sortMergeLeftKernel :: + VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) +sortMergeLeftKernel leftHashes rightHashes = runST $ do + let (leftSH, leftSI) = sortWithIndices leftHashes + (rightSH, rightSI) = sortWithIndices rightHashes + !leftN = VU.length leftHashes + !rightN = VU.length rightHashes + initCap = max 1 (min (leftN + rightN) 1_000_000) + + initLv <- VUM.unsafeNew initCap + initRv <- VUM.unsafeNew initCap + lvRef <- newSTRef initLv + rvRef <- newSTRef initRv + capRef <- newSTRef initCap + posRef <- newSTRef (0 :: Int) + + let ensureCapacity needed = do + cap <- readSTRef capRef + when (needed > cap) $ do + let newCap = max needed (cap * 2) + delta = newCap - cap + lv <- readSTRef lvRef + rv <- readSTRef rvRef + newLv <- VUM.unsafeGrow lv delta + newRv <- VUM.unsafeGrow rv delta + writeSTRef lvRef newLv + writeSTRef rvRef newRv + writeSTRef capRef newCap + + fillGroup !li !lEnd !ri !rEnd = do + let !lLen = lEnd - li + !rLen = rEnd - ri + !groupSize = lLen * rLen + !p <- readSTRef posRef + ensureCapacity (p + groupSize) + lv <- readSTRef lvRef + rv <- readSTRef rvRef + let goL !lIdx !pos + | lIdx >= lEnd = return () + | otherwise = do + let !lOrig = leftSI `VU.unsafeIndex` lIdx + goR lOrig ri pos + goL (lIdx + 1) (pos + rLen) + goR !lOrig !rIdx !pos + | rIdx >= rEnd = return () + | otherwise = do + VUM.unsafeWrite lv pos lOrig + VUM.unsafeWrite rv pos (rightSI `VU.unsafeIndex` rIdx) + goR lOrig (rIdx + 1) (pos + 1) + goL li p + writeSTRef posRef (p + groupSize) + + fill !li !ri + | li >= leftN = return () + | ri >= rightN = fillRemainingLeft li + | lh < rh = do + !p <- readSTRef posRef + ensureCapacity (p + 1) + lv <- readSTRef lvRef + rv <- readSTRef rvRef + VUM.unsafeWrite lv p (leftSI `VU.unsafeIndex` li) + VUM.unsafeWrite rv p (-1) + writeSTRef posRef (p + 1) + fill (li + 1) ri + | lh > rh = fill li (ri + 1) + | otherwise = do + let !lEnd = findGroupEnd leftSH lh (li + 1) leftN + !rEnd = findGroupEnd rightSH rh (ri + 1) rightN + fillGroup li lEnd ri rEnd + fill lEnd rEnd + where + !lh = leftSH `VU.unsafeIndex` li + !rh = rightSH `VU.unsafeIndex` ri + + fillRemainingLeft !i = do + let !remaining = leftN - i + when (remaining > 0) $ do + !p <- readSTRef posRef + ensureCapacity (p + remaining) + lv <- readSTRef lvRef + rv <- readSTRef rvRef + let go !j + | j >= remaining = return () + | otherwise = do + VUM.unsafeWrite lv (p + j) (leftSI `VU.unsafeIndex` (i + j)) + VUM.unsafeWrite rv (p + j) (-1) + go (j + 1) + go 0 + writeSTRef posRef (p + remaining) + + fill 0 0 + + !total <- readSTRef posRef + lv <- readSTRef lvRef + rv <- readSTRef rvRef + (,) + <$> VU.unsafeFreeze (VUM.slice 0 total lv) + <*> VU.unsafeFreeze (VUM.slice 0 total rv) + +{- | Assemble the result DataFrame for a left join. +Right index vectors use @-1@ sentinel, gathered via 'gatherWithSentinel'. +-} +assembleLeft :: + S.Set T.Text -> + DataFrame -> + DataFrame -> + VU.Vector Int -> + VU.Vector Int -> + DataFrame +assembleLeft csSet left right leftIxs rightIxs = + let !resultLen = VU.length leftIxs + leftColSet = S.fromList (D.columnNames left) + rightColNames = D.columnNames right + + expandedLeftCols = VB.map (D.atIndicesStable leftIxs) (D.columns left) + expandedRightCols = VB.map (D.gatherWithSentinel rightIxs) (D.columns right) + + getExpandedLeft name = do + idx <- M.lookup name (D.columnIndices left) + return (expandedLeftCols `VB.unsafeIndex` idx) + + getExpandedRight name = do + idx <- M.lookup name (D.columnIndices right) + return (expandedRightCols `VB.unsafeIndex` idx) + + baseDf = + left + { columns = expandedLeftCols + , dataframeDimensions = (resultLen, snd (D.dataframeDimensions left)) + , derivingExpressions = M.empty + } + + insertIfPresent _ Nothing df = df + insertIfPresent name (Just c) df = D.insertColumn name c df + in D.fold + ( \name df -> + if S.member name csSet + then df + else + if S.member name leftColSet + then + insertIfPresent + name + (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name) + df + else insertIfPresent name (getExpandedRight name) df + ) + rightColNames + baseDf + +{- | Performs a right join on two dataframes using the specified key columns. +Returns all rows from the right dataframe, with matching rows from the left dataframe. +Non-matching rows will have Nothing/null values for columns from the left dataframe. + +==== __Example__ +@ +ghci> df = D.fromNamedColumns [("key", D.fromList ["K0", "K1", "K2", "K3"]), ("A", D.fromList ["A0", "A1", "A2", "A3"])] +ghci> other = D.fromNamedColumns [("key", D.fromList ["K0", "K1"]), ("B", D.fromList ["B0", "B1"])] +ghci> D.rightJoin ["key"] df other + +----------------- + key | A | B +------|-----|---- + Text | Text| Text +------|-----|---- + K0 | A0 | B0 + K1 | A1 | B1 + +@ +-} +rightJoin :: + [T.Text] -> DataFrame -> DataFrame -> DataFrame +rightJoin cs left right = leftJoinWithCallPoint "rightJoin" cs right left + +fullOuterJoin :: + [T.Text] -> DataFrame -> DataFrame -> DataFrame +fullOuterJoin cs left right + | D.null right || D.nRows right == 0 = left + | D.null left || D.nRows left == 0 = right + | otherwise = fullOuterJoinNonEmpty cs left right + +fullOuterJoinNonEmpty :: [T.Text] -> DataFrame -> DataFrame -> DataFrame +fullOuterJoinNonEmpty cs left right = + let + csSet = S.fromList cs + leftRows = fst (D.dimensions left) + rightRows = fst (D.dimensions right) + + leftKeyIdxs = validatedKeyColIndices "fullOuterJoin" csSet left + rightKeyIdxs = validatedKeyColIndices "fullOuterJoin" csSet right + leftHashes = D.computeRowHashes leftKeyIdxs left + rightHashes = D.computeRowHashes rightKeyIdxs right + + -- Both sides can have nulls in full outer + (leftIxs, rightIxs) + | max leftRows rightRows > joinStrategyThreshold = + sortMergeFullOuterKernel leftHashes rightHashes + | otherwise = + hashFullOuterKernel leftHashes rightHashes + in + -- Both index vectors use -1 as sentinel + assembleFullOuter csSet left right leftIxs rightIxs + +{- | Hash-based full outer join kernel. +Builds compact indices on both sides. +Returns @(leftExpandedIndices, rightExpandedIndices)@ with @-1@ sentinels. +-} +hashFullOuterKernel :: + VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) +hashFullOuterKernel leftHashes rightHashes = runST $ do + let leftCI = buildCompactIndex leftHashes + rightCI = buildCompactIndex rightHashes + leftOff = ciOffsets leftCI + rightOff = ciOffsets rightCI + leftSI = ciSortedIndices leftCI + rightSI = ciSortedIndices rightCI + + -- Count: matched + left-only + right-only + let leftEntries = HM.toList leftOff + rightEntries = HM.toList rightOff + + !matchedCount = + foldl' + ( \acc (h, (_, ll)) -> + case HM.lookup h rightOff of + Nothing -> acc + Just (_, rl) -> acc + ll * rl + ) + 0 + leftEntries + + !leftOnlyCount = + foldl' + ( \acc (h, (_, ll)) -> + if HM.member h rightOff then acc else acc + ll + ) + 0 + leftEntries + + !rightOnlyCount = + foldl' + ( \acc (h, (_, rl)) -> + if HM.member h leftOff then acc else acc + rl + ) + 0 + rightEntries + + !totalCount = matchedCount + leftOnlyCount + rightOnlyCount + + lv <- VUM.unsafeNew totalCount + rv <- VUM.unsafeNew totalCount + posRef <- newSTRef (0 :: Int) + + -- Fill matched + left-only (iterate left keys) + forM_ leftEntries $ \(h, (lStart, lLen)) -> do + !p <- readSTRef posRef + case HM.lookup h rightOff of + Nothing -> do + -- Left-only rows + let goL !j !q + | j >= lLen = return () + | otherwise = do + VUM.unsafeWrite lv q (leftSI `VU.unsafeIndex` (lStart + j)) + VUM.unsafeWrite rv q (-1) + goL (j + 1) (q + 1) + goL 0 p + writeSTRef posRef (p + lLen) + Just (!rStart, !rLen) -> do + -- Cross product + fillCrossProduct + leftSI + rightSI + lStart + (lStart + lLen) + rStart + (rStart + rLen) + lv + rv + p + writeSTRef posRef (p + lLen * rLen) + + -- Fill right-only (iterate right keys not in left) + forM_ rightEntries $ \(h, (rStart, rLen)) -> + case HM.lookup h leftOff of + Just _ -> return () + Nothing -> do + !p <- readSTRef posRef + let goR !j !q + | j >= rLen = return () + | otherwise = do + VUM.unsafeWrite lv q (-1) + VUM.unsafeWrite rv q (rightSI `VU.unsafeIndex` (rStart + j)) + goR (j + 1) (q + 1) + goR 0 p + writeSTRef posRef (p + rLen) + + (,) <$> VU.unsafeFreeze lv <*> VU.unsafeFreeze rv + +{- | Sort-merge full outer join kernel. +Returns @(leftExpandedIndices, rightExpandedIndices)@ with @-1@ sentinels. +-} +sortMergeFullOuterKernel :: + VU.Vector Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int) +sortMergeFullOuterKernel leftHashes rightHashes = runST $ do + let (leftSH, leftSI) = sortWithIndices leftHashes + (rightSH, rightSI) = sortWithIndices rightHashes + !leftN = VU.length leftHashes + !rightN = VU.length rightHashes + + -- Pass 1: count + let countLoop !li !ri !c + | li >= leftN && ri >= rightN = c + | li >= leftN = c + (rightN - ri) + | ri >= rightN = c + (leftN - li) + | lh < rh = countLoop (li + 1) ri (c + 1) + | lh > rh = countLoop li (ri + 1) (c + 1) + | otherwise = + let !lEnd = findGroupEnd leftSH lh (li + 1) leftN + !rEnd = findGroupEnd rightSH rh (ri + 1) rightN + in countLoop lEnd rEnd (c + (lEnd - li) * (rEnd - ri)) + where + !lh = leftSH `VU.unsafeIndex` li + !rh = rightSH `VU.unsafeIndex` ri + !totalRows = countLoop 0 0 0 + + -- Pass 2: fill + lv <- VUM.unsafeNew totalRows + rv <- VUM.unsafeNew totalRows + + let fill !li !ri !pos + | li >= leftN && ri >= rightN = return () + | li >= leftN = fillRemainingRight ri pos + | ri >= rightN = fillRemainingLeft li pos + | lh < rh = do + VUM.unsafeWrite lv pos (leftSI `VU.unsafeIndex` li) + VUM.unsafeWrite rv pos (-1) + fill (li + 1) ri (pos + 1) + | lh > rh = do + VUM.unsafeWrite lv pos (-1) + VUM.unsafeWrite rv pos (rightSI `VU.unsafeIndex` ri) + fill li (ri + 1) (pos + 1) + | otherwise = do + let !lEnd = findGroupEnd leftSH lh (li + 1) leftN + !rEnd = findGroupEnd rightSH rh (ri + 1) rightN + !groupSize = (lEnd - li) * (rEnd - ri) + fillCrossProduct leftSI rightSI li lEnd ri rEnd lv rv pos + fill lEnd rEnd (pos + groupSize) + where + !lh = leftSH `VU.unsafeIndex` li + !rh = rightSH `VU.unsafeIndex` ri + + fillRemainingLeft !i !pos + | i >= leftN = return () + | otherwise = do + VUM.unsafeWrite lv pos (leftSI `VU.unsafeIndex` i) + VUM.unsafeWrite rv pos (-1) + fillRemainingLeft (i + 1) (pos + 1) + + fillRemainingRight !i !pos + | i >= rightN = return () + | otherwise = do + VUM.unsafeWrite lv pos (-1) + VUM.unsafeWrite rv pos (rightSI `VU.unsafeIndex` i) + fillRemainingRight (i + 1) (pos + 1) + + fill 0 0 0 + (,) <$> VU.unsafeFreeze lv <*> VU.unsafeFreeze rv + +{- | Assemble the result DataFrame for a full outer join. +Both index vectors use @-1@ sentinel; all columns gathered via +'gatherWithSentinel'. Key columns are coalesced (first non-null wins). +-} +assembleFullOuter :: + S.Set T.Text -> + DataFrame -> + DataFrame -> + VU.Vector Int -> + VU.Vector Int -> + DataFrame +assembleFullOuter csSet left right leftIxs rightIxs = + let !resultLen = VU.length leftIxs + leftColSet = S.fromList (D.columnNames left) + rightColNames = D.columnNames right + + expandedLeftCols = VB.map (D.gatherWithSentinel leftIxs) (D.columns left) + expandedRightCols = VB.map (D.gatherWithSentinel rightIxs) (D.columns right) + + getExpandedLeft name = do + idx <- M.lookup name (D.columnIndices left) + return (expandedLeftCols `VB.unsafeIndex` idx) + + getExpandedRight name = do + idx <- M.lookup name (D.columnIndices right) + return (expandedRightCols `VB.unsafeIndex` idx) + + baseDf = + left + { columns = expandedLeftCols + , dataframeDimensions = (resultLen, snd (D.dataframeDimensions left)) + , derivingExpressions = M.empty + } + + insertIfPresent _ Nothing df = df + insertIfPresent name (Just c) df = D.insertColumn name c df + + -- Coalesce two nullable columns: take first non-Nothing per row, + -- producing a non-optional column. + coalesceKeyColumn :: Column -> Column -> Column + coalesceKeyColumn + (BoxedColumn lBm (lCol :: VB.Vector a)) + (BoxedColumn rBm (rCol :: VB.Vector b)) = + case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> + let asMaybe bm = + VB.imap + ( \i v -> case bm of + Just bm' -> if bitmapTestBit bm' i then Just v else Nothing + Nothing -> Just v + ) + lMaybe = asMaybe lBm lCol + rMaybe = asMaybe rBm rCol + in D.fromVector $ + VB.zipWith + ( \l r -> + fromMaybe (error "fullOuterJoin: null on both sides of key column") (l <|> r) + ) + lMaybe + rMaybe + Nothing -> error "Cannot join columns of different types" + coalesceKeyColumn _ _ = error "fullOuterJoin: expected nullable column for key columns" + in D.fold + ( \name df -> + if S.member name csSet + then -- Key column: coalesce left and right + case (getExpandedLeft name, getExpandedRight name) of + (Just lc, Just rc) -> D.insertColumn name (coalesceKeyColumn lc rc) df + _ -> df + else + if S.member name leftColSet + then + insertIfPresent + name + (D.mergeColumns <$> getExpandedLeft name <*> getExpandedRight name) + df + else insertIfPresent name (getExpandedRight name) df + ) + rightColNames + baseDf diff --git a/src/DataFrame/Operations/Merge.hs b/src/DataFrame/Operations/Merge.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Merge.hs @@ -0,0 +1,73 @@ +{-# LANGUAGE InstanceSigs #-} +{-# OPTIONS_GHC -Wno-orphans #-} + +module DataFrame.Operations.Merge where + +import qualified Data.List as L +import qualified Data.Text as T +import qualified DataFrame.Internal.Column as D +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Operations.Core as D + +import Data.Maybe + +{- | Vertically merge two dataframes using shared columns. +Columns that exist in only one dataframe are padded with Nothing. +-} +instance Semigroup D.DataFrame where + (<>) :: D.DataFrame -> D.DataFrame -> D.DataFrame + (<>) a b = + let + addColumns a' b' df name + | snd (D.dimensions a') == 0 && snd (D.dimensions b') == 0 = df + | snd (D.dimensions a') == 0 = fromMaybe df $ do + col <- D.getColumn name b' + pure $ D.insertColumn name col df + | snd (D.dimensions b') == 0 = fromMaybe df $ do + col <- D.getColumn name a' + pure $ D.insertColumn name col df + | otherwise = + let + numRowsA = fst $ D.dimensions a' + numRowsB = fst $ D.dimensions b' + sumRows = numRowsA + numRowsB + + optA = D.getColumn name a' + optB = D.getColumn name b' + in + case optB of + Nothing -> case optA of + Nothing -> + -- N.B. this case should never happen, because we're dealing with columns coming from + -- union of column names of both dataframes. Nothing + Nothing would mean column + -- wasn't in either dataframe, which shouldn't happen + D.insertColumn name (D.fromList ([] :: [T.Text])) df + Just a'' -> + D.insertColumn name (D.expandColumn sumRows a'') df + Just b'' -> case optA of + Nothing -> + D.insertColumn name (D.leftExpandColumn sumRows b'') df + Just a'' -> + let concatedColumns = D.concatColumnsEither a'' b'' + in D.insertColumn name concatedColumns df + result = L.foldl' (addColumns a b) D.empty (D.columnNames a `L.union` D.columnNames b) + in + result + { D.derivingExpressions = D.derivingExpressions a <> D.derivingExpressions b + } + +instance Monoid D.DataFrame where + mempty = D.empty + +-- | Add two dataframes side by side/horizontally. +(|||) :: D.DataFrame -> D.DataFrame -> D.DataFrame +(|||) a b = + let result = + D.fold + (\name acc -> D.insertColumn name (D.unsafeGetColumn name b) acc) + (D.columnNames b) + a + in result + { D.derivingExpressions = + D.derivingExpressions result <> D.derivingExpressions b + } diff --git a/src/DataFrame/Operations/Permutation.hs b/src/DataFrame/Operations/Permutation.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Permutation.hs @@ -0,0 +1,168 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.Operations.Permutation where + +import qualified Data.List as L +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 + +import Control.Exception (throw) +import Control.Monad.ST (runST) +import Data.Type.Equality (testEquality, (:~:) (Refl)) +import Data.Vector.Internal.Check (HasCallStack) +import DataFrame.Errors (DataFrameException (..)) +import DataFrame.Internal.Column (Column (..), Columnable, atIndicesStable) +import DataFrame.Internal.DataFrame ( + DataFrame (..), + columnNames, + unsafeGetColumn, + ) +import DataFrame.Internal.Expression (Expr (Col), getColumns) +import DataFrame.Operations.Core (dimensions) +import DataFrame.Operations.Transformations (derive) +import System.Random (Random (randomR), RandomGen) +import Type.Reflection (typeRep) + +-- | Sort order taken as a parameter by the 'sortBy' function. +data SortOrder where + Asc :: (Columnable a, Ord a) => Expr a -> SortOrder + Desc :: (Columnable a, Ord a) => Expr a -> SortOrder + +instance Eq SortOrder where + (==) :: SortOrder -> SortOrder -> Bool + (==) (Asc _) (Asc _) = True + (==) (Desc _) (Desc _) = True + (==) _ _ = False + +sortOrderColumns :: SortOrder -> [T.Text] +sortOrderColumns (Asc e) = getColumns e +sortOrderColumns (Desc e) = getColumns e + +mustFlipCompare :: SortOrder -> Bool +mustFlipCompare (Asc _) = True +mustFlipCompare (Desc _) = False + +{- | Materialize any compound sort expressions into synthetic columns on +a working dataframe, returning rewritten 'SortOrder's that reference +those columns by name. +-} +prepareSortColumns :: [SortOrder] -> DataFrame -> ([SortOrder], DataFrame) +prepareSortColumns = go 0 + where + go _ [] acc = ([], acc) + go i (ord : rest) acc = + let (ord', acc') = materializeSortOrder i ord acc + (rest', acc'') = go (i + 1) rest acc' + in (ord' : rest', acc'') + +materializeSortOrder :: Int -> SortOrder -> DataFrame -> (SortOrder, DataFrame) +materializeSortOrder _ ord@(Asc (Col _)) df = (ord, df) +materializeSortOrder _ ord@(Desc (Col _)) df = (ord, df) +materializeSortOrder i (Asc (e :: Expr a)) df = + let name = syntheticName i + in (Asc (Col name :: Expr a), derive name e df) +materializeSortOrder i (Desc (e :: Expr a)) df = + let name = syntheticName i + in (Desc (Col name :: Expr a), derive name e df) + +syntheticName :: Int -> T.Text +syntheticName i = "__sortBy_synthetic_" <> T.pack (show i) <> "__" + +{- | O(k log n) Sorts the dataframe by a given row. + +> sortBy Ascending ["Age"] df +-} +sortBy :: + [SortOrder] -> + DataFrame -> + DataFrame +sortBy sortOrds df + | not (null missing) = + throw $ + ColumnsNotFoundException + missing + "sortBy" + (columnNames df) + | otherwise = + let + (sortOrds', df') = prepareSortColumns sortOrds df + comparators = map (`sortOrderComparator` df') sortOrds' + compositeCompare i j = mconcat [c i j | c <- comparators] + nRows = fst (dataframeDimensions df') + indexes = sortIndices compositeCompare nRows + in + df{columns = V.map (atIndicesStable indexes) (columns df)} + where + referenced = L.nub (concatMap sortOrderColumns sortOrds) + missing = referenced L.\\ columnNames df + +{- | Build a row-index comparator from a SortOrder and a DataFrame. +The Ord dictionary is recovered from the SortOrder GADT. +-} +sortOrderComparator :: SortOrder -> DataFrame -> Int -> Int -> Ordering +sortOrderComparator (Asc (Col name :: Expr a)) df = + case unsafeGetColumn name df of + BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> \i j -> compare (v `V.unsafeIndex` i) (v `V.unsafeIndex` j) + Nothing -> \_ _ -> EQ + UnboxedColumn _ (v :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> \i j -> compare (v `VU.unsafeIndex` i) (v `VU.unsafeIndex` j) + Nothing -> \_ _ -> EQ +sortOrderComparator (Desc (Col name :: Expr a)) df = + case unsafeGetColumn name df of + BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> \i j -> compare (v `V.unsafeIndex` j) (v `V.unsafeIndex` i) + Nothing -> \_ _ -> EQ + UnboxedColumn _ (v :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> \i j -> compare (v `VU.unsafeIndex` j) (v `VU.unsafeIndex` i) + Nothing -> \_ _ -> EQ +sortOrderComparator _ _ = error "Sorting on compound column" + +-- | Sort row indices using a comparator function. +sortIndices :: (Int -> Int -> Ordering) -> Int -> VU.Vector Int +sortIndices cmp nRows = runST $ do + withIndexes <- VG.thaw (V.generate nRows id :: V.Vector Int) + VA.sortBy cmp withIndexes + sorted <- VG.unsafeFreeze withIndexes + return (VU.convert sorted) + +shuffle :: + (RandomGen g) => + g -> + DataFrame -> + DataFrame +shuffle pureGen df = + let + indexes = shuffledIndices pureGen (fst (dimensions df)) + in + df{columns = V.map (atIndicesStable indexes) (columns df)} + +shuffledIndices :: (HasCallStack, RandomGen g) => g -> Int -> VU.Vector Int +shuffledIndices pureGen k + | k < 0 = error $ "Vector index may not be a neative number: " <> show k + | k == 0 = VU.empty + | otherwise = shuffleVec pureGen + where + shuffleVec :: (RandomGen g) => g -> VU.Vector Int + shuffleVec g = runST $ do + vm <- VUM.generate k id + let (n, nGen) = randomR (1, k - 1) g + go vm n nGen + VU.unsafeFreeze vm + + go _v (-1) _ = pure () + go _v 0 _ = pure () + go v maxInd gen = + let + (n, nextGen) = randomR (1, maxInd) gen + in + VUM.swap v 0 n *> go (VUM.tail v) (maxInd - 1) nextGen diff --git a/src/DataFrame/Operations/Statistics.hs b/src/DataFrame/Operations/Statistics.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Statistics.hs @@ -0,0 +1,399 @@ +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -Wno-orphans #-} + +module DataFrame.Operations.Statistics where + +import qualified Data.List as L +import qualified Data.Map as M +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Unboxed as VU + +import Prelude as P + +import Control.Exception (throw) +import Data.Function ((&)) +import Data.Maybe (fromMaybe, isJust) +import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) +import DataFrame.Errors (DataFrameException (..)) +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame ( + DataFrame (..), + columnNames, + empty, + fromNamedColumns, + getColumn, + ) +import DataFrame.Internal.Expression +import DataFrame.Internal.Interpreter +import DataFrame.Internal.Nullable (BaseType) +import DataFrame.Internal.Row (showValue, toAny) +import DataFrame.Internal.Statistics +import DataFrame.Internal.Types +import DataFrame.Operations.Core +import DataFrame.Operations.Subset (filterJust) +import DataFrame.Operations.Transformations (ImputeOp (..), imputeCore) +import Text.Printf (printf) +import Type.Reflection (typeRep) + +{- | Show a frequency table for a categorical feaure. + +__Examples:__ + +@ +ghci> df <- D.readCsv ".\/data\/housing.csv" + +ghci> D.frequencies "ocean_proximity" df + +--------------------------------------------------------------------- + Statistic | <1H OCEAN | INLAND | ISLAND | NEAR BAY | NEAR OCEAN +----------------|-----------|--------|--------|----------|----------- + Text | Any | Any | Any | Any | Any +----------------|-----------|--------|--------|----------|----------- + Count | 9136 | 6551 | 5 | 2290 | 2658 + Percentage (%) | 44.26% | 31.74% | 0.02% | 11.09% | 12.88% +@ +-} +frequencies :: + forall a. (Columnable a, Ord a) => Expr a -> DataFrame -> DataFrame +frequencies expr df = + let + counts = valueCounts expr df + calculatePercentage cs k = toAny $ toPct2dp (fromIntegral k / fromIntegral (P.sum $ map snd cs)) + initDf = + empty + & insertVector "Statistic" (V.fromList ["Count" :: T.Text, "Percentage (%)"]) + freqs _col' = + L.foldl' + ( \d (col'', k) -> + insertVector + (showValue @a col'') + (V.fromList [toAny k, calculatePercentage counts k]) + d + ) + initDf + counts + in + case columnAsVector expr df of + Left err -> throw err + Right column -> freqs column + +-- | Calculates the mean of a given column as a standalone value. +mean :: + forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double +mean (Col name) df = case _getColumnAsDouble name df of + Just xs -> meanDouble' xs + Nothing -> error "[INTERNAL ERROR] Column is non-numeric" +mean expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> case toUnboxedVector @a col of + Left e -> throw e + Right xs -> mean' xs + +meanMaybe :: + forall a. (Columnable a, Real a) => Expr (Maybe a) -> DataFrame -> Double +meanMaybe (Col name) df = + (mean' . optionalToDoubleVector) + (either throw id (columnAsVector (Col @(Maybe a) name) df)) +meanMaybe expr df = case interpret @(Maybe a) df expr of + Left e -> throw e + Right (TColumn col) -> case toVector @(Maybe a) col of + Left e -> throw e + Right xs -> (mean' . optionalToDoubleVector) xs + +-- | Calculates the median of a given column as a standalone value. +median :: + forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double +median (Col name) df = case columnAsUnboxedVector (Col @a name) df of + Right xs -> median' xs + Left e -> throw e +median expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> case toUnboxedVector @a col of + Left e -> throw e + Right xs -> median' xs + +-- | Calculates the median of a given column (containing optional values) as a standalone value. +medianMaybe :: + forall a. (Columnable a, Real a) => Expr (Maybe a) -> DataFrame -> Double +medianMaybe (Col name) df = + (median' . optionalToDoubleVector) + (either throw id (columnAsVector (Col @(Maybe a) name) df)) +medianMaybe expr df = case interpret @(Maybe a) df expr of + Left e -> throw e + Right (TColumn col) -> case toVector @(Maybe a) col of + Left e -> throw e + Right xs -> (median' . optionalToDoubleVector) xs + +-- | Calculates the nth percentile of a given column as a standalone value. +percentile :: + forall a. + (Columnable a, Real a, VU.Unbox a) => Int -> Expr a -> DataFrame -> Double +percentile n (Col name) df = case columnAsUnboxedVector (Col @a name) df of + Right xs -> percentile' n xs + Left e -> throw e +percentile n expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> case toUnboxedVector @a col of + Left e -> throw e + Right xs -> percentile' n xs + +-- | Calculates the nth percentile of a given column as a standalone value. +genericPercentile :: + forall a. + (Columnable a, Ord a) => Int -> Expr a -> DataFrame -> a +genericPercentile n (Col name) df = case columnAsVector (Col @a name) df of + Right xs -> percentileOrd' n xs + Left e -> throw e +genericPercentile n expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> case toVector @a col of + Left e -> throw e + Right xs -> percentileOrd' n xs + +-- | Calculates the standard deviation of a given column as a standalone value. +standardDeviation :: + forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double +standardDeviation (Col name) df = case columnAsUnboxedVector (Col @a name) df of + Right xs -> (sqrt . variance') xs + Left e -> throw e +standardDeviation expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> case toUnboxedVector @a col of + Left e -> throw e + Right xs -> (sqrt . variance') xs + +-- | Calculates the skewness of a given column as a standalone value. +skewness :: + forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double +skewness (Col name) df = case columnAsUnboxedVector (Col @a name) df of + Right xs -> skewness' xs + Left e -> throw e +skewness expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> case toUnboxedVector @a col of + Left e -> throw e + Right xs -> skewness' xs + +-- | Calculates the variance of a given column as a standalone value. +variance :: + forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double +variance (Col name) df = case _getColumnAsDouble name df of + Just xs -> varianceDouble' xs + Nothing -> error "[INTERNAL ERROR] Column is non-numeric" +variance expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> case toUnboxedVector @a col of + Left e -> throw e + Right xs -> variance' xs + +-- | Calculates the inter-quartile range of a given column as a standalone value. +interQuartileRange :: + forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> DataFrame -> Double +interQuartileRange (Col name) df = case columnAsUnboxedVector (Col @a name) df of + Right xs -> interQuartileRange' xs + Left e -> throw e +interQuartileRange expr df = case interpret df expr of + Left e -> throw e + Right (TColumn col) -> case toUnboxedVector @a col of + Left e -> throw e + Right xs -> interQuartileRange' xs + +-- | Calculates the Pearson's correlation coefficient between two given columns as a standalone value. +correlation :: T.Text -> T.Text -> DataFrame -> Maybe Double +correlation first second df = do + f <- _getColumnAsDouble first df + s <- _getColumnAsDouble second df + correlation' f s + +_getColumnAsDouble :: T.Text -> DataFrame -> Maybe (VU.Vector Double) +_getColumnAsDouble name df = case getColumn name df of + Just (UnboxedColumn _ (f :: VU.Vector a)) -> case testEquality (typeRep @a) (typeRep @Double) of + Just Refl -> Just f + Nothing -> case sIntegral @a of + STrue -> Just (VU.map fromIntegral f) + SFalse -> case sFloating @a of + STrue -> Just (VU.map realToFrac f) + SFalse -> Nothing + Nothing -> + throw $ + ColumnsNotFoundException [name] "_getColumnAsDouble" (M.keys $ columnIndices df) + _ -> Nothing -- Return a type mismatch error here. +{-# INLINE _getColumnAsDouble #-} + +optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double +optionalToDoubleVector = + VU.fromList + . V.foldl' + (\acc e -> if isJust e then realToFrac (fromMaybe 0 e) : acc else acc) + [] + +-- | Calculates the sum of a given column as a standalone value. +sum :: + forall a. (Columnable a, Num a) => Expr a -> DataFrame -> a +sum (Col name) df = case getColumn name df of + Nothing -> throw $ ColumnsNotFoundException [name] "sum" (M.keys $ columnIndices df) + Just ((UnboxedColumn _ (column :: VU.Vector a'))) -> case testEquality (typeRep @a') (typeRep @a) of + Just Refl -> VG.sum column + Nothing -> 0 + Just ((BoxedColumn _ (column :: V.Vector a'))) -> case testEquality (typeRep @a') (typeRep @a) of + Just Refl -> VG.sum column + Nothing -> 0 +sum expr df = case interpret df expr of + Left e -> throw e + Right (TColumn xs) -> case toVector @a @V.Vector xs of + Left e -> throw e + Right xs' -> VG.sum xs' + +{- | /O(n)/ Impute missing values in a column using a derived scalar. + +Given + +* an expression @f :: 'Expr' b -> 'Expr' b@ that, when interpreted over a + non-nullable column, produces the same value in every row (for example a + mean, median, or other aggregate), and +* a nullable column @'Expr' ('Maybe' b)@ + +this function: + +1. Drops all @Nothing@ values from the target column. +2. Interprets @f@ on the remaining non-null values. +3. Checks that the resulting column contains a single repeated value. +4. Uses that value to impute all @Nothing@s in the original column. + +==== __Throws__ + +* 'DataFrameException' - if the column does not exist, is empty, + +==== __Example__ +@ +>>> :set -XOverloadedStrings +>>> import qualified DataFrame as D +>>> let df = +... D.fromNamedColumns +... [ ("age", D.fromList [Just 10, Nothing, Just 20 :: Maybe Int]) ] +>>> +>>> -- Impute missing ages with the mean of the observed ages +>>> D.imputeWith F.mean "age" df +-- age +-- ---- +-- 10 +-- 15 +-- 20 +@ +-} +instance {-# OVERLAPPING #-} (Columnable b) => ImputeOp (Maybe b) where + runImpute = imputeCore + + runImputeWith f col@(Col columnName) df = + case interpret @b (filterJust columnName df) (f (Col @b columnName)) of + Left e -> throw e + Right (TColumn value) -> case headColumn @b value of + Left e -> throw e + Right h -> + if all (== h) (toList @b value) + then imputeCore col h df + else error "Impute expression returned more than one value" + runImputeWith _ _ df = df + +imputeWith :: + forall a. + (ImputeOp a, Columnable (BaseType a)) => + (Expr (BaseType a) -> Expr (BaseType a)) -> + Expr a -> + DataFrame -> + DataFrame +imputeWith = runImputeWith + +applyStatistic :: + (VU.Vector Double -> Double) -> T.Text -> DataFrame -> Maybe Double +applyStatistic f name df = apply =<< _getColumnAsDouble name (filterJust name df) + where + apply col = + let + res = f col + in + if isNaN res then Nothing else pure res +{-# INLINE applyStatistic #-} + +applyStatistics :: + (VU.Vector Double -> VU.Vector Double) -> + T.Text -> + DataFrame -> + Maybe (VU.Vector Double) +applyStatistics f name df = fmap f (_getColumnAsDouble name (filterJust name df)) + +-- | Descriptive statistics of the numeric columns. +summarize :: DataFrame -> DataFrame +summarize df = + fold + columnStats + (columnNames df) + ( fromNamedColumns + [ + ( "Statistic" + , fromList + [ "Count" :: T.Text + , "Mean" + , "Minimum" + , "25%" + , "Median" + , "75%" + , "Max" + , "StdDev" + , "IQR" + , "Skewness" + ] + ) + ] + ) + where + columnStats name d = + if all isJust (stats name) + then + insertUnboxedVector + name + (VU.fromList (map (roundTo 2 . fromMaybe 0) $ stats name)) + d + else d + stats name = + let + count = fromIntegral . numElements <$> getColumn name df + quantiles = applyStatistics (quantiles' (VU.fromList [0, 1, 2, 3, 4]) 4) name df + min' = flip (VG.!) 0 <$> quantiles + quartile1 = flip (VG.!) 1 <$> quantiles + medianVal = flip (VG.!) 2 <$> quantiles + quartile3 = flip (VG.!) 3 <$> quantiles + max' = flip (VG.!) 4 <$> quantiles + iqr = (-) <$> quartile3 <*> quartile1 + doubleColumn col = _getColumnAsDouble col (filterJust col df) + in + [ count + , mean' <$> doubleColumn name + , min' + , quartile1 + , medianVal + , quartile3 + , max' + , sqrt . variance' <$> doubleColumn name + , iqr + , skewness' <$> doubleColumn name + ] + +-- | Round a @Double@ to Specified Precision +roundTo :: Int -> Double -> Double +roundTo n x = fromInteger (round $ x * 10 ^ n) / 10.0 ^^ n + +toPct2dp :: Double -> String +toPct2dp x + | x < 0.00005 = "<0.01%" + | otherwise = printf "%.2f%%" (x * 100) diff --git a/src/DataFrame/Operations/Subset.hs b/src/DataFrame/Operations/Subset.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Subset.hs @@ -0,0 +1,542 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.Operations.Subset where + +import qualified Data.List as L +import qualified Data.Map as M +import qualified Data.Set as S +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Unboxed as VU +import qualified Prelude + +import Control.Exception (throw) +import Data.Function ((&)) +import Data.Maybe ( + fromJust, + fromMaybe, + isJust, + isNothing, + ) +import Data.Type.Equality (TestEquality (..)) +import DataFrame.Errors ( + DataFrameException (..), + TypeErrorContext (..), + ) +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame ( + DataFrame (..), + columnNames, + derivingExpressions, + empty, + getColumn, + insertColumn, + unsafeGetColumn, + ) +import DataFrame.Internal.Expression +import DataFrame.Internal.Interpreter +import DataFrame.Operations.Core +import DataFrame.Operations.Merge () +import DataFrame.Operations.Transformations (apply) +import DataFrame.Operators +import System.Random +import Type.Reflection +import Prelude hiding (filter, take) + +#if MIN_VERSION_random(1,3,0) +type SplittableGen g = (SplitGen g, RandomGen g) + +splitForStratified :: SplittableGen g => g -> (g, g) +splitForStratified = splitGen +#else +type SplittableGen g = RandomGen g + +splitForStratified :: SplittableGen g => g -> (g, g) +splitForStratified = split +#endif + +-- | O(k * n) Take the first n rows of a DataFrame. +take :: Int -> DataFrame -> DataFrame +take n d = d{columns = V.map (takeColumn n') (columns d), dataframeDimensions = (n', c)} + where + (r, c) = dataframeDimensions d + n' = clip n 0 r + +-- | O(k * n) Take the last n rows of a DataFrame. +takeLast :: Int -> DataFrame -> DataFrame +takeLast n d = + d + { columns = V.map (takeLastColumn n') (columns d) + , dataframeDimensions = (n', c) + } + where + (r, c) = dataframeDimensions d + n' = clip n 0 r + +-- | O(k * n) Drop the first n rows of a DataFrame. +drop :: Int -> DataFrame -> DataFrame +drop n d = + d + { columns = V.map (sliceColumn n' (max (r - n') 0)) (columns d) + , dataframeDimensions = (max (r - n') 0, c) + } + where + (r, c) = dataframeDimensions d + n' = clip n 0 r + +-- | O(k * n) Drop the last n rows of a DataFrame. +dropLast :: Int -> DataFrame -> DataFrame +dropLast n d = + d{columns = V.map (sliceColumn 0 n') (columns d), dataframeDimensions = (n', c)} + where + (r, c) = dataframeDimensions d + n' = clip (r - n) 0 r + +-- | O(k * n) Take a range of rows of a DataFrame. +range :: (Int, Int) -> DataFrame -> DataFrame +range (start, end) d = + d + { columns = V.map (sliceColumn (clip start 0 r) n') (columns d) + , dataframeDimensions = (n', c) + } + where + (r, c) = dataframeDimensions d + n' = clip (end - start) 0 r + +clip :: Int -> Int -> Int -> Int +clip n left right = min right $ max n left + +{- | O(n * k) Filter rows by a given condition. + +> filter "x" even df +-} +filter :: + forall a. + (Columnable a) => + -- | Column to filter by + Expr a -> + -- | Filter condition + (a -> Bool) -> + -- | Dataframe to filter + DataFrame -> + DataFrame +filter (Col filterColumnName) condition df = case getColumn filterColumnName df of + Nothing -> + throw $ + ColumnsNotFoundException [filterColumnName] "filter" (M.keys $ columnIndices df) + Just _col@(BoxedColumn bm (column :: V.Vector b)) -> + -- Check direct type match first, then try Maybe b match for nullable columns + case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> filterByVector filterColumnName column condition df + Nothing -> case (bm, typeRep @a) of + (Just bm', App tMaybe tInner) -> case eqTypeRep tMaybe (typeRep @Maybe) of + Just HRefl -> case testEquality tInner (typeRep @b) of + Just Refl -> + let maybeVec = V.imap (\i v -> if bitmapTestBit bm' i then Just v else Nothing) column + in filterByVector filterColumnName maybeVec condition df + Nothing -> filterByVector filterColumnName column condition df + Nothing -> filterByVector filterColumnName column condition df + _ -> filterByVector filterColumnName column condition df + Just _col@(UnboxedColumn bm (column :: VU.Vector b)) -> + case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> filterByVector filterColumnName column condition df + Nothing -> case (bm, typeRep @a) of + (Just bm', App tMaybe tInner) -> case eqTypeRep tMaybe (typeRep @Maybe) of + Just HRefl -> case testEquality tInner (typeRep @b) of + Just Refl -> + let maybeVec = V.generate (VU.length column) $ \i -> + if bitmapTestBit bm' i then Just (VU.unsafeIndex column i) else Nothing + in filterByVector filterColumnName maybeVec condition df + Nothing -> filterByVector filterColumnName column condition df + Nothing -> filterByVector filterColumnName column condition df + _ -> filterByVector filterColumnName column condition df +filter expr condition df = + let + (TColumn col') = case interpret @a df (normalize expr) of + Left e -> throw e + Right c -> c + indexes = case findIndices condition col' of + Right ixs -> ixs + Left e -> throw e + c' = snd $ dataframeDimensions df + in + df + { columns = V.map (atIndicesStable indexes) (columns df) + , dataframeDimensions = (VU.length indexes, c') + } + +filterByVector :: + forall a b v. + (VG.Vector v b, VG.Vector v Int, Columnable a, Columnable b) => + T.Text -> v b -> (a -> Bool) -> DataFrame -> DataFrame +filterByVector filterColumnName column condition df = case testEquality (typeRep @a) (typeRep @b) of + Nothing -> + throw $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right $ typeRep @a + , expectedType = Right $ typeRep @b + , errorColumnName = Just (T.unpack filterColumnName) + , callingFunctionName = Just "filter" + } + ) + Just Refl -> + let + ixs = VG.convert (VG.findIndices condition column) + in + df + { columns = V.map (atIndicesStable ixs) (columns df) + , dataframeDimensions = (VG.length ixs, snd (dataframeDimensions df)) + } + +{- | O(k) a version of filter where the predicate comes first. + +> filterBy even "x" df +-} +filterBy :: (Columnable a) => (a -> Bool) -> Expr a -> DataFrame -> DataFrame +filterBy = flip filter + +{- | O(k) filters the dataframe with a boolean expression. + +> filterWhere (F.col @Int x + F.col y F.> 5) df +-} +filterWhere :: Expr Bool -> DataFrame -> DataFrame +filterWhere expr df = + let + (TColumn col') = case interpret @Bool df (normalize expr) of + Left e -> throw e + Right c -> c + indexes = case findIndices id col' of + Right ixs -> ixs + Left e -> throw e + c' = snd $ dataframeDimensions df + in + df + { columns = V.map (atIndicesStable indexes) (columns df) + , dataframeDimensions = (VU.length indexes, c') + } + +{- | O(k) removes all rows with `Nothing` in a given column from the dataframe. + +> filterJust "col" df +-} +filterJust :: T.Text -> DataFrame -> DataFrame +filterJust colName df = case getColumn colName df of + Nothing -> + throw $ + ColumnsNotFoundException [colName] "filterJust" (M.keys $ columnIndices df) + Just column | hasMissing column -> case column of + BoxedColumn (Just _) (_col :: V.Vector a) -> + filter (Col @(Maybe a) colName) isJust df & apply @(Maybe a) fromJust colName + UnboxedColumn (Just _) (_col :: VU.Vector a) -> + filter (Col @(Maybe a) colName) isJust df & apply @(Maybe a) fromJust colName + _ -> df + Just _ -> df + +{- | O(k) returns all rows with `Nothing` in a give column. + +> filterNothing "col" df +-} +filterNothing :: T.Text -> DataFrame -> DataFrame +filterNothing colName df = case getColumn colName df of + Nothing -> + throw $ + ColumnsNotFoundException [colName] "filterNothing" (M.keys $ columnIndices df) + Just column | hasMissing column -> case column of + BoxedColumn (Just _) (_col :: V.Vector a) -> filter (Col @(Maybe a) colName) isNothing df + UnboxedColumn (Just _) (_col :: VU.Vector a) -> filter (Col @(Maybe a) colName) isNothing df + _ -> df + _ -> df + +{- | O(n * k) removes all rows with `Nothing` from the dataframe. + +> filterAllJust df +-} +filterAllJust :: DataFrame -> DataFrame +filterAllJust df = foldr filterJust df (columnNames df) + +{- | O(n * k) keeps any row with a null value. + +> filterAllNothing df +-} +filterAllNothing :: DataFrame -> DataFrame +filterAllNothing df = foldr filterNothing df (columnNames df) + +{- | O(k) cuts the dataframe in a cube of size (a, b) where + a is the length and b is the width. + +> cube (10, 5) df +-} +cube :: (Int, Int) -> DataFrame -> DataFrame +cube (len, width) = take len . selectBy [ColumnIndexRange (0, width - 1)] + +{- | O(n) Selects a number of columns in a given dataframe. + +> select ["name", "age"] df +-} +select :: + [T.Text] -> + DataFrame -> + DataFrame +select cs df + | L.null cs = empty + | any (`notElem` columnNames df) cs = + throw $ + ColumnsNotFoundException + (cs L.\\ columnNames df) + "select" + (columnNames df) + | otherwise = + let result = L.foldl' addKeyValue empty cs + filteredExprs = M.filterWithKey (\k _ -> k `L.elem` cs) (derivingExpressions df) + in result{derivingExpressions = filteredExprs} + where + addKeyValue d k = fromMaybe df $ do + col' <- getColumn k df + pure $ insertColumn k col' d + +data SelectionCriteria + = ColumnProperty (Column -> Bool) + | ColumnNameProperty (T.Text -> Bool) + | ColumnTextRange (T.Text, T.Text) + | ColumnIndexRange (Int, Int) + | ColumnName T.Text + +{- | Criteria for selecting a column by name. + +> selectBy [byName "Age"] df + +equivalent to: + +> select ["Age"] df +-} +byName :: T.Text -> SelectionCriteria +byName = ColumnName + +{- | Criteria for selecting columns whose property satisfies given predicate. + +> selectBy [byProperty isNumeric] df +-} +byProperty :: (Column -> Bool) -> SelectionCriteria +byProperty = ColumnProperty + +{- | Criteria for selecting columns whose name satisfies given predicate. + +> selectBy [byNameProperty (T.isPrefixOf "weight")] df +-} +byNameProperty :: (T.Text -> Bool) -> SelectionCriteria +byNameProperty = ColumnNameProperty + +{- | Criteria for selecting columns whose names are in the given lexicographic range (inclusive). + +> selectBy [byNameRange ("a", "c")] df +-} +byNameRange :: (T.Text, T.Text) -> SelectionCriteria +byNameRange = ColumnTextRange + +{- | Criteria for selecting columns whose indices are in the given (inclusive) range. + +> selectBy [byIndexRange (0, 5)] df +-} +byIndexRange :: (Int, Int) -> SelectionCriteria +byIndexRange = ColumnIndexRange + +-- | O(n) select columns by column predicate name. +selectBy :: [SelectionCriteria] -> DataFrame -> DataFrame +selectBy xs df = select finalSelection df + where + finalSelection = Prelude.filter (`S.member` columnsWithProperties) (columnNames df) + columnsWithProperties = S.fromList (L.foldl' columnWithProperty [] xs) + columnWithProperty acc (ColumnName colName) = acc ++ [colName] + columnWithProperty acc (ColumnNameProperty f) = acc ++ L.filter f (columnNames df) + columnWithProperty acc (ColumnTextRange (from, to)) = + acc + ++ reverse + (Prelude.dropWhile (to /=) $ reverse $ dropWhile (from /=) (columnNames df)) + columnWithProperty acc (ColumnIndexRange (from, to)) = acc ++ Prelude.take (to - from + 1) (Prelude.drop from (columnNames df)) + columnWithProperty acc (ColumnProperty f) = + acc + ++ map fst (L.filter (\(_k, v) -> v `elem` ixs) (M.toAscList (columnIndices df))) + where + ixs = V.ifoldl' (\acc' i c -> if f c then i : acc' else acc') [] (columns df) + +{- | O(n) inverse of select + +> exclude ["Name"] df +-} +exclude :: + [T.Text] -> + DataFrame -> + DataFrame +exclude cs df = + let keysToKeep = columnNames df L.\\ cs + in select keysToKeep df + +{- | Sample a dataframe. The double parameter must be between 0 and 1 (inclusive). + +==== __Example__ +@ +ghci> import System.Random +ghci> D.sample (mkStdGen 137) 0.1 df + +@ +-} +sample :: (RandomGen g) => g -> Double -> DataFrame -> DataFrame +sample pureGen p df = + let + rand = mkRandom pureGen (fst (dataframeDimensions df)) (0 :: Double) 1 + cRand = col @Double "__rand__" + in + df + & insertColumn (name cRand) rand + & filterWhere (cRand .>=. Lit (1 - p)) + & exclude [name cRand] + +{- | Split a dataset into two. The first in the tuple gets a sample of p (0 <= p <= 1) and the second gets (1 - p). This is useful for creating test and train splits. + +==== __Example__ +@ +ghci> import System.Random +ghci> D.randomSplit (mkStdGen 137) 0.9 df + +@ +-} +randomSplit :: + (RandomGen g) => g -> Double -> DataFrame -> (DataFrame, DataFrame) +randomSplit pureGen p df = + let + rand = mkRandom pureGen (fst (dataframeDimensions df)) (0 :: Double) 1 + cRand = col @Double "__rand__" + withRand = df & insertColumn (name cRand) rand + in + ( withRand + & filterWhere (cRand .<=. Lit p) + & exclude [name cRand] + , withRand + & filterWhere + (cRand .>. Lit p) + & exclude [name cRand] + ) + +{- | Creates n folds of a dataframe. + +==== __Example__ +@ +ghci> import System.Random +ghci> D.kFolds (mkStdGen 137) 5 df + +@ +-} +kFolds :: (RandomGen g) => g -> Int -> DataFrame -> [DataFrame] +kFolds pureGen folds df = + let + rand = mkRandom pureGen (fst (dataframeDimensions df)) (0 :: Double) 1 + cRand = col @Double "__rand__" + withRand = df & insertColumn (name cRand) rand + partitionSize = 1 / fromIntegral folds + singleFold n d = + d & filterWhere (cRand .>=. Lit (fromIntegral n * partitionSize)) + go (-1) _ = [] + go n d = + let + d' = singleFold n d + d'' = d & filterWhere (cRand .<. Lit (fromIntegral n * partitionSize)) + in + d' : go (n - 1) d'' + in + map (exclude [name cRand]) (go (folds - 1) withRand) + +-- | Convert any Column to a vector of Text labels (one per row). +columnToTextVec :: Column -> V.Vector T.Text +columnToTextVec (BoxedColumn bm (col' :: V.Vector a)) = + case bm of + Nothing -> case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> col' + Nothing -> V.map (T.pack . show) col' + Just bitmap -> + V.imap (\i x -> if bitmapTestBit bitmap i then T.pack (show x) else "null") col' +columnToTextVec (UnboxedColumn bm col') = + case bm of + Nothing -> V.map (T.pack . show) (V.convert col') + Just bitmap -> + V.generate (VU.length col') $ \i -> + if bitmapTestBit bitmap i then T.pack (show (col' VU.! i)) else "null" + +-- | Build a map from stringified label to row indices. +groupByIndices :: Column -> M.Map T.Text (VU.Vector Int) +groupByIndices col' = + let textVec = columnToTextVec col' + (grouped, _) = + V.foldl' + (\(!m, !i) key -> (M.insertWith (++) key [i] m, i + 1)) + (M.empty, 0) + textVec + in M.map (VU.fromList . L.reverse) grouped + +-- | Select rows at the given indices from all columns. +rowsAtIndices :: VU.Vector Int -> DataFrame -> DataFrame +rowsAtIndices ixs df = + df + { columns = V.map (atIndicesStable ixs) (columns df) + , dataframeDimensions = (VU.length ixs, snd (dataframeDimensions df)) + } + +{- | Sample a dataframe, preserving per-stratum proportions. + +==== __Example__ +@ +ghci> import System.Random +ghci> D.stratifiedSample (mkStdGen 42) 0.8 "label" df +@ +-} +stratifiedSample :: + forall a g. + (SplittableGen g, Columnable a) => + g -> Double -> Expr a -> DataFrame -> DataFrame +stratifiedSample gen p strataCol df = + let col' = case strataCol of + Col colName -> unsafeGetColumn colName df + _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol)) + groups = M.elems (groupByIndices col') + go _ [] = mempty + go g (ixs : rest) = + let stratum = rowsAtIndices ixs df + (g1, g2) = splitForStratified g + in sample g1 p stratum <> go g2 rest + in go gen groups + +{- | Split a dataframe into two, preserving per-stratum proportions. + +==== __Example__ +@ +ghci> import System.Random +ghci> D.stratifiedSplit (mkStdGen 42) 0.8 "label" df +@ +-} +stratifiedSplit :: + forall a g. + (SplittableGen g, Columnable a) => + g -> Double -> Expr a -> DataFrame -> (DataFrame, DataFrame) +stratifiedSplit gen p strataCol df = + let col' = case strataCol of + Col colName -> unsafeGetColumn colName df + _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol)) + groups = M.elems (groupByIndices col') + go _ [] = (mempty, mempty) + go g (ixs : rest) = + let stratum = rowsAtIndices ixs df + (g1, g2) = splitForStratified g + (tr, va) = randomSplit g1 p stratum + (trAcc, vaAcc) = go g2 rest + in (tr <> trAcc, va <> vaAcc) + in go gen groups diff --git a/src/DataFrame/Operations/Transformations.hs b/src/DataFrame/Operations/Transformations.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Transformations.hs @@ -0,0 +1,244 @@ +{-# LANGUAGE ConstrainedClassMethods #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UndecidableSuperClasses #-} + +module DataFrame.Operations.Transformations where + +import qualified Data.List as L +import qualified Data.Map as M +import qualified Data.Text as T +import qualified Data.Vector as V + +import Control.Exception (throw) +import Data.Maybe +import DataFrame.Errors (DataFrameException (..), TypeErrorContext (..)) +import DataFrame.Internal.Column ( + Columnable, + TypedColumn (..), + hasMissing, + ifoldrColumn, + imapColumn, + mapColumn, + ) +import DataFrame.Internal.DataFrame (DataFrame (..), getColumn, insertColumn) +import DataFrame.Internal.Expression +import DataFrame.Internal.Interpreter +import DataFrame.Internal.Nullable (BaseType) +import DataFrame.Operations.Core + +-- | O(k) Apply a function to a given column in a dataframe. +apply :: + forall b c. + (Columnable b, Columnable c) => + -- | function to apply + (b -> c) -> + -- | Column name + T.Text -> + -- | DataFrame to apply operation to + DataFrame -> + DataFrame +apply f columnName d = case safeApply f columnName d of + Left (TypeMismatchException context) -> + throw $ TypeMismatchException (context{callingFunctionName = Just "apply"}) + Left exception -> throw exception + Right df -> df + +-- | O(k) Safe version of the apply function. Returns (instead of throwing) the error. +safeApply :: + forall b c. + (Columnable b, Columnable c) => + -- | function to apply + (b -> c) -> + -- | Column name + T.Text -> + -- | DataFrame to apply operation to + DataFrame -> + Either DataFrameException DataFrame +safeApply f columnName d = case getColumn columnName d of + Nothing -> + Left $ ColumnsNotFoundException [columnName] "apply" (M.keys $ columnIndices d) + Just column -> do + column' <- mapColumn f column + pure $ insertColumn columnName column' d + +{- | O(k) Apply a function to an expression in a dataframe and +add the result into `alias` column. +-} +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) + { derivingExpressions = M.insert name (UExpr expr) (derivingExpressions df) + } + +{- | O(k) Apply a function to an expression in a dataframe and +add the result into `alias` column but + +==== __Examples__ + +>>> (z, df') = deriveWithExpr "z" (F.col @Int "x" + F.col "y") df +>>> filterWhere (z .>= 50) +-} +deriveWithExpr :: + forall a. (Columnable a) => T.Text -> Expr a -> DataFrame -> (Expr a, DataFrame) +deriveWithExpr name expr df = case interpret @a df (normalize expr) of + Left e -> throw e + Right (TColumn value) -> + ( Col name + , (insertColumn name value df) + { derivingExpressions = M.insert name (UExpr expr) (derivingExpressions df) + } + ) + +deriveMany :: [NamedExpr] -> DataFrame -> DataFrame +deriveMany exprs df = + let + f (name, UExpr (expr :: Expr a)) d = + case interpret @a df expr of + Left e -> throw e + Right (TColumn value) -> insertColumn name value d + in + fold f exprs df + +-- | O(k * n) Apply a function to given column names in a dataframe. +applyMany :: + (Columnable b, Columnable c) => + (b -> c) -> + [T.Text] -> + DataFrame -> + DataFrame +applyMany f names df = L.foldl' (flip (apply f)) df names + +-- | O(k) Convenience function that applies to an int column. +applyInt :: + (Columnable b) => + -- | function to apply + (Int -> b) -> + -- | Column name + T.Text -> + -- | DataFrame to apply operation to + DataFrame -> + DataFrame +applyInt = apply + +-- | O(k) Convenience function that applies to an double column. +applyDouble :: + (Columnable b) => + -- | function to apply + (Double -> b) -> + -- | Column name + T.Text -> + -- | DataFrame to apply operation to + DataFrame -> + DataFrame +applyDouble = apply + +{- | O(k * n) Apply a function to a column only if there is another column +value that matches the given criterion. + +> applyWhere (<20) "Age" (const "Gen-Z") "Generation" df +-} +applyWhere :: + forall a b. + (Columnable a, Columnable b) => + -- | Filter condition + (a -> Bool) -> + -- | Criterion Column + T.Text -> + -- | function to apply + (b -> b) -> + -- | Column name + T.Text -> + -- | DataFrame to apply operation to + DataFrame -> + DataFrame +applyWhere condition filterColumnName f columnName df = case getColumn filterColumnName df of + Nothing -> + throw $ + ColumnsNotFoundException + [filterColumnName] + "applyWhere" + (M.keys $ columnIndices df) + Just column -> case ifoldrColumn + (\i val acc -> if condition val then V.cons i acc else acc) + V.empty + column of + Left e -> throw e + Right indexes -> + if V.null indexes + then df + else L.foldl' (\d i -> applyAtIndex i f columnName d) df indexes + +-- | O(k) Apply a function to the column at a given index. +applyAtIndex :: + forall a. + (Columnable a) => + -- | Index + Int -> + -- | function to apply + (a -> a) -> + -- | Column name + T.Text -> + -- | DataFrame to apply operation to + DataFrame -> + DataFrame +applyAtIndex i f columnName df = case getColumn columnName df of + Nothing -> + throw $ + ColumnsNotFoundException [columnName] "applyAtIndex" (M.keys $ columnIndices df) + Just column -> case imapColumn (\index value -> if index == i then f value else value) column of + Left e -> throw e + Right column' -> insertColumn columnName column' df + +-- | Core impute implementation for nullable columns. Silently no-ops on non-nullable columns. +imputeCore :: + forall b. + (Columnable b) => + Expr (Maybe b) -> + b -> + DataFrame -> + DataFrame +imputeCore (Col columnName) value df = case getColumn columnName df of + Nothing -> + throw $ + ColumnsNotFoundException [columnName] "impute" (M.keys $ columnIndices df) + Just col | hasMissing col -> case safeApply (fromMaybe value) columnName df of + Left (TypeMismatchException context) -> throw $ TypeMismatchException (context{callingFunctionName = Just "impute"}) + Left exception -> throw exception + Right res -> res + _ -> df +imputeCore _ _ df = df + +class (Columnable a) => ImputeOp a where + runImpute :: Expr a -> BaseType a -> DataFrame -> DataFrame + runImputeWith :: + (Columnable (BaseType a)) => + (Expr (BaseType a) -> Expr (BaseType a)) -> + Expr a -> + DataFrame -> + DataFrame + +instance {-# OVERLAPPABLE #-} (Columnable a) => ImputeOp a where + runImpute _ _ df = df + runImputeWith _ _ df = df + +{- | Replace all instances of `Nothing` in a column with the given value. +When the column is already non-nullable, this is a silent no-op. +-} +impute :: + forall a. + (ImputeOp a) => + Expr a -> + BaseType a -> + DataFrame -> + DataFrame +impute = runImpute diff --git a/src/DataFrame/Operations/Typing.hs b/src/DataFrame/Operations/Typing.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Typing.hs @@ -0,0 +1,474 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.Operations.Typing where + +import qualified Data.Map as M +import qualified Data.Text as T +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 Control.Applicative (asum) +import Control.Monad (join) +import Control.Monad.ST (runST) +import Data.Maybe (fromMaybe) +import qualified Data.Proxy as P +import Data.Time +import Data.Type.Equality (TestEquality (..)) +import DataFrame.Internal.Column ( + Bitmap, + Column (..), + Columnable, + bitmapTestBit, + ensureOptional, + finalizeParseResult, + fromVector, + ) +import DataFrame.Internal.DataFrame ( + DataFrame (..), + insertColumn, + unsafeGetColumn, + ) +import DataFrame.Internal.Parsing +import DataFrame.Internal.Schema +import DataFrame.Operations.Core +import Text.Read +import Type.Reflection + +type DateFormat = String + +{- | How parse failures are surfaced in the resulting column. + +* 'NoSafeRead' — strict parsing: failures throw (via 'read'). +* 'MaybeRead' — failures become 'Nothing'; columns are wrapped as @Maybe a@. +* 'EitherRead' — failures become @Left rawText@; columns are wrapped as + @Either Text a@, preserving the original input so callers can inspect it. +-} +data SafeReadMode + = NoSafeRead + | MaybeRead + | EitherRead + deriving (Eq, Show, Read) + +-- | Options controlling how text columns are parsed into typed values. +data ParseOptions = ParseOptions + { missingValues :: [T.Text] + -- ^ Values to treat as @Nothing@ when the effective mode is 'MaybeRead'. + , sampleSize :: Int + -- ^ Number of rows to inspect when inferring a column's type (0 = all rows). + , parseSafe :: SafeReadMode + {- ^ Default 'SafeReadMode' applied to every column that does not have an + entry in 'parseSafeOverrides'. 'NoSafeRead' only treats empty strings as + missing; 'MaybeRead' additionally treats 'missingValues' and nullish + strings as @Nothing@; 'EitherRead' wraps the resulting column as + @Either Text a@ with the raw input preserved on failure. + -} + , parseSafeOverrides :: [(T.Text, SafeReadMode)] + {- ^ Per-column overrides. When a column name is present here, its value + takes precedence over 'parseSafe'. Typical use: strict IDs + (@NoSafeRead@) alongside lenient fields (@MaybeRead@/@EitherRead@). + -} + , parseDateFormat :: DateFormat + -- ^ Date format string as accepted by "Data.Time.Format" (e.g. @\"%Y-%m-%d\"@). + } + +{- | Sensible out-of-the-box parse options: infer from the first 100 rows, + treat common nullish strings as missing, and expect ISO 8601 dates. +-} +defaultParseOptions :: ParseOptions +defaultParseOptions = + ParseOptions + { missingValues = [] + , sampleSize = 100 + , parseSafe = MaybeRead + , parseSafeOverrides = [] + , parseDateFormat = "%Y-%m-%d" + } + +{- | Resolve a column's effective 'SafeReadMode': the override if present, +otherwise the default. +-} +effectiveSafeRead :: + SafeReadMode -> [(T.Text, SafeReadMode)] -> T.Text -> SafeReadMode +effectiveSafeRead def overrides name = fromMaybe def (lookup name overrides) + +parseDefaults :: ParseOptions -> DataFrame -> DataFrame +parseDefaults opts df = df{columns = V.imap forCol (columns df)} + where + -- Index -> column name: reverse the columnIndices map once. + nameAt = + let inverted = M.fromList [(i, n) | (n, i) <- M.toList (columnIndices df)] + in \i -> M.findWithDefault "" i inverted + forCol i col = + let mode = + effectiveSafeRead + (parseSafe opts) + (parseSafeOverrides opts) + (nameAt i) + in parseDefault opts{parseSafe = mode, parseSafeOverrides = []} col + +parseDefault :: ParseOptions -> Column -> Column +parseDefault opts (BoxedColumn Nothing (c :: V.Vector a)) = + case (typeRep @a) `testEquality` (typeRep @T.Text) of + Nothing -> case (typeRep @a) `testEquality` (typeRep @String) of + Just Refl -> parseFromExamples opts (V.map T.pack c) + Nothing -> BoxedColumn Nothing c + Just Refl -> parseFromExamples opts c +parseDefault opts (BoxedColumn (Just bm) (c :: V.Vector a)) = + case (typeRep @a) `testEquality` (typeRep @T.Text) of + Nothing -> case (typeRep @a) `testEquality` (typeRep @String) of + Just Refl -> + parseFromExamples + opts + (V.imap (\i x -> if bitmapTestBit bm i then T.pack x else "") c) + Nothing -> BoxedColumn (Just bm) c + Just Refl -> + parseFromExamples opts (V.imap (\i x -> if bitmapTestBit bm i then x else "") c) +parseDefault _ column = column + +parseFromExamples :: ParseOptions -> V.Vector T.Text -> Column +parseFromExamples opts cols = + let isNull = case parseSafe opts of + NoSafeRead -> T.null + _ -> isNullishOrMissing (missingValues opts) + -- `examples` is small (≤ sampleSize, default 100), so the + -- Maybe-wrap allocation here is ignorable. The full-column + -- equivalent (`asMaybeText = V.map ... cols`) has been removed: + -- handlers now walk `cols` directly with `isNull`. + examples = V.map (classify isNull) (V.take (sampleSize opts) cols) + dfmt = parseDateFormat opts + assumption = makeParsingAssumption dfmt examples + in case parseSafe opts of + EitherRead -> handleEitherAssumption dfmt assumption cols + mode -> + let result = case assumption of + BoolAssumption -> handleBoolAssumption isNull cols + IntAssumption -> handleIntAssumption isNull cols + DoubleAssumption -> handleDoubleAssumption isNull cols + TextAssumption -> handleTextAssumption isNull cols + DateAssumption -> handleDateAssumption dfmt isNull cols + NoAssumption -> handleNoAssumption dfmt isNull cols + in if mode == MaybeRead then ensureOptional result else result + where + classify p t = if p t then Nothing else Just t + +{- | For 'EitherRead' mode: take the chosen parsing assumption and produce an +@Either Text a@ column. Successful parses become @Right@; any row that fails +to parse as the chosen type (including null/missing cells) becomes @Left@ +carrying the raw input text verbatim. +-} +handleEitherAssumption :: + DateFormat -> ParsingAssumption -> V.Vector T.Text -> Column +handleEitherAssumption dfmt assumption raw = case assumption of + BoolAssumption -> fromVector (V.map (toEither readBool) raw) + IntAssumption -> fromVector (V.map (toEither readInt) raw) + DoubleAssumption -> fromVector (V.map (toEither readDouble) raw) + DateAssumption -> fromVector (V.map (toEither (parseTimeOpt dfmt)) raw) + -- TextAssumption and NoAssumption degenerate to Either Text Text; treat + -- empty strings as Left "" so the convention (Left = missing/failure) stays + -- consistent across column types. + TextAssumption -> fromVector (V.map textToEither raw) + NoAssumption -> fromVector (V.map textToEither raw) + where + toEither :: (T.Text -> Maybe a) -> T.Text -> Either T.Text a + toEither p t = maybe (Left t) Right (p t) + + textToEither :: T.Text -> Either T.Text T.Text + textToEither t = if T.null t then Left t else Right t + +parseUnboxedColumnWithPred :: + forall src a. + (VU.Unbox a) => + a -> + (src -> Bool) -> + (src -> Maybe a) -> + V.Vector src -> + Maybe (Maybe Bitmap, VU.Vector a) +parseUnboxedColumnWithPred nullValue isNull parser vec = runST $ do + let n = V.length vec + values <- VUM.unsafeNew n + vmask <- VUM.unsafeNew n + let go !i !anyNull + | i >= n = finalizeParseResult values vmask anyNull + | otherwise = + let !src = V.unsafeIndex vec i + in if isNull src + then do + VUM.unsafeWrite vmask i 0 + VUM.unsafeWrite values i nullValue + go (i + 1) True + else case parser src of + Just v -> do + VUM.unsafeWrite vmask i 1 + VUM.unsafeWrite values i v + go (i + 1) anyNull + Nothing -> return Nothing + go 0 False +{-# INLINE parseUnboxedColumnWithPred #-} + +-- | Wrap a successful 'parseUnboxedColumnWithPred' result as a 'Column'. +unboxedOrFallback :: + (Columnable a, VU.Unbox a) => + Maybe (Maybe Bitmap, VU.Vector a) -> + Column -> + Column +unboxedOrFallback (Just (mbm, vec)) _ = UnboxedColumn mbm vec +unboxedOrFallback Nothing fallback = fallback + +handleBoolAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column +handleBoolAssumption isNull cols = + unboxedOrFallback + (parseUnboxedColumnWithPred False isNull readBool cols) + (handleTextAssumption isNull cols) + +handleIntAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column +handleIntAssumption isNull cols = + case parseUnboxedColumnWithPred 0 isNull readInt cols of + Just (mbm, vec) -> UnboxedColumn mbm vec + Nothing -> + unboxedOrFallback + (parseUnboxedColumnWithPred 0 isNull readDouble cols) + (handleTextAssumption isNull cols) + +handleDoubleAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column +handleDoubleAssumption isNull cols = + unboxedOrFallback + (parseUnboxedColumnWithPred 0 isNull readDouble cols) + (handleTextAssumption isNull cols) + +{- | Text columns: no parse, just null-marking. When the whole column +is non-null we return a plain 'V.Vector T.Text'; otherwise we emit a +@V.Vector (Maybe T.Text)@ the same shape the old code produced. +-} +handleTextAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column +handleTextAssumption isNull cols + | V.any isNull cols = + fromVector + (V.map (\t -> if isNull t then Nothing else Just t) cols) + | otherwise = fromVector cols + +{- | Date: single parse pass, boxed because 'Day' is not unboxable. +Bails to 'handleTextAssumption' the moment a non-null cell fails to +parse as a 'Day'. Still avoids the outer @V.Vector (Maybe T.Text)@ +allocation — we walk @cols@ directly with @isNull@. +-} +handleDateAssumption :: + DateFormat -> (T.Text -> Bool) -> V.Vector T.Text -> Column +handleDateAssumption dateFormat isNull cols = + case parseBoxedMaybeColumn isNull (parseTimeOpt dateFormat) cols of + Just (anyNull, vec) + -- `vec :: V.Vector (Maybe Day)`. If no nulls, strip the + -- outer 'Maybe' (every cell is guaranteed 'Just') so the + -- column type stays 'Day' rather than becoming 'Maybe Day'. + | anyNull -> fromVector vec + | otherwise -> fromVector (V.mapMaybe id vec) + Nothing -> handleTextAssumption isNull cols + +parseBoxedMaybeColumn :: + (T.Text -> Bool) -> + (T.Text -> Maybe a) -> + V.Vector T.Text -> + Maybe (Bool, V.Vector (Maybe a)) +parseBoxedMaybeColumn isNull parser cols = runST $ do + let n = V.length cols + out <- VM.new n + let loop !i !anyNull + | i >= n = do + frozen <- V.unsafeFreeze out + return (Just (anyNull, frozen)) + | otherwise = + let !t = V.unsafeIndex cols i + in if isNull t + then do + VM.unsafeWrite out i Nothing + loop (i + 1) True + else case parser t of + Just v -> do + VM.unsafeWrite out i (Just v) + loop (i + 1) anyNull + Nothing -> return Nothing + loop 0 False + +handleNoAssumption :: + DateFormat -> (T.Text -> Bool) -> V.Vector T.Text -> Column +handleNoAssumption dateFormat isNull cols + -- Only reached when the 100-row sample was all-null. Try each + -- concrete type in turn; fall back to Text otherwise. + | V.all isNull cols = + fromVector (V.map (const (Nothing :: Maybe T.Text)) cols) + | Just (mbm, vec) <- parseUnboxedColumnWithPred False isNull readBool cols = + UnboxedColumn mbm vec + | Just (mbm, vec) <- parseUnboxedColumnWithPred 0 isNull readInt cols = + UnboxedColumn mbm vec + | Just (mbm, vec) <- parseUnboxedColumnWithPred 0 isNull readDouble cols = + UnboxedColumn mbm vec + | otherwise = case parseBoxedMaybeColumn isNull (parseTimeOpt dateFormat) cols of + Just (anyNull, vec) + -- `vec :: V.Vector (Maybe Day)`. If no nulls, strip the + -- outer 'Maybe' (every cell is guaranteed 'Just') so the + -- column type stays 'Day' rather than becoming 'Maybe Day'. + | anyNull -> fromVector vec + | otherwise -> fromVector (V.mapMaybe id vec) + Nothing -> handleTextAssumption isNull cols + +{- | Predicate matching what 'parseSafe == NoSafeRead' previously used: +only empty strings are treated as missing. + +We still expose 'convertNullish' \/ 'convertOnlyEmpty' below because +other parts of the library reference them, but neither is used by +'parseFromExamples' any longer. +-} +isNullishOrMissing :: [T.Text] -> T.Text -> Bool +isNullishOrMissing missing v = isNullish v || v `elem` missing + +convertNullish :: [T.Text] -> T.Text -> Maybe T.Text +convertNullish missing v = if isNullish v || v `elem` missing then Nothing else Just v + +convertOnlyEmpty :: T.Text -> Maybe T.Text +convertOnlyEmpty v = if v == "" then Nothing else Just v + +parseTimeOpt :: DateFormat -> T.Text -> Maybe Day +parseTimeOpt dateFormat s = + parseTimeM {- Accept leading/trailing whitespace -} + True + defaultTimeLocale + dateFormat + (T.unpack s) + +unsafeParseTime :: DateFormat -> T.Text -> Day +unsafeParseTime dateFormat s = + parseTimeOrError {- Accept leading/trailing whitespace -} + True + defaultTimeLocale + dateFormat + (T.unpack s) + +hasNullValues :: (Eq a) => V.Vector (Maybe a) -> Bool +hasNullValues = V.any (== Nothing) + +vecSameConstructor :: V.Vector (Maybe a) -> V.Vector (Maybe b) -> Bool +vecSameConstructor xs ys = (V.length xs == V.length ys) && V.and (V.zipWith hasSameConstructor xs ys) + where + hasSameConstructor :: Maybe a -> Maybe b -> Bool + hasSameConstructor (Just _) (Just _) = True + hasSameConstructor Nothing Nothing = True + hasSameConstructor _ _ = False + +makeParsingAssumption :: + DateFormat -> V.Vector (Maybe T.Text) -> ParsingAssumption +makeParsingAssumption dateFormat asMaybeText + -- All the examples are "NA", "Null", "", so we can't make any shortcut + -- assumptions and just have to go the long way. + | V.all (== Nothing) asMaybeText = NoAssumption + -- After accounting for nulls, parsing for Ints and Doubles results in the + -- same corresponding positions of Justs and Nothings, so we assume + -- that the best way to parse is Int + | vecSameConstructor asMaybeText asMaybeBool = BoolAssumption + | vecSameConstructor asMaybeText asMaybeInt + && vecSameConstructor asMaybeText asMaybeDouble = + IntAssumption + -- After accounting for nulls, the previous condition fails, so some (or none) can be parsed as Ints + -- and some can be parsed as Doubles, so we make the assumpotion of doubles. + | vecSameConstructor asMaybeText asMaybeDouble = DoubleAssumption + -- After accounting for nulls, parsing for Dates results in the same corresponding + -- positions of Justs and Nothings, so we assume that the best way to parse is Date. + | vecSameConstructor asMaybeText asMaybeDate = DateAssumption + | otherwise = TextAssumption + where + asMaybeBool = V.map (>>= readBool) asMaybeText + asMaybeInt = V.map (>>= readInt) asMaybeText + asMaybeDouble = V.map (>>= readDouble) asMaybeText + asMaybeDate = V.map (>>= parseTimeOpt dateFormat) asMaybeText + +data ParsingAssumption + = BoolAssumption + | IntAssumption + | DoubleAssumption + | DateAssumption + | NoAssumption + | TextAssumption + +{- | Re-type columns of a 'DataFrame' according to the supplied schema map. +The caller provides a @resolveMode@ function that maps a column name to its +'SafeReadMode' — typically built from a global default plus an overrides map +via 'effectiveSafeRead'. +-} +parseWithTypes :: + (T.Text -> SafeReadMode) -> + M.Map T.Text SchemaType -> + DataFrame -> + DataFrame +parseWithTypes resolveMode ts df + | M.null ts = df + | otherwise = + M.foldrWithKey + (\k v d -> insertColumn k (asType (resolveMode k) v (unsafeGetColumn k d)) d) + df + ts + where + -- \| Re-parse a plain (non-Maybe, non-Either) target type according to the + -- 'SafeReadMode'. @toStr@ converts column elements to a 'String' ready for + -- 'Read'. + plainType :: + forall a b. + (Columnable a, Read a) => + SafeReadMode -> V.Vector b -> (b -> String) -> Column + plainType mode col toStr = case mode of + NoSafeRead -> fromVector (V.map ((read @a) . toStr) col) + MaybeRead -> fromVector (V.map ((readMaybe @a) . toStr) col) + EitherRead -> fromVector (V.map ((readEitherRaw @a) . toStr) col) + + asType :: SafeReadMode -> SchemaType -> Column -> Column + asType mode (SType (_ :: P.Proxy a)) c@(BoxedColumn _ (col :: V.Vector b)) = case typeRep @a of + App t1 _t2 -> case eqTypeRep t1 (typeRep @Maybe) of + Just HRefl -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> c + Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of + Just Refl -> fromVector (V.map (join . (readAsMaybe @a) . T.unpack) col) + Nothing -> fromVector (V.map (join . (readAsMaybe @a) . show) col) + Nothing -> case t1 of + App t1' _t2' -> case eqTypeRep t1' (typeRep @Either) of + Just HRefl -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> c + Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of + Just Refl -> fromVector (V.map ((readAsEither @a) . T.unpack) col) + Nothing -> fromVector (V.map ((readAsEither @a) . show) col) + Nothing -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> c + Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of + Just Refl -> plainType @a mode col T.unpack + Nothing -> plainType @a mode col show + _ -> c + _ -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> c + Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of + Just Refl -> plainType @a mode col T.unpack + Nothing -> plainType @a mode col show + asType _ _ c = c + +readAsMaybe :: (Read a) => String -> Maybe a +readAsMaybe s + | null s = Nothing + | otherwise = readMaybe $ "Just " <> s + +readAsEither :: (Read a) => String -> a +readAsEither v = case asum [readMaybe $ "Left " <> s, readMaybe $ "Right " <> s] of + Nothing -> error $ "Couldn't read value: " <> s + Just v' -> v' + where + s = if null v then "\"\"" else v + +{- | Try 'readMaybe'; on failure return @Left raw@ where @raw@ is the original +input text. Used by 'parseWithTypes' under 'EitherRead'. +-} +readEitherRaw :: forall a. (Read a) => String -> Either T.Text a +readEitherRaw s = case readMaybe s of + Just v -> Right v + Nothing -> Left (T.pack s) diff --git a/src/DataFrame/Typed/Access.hs b/src/DataFrame/Typed/Access.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Access.hs @@ -0,0 +1,55 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +module DataFrame.Typed.Access ( + -- * Typed column access + columnAsVector, + columnAsList, +) where + +import Control.Exception (throw) +import Data.Proxy (Proxy (..)) +import qualified Data.Text as T +import qualified Data.Vector as V +import GHC.TypeLits (KnownSymbol, symbolVal) + +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Expression (Expr (Col)) +import qualified DataFrame.Operations.Core as D +import DataFrame.Typed.Schema (AssertPresent, SafeLookup) +import DataFrame.Typed.Types (TypedDataFrame (..)) + +{- | Retrieve a column as a boxed 'Vector', with the type determined by +the schema. The column must exist (enforced at compile time). +-} +columnAsVector :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , Columnable a + , AssertPresent name cols + ) => + TypedDataFrame cols -> V.Vector a +columnAsVector (TDF df) = + either throw id $ D.columnAsVector (Col @a colName) df + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Retrieve a column as a list, with the type determined by the schema. +columnAsList :: + forall name cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , Columnable a + , AssertPresent name cols + ) => + TypedDataFrame cols -> [a] +columnAsList (TDF df) = + D.columnAsList (Col @a colName) df + where + colName = T.pack (symbolVal (Proxy @name)) diff --git a/src/DataFrame/Typed/Aggregate.hs b/src/DataFrame/Typed/Aggregate.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Aggregate.hs @@ -0,0 +1,118 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +module DataFrame.Typed.Aggregate ( + -- * Typed groupBy + groupBy, + + -- * Naming an aggregation + as, + + -- * Running aggregations + aggregate, + + -- * Escape hatch + aggregateUntyped, +) where + +import Data.Proxy (Proxy (..)) +import qualified Data.Text as T +import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) + +import DataFrame.Internal.Column (Columnable) +import qualified DataFrame.Internal.DataFrame as D +import DataFrame.Internal.Expression (NamedExpr) +import qualified DataFrame.Operations.Aggregation as DA + +import DataFrame.Typed.Freeze (unsafeFreeze) +import DataFrame.Typed.Schema +import DataFrame.Typed.Types + +{- | Group a typed DataFrame by one or more key columns. + +@ +grouped = groupBy \@'[\"department\"] employees +@ +-} +groupBy :: + forall (keys :: [Symbol]) cols. + (AllKnownSymbol keys, AssertAllPresent keys cols) => + TypedDataFrame cols -> TypedGrouped keys cols +groupBy (TDF df) = TGD (DA.groupBy (symbolVals @keys) df) + +{- | Build a named aggregation entry. The result column name is supplied via +@TypeApplications@; the underlying expression is validated against the +source schema at compile time. + +@as@ produces a /transformer/ on the aggregation chain — entries compose +with plain @(.)@ from Prelude (or via @(|>)@ for SQL-like postfix +reading). 'aggregate' applies the composed transformer to the empty chain +internally, so no terminator is needed. + +==== __Prefix form__ + +@ +result = grouped |> aggregate + ( as \@\"total\" (sum (col \@\"amount\")) + . as \@\"orders\" (count (col \@\"order_id\")) + . as \@\"avg\" (mean (col \@\"amount\")) + ) +@ + +==== __Postfix form (SQL-like)__ + +@ +result = grouped |> aggregate + ( (sum (col \@\"amount\") |> as \@\"total\") + . (count (col \@\"order_id\") |> as \@\"orders\") + . (mean (col \@\"amount\") |> as \@\"avg\") + ) +@ + +Per-entry parentheses are required in the postfix form because +@(.)@ binds tighter than @(|>)@. +-} +as :: + forall name a keys cols aggs. + (KnownSymbol name, Columnable a) => + TExpr cols a -> + TAgg keys cols aggs -> + TAgg keys cols (Column name a ': aggs) +as = TAggCons (T.pack (symbolVal (Proxy @name))) + +{- | Run a typed aggregation against a grouped DataFrame. + +The first argument is a chain of 'as' entries composed with @(.)@. The +empty composition (@id@) yields just the group keys. The result schema is +the group-key columns followed by the aggregation columns in declaration +order. + +@ +result = grouped |> aggregate + ( as \@\"total\" (sum (col \@\"amount\")) + . as \@\"orders\" (count (col \@\"order_id\")) + ) +-- result :: TypedDataFrame +-- '[ Column \"region\" Text +-- , Column \"total\" Double +-- , Column \"orders\" Int +-- ] +@ +-} +aggregate :: + forall keys cols aggs. + (TAgg keys cols '[] -> TAgg keys cols aggs) -> + TypedGrouped keys cols -> + TypedDataFrame (Append (GroupKeyColumns keys cols) (Reverse aggs)) +aggregate build (TGD gdf) = + unsafeFreeze (DA.aggregate (taggToNamedExprs (build TAggNil)) gdf) + +-- | Escape hatch: run an untyped aggregation and return a raw 'DataFrame'. +aggregateUntyped :: [NamedExpr] -> TypedGrouped keys cols -> D.DataFrame +aggregateUntyped exprs (TGD gdf) = DA.aggregate exprs gdf diff --git a/src/DataFrame/Typed/Expr.hs b/src/DataFrame/Typed/Expr.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Expr.hs @@ -0,0 +1,644 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -Wno-orphans #-} + +{- | Type-safe expression construction for typed DataFrames. + +Unlike the untyped @Expr a@ where column references are unchecked strings, +'TExpr' ensures at compile time that: + +* Referenced columns exist in the schema +* Column types match the expression type + +== Example + +@ +type Schema = '[Column \"age\" Int, Column \"salary\" Double] + +-- This compiles: +goodExpr :: TExpr Schema Double +goodExpr = col \@\"salary\" + +-- This gives a compile-time error (column not found): +badExpr :: TExpr Schema Double +badExpr = col \@\"nonexistent\" + +-- This gives a compile-time error (type mismatch): +wrongType :: TExpr Schema Int +wrongType = col \@\"salary\" -- salary is Double, not Int +@ +-} +module DataFrame.Typed.Expr ( + -- * Core typed expression type (re-exported from Types) + TExpr (..), + + -- * Column reference (schema-checked) + col, + + -- * Literals + lit, + + -- * Conditional + ifThenElse, + + -- * Unary / binary lifting + lift, + lift2, + nullLift, + nullLift2, + + -- * Same-type comparison operators + (.==.), + (./=.), + (.<.), + (.<=.), + (.>=.), + (.>.), + + -- * Same-type arithmetic operators + (.+.), + (.-.), + (.*.), + (./.), + + -- * Same-type exponentiation operators + (.^^.), + (.^.), + + -- * Nullable-aware arithmetic operators + (.+), + (.-), + (.*), + (./), + + -- * Nullable-aware exponentiation operators + (.^^), + (.^), + + -- * Nullable-aware comparison operators (three-valued logic) + (.==), + (./=), + (.<), + (.<=), + (.>=), + (.>), + + -- * Logical operators + (.&&.), + (.||.), + (.&&), + (.||), + DataFrame.Typed.Expr.not, + + -- * Aggregation combinators + sum, + mean, + median, + count, + countAll, + minimum, + maximum, + collect, + over, + + -- * Cast / coercion expressions + castExpr, + castExprWithDefault, + castExprEither, + unsafeCastExpr, + toDouble, + + -- * Sort helpers + asc, + desc, +) where + +import Data.Either (fromRight) +import Data.Proxy (Proxy (..)) +import Data.String (IsString (..)) +import qualified Data.Text as T +import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) + +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Expression ( + BinaryOp (..), + Expr (..), + UnaryOp (..), + ) +import DataFrame.Internal.Nullable ( + BaseType, + DivWidenOp, + NullCmpResult, + NullLift1Op (applyNull1), + NullLift1Result, + NullLift2Op (applyNull2), + NullLift2Result, + NullableCmpOp (nullCmpOp), + NumericWidenOp, + WidenResult, + WidenResultDiv, + divArithOp, + widenArithOp, + widenCmpOp, + ) +import DataFrame.Internal.Types (Promote, PromoteDiv) + +import qualified Data.Vector.Unboxed as VU +import DataFrame.Typed.Schema ( + AllKnownSymbol, + AssertAllPresent, + AssertPresent, + SafeLookup, + symbolVals, + ) +import DataFrame.Typed.Types (TExpr (..), TSortOrder (..)) +import Prelude hiding (maximum, minimum, sum) + +{- | Create a typed column reference. This is the key type-safety entry point. + +The column name must exist in @cols@ and its type must match @a@. +Both checks happen at compile time via type families. + +@ +salary :: TExpr '[Column \"salary\" Double] Double +salary = col \@\"salary\" +@ +-} +col :: + forall (name :: Symbol) cols a. + ( KnownSymbol name + , a ~ SafeLookup name cols + , Columnable a + , AssertPresent name cols + ) => + TExpr cols a +col = TExpr (Col (T.pack (symbolVal (Proxy @name)))) + +{- | Create a literal expression. Valid for any schema since it +references no columns. +-} +lit :: (Columnable a) => a -> TExpr cols a +lit = TExpr . Lit + +-- | Conditional expression. +ifThenElse :: + (Columnable a) => + TExpr cols Bool -> TExpr cols a -> TExpr cols a -> TExpr cols a +ifThenElse (TExpr c) (TExpr t) (TExpr e) = TExpr (If c t e) + +------------------------------------------------------------------------------- +-- Numeric instances (mirror Expr's instances) +------------------------------------------------------------------------------- + +instance (Num a, Columnable a) => Num (TExpr cols a) where + (TExpr a) + (TExpr b) = TExpr (a + b) + (TExpr a) - (TExpr b) = TExpr (a - b) + (TExpr a) * (TExpr b) = TExpr (a * b) + negate (TExpr a) = TExpr (negate a) + abs (TExpr a) = TExpr (abs a) + signum (TExpr a) = TExpr (signum a) + fromInteger = TExpr . fromInteger + +instance (Fractional a, Columnable a) => Fractional (TExpr cols a) where + fromRational = TExpr . fromRational + (TExpr a) / (TExpr b) = TExpr (a / b) + +instance (Floating a, Columnable a) => Floating (TExpr cols a) where + pi = TExpr pi + exp (TExpr a) = TExpr (exp a) + sqrt (TExpr a) = TExpr (sqrt a) + log (TExpr a) = TExpr (log a) + (TExpr a) ** (TExpr b) = TExpr (a ** b) + logBase (TExpr a) (TExpr b) = TExpr (logBase a b) + sin (TExpr a) = TExpr (sin a) + cos (TExpr a) = TExpr (cos a) + tan (TExpr a) = TExpr (tan a) + asin (TExpr a) = TExpr (asin a) + acos (TExpr a) = TExpr (acos a) + atan (TExpr a) = TExpr (atan a) + sinh (TExpr a) = TExpr (sinh a) + cosh (TExpr a) = TExpr (cosh a) + asinh (TExpr a) = TExpr (asinh a) + acosh (TExpr a) = TExpr (acosh a) + atanh (TExpr a) = TExpr (atanh a) + +instance (IsString a, Columnable a) => IsString (TExpr cols a) where + fromString = TExpr . fromString + +------------------------------------------------------------------------------- +-- Lifting arbitrary functions +------------------------------------------------------------------------------- + +-- | Lift a unary function into a typed expression. +lift :: + (Columnable a, Columnable b) => (a -> b) -> TExpr cols a -> TExpr cols b +lift f (TExpr e) = TExpr (Unary (MkUnaryOp f "unaryUdf" Nothing) e) + +-- | Lift a binary function into typed expressions. +lift2 :: + (Columnable a, Columnable b, Columnable c) => + (a -> b -> c) -> TExpr cols a -> TExpr cols b -> TExpr cols c +lift2 f (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp f "binaryUdf" Nothing False 0) a b) + +{- | Typed 'nullLift': lift a unary function with nullable propagation. +When the input is @Maybe a@, 'Nothing' short-circuits; when plain @a@, applies directly. +The return type is inferred via 'NullLift1Result': no annotation needed. +-} +nullLift :: + (NullLift1Op a r (NullLift1Result a r), Columnable (NullLift1Result a r)) => + (BaseType a -> r) -> + TExpr cols a -> + TExpr cols (NullLift1Result a r) +nullLift f (TExpr e) = TExpr (Unary (MkUnaryOp (applyNull1 f) "nullLift" Nothing) e) + +{- | Typed 'nullLift2': lift a binary function with nullable propagation. +Any 'Nothing' operand short-circuits to 'Nothing' in the result. +The return type is inferred via 'NullLift2Result': no annotation needed. +-} +nullLift2 :: + (NullLift2Op a b r (NullLift2Result a b r), Columnable (NullLift2Result a b r)) => + (BaseType a -> BaseType b -> r) -> + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullLift2Result a b r) +nullLift2 f (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (applyNull2 f) "nullLift2" Nothing False 0) a b) + +infixl 4 .==., ./=., .<., .<=., .>=., .>. +infix 4 .==, ./=, .<, .<=, .>=, .> +infixr 3 .&&., .&& +infixr 2 .||., .|| +infixl 6 .+., .-. +infixl 7 .*., ./. +infixr 8 .^^., .^^, .^., .^ + +(.==.) :: + (Columnable a, Eq a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool +(.==.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (==) "eq" (Just "==") True 4) a b) + +(./=.) :: + (Columnable a, Eq a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool +(./=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (/=) "neq" (Just "/=") True 4) a b) + +(.<.) :: + (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool +(.<.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<) "lt" (Just "<") False 4) a b) + +(.<=.) :: + (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool +(.<=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (<=) "leq" (Just "<=") False 4) a b) + +(.>=.) :: + (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool +(.>=.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>=) "geq" (Just ">=") False 4) a b) + +(.>.) :: + (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool +(.>.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>) "gt" (Just ">") False 4) a b) + +-- Same-type arithmetic operators + +(.+.) :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a -> TExpr cols a +(.+.) = (+) + +(.-.) :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a -> TExpr cols a +(.-.) = (-) + +(.*.) :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a -> TExpr cols a +(.*.) = (*) + +(./.) :: + (Columnable a, Fractional a) => TExpr cols a -> TExpr cols a -> TExpr cols a +(./.) = (/) + +-- Same-type exponentiation operators + +(.^^.) :: + (Columnable a, Columnable b, Fractional a, Integral b) => + TExpr cols a -> TExpr cols b -> TExpr cols a +(.^^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^^) "pow" (Just ".^^.") False 8) a b) + +(.^.) :: + (Columnable a, Columnable b, Num a, Integral b) => + TExpr cols a -> TExpr cols b -> TExpr cols a +(.^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^) "pow" (Just ".^.") False 8) a b) + +(.&&.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool +(.&&.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (&&) "and" (Just ".&&.") True 3) a b) + +(.||.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool +(.||.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (||) "or" (Just ".||.") True 2) a b) + +-- | Nullable-aware logical AND. Returns @Maybe Bool@ when either operand is nullable. +(.&&) :: + (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullCmpResult a b) +(.&&) (TExpr a) (TExpr b) = + TExpr (Binary (MkBinaryOp (nullCmpOp (&&)) "nulland" (Just ".&&") True 3) a b) + +-- | Nullable-aware logical OR. Returns @Maybe Bool@ when either operand is nullable. +(.||) :: + (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullCmpResult a b) +(.||) (TExpr a) (TExpr b) = + TExpr (Binary (MkBinaryOp (nullCmpOp (||)) "nullor" (Just ".||") True 2) a b) + +------------------------------------------------------------------------------- +-- Nullable-aware arithmetic operators +------------------------------------------------------------------------------- + +infixl 6 .+, .- +infixl 7 .*, ./ + +{- | Nullable-aware addition. Works for all combinations of nullable\/non-nullable operands. +@col \@\"x\" '.+' col \@\"y\" -- :: TExpr cols (Maybe Int) when y :: Maybe Int@ +-} +(.+) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b) + , Num (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (WidenResult a b) +(.+) (TExpr a) (TExpr b) = + TExpr + ( Binary + (MkBinaryOp (applyNull2 (widenArithOp (+))) "nulladd" (Just "+") True 6) + a + b + ) + +-- | Nullable-aware subtraction. +(.-) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b) + , Num (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (WidenResult a b) +(.-) (TExpr a) (TExpr b) = + TExpr + ( Binary + (MkBinaryOp (applyNull2 (widenArithOp (-))) "nullsub" (Just "-") False 6) + a + b + ) + +-- | Nullable-aware multiplication. +(.*) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b) + , Num (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (WidenResult a b) +(.*) (TExpr a) (TExpr b) = + TExpr + ( Binary + (MkBinaryOp (applyNull2 (widenArithOp (*))) "nullmul" (Just "*") True 7) + a + b + ) + +-- | Nullable-aware division. Integral operands are promoted to Double. +(./) :: + ( DivWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (PromoteDiv (BaseType a) (BaseType b)) (WidenResultDiv a b) + , Fractional (PromoteDiv (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (WidenResultDiv a b) +(./) (TExpr a) (TExpr b) = + TExpr + ( Binary + (MkBinaryOp (applyNull2 (divArithOp (/))) "nulldiv" (Just "/") False 7) + a + b + ) + +-- | Nullable-aware exponentiation (fractional base, integral exponent). +(.^^) :: + ( Columnable (BaseType a) + , Columnable (BaseType b) + , Fractional (BaseType a) + , Integral (BaseType b) + , NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (BaseType a) a + , Num (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> TExpr cols b -> TExpr cols a +(.^^) (TExpr a) (TExpr b) = + TExpr (Binary (MkBinaryOp (applyNull2 (^^)) "pow" (Just ".^^") False 8) a b) + +-- | Nullable-aware exponentiation (num base, integral exponent). +(.^) :: + ( Columnable (BaseType a) + , Columnable (BaseType b) + , Num (BaseType a) + , Integral (BaseType b) + , NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (BaseType a) a + , Num (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> TExpr cols b -> TExpr cols a +(.^) (TExpr a) (TExpr b) = + TExpr (Binary (MkBinaryOp (applyNull2 (^)) "pow" (Just ".^") False 8) a b) + +------------------------------------------------------------------------------- +-- Nullable-aware comparison operators (three-valued logic) +------------------------------------------------------------------------------- + +{- | Nullable-aware equality. Widens numeric operands to their common type, +so @TExpr cols Double .== TExpr cols Int@ typechecks. Returns @Maybe Bool@ +when either operand is nullable. +-} +(.==) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Eq (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullCmpResult a b) +(.==) (TExpr a) (TExpr b) = + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (==))) "eq" (Just "==") True 4) a b) + +-- | Nullable-aware inequality. Widens numeric operands to their common type. +(./=) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Eq (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullCmpResult a b) +(./=) (TExpr a) (TExpr b) = + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (/=))) "neq" (Just "/=") True 4) a b) + +-- | Nullable-aware less-than. Widens numeric operands to their common type. +(.<) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullCmpResult a b) +(.<) (TExpr a) (TExpr b) = + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<))) "lt" (Just "<") False 4) a b) + +-- | Nullable-aware less-than-or-equal. Widens numeric operands to their common type. +(.<=) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullCmpResult a b) +(.<=) (TExpr a) (TExpr b) = + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<=))) "leq" (Just "<=") False 4) a b) + +-- | Nullable-aware greater-than-or-equal. Widens numeric operands to their common type. +(.>=) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullCmpResult a b) +(.>=) (TExpr a) (TExpr b) = + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (>=))) "geq" (Just ">=") False 4) a b) + +-- | Nullable-aware greater-than. Widens numeric operands to their common type. +(.>) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullCmpResult a b) +(.>) (TExpr a) (TExpr b) = + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (>))) "gt" (Just ">") False 4) a b) + +not :: TExpr cols Bool -> TExpr cols Bool +not (TExpr e) = TExpr (Unary (MkUnaryOp Prelude.not "not" (Just "!")) e) + +------------------------------------------------------------------------------- +-- Aggregation combinators +------------------------------------------------------------------------------- + +sum :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a +sum (TExpr e) = TExpr (F.sum e) + +mean :: (Columnable a, Real a) => TExpr cols a -> TExpr cols Double +mean (TExpr e) = TExpr (F.mean e) + +median :: + (Columnable a, Real a, VU.Unbox a) => TExpr cols a -> TExpr cols Double +median (TExpr e) = TExpr (F.median e) + +count :: (Columnable a) => TExpr cols a -> TExpr cols Int +count (TExpr e) = TExpr (F.count e) + +-- | Row count, the equivalent of SQL's @COUNT(*)@. +countAll :: TExpr cols Int +countAll = TExpr F.countAll + +minimum :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a +minimum (TExpr e) = TExpr (F.minimum e) + +maximum :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a +maximum (TExpr e) = TExpr (F.maximum e) + +collect :: (Columnable a) => TExpr cols a -> TExpr cols [a] +collect (TExpr e) = TExpr (F.collect e) + +over :: + forall (names :: [Symbol]) cols a. + (Columnable a, AllKnownSymbol names, AssertAllPresent names cols) => + TExpr cols a -> TExpr cols a +over (TExpr e) = TExpr{unTExpr = F.over (symbolVals @names) e} + +------------------------------------------------------------------------------- +-- Cast / coercion expressions +------------------------------------------------------------------------------- + +castExpr :: + forall b cols src. + (Columnable b, Columnable src, Read b) => TExpr cols src -> TExpr cols (Maybe b) +castExpr (TExpr e) = + TExpr + (CastExprWith @b @(Maybe b) @src "castExpr" (either (const Nothing) Just) e) + +castExprWithDefault :: + forall b cols src. + (Columnable b, Columnable src, Read b) => b -> TExpr cols src -> TExpr cols b +castExprWithDefault def (TExpr e) = + TExpr + ( CastExprWith @b @b @src + ("castExprWithDefault:" <> T.pack (show def)) + (fromRight def) + e + ) + +castExprEither :: + forall b cols src. + (Columnable b, Columnable src, Read b) => + TExpr cols src -> TExpr cols (Either T.Text b) +castExprEither (TExpr e) = + TExpr + ( CastExprWith @b @(Either T.Text b) @src + "castExprEither" + (either (Left . T.pack) Right) + e + ) + +unsafeCastExpr :: + forall b cols src. + (Columnable b, Columnable src, Read b) => TExpr cols src -> TExpr cols b +unsafeCastExpr (TExpr e) = + TExpr + ( CastExprWith @b @b @src + "unsafeCastExpr" + (fromRight (error "unsafeCastExpr: unexpected Nothing in column")) + e + ) + +toDouble :: (Columnable a, Real a) => TExpr cols a -> TExpr cols Double +toDouble (TExpr e) = TExpr (F.toDouble e) + +-- | Create an ascending sort order from a typed expression. +asc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols +asc = Asc + +-- | Create a descending sort order from a typed expression. +desc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols +desc = Desc diff --git a/src/DataFrame/Typed/Join.hs b/src/DataFrame/Typed/Join.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Join.hs @@ -0,0 +1,72 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} + +module DataFrame.Typed.Join ( + -- * Typed joins + innerJoin, + leftJoin, + rightJoin, + fullOuterJoin, +) where + +import GHC.TypeLits (Symbol) + +import qualified DataFrame.Operations.Join as DJ + +import DataFrame.Typed.Freeze (unsafeFreeze) +import DataFrame.Typed.Schema +import DataFrame.Typed.Types (TypedDataFrame (..)) + +-- | Typed inner join on one or more key columns. +innerJoin :: + forall (keys :: [Symbol]) left right. + (AllKnownSymbol keys) => + TypedDataFrame left -> + TypedDataFrame right -> + TypedDataFrame (InnerJoinSchema keys left right) +innerJoin (TDF l) (TDF r) = + unsafeFreeze (DJ.innerJoin keyNames r l) + where + keyNames = symbolVals @keys + +-- | Typed left join. +leftJoin :: + forall (keys :: [Symbol]) left right. + (AllKnownSymbol keys) => + TypedDataFrame left -> + TypedDataFrame right -> + TypedDataFrame (LeftJoinSchema keys left right) +leftJoin (TDF l) (TDF r) = + unsafeFreeze (DJ.leftJoin keyNames l r) + where + keyNames = symbolVals @keys + +-- | Typed right join. +rightJoin :: + forall (keys :: [Symbol]) left right. + (AllKnownSymbol keys) => + TypedDataFrame left -> + TypedDataFrame right -> + TypedDataFrame (RightJoinSchema keys left right) +rightJoin (TDF l) (TDF r) = + unsafeFreeze (DJ.rightJoin keyNames l r) + where + keyNames = symbolVals @keys + +-- | Typed full outer join. +fullOuterJoin :: + forall (keys :: [Symbol]) left right. + (AllKnownSymbol keys) => + TypedDataFrame left -> + TypedDataFrame right -> + TypedDataFrame (FullOuterJoinSchema keys left right) +fullOuterJoin (TDF l) (TDF r) = + unsafeFreeze (DJ.fullOuterJoin keyNames r l) + where + keyNames = symbolVals @keys diff --git a/src/DataFrame/Typed/Operations.hs b/src/DataFrame/Typed/Operations.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Operations.hs @@ -0,0 +1,379 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +module DataFrame.Typed.Operations ( + -- * Schema-preserving operations + filterWhere, + filter, + filterBy, + filterAllJust, + filterJust, + filterNothing, + sortBy, + take, + takeLast, + drop, + dropLast, + range, + cube, + distinct, + sample, + shuffle, + + -- * Schema-modifying operations + derive, + impute, + select, + exclude, + rename, + renameMany, + insert, + insertColumn, + insertVector, + cloneColumn, + dropColumn, + replaceColumn, + + -- * Metadata + dimensions, + nRows, + nColumns, + columnNames, + + -- * Vertical merge + append, +) where + +import Data.Proxy (Proxy (..)) +import qualified Data.Text as T +import qualified Data.Vector as V +import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) +import System.Random (RandomGen) +import Prelude hiding (drop, filter, take) + +import qualified DataFrame.Functions as DF +import DataFrame.Internal.Column (Columnable) +import qualified DataFrame.Internal.Column as C +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Operations.Aggregation as DA +import qualified DataFrame.Operations.Core as D +import DataFrame.Operations.Merge () +import qualified DataFrame.Operations.Permutation as D +import qualified DataFrame.Operations.Subset as D +import qualified DataFrame.Operations.Transformations as D + +import DataFrame.Typed.Freeze (unsafeFreeze) +import DataFrame.Typed.Schema +import DataFrame.Typed.Types (TExpr (..), TSortOrder (..), TypedDataFrame (..)) +import qualified DataFrame.Typed.Types as T + +------------------------------------------------------------------------------- +-- Schema-preserving operations +------------------------------------------------------------------------------- + +{- | Filter rows where a boolean expression evaluates to True. +The expression is validated against the schema at compile time. +-} +filterWhere :: TExpr cols Bool -> TypedDataFrame cols -> TypedDataFrame cols +filterWhere (TExpr expr) (TDF df) = TDF (D.filterWhere expr df) + +-- | Filter rows by applying a predicate to a typed expression. +filter :: + (Columnable a) => + TExpr cols a -> (a -> Bool) -> TypedDataFrame cols -> TypedDataFrame cols +filter (TExpr expr) pred' (TDF df) = TDF (D.filter expr pred' df) + +-- | Filter rows by a predicate on a column expression (flipped argument order). +filterBy :: + (Columnable a) => + (a -> Bool) -> TExpr cols a -> TypedDataFrame cols -> TypedDataFrame cols +filterBy pred' (TExpr expr) (TDF df) = TDF (D.filterBy pred' expr df) + +{- | Keep only rows where ALL Optional columns have Just values. +Strips 'Maybe' from all column types in the result schema. + +@ +df :: TDF '[Column \"x\" (Maybe Double), Column \"y\" Int] +filterAllJust df :: TDF '[Column \"x\" Double, Column \"y\" Int] +@ +-} +filterAllJust :: TypedDataFrame cols -> TypedDataFrame (StripAllMaybe cols) +filterAllJust (TDF df) = unsafeFreeze (D.filterAllJust df) + +{- | Keep only rows where the named column has Just values. +Strips 'Maybe' from that column's type in the result schema. + +@ +filterJust \@\"x\" df +@ +-} +filterJust :: + forall name cols. + ( KnownSymbol name + , AssertPresent name cols + ) => + TypedDataFrame cols -> TypedDataFrame (StripMaybeAt name cols) +filterJust (TDF df) = unsafeFreeze (D.filterJust colName df) + where + colName = T.pack (symbolVal (Proxy @name)) + +{- | Keep only rows where the named column has Nothing. +Schema is preserved (column types unchanged, just fewer rows). +-} +filterNothing :: + forall name cols. + ( KnownSymbol name + , AssertPresent name cols + ) => + TypedDataFrame cols -> TypedDataFrame cols +filterNothing (TDF df) = TDF (D.filterNothing colName df) + where + colName = T.pack (symbolVal (Proxy @name)) + +{- | Sort by the given typed sort orders. +Sort orders reference columns that are validated against the schema. +-} +sortBy :: [TSortOrder cols] -> TypedDataFrame cols -> TypedDataFrame cols +sortBy ords (TDF df) = TDF (D.sortBy (map toUntypedSort ords) df) + where + toUntypedSort :: TSortOrder cols -> D.SortOrder + toUntypedSort (Asc (TExpr e)) = D.Asc e + toUntypedSort (Desc (TExpr e)) = D.Desc e + +-- | Take the first @n@ rows. +take :: Int -> TypedDataFrame cols -> TypedDataFrame cols +take n (TDF df) = TDF (D.take n df) + +-- | Take the last @n@ rows. +takeLast :: Int -> TypedDataFrame cols -> TypedDataFrame cols +takeLast n (TDF df) = TDF (D.takeLast n df) + +-- | Drop the first @n@ rows. +drop :: Int -> TypedDataFrame cols -> TypedDataFrame cols +drop n (TDF df) = TDF (D.drop n df) + +-- | Drop the last @n@ rows. +dropLast :: Int -> TypedDataFrame cols -> TypedDataFrame cols +dropLast n (TDF df) = TDF (D.dropLast n df) + +-- | Take rows in the given range (start, end). +range :: (Int, Int) -> TypedDataFrame cols -> TypedDataFrame cols +range r (TDF df) = TDF (D.range r df) + +-- | Take a sub-cube of the DataFrame. +cube :: (Int, Int) -> TypedDataFrame cols -> TypedDataFrame cols +cube c (TDF df) = TDF (D.cube c df) + +-- | Remove duplicate rows. +distinct :: TypedDataFrame cols -> TypedDataFrame cols +distinct (TDF df) = TDF (DA.distinct df) + +-- | Randomly sample a fraction of rows. +sample :: + (RandomGen g) => g -> Double -> TypedDataFrame cols -> TypedDataFrame cols +sample g frac (TDF df) = TDF (D.sample g frac df) + +-- | Shuffle all rows randomly. +shuffle :: (RandomGen g) => g -> TypedDataFrame cols -> TypedDataFrame cols +shuffle g (TDF df) = TDF (D.shuffle g df) + +------------------------------------------------------------------------------- +-- Schema-modifying operations +------------------------------------------------------------------------------- + +{- | Derive a new column from a typed expression. The column name must NOT +already exist in the schema (enforced at compile time via 'AssertAbsent'). +The expression is validated against the current schema. + +@ +df' = derive \@\"total\" (col \@\"price\" * col \@\"qty\") df +-- df' :: TDF (Column \"total\" Double ': originalCols) +@ +-} +derive :: + forall name a cols. + ( KnownSymbol name + , Columnable a + , AssertAbsent name cols + ) => + TExpr cols a -> + TypedDataFrame cols -> + TypedDataFrame (Snoc cols (T.Column name a)) +derive (TExpr expr) (TDF df) = unsafeFreeze (D.derive colName expr df) + where + colName = T.pack (symbolVal (Proxy @name)) + +impute :: + forall name a cols. + ( KnownSymbol name + , Columnable a + , Maybe a ~ Lookup name cols + ) => + a -> + TypedDataFrame cols -> + TypedDataFrame (Impute name cols) +impute value (TDF df) = + unsafeFreeze + (D.derive colName (DF.fromMaybe value (DF.col @(Maybe a) colName)) df) + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Select a subset of columns by name. +select :: + forall (names :: [Symbol]) cols. + (AllKnownSymbol names, AssertAllPresent names cols) => + TypedDataFrame cols -> TypedDataFrame (SubsetSchema names cols) +select (TDF df) = unsafeFreeze (D.select (symbolVals @names) df) + +-- | Exclude columns by name. +exclude :: + forall (names :: [Symbol]) cols. + (AllKnownSymbol names) => + TypedDataFrame cols -> TypedDataFrame (ExcludeSchema names cols) +exclude (TDF df) = unsafeFreeze (D.exclude (symbolVals @names) df) + +-- | Rename a column. +rename :: + forall old new cols. + (KnownSymbol old, KnownSymbol new) => + TypedDataFrame cols -> TypedDataFrame (RenameInSchema old new cols) +rename (TDF df) = unsafeFreeze (D.rename oldName newName df) + where + oldName = T.pack (symbolVal (Proxy @old)) + newName = T.pack (symbolVal (Proxy @new)) + +-- | Rename multiple columns from a type-level list of pairs. +renameMany :: + forall (pairs :: [(Symbol, Symbol)]) cols. + (AllKnownPairs pairs) => + TypedDataFrame cols -> TypedDataFrame (RenameManyInSchema pairs cols) +renameMany (TDF df) = unsafeFreeze (foldRenames (pairVals @pairs) df) + where + foldRenames [] df' = df' + foldRenames ((old, new) : rest) df' = foldRenames rest (D.rename old new df') + +-- | Insert a new column from a Foldable container. +insert :: + forall name a cols t. + ( KnownSymbol name + , Columnable a + , Foldable t + , AssertAbsent name cols + ) => + t a -> TypedDataFrame cols -> TypedDataFrame (T.Column name a ': cols) +insert xs (TDF df) = unsafeFreeze (D.insert colName xs df) + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Insert a raw 'Column' value. +insertColumn :: + forall name a cols. + ( KnownSymbol name + , Columnable a + , AssertAbsent name cols + ) => + C.Column -> TypedDataFrame cols -> TypedDataFrame (T.Column name a ': cols) +insertColumn col (TDF df) = unsafeFreeze (D.insertColumn colName col df) + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Insert a boxed 'Vector'. +insertVector :: + forall name a cols. + ( KnownSymbol name + , Columnable a + , AssertAbsent name cols + ) => + V.Vector a -> TypedDataFrame cols -> TypedDataFrame (T.Column name a ': cols) +insertVector vec (TDF df) = unsafeFreeze (D.insertVector colName vec df) + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Clone an existing column under a new name. +cloneColumn :: + forall old new cols. + ( KnownSymbol old + , KnownSymbol new + , AssertPresent old cols + , AssertAbsent new cols + ) => + TypedDataFrame cols -> TypedDataFrame (T.Column new (Lookup old cols) ': cols) +cloneColumn (TDF df) = unsafeFreeze (D.cloneColumn oldName newName df) + where + oldName = T.pack (symbolVal (Proxy @old)) + newName = T.pack (symbolVal (Proxy @new)) + +-- | Drop a column by name. +dropColumn :: + forall name cols. + ( KnownSymbol name + , AssertPresent name cols + ) => + TypedDataFrame cols -> TypedDataFrame (RemoveColumn name cols) +dropColumn (TDF df) = unsafeFreeze (D.exclude [colName] df) + where + colName = T.pack (symbolVal (Proxy @name)) + +{- | Replace an existing column with new values derived from a typed expression. +The column must already exist and the new type must match. +-} +replaceColumn :: + forall name a cols. + ( KnownSymbol name + , Columnable a + , a ~ SafeLookup name cols + , AssertPresent name cols + ) => + TExpr cols a -> TypedDataFrame cols -> TypedDataFrame cols +replaceColumn (TExpr expr) (TDF df) = unsafeFreeze (D.derive colName expr df) + where + colName = T.pack (symbolVal (Proxy @name)) + +-- | Vertically merge two DataFrames with the same schema. +append :: TypedDataFrame cols -> TypedDataFrame cols -> TypedDataFrame cols +append (TDF a) (TDF b) = TDF (a <> b) + +------------------------------------------------------------------------------- +-- Metadata (pass-through) +------------------------------------------------------------------------------- + +dimensions :: TypedDataFrame cols -> (Int, Int) +dimensions (TDF df) = D.dimensions df + +nRows :: TypedDataFrame cols -> Int +nRows (TDF df) = D.nRows df + +nColumns :: TypedDataFrame cols -> Int +nColumns (TDF df) = D.nColumns df + +columnNames :: TypedDataFrame cols -> [T.Text] +columnNames (TDF df) = D.columnNames df + +------------------------------------------------------------------------------- +-- Internal helpers +------------------------------------------------------------------------------- + +-- | Helper class for extracting [(Text, Text)] from type-level pairs. +class AllKnownPairs (pairs :: [(Symbol, Symbol)]) where + pairVals :: [(T.Text, T.Text)] + +instance AllKnownPairs '[] where + pairVals = [] + +instance + (KnownSymbol a, KnownSymbol b, AllKnownPairs rest) => + AllKnownPairs ('(a, b) ': rest) + where + pairVals = + ( T.pack (symbolVal (Proxy @a)) + , T.pack (symbolVal (Proxy @b)) + ) + : pairVals @rest