dataframe-operations (empty) → 1.0.0.0
raw patch · 19 files changed
+6537/−0 lines, 19 filesdep +basedep +bytestringdep +containers
Dependencies added: base, bytestring, containers, dataframe-core, dataframe-parsing, deepseq, hashable, random, regex-tdfa, text, time, unordered-containers, vector, vector-algorithms
Files
- LICENSE +20/−0
- dataframe-operations.cabal +65/−0
- src/DataFrame/Functions.hs +676/−0
- src/DataFrame/Internal/Statistics.hs +285/−0
- src/DataFrame/Monad.hs +101/−0
- src/DataFrame/Operations/Aggregation.hs +165/−0
- src/DataFrame/Operations/Core.hs +955/−0
- src/DataFrame/Operations/Join.hs +1102/−0
- src/DataFrame/Operations/Merge.hs +73/−0
- src/DataFrame/Operations/Permutation.hs +168/−0
- src/DataFrame/Operations/Statistics.hs +399/−0
- src/DataFrame/Operations/Subset.hs +542/−0
- src/DataFrame/Operations/Transformations.hs +244/−0
- src/DataFrame/Operations/Typing.hs +474/−0
- src/DataFrame/Typed/Access.hs +55/−0
- src/DataFrame/Typed/Aggregate.hs +118/−0
- src/DataFrame/Typed/Expr.hs +644/−0
- src/DataFrame/Typed/Join.hs +72/−0
- src/DataFrame/Typed/Operations.hs +379/−0
+ LICENSE view
@@ -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.
+ dataframe-operations.cabal view
@@ -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
+ src/DataFrame/Functions.hs view
@@ -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
+ src/DataFrame/Internal/Statistics.hs view
@@ -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+ #-}
+ src/DataFrame/Monad.hs view
@@ -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)
+ src/DataFrame/Operations/Aggregation.hs view
@@ -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
+ src/DataFrame/Operations/Core.hs view
@@ -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 <Type>` 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 <Type>` 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
+ src/DataFrame/Operations/Join.hs view
@@ -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
+ src/DataFrame/Operations/Merge.hs view
@@ -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+ }
+ src/DataFrame/Operations/Permutation.hs view
@@ -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
+ src/DataFrame/Operations/Statistics.hs view
@@ -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)
+ src/DataFrame/Operations/Subset.hs view
@@ -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
+ src/DataFrame/Operations/Transformations.hs view
@@ -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
+ src/DataFrame/Operations/Typing.hs view
@@ -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)
+ src/DataFrame/Typed/Access.hs view
@@ -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))
+ src/DataFrame/Typed/Aggregate.hs view
@@ -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
+ src/DataFrame/Typed/Expr.hs view
@@ -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
+ src/DataFrame/Typed/Join.hs view
@@ -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
+ src/DataFrame/Typed/Operations.hs view
@@ -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