diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,20 @@ # Revision history for dataframe +## 1.1.1.0 +### New features +* Add `DataFrame.Typed.Lazy` module — a type-safe lazy query pipeline combining compile-time schema tracking with deferred execution. +* Add `fromCsv` function for parsing a CSV string directly into a DataFrame. +* Add `DataKinds` extension and `DataFrame.Typed` import to the GHCi file for easier interactive typed dataframe workflows. + +### Performance +* Specialize and inline aggregation functions (`sum`, `mean`, `variance`, `median`, `stddev`, etc.) to avoid expensive numeric conversions at runtime. +* Remove `Ord` constraint from `Columnable'` and move it to call sites, reducing unnecessary constraint propagation. +* Replace exponential type-level `If` nesting in typed schema families with linear helper type families, fixing slow compilation times. + +### Bug fixes +* Fix Functions module compilation under GHC 9.10 (#194). +* Document and test `safeColumns` option in `ParquetReadOptions` (#190). + ## 1.1.0.0 ### Breaking changes * Remove `OptionalColumn` constructor; fold nullability into `BoxedColumn`/`UnboxedColumn` via bit-packed bitmap. diff --git a/dataframe.cabal b/dataframe.cabal --- a/dataframe.cabal +++ b/dataframe.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: dataframe -version: 1.1.0.0 +version: 1.1.1.0 synopsis: A fast, safe, and intuitive DataFrame library. @@ -114,6 +114,7 @@ DataFrame.Typed.Aggregate, DataFrame.Typed.TH, DataFrame.Typed.Expr, + DataFrame.Typed.Lazy, DataFrame.Typed build-depends: base >= 4 && <5, deepseq >= 1 && < 2, @@ -248,6 +249,7 @@ Functions, GenDataFrame, Internal.Parsing, + IO.CSV, IO.JSON, Operations.Aggregations, Operations.Apply, diff --git a/ffi/DataFrame/IR.hs b/ffi/DataFrame/IR.hs --- a/ffi/DataFrame/IR.hs +++ b/ffi/DataFrame/IR.hs @@ -1,4 +1,5 @@ {-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -15,16 +16,19 @@ import Data.Aeson (FromJSON (..), withObject, (.:)) import Data.Aeson.Types (Parser) +import qualified Data.ByteString as BS +import Data.Int (Int16, Int32, Int64, Int8) import qualified Data.Text as T import Data.Type.Equality ( TestEquality (testEquality), type (:~:) (Refl), + type (:~~:) (HRefl), ) import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU -import Data.Word (Word64) +import Data.Word (Word, Word16, Word32, Word64, Word8) import Foreign (wordPtrToPtr) -import Type.Reflection (typeRep) +import Type.Reflection (SomeTypeRep (..), eqTypeRep, typeRep) import DataFrame.Functions (count, mean, meanMaybe, sumMaybe) import qualified DataFrame.Functions as Functions @@ -34,7 +38,7 @@ readSeparated, readTsv, ) -import DataFrame.Internal.Column (Column (..)) +import DataFrame.Internal.Column (Column (..), Columnable) import DataFrame.Internal.DataFrame (DataFrame, unsafeGetColumn) import DataFrame.Internal.Expression (Expr (..), NamedExpr) import DataFrame.Operations.Aggregation (aggregate, groupBy) @@ -111,16 +115,38 @@ executePlan (Limit k node) = Subset.take k <$> executePlan node --- | Build a SortOrder from a column's runtime type. +{- | Build a SortOrder from a column's runtime type. +Uses type dispatch to recover Ord for known column types. +-} mkSortOrder :: Bool -> T.Text -> Column -> SortOrder -mkSortOrder True n (UnboxedColumn Nothing (_ :: VU.Vector a)) = Asc (Col @a n) -mkSortOrder False n (UnboxedColumn Nothing (_ :: VU.Vector a)) = Desc (Col @a n) -mkSortOrder True n (UnboxedColumn (Just _) (_ :: VU.Vector a)) = Asc (Col @(Maybe a) n) -mkSortOrder False n (UnboxedColumn (Just _) (_ :: VU.Vector a)) = Desc (Col @(Maybe a) n) -mkSortOrder True n (BoxedColumn Nothing (_ :: V.Vector a)) = Asc (Col @a n) -mkSortOrder False n (BoxedColumn Nothing (_ :: V.Vector a)) = Desc (Col @a n) -mkSortOrder True n (BoxedColumn (Just _) (_ :: V.Vector a)) = Asc (Col @(Maybe a) n) -mkSortOrder False n (BoxedColumn (Just _) (_ :: V.Vector a)) = Desc (Col @(Maybe a) n) +mkSortOrder isAsc name col = dispatchType (columnTypeRep col) + where + columnTypeRep :: Column -> SomeTypeRep + columnTypeRep (UnboxedColumn _ (_ :: VU.Vector a)) = SomeTypeRep (typeRep @a) + columnTypeRep (BoxedColumn _ (_ :: V.Vector a)) = SomeTypeRep (typeRep @a) + mk :: (Columnable a, Ord a) => Expr a -> SortOrder + mk = if isAsc then Asc else Desc + dispatchType (SomeTypeRep tr) + | Just HRefl <- eqTypeRep tr (typeRep @Int) = mk (Col @Int name) + | Just HRefl <- eqTypeRep tr (typeRep @Int8) = mk (Col @Int8 name) + | Just HRefl <- eqTypeRep tr (typeRep @Int16) = mk (Col @Int16 name) + | Just HRefl <- eqTypeRep tr (typeRep @Int32) = mk (Col @Int32 name) + | Just HRefl <- eqTypeRep tr (typeRep @Int64) = mk (Col @Int64 name) + | Just HRefl <- eqTypeRep tr (typeRep @Word) = mk (Col @Word name) + | Just HRefl <- eqTypeRep tr (typeRep @Word8) = mk (Col @Word8 name) + | Just HRefl <- eqTypeRep tr (typeRep @Word16) = mk (Col @Word16 name) + | Just HRefl <- eqTypeRep tr (typeRep @Word32) = mk (Col @Word32 name) + | Just HRefl <- eqTypeRep tr (typeRep @Word64) = mk (Col @Word64 name) + | Just HRefl <- eqTypeRep tr (typeRep @Integer) = mk (Col @Integer name) + | Just HRefl <- eqTypeRep tr (typeRep @Double) = mk (Col @Double name) + | Just HRefl <- eqTypeRep tr (typeRep @Float) = mk (Col @Float name) + | Just HRefl <- eqTypeRep tr (typeRep @Bool) = mk (Col @Bool name) + | Just HRefl <- eqTypeRep tr (typeRep @Char) = mk (Col @Char name) + | Just HRefl <- eqTypeRep tr (typeRep @T.Text) = mk (Col @T.Text name) + | Just HRefl <- eqTypeRep tr (typeRep @String) = mk (Col @String name) + | Just HRefl <- eqTypeRep tr (typeRep @BS.ByteString) = + mk (Col @BS.ByteString name) + | otherwise = error $ "mkSortOrder: unsupported column type: " ++ show tr -- | Dispatch aggregation by fn name and runtime column type. buildNamedExpr :: DataFrame -> AggSpec -> IO NamedExpr diff --git a/src/DataFrame.hs b/src/DataFrame.hs --- a/src/DataFrame.hs +++ b/src/DataFrame.hs @@ -250,6 +250,8 @@ ReadOptions (..), TypeSpec (..), defaultReadOptions, + fromCsv, + fromCsvBytes, readCsv, readCsvWithOpts, readSeparated, @@ -282,6 +284,7 @@ empty, null, toCsv, + toCsv', toMarkdown, toMarkdown', toSeparated, diff --git a/src/DataFrame/DecisionTree.hs b/src/DataFrame/DecisionTree.hs --- a/src/DataFrame/DecisionTree.hs +++ b/src/DataFrame/DecisionTree.hs @@ -3,6 +3,7 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} @@ -20,22 +21,82 @@ import Control.Exception (throw) import Control.Monad (guard) -import Data.Containers.ListUtils (nubOrd) import Data.Function (on) #if MIN_VERSION_base(4,20,0) -import Data.List (maximumBy, minimumBy, sort, sortBy) +import Data.List (maximumBy, minimumBy, nub, sort, sortBy) #else -import Data.List (foldl', maximumBy, minimumBy, sort, sortBy) +import Data.List (foldl', maximumBy, minimumBy, nub, sort, sortBy) #endif +import Data.Int (Int16, Int32, Int64, Int8) import qualified Data.Map.Strict as M +import Data.Proxy (Proxy (..)) import qualified Data.Text as T import Data.Type.Equality import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU -import Type.Reflection (typeRep) +import Data.Word (Word16, Word32, Word64, Word8) +import Type.Reflection (SomeTypeRep (..), typeRep) import DataFrame.Operators +{- | Declares which column types support ordering for decision tree splits. + +Use 'orderable' to register a type, and '<>' to combine: + +@ +defaultTreeConfig + { columnOrdering = defaultColumnOrdering <> orderable \@MyCustomType + } +@ +-} +newtype ColumnOrdering = ColumnOrdering (M.Map SomeTypeRep OrdDict) + +instance Semigroup ColumnOrdering where + ColumnOrdering a <> ColumnOrdering b = ColumnOrdering (a <> b) + +instance Monoid ColumnOrdering where + mempty = ColumnOrdering M.empty + +-- | Register a type as orderable for decision tree splits. +orderable :: forall a. (Columnable a, Ord a) => ColumnOrdering +orderable = ColumnOrdering (M.singleton (SomeTypeRep (typeRep @a)) (OrdDict (Proxy @a))) + +-- | All standard numeric, text, and primitive types. +defaultColumnOrdering :: ColumnOrdering +defaultColumnOrdering = + mconcat + [ orderable @Int + , orderable @Int8 + , orderable @Int16 + , orderable @Int32 + , orderable @Int64 + , orderable @Word + , orderable @Word8 + , orderable @Word16 + , orderable @Word32 + , orderable @Word64 + , orderable @Integer + , orderable @Double + , orderable @Float + , orderable @Bool + , orderable @Char + , orderable @T.Text + , orderable @String + ] + +-- Internal: existential Ord dictionary. +data OrdDict where + OrdDict :: (Columnable a, Ord a) => Proxy a -> OrdDict + +-- Internal: look up Ord for type @a@. +withOrdFrom :: + forall a r. (Columnable a) => ColumnOrdering -> ((Ord a) => r) -> Maybe r +withOrdFrom (ColumnOrdering m) k = case M.lookup (SomeTypeRep (typeRep @a)) m of + Just (OrdDict (_ :: Proxy b)) -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> Just k + Nothing -> Nothing + Nothing -> Nothing + data TreeConfig = TreeConfig { maxTreeDepth :: Int , minSamplesSplit :: Int @@ -45,8 +106,8 @@ , synthConfig :: SynthConfig , taoIterations :: Int , taoConvergenceTol :: Double + , columnOrdering :: ColumnOrdering } - deriving (Eq, Show) data SynthConfig = SynthConfig { maxExprDepth :: Int @@ -82,6 +143,7 @@ , synthConfig = defaultSynthConfig , taoIterations = 10 , taoConvergenceTol = 1e-6 + , columnOrdering = defaultColumnOrdering } data Tree a @@ -101,7 +163,7 @@ -- | Fit a TAO decision tree fitDecisionTree :: forall a. - (Columnable a) => + (Columnable a, Ord a) => TreeConfig -> Expr a -> DataFrame -> @@ -109,7 +171,7 @@ fitDecisionTree cfg (Col target) df = let conds = - nubOrd $ + nub $ numericConditions cfg (exclude [target] df) ++ generateConditionsOld cfg (exclude [target] df) @@ -124,7 +186,7 @@ taoOptimize :: forall a. - (Columnable a) => + (Columnable a, Ord a) => TreeConfig -> T.Text -> -- Target column name [Expr Bool] -> -- Candidate conditions @@ -151,7 +213,7 @@ taoIteration :: forall a. - (Columnable a) => + (Columnable a, Ord a) => TreeConfig -> T.Text -> [Expr Bool] -> @@ -168,7 +230,7 @@ optimizeDepthLevel :: forall a. - (Columnable a) => + (Columnable a, Ord a) => TreeConfig -> T.Text -> [Expr Bool] -> @@ -181,7 +243,7 @@ optimizeAtDepth :: forall a. - (Columnable a) => + (Columnable a, Ord a) => TreeConfig -> T.Text -> [Expr Bool] -> @@ -224,7 +286,7 @@ optimizeNode :: forall a. - (Columnable a) => + (Columnable a, Ord a) => TreeConfig -> T.Text -> [Expr Bool] -> @@ -333,8 +395,8 @@ identifyCarePoints target df indices leftTree rightTree = case interpret @a df (Col target) of Left _ -> [] - Right (TColumn col) -> - case toVector @a col of + Right (TColumn column) -> + case toVector @a column of Left _ -> [] Right targetVals -> V.toList $ V.mapMaybe (checkPoint targetVals) indices @@ -366,8 +428,8 @@ predictWithTree target df idx (Branch cond left right) = case interpret @Bool df cond of Left _ -> predictWithTree @a target df idx left -- Default to left on error - Right (TColumn col) -> - case toVector @Bool col of + Right (TColumn column) -> + case toVector @Bool column of Left _ -> predictWithTree @a target df idx left Right boolVals -> if boolVals V.! idx @@ -378,8 +440,8 @@ countCarePointErrors cond df carePoints = case interpret @Bool df cond of Left _ -> length carePoints - Right (TColumn col) -> - case toVector @Bool col of + Right (TColumn column) -> + case toVector @Bool column of Left _ -> length carePoints Right boolVals -> length $ filter (isMisclassified boolVals) carePoints @@ -395,15 +457,15 @@ partitionIndices cond df indices = case interpret @Bool df cond of Left _ -> (indices, V.empty) - Right (TColumn col) -> - case toVector @Bool col of + Right (TColumn column) -> + case toVector @Bool column of Left _ -> (indices, V.empty) Right boolVals -> V.partition (boolVals V.!) indices majorityValueFromIndices :: forall a. - (Columnable a) => + (Columnable a, Ord a) => T.Text -> DataFrame -> V.Vector Int -> @@ -437,8 +499,8 @@ | otherwise = case interpret @a df (Col target) of Left _ -> 1.0 - Right (TColumn col) -> - case toVector @a col of + Right (TColumn column) -> + case toVector @a column of Left _ -> 1.0 Right targetVals -> let @@ -476,7 +538,7 @@ buildGreedyTree :: forall a. - (Columnable a) => + (Columnable a, Ord a) => TreeConfig -> Int -> T.Text -> @@ -501,7 +563,7 @@ findBestGreedySplit :: forall a. - (Columnable a) => + (Columnable a, Ord a) => TreeConfig -> T.Text -> [Expr Bool] -> DataFrame -> Maybe (Expr Bool) findBestGreedySplit cfg target conds df = let @@ -637,24 +699,24 @@ numericCols :: DataFrame -> [NumExpr] numericCols df = concatMap extract (columnNames df) where - extract col = case unsafeGetColumn col df of + extract colName = case unsafeGetColumn colName df of UnboxedColumn Nothing (_ :: VU.Vector b) -> case testEquality (typeRep @b) (typeRep @Double) of - Just Refl -> [NDouble (Col col)] + Just Refl -> [NDouble (Col colName)] Nothing -> case sIntegral @b of - STrue -> [NDouble (F.toDouble (Col @b col))] + STrue -> [NDouble (F.toDouble (Col @b colName))] SFalse -> [] BoxedColumn (Just _) (_ :: V.Vector b) -> case testEquality (typeRep @b) (typeRep @Double) of - Just Refl -> [NMaybeDouble (Col @(Maybe b) col)] + Just Refl -> [NMaybeDouble (Col @(Maybe b) colName)] Nothing -> case sIntegral @b of - STrue -> [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) col))] + STrue -> [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) colName))] SFalse -> [] UnboxedColumn (Just _) (_ :: VU.Vector b) -> case testEquality (typeRep @b) (typeRep @Double) of - Just Refl -> [NMaybeDouble (Col @(Maybe b) col)] + Just Refl -> [NMaybeDouble (Col @(Maybe b) colName)] Nothing -> case sIntegral @b of - STrue -> [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) col))] + STrue -> [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) colName))] SFalse -> [] _ -> [] @@ -701,19 +763,21 @@ generateConditionsOld :: TreeConfig -> DataFrame -> [Expr Bool] generateConditionsOld cfg df = let + ords = columnOrdering cfg genConds :: T.Text -> [Expr Bool] genConds colName = case unsafeGetColumn colName df of - (BoxedColumn Nothing (col :: V.Vector a)) -> - let ps = map (Lit . (`percentileOrd'` col)) [1, 25, 75, 99] - in map (F.lift2 (==) (Col @a colName)) ps - (BoxedColumn (Just _) (col :: V.Vector a)) -> case sFloating @a of + (BoxedColumn Nothing (column :: V.Vector a)) -> + case withOrdFrom @a ords (map (Lit . (`percentileOrd'` column)) [1, 25, 75, 99]) of + Just ps -> map (F.lift2 (==) (Col @a colName)) ps + Nothing -> [] + (BoxedColumn (Just _) (column :: V.Vector a)) -> case sFloating @a of STrue -> [] -- handled by numericCols / numericExprs SFalse -> case sIntegral @a of STrue -> [] -- handled by numericCols / numericExprs SFalse -> - map - (F.lift2 (==) (Col @(Maybe a) colName) . Lit . Just . (`percentileOrd'` col)) - [1, 25, 75, 99] + case withOrdFrom @a ords (map (Lit . Just . (`percentileOrd'` column)) [1, 25, 75, 99]) of + Just ps -> map (F.lift2 (==) (Col @(Maybe a) colName)) ps + Nothing -> [] (UnboxedColumn _ (_ :: VU.Vector a)) -> [] columnConds = @@ -743,9 +807,10 @@ Nothing -> [] Just Refl -> case testEquality (typeRep @a) (typeRep @T.Text) of Nothing -> - [ F.lift2 (<=) (Col @(Maybe a) l) (Col r) - , F.lift2 (==) (Col @(Maybe a) l) (Col r) - ] + case withOrdFrom @a ords [F.lift2 (<=) (Col @(Maybe a) l) (Col r)] of + Just leExprs -> + leExprs ++ [F.lift2 (==) (Col @(Maybe a) l) (Col r)] + Nothing -> [F.lift2 (==) (Col @(Maybe a) l) (Col r)] Just Refl -> [F.lift2 (==) (Col @(Maybe a) l) (Col r)] _ -> [] in @@ -754,7 +819,8 @@ partitionDataFrame :: Expr Bool -> DataFrame -> (DataFrame, DataFrame) partitionDataFrame cond df = (filterWhere cond df, filterWhere (F.not cond) df) -calculateGini :: forall a. (Columnable a) => T.Text -> DataFrame -> Double +calculateGini :: + forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> Double calculateGini target df = let n = fromIntegral $ nRows df counts = getCounts @a target df @@ -762,19 +828,20 @@ probs = map (\c -> (fromIntegral c + 1) / (n + numClasses)) (M.elems counts) in if n == 0 then 0 else 1 - sum (map (^ 2) probs) -majorityValue :: forall a. (Columnable a) => T.Text -> DataFrame -> a +majorityValue :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> a majorityValue target df = let counts = getCounts @a target df in if M.null counts then error "Empty DataFrame in leaf" else fst $ maximumBy (compare `on` snd) (M.toList counts) -getCounts :: forall a. (Columnable a) => T.Text -> DataFrame -> M.Map a Int +getCounts :: + forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> M.Map a Int getCounts target df = case interpret @a df (Col target) of Left e -> throw e - Right (TColumn col) -> - case toVector @a col of + Right (TColumn column) -> + case toVector @a column of Left e -> throw e Right vals -> foldl' (\acc x -> M.insertWith (+) x 1 acc) M.empty (V.toList vals) @@ -782,8 +849,8 @@ percentile p expr df = case interpret @Double df expr of Left _ -> 0 - Right (TColumn col) -> - case toVector @Double col of + Right (TColumn column) -> + case toVector @Double column of Left _ -> 0 Right vals -> let sorted = V.fromList $ sort $ V.toList vals @@ -793,7 +860,7 @@ buildTree :: forall a. - (Columnable a) => + (Columnable a, Ord a) => TreeConfig -> Int -> T.Text -> @@ -810,7 +877,7 @@ findBestSplit :: forall a. - (Columnable a) => + (Columnable a, Ord a) => TreeConfig -> T.Text -> [Expr Bool] -> DataFrame -> Maybe (Expr Bool) findBestSplit = findBestGreedySplit @a @@ -823,7 +890,7 @@ -- | Compute normalised class probabilities from a subset of training rows. probsFromIndices :: forall a. - (Columnable a) => + (Columnable a, Ord a) => T.Text -> DataFrame -> V.Vector Int -> @@ -831,8 +898,8 @@ probsFromIndices target df indices = case interpret @a df (Col target) of Left _ -> M.empty - Right (TColumn col) -> - case toVector @a col of + Right (TColumn column) -> + case toVector @a column of Left _ -> M.empty Right vals -> let counts = @@ -849,7 +916,7 @@ -} buildProbTree :: forall a. - (Columnable a) => + (Columnable a, Ord a) => Tree a -> T.Text -> DataFrame -> @@ -879,7 +946,7 @@ -} fitProbTree :: forall a. - (Columnable a) => + (Columnable a, Ord a) => TreeConfig -> Expr a -> -- target column, e.g. @Col \"label\"@ DataFrame -> @@ -887,7 +954,7 @@ fitProbTree cfg (Col target) df = let conds = - nubOrd $ + nub $ numericConditions cfg (exclude [target] df) ++ generateConditionsOld cfg (exclude [target] df) initialTree = buildGreedyTree @a cfg (maxTreeDepth cfg) target conds df @@ -915,11 +982,11 @@ -} probExprs :: forall a. - (Columnable a) => + (Columnable a, Ord a) => ProbTree a -> M.Map a (Expr Double) probExprs tree = - let classes = nubOrd (allClasses tree) + let classes = nub (allClasses tree) in M.fromList [(c, classExpr c tree) | c <- classes] where allClasses :: ProbTree a -> [a] diff --git a/src/DataFrame/Display/Terminal/Plot.hs b/src/DataFrame/Display/Terminal/Plot.hs --- a/src/DataFrame/Display/Terminal/Plot.hs +++ b/src/DataFrame/Display/Terminal/Plot.hs @@ -10,6 +10,7 @@ module DataFrame.Display.Terminal.Plot where import Control.Monad +import qualified Data.Bifunctor import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T @@ -18,9 +19,10 @@ import qualified Data.Vector as V import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU +import Data.Word (Word8) import DataFrame.Internal.Types import GHC.Stack (HasCallStack) -import Type.Reflection (typeRep) +import Type.Reflection (TypeRep, typeRep) import DataFrame.Internal.Column (Column (..), Columnable, isNumeric) import qualified DataFrame.Internal.Column as D @@ -339,18 +341,54 @@ Just idx -> let col = columns df V.! idx in case col of - BoxedColumn _ vec -> - let counts = countValues vec - in Just [(T.pack (show k), fromIntegral v) | (k, v) <- counts] - UnboxedColumn _ vec -> - let counts = countValuesUnboxed vec - in Just [(T.pack (show k), fromIntegral v) | (k, v) <- counts] + BoxedColumn _ (vec :: V.Vector a) -> + Just (countBoxed (typeRep @a) vec) + UnboxedColumn _ (vec :: VU.Vector a) -> + Just (countUnboxed (typeRep @a) vec) where - countValues :: (Ord a, Show a) => V.Vector a -> [(a, Int)] + countBoxed :: + forall a. (Show a) => TypeRep a -> V.Vector a -> [(T.Text, Double)] + countBoxed tr vec + | Just Refl <- testEquality tr (typeRep @T.Text) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @String) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Integer) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Int) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Double) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Float) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Bool) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Char) = toPairs $ countValues vec + | otherwise = countByShow $ V.toList vec + + countUnboxed :: + forall a. (Show a, VU.Unbox a) => TypeRep a -> VU.Vector a -> [(T.Text, Double)] + countUnboxed tr vec + | Just Refl <- testEquality tr (typeRep @Int) = toPairs $ countValuesUnboxed vec + | Just Refl <- testEquality tr (typeRep @Double) = + toPairs $ countValuesUnboxed vec + | Just Refl <- testEquality tr (typeRep @Float) = + toPairs $ countValuesUnboxed vec + | Just Refl <- testEquality tr (typeRep @Bool) = + toPairs $ countValuesUnboxed vec + | Just Refl <- testEquality tr (typeRep @Char) = + toPairs $ countValuesUnboxed vec + | Just Refl <- testEquality tr (typeRep @Word8) = + toPairs $ countValuesUnboxed vec + | otherwise = countByShow $ VU.toList vec + + toPairs :: (Show a) => [(a, Int)] -> [(T.Text, Double)] + toPairs = map (\(k, v) -> (T.pack (show k), fromIntegral v)) + + countValues :: (Ord a) => V.Vector a -> [(a, Int)] countValues vec = M.toList $ V.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec - countValuesUnboxed :: (Ord a, Show a, VU.Unbox a) => VU.Vector a -> [(a, Int)] + countValuesUnboxed :: (Ord a, VU.Unbox a) => VU.Vector a -> [(a, Int)] countValuesUnboxed vec = M.toList $ VU.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec + + countByShow :: (Show a) => [a] -> [(T.Text, Double)] + countByShow xs = + map (Data.Bifunctor.bimap T.pack fromIntegral) $ + M.toList $ + L.foldl' (\acc x -> M.insertWith (+) (show x) 1 acc) M.empty xs isNumericColumnCheck :: T.Text -> DataFrame -> Bool isNumericColumnCheck colName df = isNumericColumn df colName diff --git a/src/DataFrame/Display/Terminal/PrettyPrint.hs b/src/DataFrame/Display/Terminal/PrettyPrint.hs --- a/src/DataFrame/Display/Terminal/PrettyPrint.hs +++ b/src/DataFrame/Display/Terminal/PrettyPrint.hs @@ -62,7 +62,7 @@ separator = T.intercalate "-|-" [T.replicate width (T.singleton '-') | width <- widths] fillCols fill cols = T.intercalate " | " [fill c width col | (c, width, col) <- zip3 cs widths cols] - lines = + outputLines = if properMarkdown then T.concat ["| ", fillCols colTitleFill consolidatedHeader, " |"] @@ -75,4 +75,4 @@ : fillCols colTitleFill types : separator : map (fillCols colValueFill) rows - in T.unlines lines + in T.unlines outputLines diff --git a/src/DataFrame/Display/Web/Plot.hs b/src/DataFrame/Display/Web/Plot.hs --- a/src/DataFrame/Display/Web/Plot.hs +++ b/src/DataFrame/Display/Web/Plot.hs @@ -10,6 +10,7 @@ module DataFrame.Display.Web.Plot where import Control.Monad +import qualified Data.Bifunctor import Data.Char import qualified Data.List as L import qualified Data.Map as M @@ -19,9 +20,10 @@ import qualified Data.Vector as V import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU +import Data.Word (Word8) import GHC.Stack (HasCallStack) import System.Random (newStdGen, randomRs) -import Type.Reflection (typeRep) +import Type.Reflection (TypeRep, typeRep) import DataFrame.Internal.Column (Column (..), Columnable, isNumeric) import qualified DataFrame.Internal.Column as D @@ -895,19 +897,56 @@ let col = columns df V.! idx in case col of BoxedColumn _ (vec :: V.Vector a) -> - let counts = countValues vec - in case testEquality (typeRep @a) (typeRep @T.Text) of - Nothing -> Just [(T.pack (show k), fromIntegral v) | (k, v) <- counts] - Just Refl -> Just [(k, fromIntegral v) | (k, v) <- counts] - UnboxedColumn _ vec -> - let counts = countValuesUnboxed vec - in Just [(T.pack (show k), fromIntegral v) | (k, v) <- counts] + Just (countBoxed (typeRep @a) vec) + UnboxedColumn _ (vec :: VU.Vector a) -> + Just (countUnboxed (typeRep @a) vec) where - countValues :: (Ord a, Show a) => V.Vector a -> [(a, Int)] + countBoxed :: + forall a. (Show a) => TypeRep a -> V.Vector a -> [(T.Text, Double)] + countBoxed tr vec + | Just Refl <- testEquality tr (typeRep @T.Text) = toPairsText $ countValues vec + | Just Refl <- testEquality tr (typeRep @String) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Integer) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Int) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Double) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Float) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Bool) = toPairs $ countValues vec + | Just Refl <- testEquality tr (typeRep @Char) = toPairs $ countValues vec + | otherwise = countByShow $ V.toList vec + + countUnboxed :: + forall a. (Show a, VU.Unbox a) => TypeRep a -> VU.Vector a -> [(T.Text, Double)] + countUnboxed tr vec + | Just Refl <- testEquality tr (typeRep @Int) = toPairs $ countValuesUnboxed vec + | Just Refl <- testEquality tr (typeRep @Double) = + toPairs $ countValuesUnboxed vec + | Just Refl <- testEquality tr (typeRep @Float) = + toPairs $ countValuesUnboxed vec + | Just Refl <- testEquality tr (typeRep @Bool) = + toPairs $ countValuesUnboxed vec + | Just Refl <- testEquality tr (typeRep @Char) = + toPairs $ countValuesUnboxed vec + | Just Refl <- testEquality tr (typeRep @Word8) = + toPairs $ countValuesUnboxed vec + | otherwise = countByShow $ VU.toList vec + + toPairs :: (Show a) => [(a, Int)] -> [(T.Text, Double)] + toPairs = map (\(k, v) -> (T.pack (show k), fromIntegral v)) + + toPairsText :: [(T.Text, Int)] -> [(T.Text, Double)] + toPairsText = map (Data.Bifunctor.second fromIntegral) + + countValues :: (Ord a) => V.Vector a -> [(a, Int)] countValues vec = M.toList $ V.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec - countValuesUnboxed :: (Ord a, Show a, VU.Unbox a) => VU.Vector a -> [(a, Int)] + countValuesUnboxed :: (Ord a, VU.Unbox a) => VU.Vector a -> [(a, Int)] countValuesUnboxed vec = M.toList $ VU.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec + + countByShow :: (Show a) => [a] -> [(T.Text, Double)] + countByShow xs = + map (Data.Bifunctor.bimap T.pack fromIntegral) $ + M.toList $ + L.foldl' (\acc x -> M.insertWith (+) (show x) 1 acc) M.empty xs groupWithOther :: Int -> [(T.Text, Double)] -> [(T.Text, Double)] groupWithOther n items = diff --git a/src/DataFrame/Errors.hs b/src/DataFrame/Errors.hs --- a/src/DataFrame/Errors.hs +++ b/src/DataFrame/Errors.hs @@ -108,14 +108,14 @@ ++ "]" typeMismatchError :: String -> String -> String -typeMismatchError givenType expectedType = +typeMismatchError givenType expType = red $ red "\n\n[Error]: Type Mismatch" ++ "\n\tWhile running your code I tried to " ++ "get a column of type: " ++ red (show givenType) ++ " but the column in the dataframe was actually of type: " - ++ green (show expectedType) + ++ green (show expType) emptyDataSetError :: T.Text -> String emptyDataSetError callPoint = diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs --- a/src/DataFrame/Functions.hs +++ b/src/DataFrame/Functions.hs @@ -2,6 +2,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} +{-# LANGUAGE IncoherentInstances #-} {-# LANGUAGE InstanceSigs #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} @@ -125,16 +126,16 @@ values are wrapped in 'Just'. Existing 'Nothing' in optional source columns stays as 'Nothing'. -} -cast :: forall a. (Columnable a) => T.Text -> Expr (Maybe a) -cast name = CastWith name "cast" (either (const Nothing) Just) +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) => a -> T.Text -> Expr a -castWithDefault def name = - CastWith name ("castWithDefault:" <> T.pack (show def)) (fromRight def) +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 @@ -142,27 +143,29 @@ them downstream. Existing 'Nothing' in optional source columns becomes @Left \"null\"@. -} -castEither :: forall a. (Columnable a) => T.Text -> Expr (Either T.Text a) -castEither name = CastWith name "castEither" (either (Left . T.pack) Right) +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) => T.Text -> Expr a -unsafeCast name = +unsafeCast :: forall a. (Columnable a, Read a) => T.Text -> Expr a +unsafeCast colName = CastWith - name + colName "unsafeCast" (fromRight (error "unsafeCast: unexpected Nothing in column")) castExpr :: - forall b src. (Columnable b, Columnable src) => Expr src -> Expr (Maybe b) + 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) => b -> Expr src -> Expr b + 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)) @@ -170,14 +173,14 @@ castExprEither :: forall b src. - (Columnable b, Columnable src) => Expr src -> Expr (Either T.Text b) + (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) => Expr src -> Expr b + forall b src. (Columnable b, Columnable src, Read b) => Expr src -> Expr b unsafeCastExpr = CastExprWith @b @b @src "unsafeCastExpr" @@ -230,9 +233,21 @@ 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 #-} 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 = @@ -245,21 +260,58 @@ . V.foldl' (\m e -> M.insertWith (+) e 1 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 = @@ -271,18 +323,58 @@ (\(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 = @@ -301,59 +393,120 @@ 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) => [(a, b)] -> Expr a -> Expr (Maybe b) + (Columnable a, Columnable b, Show (a, b)) => + [(a, b)] -> Expr a -> Expr (Maybe b) recode mapping = Unary ( MkUnaryOp @@ -367,12 +520,12 @@ forall a b. (Columnable a, Columnable b) => Expr b -> [(Expr a -> Expr Bool, b)] -> Expr a -> Expr b -recodeWithCondition fallback [] value = fallback -recodeWithCondition fallback ((cond, value) : rest) expr = ifThenElse (cond expr) (lit value) (recodeWithCondition fallback rest expr) +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) => b -> [(a, b)] -> Expr a -> Expr b + (Columnable a, Columnable b, Show (a, b)) => b -> [(a, b)] -> Expr a -> Expr b recodeWithDefault d mapping = Unary ( MkUnaryOp @@ -508,7 +661,7 @@ | otherwise -> do m <- lookupTypeName t case m of - Just name -> pure (ConT name) + Just tyName -> pure (ConT tyName) Nothing -> fail $ "Unsupported type: " ++ t0 typeFromString [tycon, t1] = AppT <$> typeFromString [tycon] <*> typeFromString [t1] typeFromString [tycon, t1, t2] = @@ -568,14 +721,14 @@ in fromNamedColumns (map (schemaElemToColumn nullableCols) leafElems) schemaElemToColumn :: S.Set T.Text -> SchemaElement -> (T.Text, Column) -schemaElemToColumn nullableCols elem = - let name = elementName elem - isNull = name `S.member` nullableCols - col = +schemaElemToColumn nullableCols element = + let colName = elementName element + isNull = colName `S.member` nullableCols + column = if isNull - then emptyNullableColumnForType (elementType elem) - else emptyColumnForType (elementType elem) - in (name, col) + then emptyNullableColumnForType (elementType element) + else emptyColumnForType (elementType element) + in (colName, column) emptyColumnForType :: TType -> Column emptyColumnForType = \case @@ -621,7 +774,7 @@ types = map (columnTypeString . (`unsafeGetColumn` df)) names specs = zipWith - ( \name type_ -> (name, maybe "" (sanitize . (<> "_")) prefix <> sanitize name, type_) + ( \colName type_ -> (colName, maybe "" (sanitize . (<> "_")) prefix <> sanitize colName, type_) ) names types diff --git a/src/DataFrame/IO/CSV.hs b/src/DataFrame/IO/CSV.hs --- a/src/DataFrame/IO/CSV.hs +++ b/src/DataFrame/IO/CSV.hs @@ -13,7 +13,6 @@ import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as C import qualified Data.ByteString.Lazy as BL -import qualified Data.List as L import qualified Data.Map.Strict as M import qualified Data.Proxy as P import qualified Data.Text as T @@ -28,11 +27,11 @@ import qualified Data.Csv.Streaming as CsvStream import Control.DeepSeq +import Control.Exception (SomeException, catch) import Control.Monad import Data.Char import qualified Data.Csv as Csv import Data.Either -import Data.Function (on) import Data.Functor import Data.IORef import Data.Maybe @@ -548,6 +547,17 @@ DataFrame -> IO () writeSeparated c filepath df = TIO.writeFile filepath (toSeparated c df) + +-- | Parse a CSV string into a DataFrame using default options. +fromCsv :: String -> IO (Either String DataFrame) +fromCsv s = do + let bs = BL.fromStrict (TE.encodeUtf8 (T.pack s)) + (Right <$> decodeSeparated defaultReadOptions bs) + `catch` (\(e :: SomeException) -> pure (Left (show e))) + +-- | Parse a lazy 'ByteString' containing CSV data into a DataFrame using default options. +fromCsvBytes :: BL.ByteString -> IO DataFrame +fromCsvBytes = decodeSeparated defaultReadOptions stripQuotes :: T.Text -> T.Text stripQuotes txt = diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs --- a/src/DataFrame/IO/Parquet.hs +++ b/src/DataFrame/IO/Parquet.hs @@ -64,6 +64,7 @@ 1. predicate filtering 2. column projection 3. row range +4. safe column promotion Column selection for @selectedColumns@ uses leaf column names only. -} @@ -198,7 +199,7 @@ let schemaElements = schema fileMetadata let sNodes = parseAll (drop 1 schemaElements) let getTypeLength :: [String] -> Maybe Int32 - getTypeLength path = findTypeLength schemaElements path 0 + getTypeLength colPath = findTypeLength schemaElements colPath 0 where findTypeLength [] _ _ = Nothing findTypeLength (s : ss) targetPath depth diff --git a/src/DataFrame/IO/Parquet/Binary.hs b/src/DataFrame/IO/Parquet/Binary.hs --- a/src/DataFrame/IO/Parquet/Binary.hs +++ b/src/DataFrame/IO/Parquet/Binary.hs @@ -26,37 +26,37 @@ -} loop :: BS.ByteString -> Word64 -> Int -> Int -> (Word64, BS.ByteString) loop bs result _ 10 = (result, bs) - loop xs result shift i = case BS.uncons xs of + loop xs' result shiftAmt i = case BS.uncons xs' of Nothing -> error "readUVarInt: not enough input bytes" - Just (b, bs) -> + Just (b, bs') -> if b < 0x80 - then (result .|. (fromIntegral b `shiftL` shift), bs) + then (result .|. (fromIntegral b `shiftL` shiftAmt), bs') else let payloadBits = fromIntegral (b .&. 0x7f) :: Word64 - in loop bs (result .|. (payloadBits `shiftL` shift)) (shift + 7) (i + 1) + in loop bs' (result .|. (payloadBits `shiftL` shiftAmt)) (shiftAmt + 7) (i + 1) readVarIntFromBytes :: (Integral a) => BS.ByteString -> (a, BS.ByteString) -readVarIntFromBytes bs = (fromIntegral n, rem) +readVarIntFromBytes bs = (fromIntegral n, remainder) where - (n, rem) = loop 0 0 bs - loop shift result bs = case BS.uncons bs of + (n, remainder) = loop 0 0 bs + loop shiftAmt result bs' = case BS.uncons bs' of Nothing -> (result, BS.empty) Just (x, xs) -> - let res = result .|. (fromIntegral (x .&. 0x7f) :: Integer) `shiftL` shift - in if x .&. 0x80 /= 0x80 then (res, xs) else loop (shift + 7) res xs + let res = result .|. (fromIntegral (x .&. 0x7f) :: Integer) `shiftL` shiftAmt + in if x .&. 0x80 /= 0x80 then (res, xs) else loop (shiftAmt + 7) res xs readIntFromBytes :: (Integral a) => BS.ByteString -> (a, BS.ByteString) readIntFromBytes bs = - let (n, rem) = readVarIntFromBytes bs + let (n, remainder) = readVarIntFromBytes bs u = fromIntegral n :: Word32 - in (fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), rem) + in (fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), remainder) readInt32FromBytes :: BS.ByteString -> (Int32, BS.ByteString) readInt32FromBytes bs = - let (n', rem) = readVarIntFromBytes @Int64 bs + let (n', remainder) = readVarIntFromBytes @Int64 bs n = fromIntegral n' :: Int32 u = fromIntegral n :: Word32 - in ((fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), rem) + in ((fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), remainder) readAndAdvance :: IORef Int -> BS.ByteString -> IO Word8 readAndAdvance bufferPos buffer = do @@ -68,12 +68,12 @@ readVarIntFromBuffer :: (Integral a) => BS.ByteString -> IORef Int -> IO a readVarIntFromBuffer buf bufferPos = do start <- readIORef bufferPos - let loop i shift result = do + let loop i shiftAmt result = do b <- readAndAdvance bufferPos buf - let res = result .|. (fromIntegral (b .&. 0x7f) :: Integer) `shiftL` shift + let res = result .|. (fromIntegral (b .&. 0x7f) :: Integer) `shiftL` shiftAmt if b .&. 0x80 /= 0x80 then return res - else loop (i + 1) (shift + 7) res + else loop (i + 1) (shiftAmt + 7) res fromIntegral <$> loop start 0 0 readIntFromBuffer :: (Integral a) => BS.ByteString -> IORef Int -> IO a @@ -96,9 +96,9 @@ readByteStringFromBytes :: BS.ByteString -> (BS.ByteString, BS.ByteString) readByteStringFromBytes xs = let - (size, rem) = readVarIntFromBytes @Int xs + (size, remainder) = readVarIntFromBytes @Int xs in - BS.splitAt size rem + BS.splitAt size remainder readByteString :: BS.ByteString -> IORef Int -> IO BS.ByteString readByteString buf pos = do diff --git a/src/DataFrame/IO/Parquet/Page.hs b/src/DataFrame/IO/Parquet/Page.hs --- a/src/DataFrame/IO/Parquet/Page.hs +++ b/src/DataFrame/IO/Parquet/Page.hs @@ -38,9 +38,9 @@ if BS.null columnBytes then pure (Nothing, BS.empty) else do - let (hdr, rem) = readPageHeader emptyPageHeader columnBytes 0 + let (hdr, remainder) = readPageHeader emptyPageHeader columnBytes 0 - let compressed = BS.take (fromIntegral $ compressedPageSize hdr) rem + let compressed = BS.take (fromIntegral $ compressedPageSize hdr) remainder fullData <- case c of ZSTD -> do @@ -65,7 +65,7 @@ other -> error ("Unsupported compression type: " ++ show other) pure ( Just $ Page hdr fullData - , BS.drop (fromIntegral $ compressedPageSize hdr) rem + , BS.drop (fromIntegral $ compressedPageSize hdr) remainder ) readPageHeader :: @@ -79,46 +79,46 @@ in case fieldContents of Nothing -> (hdr, BS.drop 1 xs) - Just (rem, elemType, identifier) -> case identifier of + Just (remainder, elemType, identifier) -> case identifier of 1 -> let - (pType, rem') = readInt32FromBytes rem + (pType, remainder') = readInt32FromBytes remainder in - readPageHeader (hdr{pageHeaderPageType = pageTypeFromInt pType}) rem' identifier + readPageHeader (hdr{pageHeaderPageType = pageTypeFromInt pType}) remainder' identifier 2 -> let - (uncompressedPageSize, rem') = readInt32FromBytes rem + (uncompressedPageSize, remainder') = readInt32FromBytes remainder in readPageHeader (hdr{uncompressedPageSize = uncompressedPageSize}) - rem' + remainder' identifier 3 -> let - (compressedPageSize, rem') = readInt32FromBytes rem + (compressedPageSize, remainder') = readInt32FromBytes remainder in - readPageHeader (hdr{compressedPageSize = compressedPageSize}) rem' identifier + readPageHeader (hdr{compressedPageSize = compressedPageSize}) remainder' identifier 4 -> let - (crc, rem') = readInt32FromBytes rem + (crc, remainder') = readInt32FromBytes remainder in - readPageHeader (hdr{pageHeaderCrcChecksum = crc}) rem' identifier + readPageHeader (hdr{pageHeaderCrcChecksum = crc}) remainder' identifier 5 -> let - (dataPageHeader, rem') = readPageTypeHeader emptyDataPageHeader rem 0 + (dataPageHeader, remainder') = readPageTypeHeader emptyDataPageHeader remainder 0 in - readPageHeader (hdr{pageTypeHeader = dataPageHeader}) rem' identifier + readPageHeader (hdr{pageTypeHeader = dataPageHeader}) remainder' identifier 6 -> error "Index page header not supported" 7 -> let - (dictionaryPageHeader, rem') = readPageTypeHeader emptyDictionaryPageHeader rem 0 + (dictionaryPageHeader, remainder') = readPageTypeHeader emptyDictionaryPageHeader remainder 0 in - readPageHeader (hdr{pageTypeHeader = dictionaryPageHeader}) rem' identifier + readPageHeader (hdr{pageTypeHeader = dictionaryPageHeader}) remainder' identifier 8 -> let - (dataPageHeaderV2, rem') = readPageTypeHeader emptyDataPageHeaderV2 rem 0 + (dataPageHeaderV2, remainder') = readPageTypeHeader emptyDataPageHeaderV2 remainder 0 in - readPageHeader (hdr{pageTypeHeader = dataPageHeaderV2}) rem' identifier + readPageHeader (hdr{pageTypeHeader = dataPageHeaderV2}) remainder' identifier n -> error $ "Unknown page header field " ++ show n readPageTypeHeader :: @@ -134,26 +134,26 @@ in case fieldContents of Nothing -> (hdr, BS.drop 1 xs) - Just (rem, elemType, identifier) -> case identifier of + Just (remainder, elemType, identifier) -> case identifier of 1 -> let - (numValues, rem') = readInt32FromBytes rem + (numValues, remainder') = readInt32FromBytes remainder in readPageTypeHeader (hdr{dictionaryPageHeaderNumValues = numValues}) - rem' + remainder' identifier 2 -> let - (enc, rem') = readInt32FromBytes rem + (enc, remainder') = readInt32FromBytes remainder in readPageTypeHeader (hdr{dictionaryPageHeaderEncoding = parquetEncodingFromInt enc}) - rem' + remainder' identifier 3 -> let - isSorted = fromMaybe (error "readPageTypeHeader: not enough bytes") (rem BS.!? 0) + isSorted = fromMaybe (error "readPageTypeHeader: not enough bytes") (remainder BS.!? 0) in readPageTypeHeader (hdr{dictionaryPageIsSorted = isSorted == compactBooleanTrue}) @@ -163,7 +163,7 @@ -- But in other cases you do. -- This might become a problem later but in the mean -- time I'm not dropping (this assumes this is the common case). - rem + remainder identifier n -> error $ "readPageTypeHeader: unsupported identifier " ++ show n @@ -176,41 +176,41 @@ in case fieldContents of Nothing -> (hdr, BS.drop 1 xs) - Just (rem, elemType, identifier) -> case identifier of + Just (remainder, elemType, identifier) -> case identifier of 1 -> let - (numValues, rem') = readInt32FromBytes rem + (numValues, remainder') = readInt32FromBytes remainder in - readPageTypeHeader (hdr{dataPageHeaderNumValues = numValues}) rem' identifier + readPageTypeHeader (hdr{dataPageHeaderNumValues = numValues}) remainder' identifier 2 -> let - (enc, rem') = readInt32FromBytes rem + (enc, remainder') = readInt32FromBytes remainder in readPageTypeHeader (hdr{dataPageHeaderEncoding = parquetEncodingFromInt enc}) - rem' + remainder' identifier 3 -> let - (enc, rem') = readInt32FromBytes rem + (enc, remainder') = readInt32FromBytes remainder in readPageTypeHeader (hdr{definitionLevelEncoding = parquetEncodingFromInt enc}) - rem' + remainder' identifier 4 -> let - (enc, rem') = readInt32FromBytes rem + (enc, remainder') = readInt32FromBytes remainder in readPageTypeHeader (hdr{repetitionLevelEncoding = parquetEncodingFromInt enc}) - rem' + remainder' identifier 5 -> let - (stats, rem') = readStatisticsFromBytes emptyColumnStatistics rem 0 + (stats, remainder') = readStatisticsFromBytes emptyColumnStatistics remainder 0 in - readPageTypeHeader (hdr{dataPageHeaderStatistics = stats}) rem' identifier + readPageTypeHeader (hdr{dataPageHeaderStatistics = stats}) remainder' identifier n -> error $ show n readPageTypeHeader hdr@(DataPageHeaderV2{..}) xs lastFieldId = if BS.null xs @@ -221,57 +221,57 @@ in case fieldContents of Nothing -> (hdr, BS.drop 1 xs) - Just (rem, elemType, identifier) -> case identifier of + Just (remainder, elemType, identifier) -> case identifier of 1 -> let - (numValues, rem') = readInt32FromBytes rem + (numValues, remainder') = readInt32FromBytes remainder in - readPageTypeHeader (hdr{dataPageHeaderV2NumValues = numValues}) rem' identifier + readPageTypeHeader (hdr{dataPageHeaderV2NumValues = numValues}) remainder' identifier 2 -> let - (numNulls, rem') = readInt32FromBytes rem + (numNulls, remainder') = readInt32FromBytes remainder in - readPageTypeHeader (hdr{dataPageHeaderV2NumNulls = numNulls}) rem' identifier + readPageTypeHeader (hdr{dataPageHeaderV2NumNulls = numNulls}) remainder' identifier 3 -> let - (numRows, rem') = readInt32FromBytes rem + (numRows, remainder') = readInt32FromBytes remainder in - readPageTypeHeader (hdr{dataPageHeaderV2NumRows = numRows}) rem' identifier + readPageTypeHeader (hdr{dataPageHeaderV2NumRows = numRows}) remainder' identifier 4 -> let - (enc, rem') = readInt32FromBytes rem + (enc, remainder') = readInt32FromBytes remainder in readPageTypeHeader (hdr{dataPageHeaderV2Encoding = parquetEncodingFromInt enc}) - rem' + remainder' identifier 5 -> let - (n, rem') = readInt32FromBytes rem + (n, remainder') = readInt32FromBytes remainder in - readPageTypeHeader (hdr{definitionLevelByteLength = n}) rem' identifier + readPageTypeHeader (hdr{definitionLevelByteLength = n}) remainder' identifier 6 -> let - (n, rem') = readInt32FromBytes rem + (n, remainder') = readInt32FromBytes remainder in - readPageTypeHeader (hdr{repetitionLevelByteLength = n}) rem' identifier + readPageTypeHeader (hdr{repetitionLevelByteLength = n}) remainder' identifier 7 -> let - (isCompressed, rem') = case BS.uncons rem of + (isCompressed, remainder') = case BS.uncons remainder of Just (b, bytes) -> ((b .&. 0x0f) == compactBooleanTrue, bytes) Nothing -> (True, BS.empty) in readPageTypeHeader (hdr{dataPageHeaderV2IsCompressed = isCompressed}) - rem' + remainder' identifier 8 -> let - (stats, rem') = readStatisticsFromBytes emptyColumnStatistics rem 0 + (stats, remainder') = readStatisticsFromBytes emptyColumnStatistics remainder 0 in readPageTypeHeader (hdr{dataPageHeaderV2Statistics = stats}) - rem' + remainder' identifier n -> error $ show n @@ -283,12 +283,12 @@ then Nothing else let modifier = fromIntegral ((x .&. 0xf0) `shiftR` 4) :: Int16 - (identifier, rem) = + (identifier, remainder) = if modifier == 0 then readIntFromBytes @Int16 xs else (lastFieldId + modifier, xs) elemType = toTType (x .&. 0x0f) - in Just (rem, elemType, identifier) + in Just (remainder, elemType, identifier) readAllPages :: CompressionCodec -> BS.ByteString -> IO [Page] readAllPages codec bytes = go bytes [] @@ -297,10 +297,10 @@ if BS.null bs then return (reverse acc) else do - (maybePage, remaining) <- readPage codec bs + (maybePage, remainderaining) <- readPage codec bs case maybePage of Nothing -> return (reverse acc) - Just page -> go remaining (page : acc) + Just page -> go remainderaining (page : acc) -- | Read n Int32 values directly into an unboxed vector (no intermediate list). readNInt32Vec :: Int -> BS.ByteString -> VU.Vector Int32 @@ -394,53 +394,53 @@ in case fieldContents of Nothing -> (cs, BS.drop 1 xs) - Just (rem, elemType, identifier) -> case identifier of + Just (remainder, elemType, identifier) -> case identifier of 1 -> let - (maxInBytes, rem') = readByteStringFromBytes rem + (maxInBytes, remainder') = readByteStringFromBytes remainder in - readStatisticsFromBytes (cs{columnMax = maxInBytes}) rem' identifier + readStatisticsFromBytes (cs{columnMax = maxInBytes}) remainder' identifier 2 -> let - (minInBytes, rem') = readByteStringFromBytes rem + (minInBytes, remainder') = readByteStringFromBytes remainder in - readStatisticsFromBytes (cs{columnMin = minInBytes}) rem' identifier + readStatisticsFromBytes (cs{columnMin = minInBytes}) remainder' identifier 3 -> let - (nullCount, rem') = readIntFromBytes @Int64 rem + (nullCount, remainder') = readIntFromBytes @Int64 remainder in - readStatisticsFromBytes (cs{columnNullCount = nullCount}) rem' identifier + readStatisticsFromBytes (cs{columnNullCount = nullCount}) remainder' identifier 4 -> let - (distinctCount, rem') = readIntFromBytes @Int64 rem + (distinctCount, remainder') = readIntFromBytes @Int64 remainder in - readStatisticsFromBytes (cs{columnDistictCount = distinctCount}) rem' identifier + readStatisticsFromBytes (cs{columnDistictCount = distinctCount}) remainder' identifier 5 -> let - (maxInBytes, rem') = readByteStringFromBytes rem + (maxInBytes, remainder') = readByteStringFromBytes remainder in - readStatisticsFromBytes (cs{columnMaxValue = maxInBytes}) rem' identifier + readStatisticsFromBytes (cs{columnMaxValue = maxInBytes}) remainder' identifier 6 -> let - (minInBytes, rem') = readByteStringFromBytes rem + (minInBytes, remainder') = readByteStringFromBytes remainder in - readStatisticsFromBytes (cs{columnMinValue = minInBytes}) rem' identifier + readStatisticsFromBytes (cs{columnMinValue = minInBytes}) remainder' identifier 7 -> - case BS.uncons rem of + case BS.uncons remainder of Nothing -> error "readStatisticsFromBytes: not enough bytes" - Just (isMaxValueExact, rem') -> + Just (isMaxValueExact, remainder') -> readStatisticsFromBytes (cs{isColumnMaxValueExact = isMaxValueExact == compactBooleanTrue}) - rem' + remainder' identifier 8 -> - case BS.uncons rem of + case BS.uncons remainder of Nothing -> error "readStatisticsFromBytes: not enough bytes" - Just (isMinValueExact, rem') -> + Just (isMinValueExact, remainder') -> readStatisticsFromBytes (cs{isColumnMinValueExact = isMinValueExact == compactBooleanTrue}) - rem' + remainder' identifier n -> error $ show n diff --git a/src/DataFrame/IO/Parquet/Thrift.hs b/src/DataFrame/IO/Parquet/Thrift.hs --- a/src/DataFrame/IO/Parquet/Thrift.hs +++ b/src/DataFrame/IO/Parquet/Thrift.hs @@ -362,9 +362,9 @@ Nothing -> return metadata Just (elemType, identifier) -> case identifier of 1 -> do - version <- readIntFromBuffer @Int32 metaDataBuf bufferPos + parsedVersion <- readIntFromBuffer @Int32 metaDataBuf bufferPos readFileMetaData - (metadata{version = version}) + (metadata{version = parsedVersion}) metaDataBuf bufferPos identifier @@ -385,9 +385,9 @@ bufferPos identifier 3 -> do - numRows <- readIntFromBuffer @Int64 metaDataBuf bufferPos + parsedNumRows <- readIntFromBuffer @Int64 metaDataBuf bufferPos readFileMetaData - (metadata{numRows = fromIntegral numRows}) + (metadata{numRows = fromIntegral parsedNumRows}) metaDataBuf bufferPos identifier @@ -400,10 +400,10 @@ -- TODO actually check elemType agrees (also for all the other underscored _elemType in this module) let _elemType = toTType sizeAndType - rowGroups <- + parsedRowGroups <- replicateM listSize (readRowGroup emptyRowGroup metaDataBuf bufferPos 0) readFileMetaData - (metadata{rowGroups = rowGroups}) + (metadata{rowGroups = parsedRowGroups}) metaDataBuf bufferPos identifier @@ -415,17 +415,17 @@ else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) let _elemType = toTType sizeAndType - keyValueMetadata <- + parsedKeyValueMetadata <- replicateM listSize (readKeyValue emptyKeyValue metaDataBuf bufferPos 0) readFileMetaData - (metadata{keyValueMetadata = keyValueMetadata}) + (metadata{keyValueMetadata = parsedKeyValueMetadata}) metaDataBuf bufferPos identifier 6 -> do - createdBy <- readString metaDataBuf bufferPos + parsedCreatedBy <- readString metaDataBuf bufferPos readFileMetaData - (metadata{createdBy = Just createdBy}) + (metadata{createdBy = Just parsedCreatedBy}) metaDataBuf bufferPos identifier @@ -437,23 +437,23 @@ else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) let _elemType = toTType sizeAndType - columnOrders <- replicateM listSize (readColumnOrder metaDataBuf bufferPos 0) + parsedColumnOrders <- replicateM listSize (readColumnOrder metaDataBuf bufferPos 0) readFileMetaData - (metadata{columnOrders = columnOrders}) + (metadata{columnOrders = parsedColumnOrders}) metaDataBuf bufferPos identifier 8 -> do - encryptionAlgorithm <- readEncryptionAlgorithm metaDataBuf bufferPos 0 + parsedEncryptionAlgorithm <- readEncryptionAlgorithm metaDataBuf bufferPos 0 readFileMetaData - (metadata{encryptionAlgorithm = encryptionAlgorithm}) + (metadata{encryptionAlgorithm = parsedEncryptionAlgorithm}) metaDataBuf bufferPos identifier 9 -> do - footerSigningKeyMetadata <- readByteString metaDataBuf bufferPos + parsedFooterSigningKeyMetadata <- readByteString metaDataBuf bufferPos readFileMetaData - (metadata{footerSigningKeyMetadata = footerSigningKeyMetadata}) + (metadata{footerSigningKeyMetadata = parsedFooterSigningKeyMetadata}) metaDataBuf bufferPos identifier @@ -478,9 +478,9 @@ pos identifier 2 -> do - typeLength <- readInt32FromBuffer buf pos + parsedTypeLength <- readInt32FromBuffer buf pos readSchemaElement - (schemaElement{typeLength = typeLength}) + (schemaElement{typeLength = parsedTypeLength}) buf pos identifier @@ -503,40 +503,40 @@ pos identifier 5 -> do - numChildren <- readInt32FromBuffer buf pos + parsedNumChildren <- readInt32FromBuffer buf pos readSchemaElement - (schemaElement{numChildren = numChildren}) + (schemaElement{numChildren = parsedNumChildren}) buf pos identifier 6 -> do - convertedType <- readInt32FromBuffer buf pos + parsedConvertedType <- readInt32FromBuffer buf pos readSchemaElement - (schemaElement{convertedType = convertedType}) + (schemaElement{convertedType = parsedConvertedType}) buf pos identifier 7 -> do - scale <- readInt32FromBuffer buf pos - readSchemaElement (schemaElement{scale = scale}) buf pos identifier + parsedScale <- readInt32FromBuffer buf pos + readSchemaElement (schemaElement{scale = parsedScale}) buf pos identifier 8 -> do - precision <- readInt32FromBuffer buf pos + parsedPrecision <- readInt32FromBuffer buf pos readSchemaElement - (schemaElement{precision = precision}) + (schemaElement{precision = parsedPrecision}) buf pos identifier 9 -> do - fieldId <- readInt32FromBuffer buf pos + parsedFieldId <- readInt32FromBuffer buf pos readSchemaElement - (schemaElement{fieldId = fieldId}) + (schemaElement{fieldId = parsedFieldId}) buf pos identifier 10 -> do - logicalType <- readLogicalType LOGICAL_TYPE_UNKNOWN buf pos 0 + parsedLogicalType <- readLogicalType LOGICAL_TYPE_UNKNOWN buf pos 0 readSchemaElement - (schemaElement{logicalType = logicalType}) + (schemaElement{logicalType = parsedLogicalType}) buf pos identifier @@ -577,8 +577,8 @@ pos identifier 7 -> do - ordinal <- readIntFromBuffer @Int16 buf pos - readRowGroup (r{ordinal = ordinal}) buf pos identifier + parsedOrdinal <- readIntFromBuffer @Int16 buf pos + readRowGroup (r{ordinal = parsedOrdinal}) buf pos identifier _ -> error $ "Unknown row group field: " ++ show identifier readColumnChunk :: @@ -598,9 +598,9 @@ pos identifier 2 -> do - columnChunkMetadataFileOffset <- readIntFromBuffer @Int64 buf pos + parsedMetadataFileOffset <- readIntFromBuffer @Int64 buf pos readColumnChunk - (c{columnChunkMetadataFileOffset = columnChunkMetadataFileOffset}) + (c{columnChunkMetadataFileOffset = parsedMetadataFileOffset}) buf pos identifier @@ -626,16 +626,16 @@ pos identifier 6 -> do - columnChunkColumnIndexOffset <- readIntFromBuffer @Int64 buf pos + parsedColumnIndexOffset <- readIntFromBuffer @Int64 buf pos readColumnChunk - (c{columnChunkColumnIndexOffset = columnChunkColumnIndexOffset}) + (c{columnChunkColumnIndexOffset = parsedColumnIndexOffset}) buf pos identifier 7 -> do - columnChunkColumnIndexLength <- readInt32FromBuffer buf pos + parsedColumnIndexLength <- readInt32FromBuffer buf pos readColumnChunk - (c{columnChunkColumnIndexLength = columnChunkColumnIndexLength}) + (c{columnChunkColumnIndexLength = parsedColumnIndexLength}) buf pos identifier @@ -682,16 +682,16 @@ numValues <- readIntFromBuffer @Int64 buf pos readColumnMetadata (cm{columnNumValues = numValues}) buf pos identifier 6 -> do - columnTotalUncompressedSize <- readIntFromBuffer @Int64 buf pos + parsedTotalUncompressedSize <- readIntFromBuffer @Int64 buf pos readColumnMetadata - (cm{columnTotalUncompressedSize = columnTotalUncompressedSize}) + (cm{columnTotalUncompressedSize = parsedTotalUncompressedSize}) buf pos identifier 7 -> do - columnTotalCompressedSize <- readIntFromBuffer @Int64 buf pos + parsedTotalCompressedSize <- readIntFromBuffer @Int64 buf pos readColumnMetadata - (cm{columnTotalCompressedSize = columnTotalCompressedSize}) + (cm{columnTotalCompressedSize = parsedTotalCompressedSize}) buf pos identifier @@ -699,31 +699,31 @@ sizeAndType <- readAndAdvance pos buf let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int let _elemType = toTType sizeAndType - columnKeyValueMetadata <- + parsedKeyValueMeta <- replicateM sizeOnly (readKeyValue emptyKeyValue buf pos 0) readColumnMetadata - (cm{columnKeyValueMetadata = columnKeyValueMetadata}) + (cm{columnKeyValueMetadata = parsedKeyValueMeta}) buf pos identifier 9 -> do - columnDataPageOffset <- readIntFromBuffer @Int64 buf pos + parsedDataPageOffset <- readIntFromBuffer @Int64 buf pos readColumnMetadata - (cm{columnDataPageOffset = columnDataPageOffset}) + (cm{columnDataPageOffset = parsedDataPageOffset}) buf pos identifier 10 -> do - columnIndexPageOffset <- readIntFromBuffer @Int64 buf pos + parsedIndexPageOffset <- readIntFromBuffer @Int64 buf pos readColumnMetadata - (cm{columnIndexPageOffset = columnIndexPageOffset}) + (cm{columnIndexPageOffset = parsedIndexPageOffset}) buf pos identifier 11 -> do - columnDictionaryPageOffset <- readIntFromBuffer @Int64 buf pos + parsedDictionaryPageOffset <- readIntFromBuffer @Int64 buf pos readColumnMetadata - (cm{columnDictionaryPageOffset = columnDictionaryPageOffset}) + (cm{columnDictionaryPageOffset = parsedDictionaryPageOffset}) buf pos identifier @@ -742,16 +742,16 @@ pos identifier 14 -> do - bloomFilterOffset <- readIntFromBuffer @Int64 buf pos + parsedBloomFilterOffset <- readIntFromBuffer @Int64 buf pos readColumnMetadata - (cm{bloomFilterOffset = bloomFilterOffset}) + (cm{bloomFilterOffset = parsedBloomFilterOffset}) buf pos identifier 15 -> do - bloomFilterLength <- readInt32FromBuffer buf pos + parsedBloomFilterLength <- readInt32FromBuffer buf pos readColumnMetadata - (cm{bloomFilterLength = bloomFilterLength}) + (cm{bloomFilterLength = parsedBloomFilterLength}) buf pos identifier @@ -815,25 +815,25 @@ IORef Int -> Int16 -> IO EncryptionAlgorithm -readAesGcmCtrV1 v@(AesGcmCtrV1 aadPrefix aadFileUnique supplyAadPrefix) buf pos lastFieldId = do +readAesGcmCtrV1 v@(AesGcmCtrV1 _aadPrefix _aadFileUnique _supplyAadPrefix) buf pos lastFieldId = do fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return v Just (elemType, identifier) -> case identifier of 1 -> do - aadPrefix <- readByteString buf pos - readAesGcmCtrV1 (v{aadPrefix = aadPrefix}) buf pos identifier + parsedAadPrefix <- readByteString buf pos + readAesGcmCtrV1 (v{aadPrefix = parsedAadPrefix}) buf pos identifier 2 -> do - aadFileUnique <- readByteString buf pos + parsedAadFileUnique <- readByteString buf pos readAesGcmCtrV1 - (v{aadFileUnique = aadFileUnique}) + (v{aadFileUnique = parsedAadFileUnique}) buf pos identifier 3 -> do - supplyAadPrefix <- readAndAdvance pos buf + parsedSupplyAadPrefix <- readAndAdvance pos buf readAesGcmCtrV1 - (v{supplyAadPrefix = supplyAadPrefix == compactBooleanTrue}) + (v{supplyAadPrefix = parsedSupplyAadPrefix == compactBooleanTrue}) buf pos identifier @@ -847,21 +847,21 @@ IORef Int -> Int16 -> IO EncryptionAlgorithm -readAesGcmV1 v@(AesGcmV1 aadPrefix aadFileUnique supplyAadPrefix) buf pos lastFieldId = do +readAesGcmV1 v@(AesGcmV1 _aadPrefix _aadFileUnique _supplyAadPrefix) buf pos lastFieldId = do fieldContents <- readField buf pos lastFieldId case fieldContents of Nothing -> return v Just (elemType, identifier) -> case identifier of 1 -> do - aadPrefix <- readByteString buf pos - readAesGcmV1 (v{aadPrefix = aadPrefix}) buf pos identifier + parsedAadPrefix <- readByteString buf pos + readAesGcmV1 (v{aadPrefix = parsedAadPrefix}) buf pos identifier 2 -> do - aadFileUnique <- readByteString buf pos - readAesGcmV1 (v{aadFileUnique = aadFileUnique}) buf pos identifier + parsedAadFileUnique <- readByteString buf pos + readAesGcmV1 (v{aadFileUnique = parsedAadFileUnique}) buf pos identifier 3 -> do - supplyAadPrefix <- readAndAdvance pos buf + parsedSupplyAadPrefix <- readAndAdvance pos buf readAesGcmV1 - (v{supplyAadPrefix = supplyAadPrefix == compactBooleanTrue}) + (v{supplyAadPrefix = parsedSupplyAadPrefix == compactBooleanTrue}) buf pos identifier @@ -986,9 +986,9 @@ Nothing -> return ss Just (elemType, identifier) -> case identifier of 1 -> do - unencodedByteArrayDataTypes <- readIntFromBuffer @Int64 buf pos + parsedUnencodedByteArrayDataTypes <- readIntFromBuffer @Int64 buf pos readSizeStatistics - (ss{unencodedByteArrayDataTypes = unencodedByteArrayDataTypes}) + (ss{unencodedByteArrayDataTypes = parsedUnencodedByteArrayDataTypes}) buf pos identifier @@ -996,10 +996,10 @@ sizeAndType <- readAndAdvance pos buf let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int let _elemType = toTType sizeAndType - repetitionLevelHistogram <- + parsedRepetitionLevelHistogram <- replicateM sizeOnly (readIntFromBuffer @Int64 buf pos) readSizeStatistics - (ss{repetitionLevelHistogram = repetitionLevelHistogram}) + (ss{repetitionLevelHistogram = parsedRepetitionLevelHistogram}) buf pos identifier @@ -1007,10 +1007,10 @@ sizeAndType <- readAndAdvance pos buf let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int let _elemType = toTType sizeAndType - definitionLevelHistogram <- + parsedDefinitionLevelHistogram <- replicateM sizeOnly (readIntFromBuffer @Int64 buf pos) readSizeStatistics - (ss{definitionLevelHistogram = definitionLevelHistogram}) + (ss{definitionLevelHistogram = parsedDefinitionLevelHistogram}) buf pos identifier @@ -1033,10 +1033,10 @@ readLogicalType :: LogicalType -> BS.ByteString -> IORef Int -> Int16 -> IO LogicalType -readLogicalType logicalType buf pos lastFieldId = do +readLogicalType parsedLogicalType buf pos lastFieldId = do fieldContents <- readField buf pos lastFieldId case fieldContents of - Nothing -> pure logicalType + Nothing -> pure parsedLogicalType Just (elemType, identifier) -> case identifier of 1 -> do -- This is an empty enum and is read as a field. @@ -1097,10 +1097,10 @@ IORef Int -> Int16 -> IO LogicalType -readIntType bitWidth intIsSigned buf pos lastFieldId = do +readIntType parsedBitWidth parsedIntIsSigned buf pos lastFieldId = do t <- readAndAdvance pos buf if t .&. 0x0f == 0 - then return (IntType bitWidth intIsSigned) + then return (IntType parsedBitWidth parsedIntIsSigned) else do let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16 identifier <- @@ -1111,10 +1111,10 @@ case identifier of 1 -> do bitWidth' <- readAndAdvance pos buf - readIntType (fromIntegral bitWidth') intIsSigned buf pos identifier + readIntType (fromIntegral bitWidth') parsedIntIsSigned buf pos identifier 2 -> do let intIsSigned' = (t .&. 0x0f) == compactBooleanTrue - readIntType bitWidth intIsSigned' buf pos identifier + readIntType parsedBitWidth intIsSigned' buf pos identifier _ -> error $ "UNKNOWN field ID for IntType: " ++ show identifier readDecimalType :: @@ -1124,17 +1124,17 @@ IORef Int -> Int16 -> IO LogicalType -readDecimalType precision scale buf pos lastFieldId = do +readDecimalType parsedPrecision parsedScale buf pos lastFieldId = do fieldContents <- readField buf pos lastFieldId case fieldContents of - Nothing -> return (DecimalType precision scale) + Nothing -> return (DecimalType parsedPrecision parsedScale) Just (elemType, identifier) -> case identifier of 1 -> do scale' <- readInt32FromBuffer buf pos - readDecimalType precision scale' buf pos identifier + readDecimalType parsedPrecision scale' buf pos identifier 2 -> do precision' <- readInt32FromBuffer buf pos - readDecimalType precision' scale buf pos identifier + readDecimalType precision' parsedScale buf pos identifier _ -> error $ "UNKNOWN field ID for DecimalType" ++ show identifier readTimeType :: @@ -1144,17 +1144,17 @@ IORef Int -> Int16 -> IO LogicalType -readTimeType isAdjustedToUTC unit buf pos lastFieldId = do +readTimeType parsedIsAdjustedToUTC parsedUnit buf pos lastFieldId = do fieldContents <- readField buf pos lastFieldId case fieldContents of - Nothing -> return (TimeType{isAdjustedToUTC = isAdjustedToUTC, unit = unit}) + Nothing -> return (TimeType{isAdjustedToUTC = parsedIsAdjustedToUTC, unit = parsedUnit}) Just (elemType, identifier) -> case identifier of 1 -> do let isAdjustedToUTC' = elemType == toTType compactBooleanTrue - readTimeType isAdjustedToUTC' unit buf pos identifier + readTimeType isAdjustedToUTC' parsedUnit buf pos identifier 2 -> do unit' <- readUnit TIME_UNIT_UNKNOWN buf pos 0 - readTimeType isAdjustedToUTC unit' buf pos identifier + readTimeType parsedIsAdjustedToUTC unit' buf pos identifier _ -> error $ "UNKNOWN field ID for TimeType" ++ show identifier readTimestampType :: @@ -1164,24 +1164,24 @@ IORef Int -> Int16 -> IO LogicalType -readTimestampType isAdjustedToUTC unit buf pos lastFieldId = do +readTimestampType parsedIsAdjustedToUTC parsedUnit buf pos lastFieldId = do fieldContents <- readField buf pos lastFieldId case fieldContents of - Nothing -> return (TimestampType{isAdjustedToUTC = isAdjustedToUTC, unit = unit}) + Nothing -> return (TimestampType{isAdjustedToUTC = parsedIsAdjustedToUTC, unit = parsedUnit}) Just (elemType, identifier) -> case identifier of 1 -> do let isAdjustedToUTC' = elemType == toTType compactBooleanTrue - readTimestampType isAdjustedToUTC' unit buf pos identifier + readTimestampType isAdjustedToUTC' parsedUnit buf pos identifier 2 -> do unit' <- readUnit TIME_UNIT_UNKNOWN buf pos 0 - readTimestampType isAdjustedToUTC unit' buf pos identifier + readTimestampType parsedIsAdjustedToUTC unit' buf pos identifier _ -> error $ "UNKNOWN field ID for TimestampType " ++ show identifier readUnit :: TimeUnit -> BS.ByteString -> IORef Int -> Int16 -> IO TimeUnit -readUnit unit buf pos lastFieldId = do +readUnit parsedUnit buf pos lastFieldId = do fieldContents <- readField buf pos lastFieldId case fieldContents of - Nothing -> return unit + Nothing -> return parsedUnit Just (elemType, identifier) -> case identifier of 1 -> do _ <- readField buf pos 0 diff --git a/src/DataFrame/Internal/Column.hs b/src/DataFrame/Internal/Column.hs --- a/src/DataFrame/Internal/Column.hs +++ b/src/DataFrame/Internal/Column.hs @@ -20,7 +20,6 @@ import qualified Data.Text as T import qualified Data.Vector as VB -import qualified Data.Vector.Algorithms.Merge as VA import qualified Data.Vector.Generic as VG import qualified Data.Vector.Mutable as VBM import qualified Data.Vector.Unboxed as VU @@ -214,10 +213,6 @@ (==) :: (Eq a) => TypedColumn a -> TypedColumn a -> Bool (==) (TColumn a) (TColumn b) = a == b -instance (Ord a) => Ord (TypedColumn a) where - compare :: (Ord a) => TypedColumn a -> TypedColumn a -> Ordering - compare (TColumn a) (TColumn b) = compare a b - -- | Gets the underlying value from a TypedColumn. unwrapTypedColumn :: TypedColumn a -> Column unwrapTypedColumn (TColumn value) = value @@ -370,19 +365,6 @@ ) (==) _ _ = False --- Generalised LEQ that does reflection. -generalLEQ :: - forall a b. (Typeable a, Typeable b, Ord a, Ord b) => a -> b -> Bool -generalLEQ x y = case testEquality (typeRep @a) (typeRep @b) of - Nothing -> False - Just Refl -> x <= y - -instance Ord Column where - (<=) :: Column -> Column -> Bool - (<=) (BoxedColumn _ (a :: VB.Vector t1)) (BoxedColumn _ (b :: VB.Vector t2)) = generalLEQ a b - (<=) (UnboxedColumn _ (a :: VU.Vector t1)) (UnboxedColumn _ (b :: VU.Vector t2)) = generalLEQ a b - (<=) _ _ = False - {- | A class for converting a vector to a column of the appropriate type. Given each Rep we tell the `toColumnRep` function which Column type to pick. -} @@ -708,7 +690,7 @@ (a -> Bool) -> Column -> Either DataFrameException (VU.Vector Int) -findIndices pred = \case +findIndices predicate = \case BoxedColumn _ (v :: VB.Vector b) -> run v VG.convert UnboxedColumn _ (v :: VU.Vector b) -> run v id where @@ -719,7 +701,7 @@ (v Int -> VU.Vector Int) -> Either DataFrameException (VU.Vector Int) run column finalize = case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> Right . finalize $ VG.findIndices pred column + Just Refl -> Right . finalize $ VG.findIndices predicate column Nothing -> Left $ TypeMismatchException @@ -729,23 +711,6 @@ , callingFunctionName = Just "findIndices" , errorColumnName = Nothing } - --- | An internal function that returns a vector of how indexes change after a column is sorted. -sortedIndexes :: Bool -> Column -> VU.Vector Int -sortedIndexes asc = \case - BoxedColumn _ column -> sortWorker VG.convert column - UnboxedColumn _ column -> sortWorker id column - where - sortWorker :: - (VG.Vector v a, Ord a, VG.Vector v (Int, a), VG.Vector v Int) => - (v Int -> VU.Vector Int) -> v a -> VU.Vector Int - sortWorker finalize column = runST $ do - withIndexes <- VG.thaw $ VG.indexed column - let cmp = if asc then compare else flip compare - VA.sortBy (\(_, b) (_, b') -> cmp b b') withIndexes - sorted <- VG.unsafeFreeze withIndexes - return $ finalize $ VG.map fst sorted -{-# INLINE sortedIndexes #-} -- | Fold (right) column with index. ifoldrColumn :: diff --git a/src/DataFrame/Internal/DataFrame.hs b/src/DataFrame/Internal/DataFrame.hs --- a/src/DataFrame/Internal/DataFrame.hs +++ b/src/DataFrame/Internal/DataFrame.hs @@ -207,6 +207,10 @@ toCsv :: DataFrame -> T.Text toCsv = toSeparated ',' +-- | Convert a DataFrame to a CSV (comma-separated) string. +toCsv' :: DataFrame -> String +toCsv' = T.unpack . toSeparated ',' + -- | Convert a DataFrame to a text representation with a custom separator. toSeparated :: Char -> DataFrame -> T.Text toSeparated sep df diff --git a/src/DataFrame/Internal/Expression.hs b/src/DataFrame/Internal/Expression.hs --- a/src/DataFrame/Internal/Expression.hs +++ b/src/DataFrame/Internal/Expression.hs @@ -53,13 +53,13 @@ data Expr a where Col :: (Columnable a) => T.Text -> Expr a CastWith :: - (Columnable a, Columnable b) => + (Columnable a, Columnable b, Read a) => T.Text -> T.Text -> (Either String a -> b) -> Expr b CastExprWith :: - (Columnable a, Columnable b, Columnable src) => + (Columnable a, Columnable b, Columnable src, Read a) => T.Text -> (Either String a -> b) -> Expr src -> @@ -234,7 +234,8 @@ (MkUnaryOp{unaryFn = atanh, unaryName = "atanh", unarySymbol = Nothing}) instance (Show a) => Show (Expr a) where - show :: forall a. (Show a) => Expr a -> String + {-# HLINT ignore "Use show" #-} + show :: Expr a -> String show (Col name) = "(col @" ++ show (typeRep @a) ++ " " ++ show name ++ ")" show (CastWith name tag _) = "(castWith " ++ show tag ++ " " ++ show name ++ ")" show (CastExprWith tag _ inner) = "(castExprWith " ++ show tag ++ " " ++ show inner ++ ")" @@ -246,7 +247,7 @@ show (Agg (FoldAgg op _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")" show (Agg (MergeAgg op _ _ _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")" -normalize :: (Eq a, Ord a, Show a, Typeable a) => Expr a -> Expr a +normalize :: (Show a, Typeable a) => Expr a -> Expr a normalize expr = case expr of Col name -> Col name CastWith n t f -> CastWith n t f @@ -278,11 +279,14 @@ exprKey (Lit val) = "1:" ++ show val exprKey (If c t e) = "2:" ++ exprKey c ++ exprKey t ++ exprKey e exprKey (Unary op e) = "3:" ++ T.unpack (unaryName op) ++ exprKey e - exprKey (Binary op e1 e2) = "4:" ++ T.unpack (binaryName op) ++ exprKey e1 ++ exprKey e2 + exprKey (Binary op e1' e2') = "4:" ++ T.unpack (binaryName op) ++ exprKey e1' ++ exprKey e2' exprKey (Agg (CollectAgg name _) e) = "5:" ++ T.unpack name ++ exprKey e exprKey (Agg (FoldAgg name _ _) e) = "5:" ++ T.unpack name ++ exprKey e exprKey (Agg (MergeAgg name _ _ _ _) e) = "5:" ++ T.unpack name ++ exprKey e +instance (Ord a, Columnable a) => Ord (Expr a) where + compare l r = compareExpr (normalize l) (normalize r) + instance (Eq a, Columnable a) => Eq (Expr a) where (==) l r = eqNormalized (normalize l) (normalize r) where @@ -307,43 +311,6 @@ n1 == n2 && e1 `exprEq` e2 eqNormalized _ _ = False -instance (Ord a, Columnable a) => Ord (Expr a) where - compare :: Expr a -> Expr a -> Ordering - compare e1 e2 = case (e1, e2) of - (Col n1, Col n2) -> compare n1 n2 - (CastWith n1 t1 _, CastWith n2 t2 _) -> compare n1 n2 <> compare t1 t2 - (CastExprWith t1 _ _, CastExprWith t2 _ _) -> compare t1 t2 - (Lit v1, Lit v2) -> compare v1 v2 - (If c1 t1 e1', If c2 t2 e2') -> - compare c1 c2 <> exprComp t1 t2 <> exprComp e1' e2' - (Unary op1 e1', Unary op2 e2') -> compare (unaryName op1) (unaryName op2) <> exprComp e1' e2' - (Binary op1 a1 b1, Binary op2 a2 b2) -> - compare (binaryName op1) (binaryName op2) <> exprComp a1 a2 <> exprComp b1 b2 - (Agg (CollectAgg n1 _) e1', Agg (CollectAgg n2 _) e2') -> compare n1 n2 <> exprComp e1' e2' - (Agg (FoldAgg n1 _ _) e1', Agg (FoldAgg n2 _ _) e2') -> compare n1 n2 <> exprComp e1' e2' - (Agg (MergeAgg n1 _ _ _ _) e1', Agg (MergeAgg n2 _ _ _ _) e2') -> compare n1 n2 <> exprComp e1' e2' - -- Different constructors - compare by priority - (Col _, _) -> LT - (_, Col _) -> GT - (CastWith{}, _) -> LT - (_, CastWith{}) -> GT - (CastExprWith{}, _) -> LT - (_, CastExprWith{}) -> GT - (Lit _, _) -> LT - (_, Lit _) -> GT - (Unary{}, _) -> LT - (_, Unary{}) -> GT - (Binary{}, _) -> LT - (_, Binary{}) -> GT - (If{}, _) -> LT - (_, If{}) -> GT - (Agg{}, _) -> LT - -exprComp :: (Columnable b, Columnable c) => Expr b -> Expr c -> Ordering -exprComp e1 e2 = case testEquality (typeOf e1) (typeOf e2) of - Just Refl -> e1 `compare` e2 - Nothing -> LT - replaceExpr :: forall a b c. (Columnable a, Columnable b, Columnable c) => @@ -363,7 +330,7 @@ If (replaceExpr new old cond) (replaceExpr new old l) (replaceExpr new old r) (Unary op value) -> Unary op (replaceExpr new old value) (Binary op l r) -> Binary op (replaceExpr new old l) (replaceExpr new old r) - (Agg op expr) -> Agg op (replaceExpr new old expr) + (Agg op inner) -> Agg op (replaceExpr new old inner) eSize :: Expr a -> Int eSize (Col _) = 1 diff --git a/src/DataFrame/Internal/Interpreter.hs b/src/DataFrame/Internal/Interpreter.hs --- a/src/DataFrame/Internal/Interpreter.hs +++ b/src/DataFrame/Internal/Interpreter.hs @@ -41,6 +41,225 @@ typeRep, ) +import Data.Int (Int16, Int32, Int64, Int8) + +-- Specializations for common aggregation types to avoid dictionary overhead. +-- foldLinearGroups: mean accumulator +{-# SPECIALIZE foldLinearGroups :: + (MeanAcc -> Double -> MeanAcc) -> + MeanAcc -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (MeanAcc -> Float -> MeanAcc) -> + MeanAcc -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (MeanAcc -> Int -> MeanAcc) -> + MeanAcc -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (MeanAcc -> Int8 -> MeanAcc) -> + MeanAcc -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (MeanAcc -> Int16 -> MeanAcc) -> + MeanAcc -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (MeanAcc -> Int32 -> MeanAcc) -> + MeanAcc -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (MeanAcc -> Int64 -> MeanAcc) -> + MeanAcc -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +-- foldLinearGroups: count accumulator +{-# SPECIALIZE foldLinearGroups :: + (Int -> Double -> Int) -> + Int -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (Int -> Float -> Int) -> + Int -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (Int -> Int -> Int) -> + Int -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (Int -> Int8 -> Int) -> + Int -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (Int -> Int16 -> Int) -> + Int -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (Int -> Int32 -> Int) -> + Int -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (Int -> Int64 -> Int) -> + Int -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +-- foldLinearGroups: sum/min/max (acc == elem) +{-# SPECIALIZE foldLinearGroups :: + (Double -> Double -> Double) -> + Double -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (Float -> Float -> Float) -> + Float -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (Int8 -> Int8 -> Int8) -> + Int8 -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (Int16 -> Int16 -> Int16) -> + Int16 -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (Int32 -> Int32 -> Int32) -> + Int32 -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} +{-# SPECIALIZE foldLinearGroups :: + (Int64 -> Int64 -> Int64) -> + Int64 -> + Column -> + VU.Vector Int -> + Int -> + Either DataFrameException Column + #-} + +-- mapColumn: finalize +{-# SPECIALIZE mapColumn :: + (MeanAcc -> Double) -> Column -> Either DataFrameException Column + #-} +{-# SPECIALIZE mapColumn :: + (Double -> Double) -> Column -> Either DataFrameException Column + #-} +{-# SPECIALIZE mapColumn :: + (Float -> Float) -> Column -> Either DataFrameException Column + #-} +{-# SPECIALIZE mapColumn :: + (Int -> Int) -> Column -> Either DataFrameException Column + #-} + +-- zipWithColumns: binary ops +{-# SPECIALIZE zipWithColumns :: + (Double -> Double -> Double) -> + Column -> + Column -> + Either DataFrameException Column + #-} +{-# SPECIALIZE zipWithColumns :: + (Float -> Float -> Float) -> + Column -> + Column -> + Either DataFrameException Column + #-} +{-# SPECIALIZE zipWithColumns :: + (Int -> Int -> Int) -> Column -> Column -> Either DataFrameException Column + #-} +{-# SPECIALIZE zipWithColumns :: + (Int8 -> Int8 -> Int8) -> Column -> Column -> Either DataFrameException Column + #-} +{-# SPECIALIZE zipWithColumns :: + (Int16 -> Int16 -> Int16) -> + Column -> + Column -> + Either DataFrameException Column + #-} +{-# SPECIALIZE zipWithColumns :: + (Int32 -> Int32 -> Int32) -> + Column -> + Column -> + Either DataFrameException Column + #-} +{-# SPECIALIZE zipWithColumns :: + (Int64 -> Int64 -> Int64) -> + Column -> + Column -> + Either DataFrameException Column + #-} + ------------------------------------------------------------------------------- -- Value: the unified result type ------------------------------------------------------------------------------- @@ -253,7 +472,7 @@ -} promoteColumnWith :: forall a b. - (Columnable a, Columnable b) => + (Columnable a, Columnable b, Read a) => (Either String a -> b) -> Column -> Either DataFrameException Column promoteColumnWith onResult col | hasElemType @b col = Right col @@ -395,7 +614,7 @@ tryParseWith :: forall a b. - (Columnable a, Columnable b) => + (Columnable a, Columnable b, Read a) => (Either String a -> b) -> Column -> Either DataFrameException Column tryParseWith onResult col = case col of BoxedColumn bm (v :: V.Vector c) -> diff --git a/src/DataFrame/Internal/Parsing.hs b/src/DataFrame/Internal/Parsing.hs --- a/src/DataFrame/Internal/Parsing.hs +++ b/src/DataFrame/Internal/Parsing.hs @@ -13,12 +13,10 @@ import Data.Attoparsec.Text hiding (decimal, double, signed) import Data.ByteString.Lex.Fractional import Data.Foldable (fold) -import Data.Maybe (fromMaybe) import Data.Text.Read (decimal, double, signed) import Data.Time (Day, defaultTimeLocale, parseTimeM) import GHC.Stack (HasCallStack) import System.IO (Handle, IOMode (..), hIsEOF, hTell, withFile) -import Text.Read (readMaybe) import Prelude hiding (takeWhile) isNullish :: T.Text -> Bool @@ -41,11 +39,6 @@ isFalseish :: T.Text -> Bool isFalseish t = t `elem` ["False", "false", "FALSE"] -readValue :: (HasCallStack, Read a) => T.Text -> a -readValue s = case readMaybe (T.unpack s) of - Nothing -> error ("Could not read value: " <> T.unpack s) - Just value -> value - readBool :: (HasCallStack) => T.Text -> Maybe Bool readBool s | isTrueish s = Just True @@ -121,12 +114,6 @@ Right (value, "") -> Right value Right (value, _) -> Left s {-# INLINE readDoubleEither #-} - -safeReadValue :: (Read a) => T.Text -> Maybe a -safeReadValue s = readMaybe (T.unpack s) - -readWithDefault :: (HasCallStack, Read a) => a -> T.Text -> a -readWithDefault v s = fromMaybe v (readMaybe (T.unpack s)) -- --------------------------------------------------------------------------- -- Attoparsec CSV parser combinators (shared between Lazy.IO.CSV and others) diff --git a/src/DataFrame/Internal/Row.hs b/src/DataFrame/Internal/Row.hs --- a/src/DataFrame/Internal/Row.hs +++ b/src/DataFrame/Internal/Row.hs @@ -13,12 +13,9 @@ import qualified Data.Map as M 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 Control.Exception (throw) -import Control.Monad.ST (runST) import Data.Function (on) import Data.Maybe (fromMaybe) import Data.Type.Equality (TestEquality (..)) @@ -27,15 +24,6 @@ import DataFrame.Internal.Column import DataFrame.Internal.DataFrame import DataFrame.Internal.Expression (Expr (..)) -import Text.ParserCombinators.ReadPrec (ReadPrec) -import Text.Read ( - Lexeme (Ident), - lexP, - parens, - readListPrec, - readListPrecDefault, - readPrec, - ) import Type.Reflection (typeOf, typeRep) data Any where @@ -47,12 +35,6 @@ Refl <- testEquality (typeOf a) (typeOf b) return $ a == b -instance Ord Any where - (<=) :: Any -> Any -> Bool - (Value a) <= (Value b) = fromMaybe False $ do - Refl <- testEquality (typeOf a) (typeOf b) - return $ a <= b - instance Show Any where show :: Any -> String show (Value a) = T.unpack (showValue a) @@ -64,15 +46,6 @@ Just Refl -> T.pack v Nothing -> (T.pack . show) v -instance Read Any where - readListPrec :: ReadPrec [Any] - readListPrec = readListPrecDefault - - readPrec :: ReadPrec Any - readPrec = parens $ do - Ident "Value" <- lexP - readPrec - -- | Wraps a value into an \Any\ type. This helps up represent rows as heterogenous lists. toAny :: forall a. (Columnable a) => a -> Any toAny = Value @@ -200,17 +173,3 @@ Nothing -> throwError name Nothing -> throw $ ColumnsNotFoundException [name] "mkRowRep" (M.keys $ columnIndices df) - -sortedIndexes' :: [Bool] -> V.Vector Row -> VU.Vector Int -sortedIndexes' flipCompare rows = runST $ do - withIndexes <- VG.thaw (V.indexed rows) - VA.sortBy (produceOrderingFromRow flipCompare `on` snd) withIndexes - sorted <- VG.unsafeFreeze withIndexes - return $ VU.generate (VG.length rows) (\i -> fst (sorted VG.! i)) - -produceOrderingFromRow :: [Bool] -> Row -> Row -> Ordering -produceOrderingFromRow mustFlips v1 v2 = V.foldr (<>) mempty vZipped - where - vFlip = V.fromList mustFlips - vZipped = - V.zipWith3 (\b e1 e2 -> if b then compare e1 e2 else compare e2 e1) vFlip v1 v2 diff --git a/src/DataFrame/Internal/Schema.hs b/src/DataFrame/Internal/Schema.hs --- a/src/DataFrame/Internal/Schema.hs +++ b/src/DataFrame/Internal/Schema.hs @@ -21,7 +21,7 @@ -- | A runtime tag for a column’s element type. data SchemaType where -- | Constructor carrying a 'Proxy' of the element type. - SType :: (Columnable a) => P.Proxy a -> SchemaType + SType :: (Columnable a, Read a) => P.Proxy a -> SchemaType {- | Show the underlying element type using 'typeRep'. @@ -59,7 +59,7 @@ >>> show (schemaType @Double) "Double" -} -schemaType :: forall a. (Columnable a) => SchemaType +schemaType :: forall a. (Columnable a, Read a) => SchemaType schemaType = SType (P.Proxy @a) {- | Logical schema of a 'DataFrame': a mapping from column names to their diff --git a/src/DataFrame/Internal/Statistics.hs b/src/DataFrame/Internal/Statistics.hs --- a/src/DataFrame/Internal/Statistics.hs +++ b/src/DataFrame/Internal/Statistics.hs @@ -63,12 +63,12 @@ deriving (Show) varianceStep :: VarAcc -> Double -> VarAcc -varianceStep (VarAcc !n !mean !m2) !x = +varianceStep (VarAcc !n !meanVal !m2) !x = let !n' = n + 1 - !delta = x - mean - !mean' = mean + delta / fromIntegral n' - !m2' = m2 + delta * (x - mean') - in VarAcc n' mean' m2' + !delta = x - meanVal + !meanVal' = meanVal + delta / fromIntegral n' + !m2' = m2 + delta * (x - meanVal') + in VarAcc n' meanVal' m2' {-# INLINE varianceStep #-} computeVariance :: VarAcc -> Double @@ -89,15 +89,15 @@ data SkewAcc = SkewAcc !Int !Double !Double !Double deriving (Show) skewnessStep :: (VU.Unbox a, Num a, Real a) => SkewAcc -> a -> SkewAcc -skewnessStep (SkewAcc !n !mean !m2 !m3) !x' = +skewnessStep (SkewAcc !n !meanVal !m2 !m3) !x' = let !n' = n + 1 x = rtf x' !k = fromIntegral n' - !delta = x - mean - !mean' = mean + delta / k + !delta = x - meanVal + !meanVal' = meanVal + delta / k !m2' = m2 + (delta ^ 2 * (k - 1)) / k !m3' = m3 + (delta ^ 3 * (k - 1) * (k - 2)) / k ^ 2 - (3 * delta * m2) / k - in SkewAcc n' mean' m2' m3' + in SkewAcc n' meanVal' m2' m3' {-# INLINE skewnessStep #-} computeSkewness :: SkewAcc -> Double diff --git a/src/DataFrame/Internal/Types.hs b/src/DataFrame/Internal/Types.hs --- a/src/DataFrame/Internal/Types.hs +++ b/src/DataFrame/Internal/Types.hs @@ -19,7 +19,7 @@ import qualified Data.Vector.Unboxed as VU import Data.Word (Word16, Word32, Word64, Word8) -type Columnable' a = (Typeable a, Show a, Ord a, Eq a, Read a) +type Columnable' a = (Typeable a, Show a, Eq a) {- | A type with column representations used to select the "right" representation when specializing the `toColumn` function. diff --git a/src/DataFrame/Lazy/IO/CSV.hs b/src/DataFrame/Lazy/IO/CSV.hs --- a/src/DataFrame/Lazy/IO/CSV.hs +++ b/src/DataFrame/Lazy/IO/CSV.hs @@ -87,13 +87,13 @@ readSeparated :: Char -> ReadOptions -> FilePath -> IO (DataFrame, (Integer, T.Text, Int)) readSeparated c opts path = do - totalRows <- case totalRows opts of + totalRows' <- case totalRows opts of Nothing -> countRows c path >>= \total -> if hasHeader opts then return (total - 1) else return total Just n -> if hasHeader opts then return (n - 1) else return n - let (_, len) = case rowRange opts of - Nothing -> (0, totalRows) - Just (start, len) -> (start, min len (totalRows - rowsRead opts)) + let (_, len') = case rowRange opts of + Nothing -> (0, totalRows') + Just (start, len'') -> (start, min len'' (totalRows' - rowsRead opts)) withFile path ReadMode $ \handle -> do firstRow <- fmap T.strip . parseSep c <$> TIO.hGetLine handle let columnNames = @@ -109,7 +109,7 @@ -- Initialize mutable vectors for each column let numColumns = length columnNames - let numRows = len + let numRows = len' -- Use this row to infer the types of the rest of the column. (dataRow, remainder) <- readSingleLine c (leftOver opts) handle @@ -169,11 +169,11 @@ IO (T.Text, Int) fillColumns n c mutableCols nullIndices unused handle = do input <- newIORef unused - rowsRead <- newIORef (0 :: Int) + rowsRead' <- newIORef (0 :: Int) forM_ [1 .. (n - 1)] $ \i -> do - isEOF <- hIsEOF handle + atEOF <- hIsEOF handle input' <- readIORef input - unless (isEOF && input' == mempty) $ do + unless (atEOF && input' == mempty) $ do parseWith (TIO.hGetChunk handle) (parseRow c) input' >>= \case Fail unconsumed ctx er -> do erpos <- hTell handle @@ -188,10 +188,10 @@ fail "Partial handler is called" Done (unconsumed :: T.Text) (row :: [T.Text]) -> do writeIORef input unconsumed - modifyIORef rowsRead (+ 1) + modifyIORef rowsRead' (+ 1) zipWithM_ (writeValue mutableCols nullIndices i) [0 ..] row l <- readIORef input - r <- readIORef rowsRead + r <- readIORef rowsRead' pure (l, r) {-# INLINE fillColumns #-} @@ -206,7 +206,7 @@ writeValue mutableCols nullIndices count colIndex value = do col <- VM.unsafeRead mutableCols colIndex res <- writeColumn count value col - let modify value = VM.unsafeModify nullIndices ((count, value) :) colIndex + let modify val = VM.unsafeModify nullIndices ((count, val) :) colIndex either modify (const (return ())) res {-# INLINE writeValue #-} @@ -392,12 +392,12 @@ else go idx start inQ (pos + 1) extract s e = - let field = BS.take (e - s) (BS.drop s bs) - in if BS.length field >= 2 - && BS.head field == quoteChar - && BS.last field == quoteChar - then BS.init (BS.tail field) - else field + let fieldVal = BS.take (e - s) (BS.drop s bs) + in if BS.length fieldVal >= 2 + && BS.head fieldVal == quoteChar + && BS.last fieldVal == quoteChar + then BS.init (BS.tail fieldVal) + else fieldVal {-# INLINE getNthFieldBs #-} -- | Allocate a fresh 'MutableColumn' for @n@ slots based on a 'SchemaType'. diff --git a/src/DataFrame/Operations/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs --- a/src/DataFrame/Operations/Aggregation.hs +++ b/src/DataFrame/Operations/Aggregation.hs @@ -62,8 +62,8 @@ (VU.fromList [0]) VU.empty | otherwise = - let !vis = VU.map fst valueIndices - !os = changingPoints valueIndices + let !vis = VU.map fst valIndices + !os = changingPoints valIndices !n = fst (dimensions df) in Grouped df @@ -75,7 +75,7 @@ indicesToGroup = M.elems $ M.filterWithKey (\k _ -> k `elem` names) (columnIndices df) doubleToInt :: Double -> Int doubleToInt = floor . (* 1000) - valueIndices = runST $ do + valIndices = runST $ do let n = fst (dimensions df) mv <- VUM.new n @@ -183,9 +183,9 @@ (VU.fromList (VU.length vs : fst (VU.ifoldl' findChangePoints initialState vs))) where initialState = ([0], snd (VU.head vs)) - findChangePoints (!offsets, !currentVal) index (_, !newVal) - | currentVal == newVal = (offsets, currentVal) - | otherwise = (index : offsets, newVal) + findChangePoints (!offs, !currentVal) index (_, !newVal) + | currentVal == newVal = (offs, currentVal) + | otherwise = (index : offs, newVal) computeRowHashes :: [Int] -> DataFrame -> VU.Vector Int computeRowHashes indices df = runST $ do @@ -275,11 +275,11 @@ All ungrouped columns will be dropped. -} aggregate :: [NamedExpr] -> GroupedDataFrame -> DataFrame -aggregate aggs gdf@(Grouped df groupingColumns valueIndices offsets _rowToGroup) = +aggregate aggs gdf@(Grouped df groupingColumns valIndices offs _rowToGroup) = let df' = selectIndices - (VU.map (valueIndices VU.!) (VU.init offsets)) + (VU.map (valIndices VU.!) (VU.init offs)) (select groupingColumns df) f (name, UExpr (expr :: Expr a)) d = diff --git a/src/DataFrame/Operations/Join.hs b/src/DataFrame/Operations/Join.hs --- a/src/DataFrame/Operations/Join.hs +++ b/src/DataFrame/Operations/Join.hs @@ -81,8 +81,8 @@ buildCompactIndex hashes = let n = VU.length hashes (sortedHashes, sortedIndices) = sortWithIndices hashes - !offsets = buildOffsets sortedHashes n 0 HM.empty - in CompactIndex sortedIndices offsets + !offs = buildOffsets sortedHashes n 0 HM.empty + in CompactIndex sortedIndices offs where buildOffsets :: VU.Vector Int -> diff --git a/src/DataFrame/Operations/Permutation.hs b/src/DataFrame/Operations/Permutation.hs --- a/src/DataFrame/Operations/Permutation.hs +++ b/src/DataFrame/Operations/Permutation.hs @@ -2,30 +2,35 @@ {-# 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 (Columnable, atIndicesStable) -import DataFrame.Internal.DataFrame (DataFrame (..)) +import DataFrame.Internal.Column (Column (..), Columnable, atIndicesStable) +import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn) import DataFrame.Internal.Expression (Expr (Col)) -import DataFrame.Internal.Row (sortedIndexes', toRowVector) import DataFrame.Operations.Core (columnNames, dimensions) 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) => Expr a -> SortOrder - Desc :: (Columnable a) => Expr a -> SortOrder + Asc :: (Columnable a, Ord a) => Expr a -> SortOrder + Desc :: (Columnable a, Ord a) => Expr a -> SortOrder instance Eq SortOrder where (==) :: SortOrder -> SortOrder -> Bool @@ -59,12 +64,44 @@ (columnNames df) | otherwise = let - indexes = sortedIndexes' mustFlips (toRowVector names 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 names = map getSortColumnName sortOrds - mustFlips = map mustFlipCompare sortOrds + +{- | 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) => diff --git a/src/DataFrame/Operations/Statistics.hs b/src/DataFrame/Operations/Statistics.hs --- a/src/DataFrame/Operations/Statistics.hs +++ b/src/DataFrame/Operations/Statistics.hs @@ -60,7 +60,8 @@ Percentage (%) | 44.26% | 31.74% | 0.02% | 11.09% | 12.88% @ -} -frequencies :: forall a. (Columnable a) => Expr a -> DataFrame -> DataFrame +frequencies :: + forall a. (Columnable a, Ord a) => Expr a -> DataFrame -> DataFrame frequencies expr df = let counts = valueCounts expr df @@ -68,11 +69,11 @@ initDf = empty & insertVector "Statistic" (V.fromList ["Count" :: T.Text, "Percentage (%)"]) - freqs col = + freqs col' = L.foldl' - ( \d (col, k) -> + ( \d (col'', k) -> insertVector - (showValue @a col) + (showValue @a col'') (V.fromList [toAny k, calculatePercentage counts k]) d ) @@ -248,7 +249,7 @@ Left e -> throw e Right (TColumn xs) -> case toVector @a @V.Vector xs of Left e -> throw e - Right xs -> VG.sum xs + Right xs' -> VG.sum xs' {- | /O(n)/ Impute missing values in a column using a derived scalar. @@ -367,7 +368,7 @@ quantiles = applyStatistics (quantiles' (VU.fromList [0, 1, 2, 3, 4]) 4) name df min' = flip (VG.!) 0 <$> quantiles quartile1 = flip (VG.!) 1 <$> quantiles - median' = flip (VG.!) 2 <$> quantiles + medianVal = flip (VG.!) 2 <$> quantiles quartile3 = flip (VG.!) 3 <$> quantiles max' = flip (VG.!) 4 <$> quantiles iqr = (-) <$> quartile3 <*> quartile1 @@ -377,7 +378,7 @@ , mean' <$> doubleColumn name , min' , quartile1 - , median' + , medianVal , quartile3 , max' , sqrt . variance' <$> doubleColumn name diff --git a/src/DataFrame/Operations/Subset.hs b/src/DataFrame/Operations/Subset.hs --- a/src/DataFrame/Operations/Subset.hs +++ b/src/DataFrame/Operations/Subset.hs @@ -132,7 +132,7 @@ Nothing -> throw $ ColumnsNotFoundException [filterColumnName] "filter" (M.keys $ columnIndices df) - Just col@(BoxedColumn bm (column :: V.Vector b)) -> + 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 @@ -145,7 +145,7 @@ 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)) -> + 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 @@ -160,10 +160,10 @@ _ -> filterByVector filterColumnName column condition df filter expr condition df = let - (TColumn col) = case interpret @a df (normalize expr) of + (TColumn col') = case interpret @a df (normalize expr) of Left e -> throw e Right c -> c - indexes = case findIndices condition col of + indexes = case findIndices condition col' of Right ixs -> ixs Left e -> throw e c' = snd $ dataframeDimensions df @@ -211,10 +211,10 @@ filterWhere :: Expr Bool -> DataFrame -> DataFrame filterWhere expr df = let - (TColumn col) = case interpret @Bool df (normalize expr) of + (TColumn col') = case interpret @Bool df (normalize expr) of Left e -> throw e Right c -> c - indexes = case findIndices id col of + indexes = case findIndices id col' of Right ixs -> ixs Left e -> throw e c' = snd $ dataframeDimensions df @@ -229,12 +229,12 @@ > filterJust "col" df -} filterJust :: T.Text -> DataFrame -> DataFrame -filterJust name df = case getColumn name df of +filterJust colName df = case getColumn colName df of Nothing -> - throw $ ColumnsNotFoundException [name] "filterJust" (M.keys $ columnIndices df) + 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) name) isJust df & apply @(Maybe a) fromJust name - UnboxedColumn (Just _) (col :: VU.Vector a) -> filter (Col @(Maybe a) name) isJust df & apply @(Maybe a) fromJust name + 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 @@ -243,13 +243,13 @@ > filterNothing "col" df -} filterNothing :: T.Text -> DataFrame -> DataFrame -filterNothing name df = case getColumn name df of +filterNothing colName df = case getColumn colName df of Nothing -> throw $ - ColumnsNotFoundException [name] "filterNothing" (M.keys $ columnIndices df) + ColumnsNotFoundException [colName] "filterNothing" (M.keys $ columnIndices df) Just column | hasMissing column -> case column of - BoxedColumn (Just _) (col :: V.Vector a) -> filter (Col @(Maybe a) name) isNothing df - UnboxedColumn (Just _) (col :: VU.Vector a) -> filter (Col @(Maybe a) name) isNothing df + 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 @@ -273,7 +273,7 @@ > cube (10, 5) df -} cube :: (Int, Int) -> DataFrame -> DataFrame -cube (length, width) = take length . selectBy [ColumnIndexRange (0, width - 1)] +cube (len, width) = take len . selectBy [ColumnIndexRange (0, width - 1)] {- | O(n) Selects a number of columns in a given dataframe. @@ -297,8 +297,8 @@ in result{derivingExpressions = filteredExprs} where addKeyValue d k = fromMaybe df $ do - col <- getColumn k df - pure $ insertColumn k col d + col' <- getColumn k df + pure $ insertColumn k col' d data SelectionCriteria = ColumnProperty (Column -> Bool) @@ -352,7 +352,7 @@ where finalSelection = Prelude.filter (`S.member` columnsWithProperties) (columnNames df) columnsWithProperties = S.fromList (L.foldl' columnWithProperty [] xs) - columnWithProperty acc (ColumnName name) = acc ++ [name] + columnWithProperty acc (ColumnName colName) = acc ++ [colName] columnWithProperty acc (ColumnNameProperty f) = acc ++ L.filter f (columnNames df) columnWithProperty acc (ColumnTextRange (from, to)) = acc @@ -363,7 +363,7 @@ 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) + ixs = V.ifoldl' (\acc' i c -> if f c then i : acc' else acc') [] (columns df) {- | O(n) inverse of select @@ -463,24 +463,24 @@ -- | 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)) = +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 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) = + 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) + 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" + 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 +groupByIndices col' = + let textVec = columnToTextVec col' (grouped, _) = V.foldl' (\(!m, !i) key -> (M.insertWith (++) key [i] m, i + 1)) @@ -509,10 +509,10 @@ (SplittableGen g, Columnable a) => g -> Double -> Expr a -> DataFrame -> DataFrame stratifiedSample gen p strataCol df = - let col = case strataCol of - Col name -> unsafeGetColumn name df + let col' = case strataCol of + Col colName -> unsafeGetColumn colName df _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol)) - groups = M.elems (groupByIndices col) + groups = M.elems (groupByIndices col') go _ [] = mempty go g (ixs : rest) = let stratum = rowsAtIndices ixs df @@ -533,10 +533,10 @@ (SplittableGen g, Columnable a) => g -> Double -> Expr a -> DataFrame -> (DataFrame, DataFrame) stratifiedSplit gen p strataCol df = - let col = case strataCol of - Col name -> unsafeGetColumn name df + let col' = case strataCol of + Col colName -> unsafeGetColumn colName df _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol)) - groups = M.elems (groupByIndices col) + groups = M.elems (groupByIndices col') go _ [] = (mempty, mempty) go g (ixs : rest) = let stratum = rowsAtIndices ixs df diff --git a/src/DataFrame/Operations/Typing.hs b/src/DataFrame/Operations/Typing.hs --- a/src/DataFrame/Operations/Typing.hs +++ b/src/DataFrame/Operations/Typing.hs @@ -288,6 +288,6 @@ 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 + Just v' -> v' where s = if null v then "\"\"" else v diff --git a/src/DataFrame/Operators.hs b/src/DataFrame/Operators.hs --- a/src/DataFrame/Operators.hs +++ b/src/DataFrame/Operators.hs @@ -48,7 +48,7 @@ (|>) = (&) as :: (Columnable a) => Expr a -> T.Text -> NamedExpr -as expr name = (name, UExpr expr) +as expr colName = (colName, UExpr expr) name :: (Show a) => Expr a -> T.Text name (Col n) = n @@ -71,7 +71,7 @@ liftDecorated :: (Columnable a, Columnable b) => (a -> b) -> T.Text -> Maybe T.Text -> Expr a -> Expr b -liftDecorated f name rep = Unary (MkUnaryOp{unaryFn = f, unaryName = name, unarySymbol = rep}) +liftDecorated f opName rep = Unary (MkUnaryOp{unaryFn = f, unaryName = opName, unarySymbol = rep}) lift2Decorated :: (Columnable c, Columnable b, Columnable a) => @@ -83,11 +83,11 @@ Expr c -> Expr b -> Expr a -lift2Decorated f name rep comm prec = +lift2Decorated f opName rep comm prec = Binary ( MkBinaryOp { binaryFn = f - , binaryName = name + , binaryName = opName , binarySymbol = rep , binaryCommutative = comm , binaryPrecedence = prec diff --git a/src/DataFrame/Synthesis.hs b/src/DataFrame/Synthesis.hs --- a/src/DataFrame/Synthesis.hs +++ b/src/DataFrame/Synthesis.hs @@ -28,12 +28,10 @@ import DataFrame.Operations.Subset (exclude) import Control.Exception (throw) -import Data.Containers.ListUtils import Data.Function import qualified Data.List as L import qualified Data.Map as M import Data.Maybe (listToMaybe) -import qualified Data.Set as S import qualified Data.Text as T import Data.Type.Equality import qualified Data.Vector.Unboxed as VU @@ -165,20 +163,20 @@ DataFrame -> [Expr a] -> [(Expr a, TypedColumn a)] -deduplicate df = go S.empty . nubOrd . L.sortBy (\e1 e2 -> compare (eSize e1) (eSize e2)) +deduplicate df = go [] . L.nub . L.sortBy (\e1 e2 -> compare (eSize e1) (eSize e2)) where go _ [] = [] go seen (x : xs) | hasInvalid = go seen xs - | S.member res seen = go seen xs - | otherwise = (x, res) : go (S.insert res seen) xs + | res `elem` seen = go seen xs + | otherwise = (x, res) : go (res : seen) xs where res = case interpret @a df x of Left e -> throw e Right v -> v hasInvalid = case res of - (TColumn (UnboxedColumn _ (col :: VU.Vector b))) -> case testEquality (typeRep @Double) (typeRep @b) of - Just Refl -> VU.any (\n -> isNaN n || isInfinite n) col + (TColumn (UnboxedColumn _ (column :: VU.Vector b))) -> case testEquality (typeRep @Double) (typeRep @b) of + Just Refl -> VU.any (\n -> isNaN n || isInfinite n) column Nothing -> False _ -> False @@ -384,9 +382,9 @@ [(Expr Double, TypedColumn a)] -> [Expr Double] pickTopN _ _ _ [] = [] -pickTopN df (TColumn col) cfg ps = +pickTopN df (TColumn column) cfg ps = let - l = case toVector @Double @VU.Vector col of + l = case toVector @Double @VU.Vector column of Left e -> throw e Right v -> v ordered = @@ -414,14 +412,14 @@ case toVector @Double @VU.Vector col' of Left e -> throw e Right v -> VU.convert v - interpretDoubleVector e = + interpretDoubleVector e' = let - (TColumn col') = case interpret df e of - Left e -> throw e + (TColumn col') = case interpret df e' of + Left err -> throw err Right v -> v in case toVector @Double @VU.Vector col' of - Left e -> throw e + Left err -> throw err Right v -> VU.convert v in trace @@ -441,9 +439,9 @@ [(Expr Bool, TypedColumn Bool)] -> [Expr Bool] pickTopNBool _ _ [] = [] -pickTopNBool df (TColumn col) ps = +pickTopNBool df (TColumn column) ps = let - l = case toVector @Double @VU.Vector col of + l = case toVector @Double @VU.Vector column of Left e -> throw e Right v -> v ordered = @@ -475,10 +473,10 @@ ordered satisfiesExamples :: DataFrame -> TypedColumn Double -> Expr Double -> Bool -satisfiesExamples df col expr = +satisfiesExamples df column expr = let result = case interpret df expr of Left e -> throw e Right v -> v in - result == col + result == column diff --git a/src/DataFrame/Typed/Expr.hs b/src/DataFrame/Typed/Expr.hs --- a/src/DataFrame/Typed/Expr.hs +++ b/src/DataFrame/Typed/Expr.hs @@ -128,12 +128,11 @@ 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 ( - AggStrategy (..), BinaryOp (..), Expr (..), - MeanAcc (..), NamedExpr, UExpr (..), UnaryOp (..), @@ -522,33 +521,22 @@ ------------------------------------------------------------------------------- sum :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a -sum (TExpr e) = TExpr (Agg (FoldAgg "sum" Nothing (+)) e) +sum (TExpr e) = TExpr (F.sum e) mean :: (Columnable a, Real a) => TExpr cols a -> TExpr cols Double -mean (TExpr e) = - TExpr - ( 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) - ) - e - ) +mean (TExpr e) = TExpr (F.mean e) count :: (Columnable a) => TExpr cols a -> TExpr cols Int -count (TExpr e) = TExpr (Agg (MergeAgg "count" (0 :: Int) (\c _ -> c + 1) (+) id) e) +count (TExpr e) = TExpr (F.count e) minimum :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -minimum (TExpr e) = TExpr (Agg (FoldAgg "minimum" Nothing min) e) +minimum (TExpr e) = TExpr (F.minimum e) maximum :: (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -maximum (TExpr e) = TExpr (Agg (FoldAgg "maximum" Nothing max) e) +maximum (TExpr e) = TExpr (F.maximum e) collect :: (Columnable a) => TExpr cols a -> TExpr cols [a] -collect (TExpr e) = TExpr (Agg (FoldAgg "collect" (Just []) (flip (:))) e) +collect (TExpr e) = TExpr (F.collect e) ------------------------------------------------------------------------------- -- Cast / coercion expressions @@ -556,14 +544,14 @@ castExpr :: forall b cols src. - (Columnable b, Columnable src) => TExpr cols src -> TExpr cols (Maybe b) + (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) => b -> TExpr cols src -> TExpr cols b + (Columnable b, Columnable src, Read b) => b -> TExpr cols src -> TExpr cols b castExprWithDefault def (TExpr e) = TExpr ( CastExprWith @b @b @src @@ -574,7 +562,8 @@ castExprEither :: forall b cols src. - (Columnable b, Columnable src) => TExpr cols src -> TExpr cols (Either T.Text b) + (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 @@ -585,7 +574,7 @@ unsafeCastExpr :: forall b cols src. - (Columnable b, Columnable src) => TExpr cols src -> TExpr cols b + (Columnable b, Columnable src, Read b) => TExpr cols src -> TExpr cols b unsafeCastExpr (TExpr e) = TExpr ( CastExprWith @b @b @src @@ -603,9 +592,9 @@ as (TExpr e) name = (name, UExpr e) -- | Create an ascending sort order from a typed expression. -asc :: (Columnable a) => TExpr cols a -> TSortOrder cols +asc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols asc = Asc -- | Create a descending sort order from a typed expression. -desc :: (Columnable a) => TExpr cols a -> TSortOrder cols +desc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols desc = Desc diff --git a/src/DataFrame/Typed/Lazy.hs b/src/DataFrame/Typed/Lazy.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Typed/Lazy.hs @@ -0,0 +1,207 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | +Module : DataFrame.Typed.Lazy +Copyright : (c) 2025 +License : MIT +Stability : experimental + +Type-safe lazy query pipelines. + +This module combines the compile-time schema tracking of 'TypedDataFrame' +with the deferred execution of 'LazyDataFrame'. Queries are built as a +logical plan tree with phantom-typed schema tracking; execution is deferred +until 'run' is called. + +@ +{\-\# LANGUAGE DataKinds, TypeApplications, TypeOperators \#-\} +import qualified DataFrame.Typed.Lazy as TL +import DataFrame.Typed (Column) + +type Schema = '[Column \"id\" Int, Column \"name\" Text, Column \"score\" Double] + +main = do + let query = TL.scanCsv \@Schema \"data.csv\" + & TL.filter (TL.col \@\"score\" TL..>. TL.lit 0.5) + & TL.select \@'[\"id\", \"name\"] + df <- TL.run query -- TypedDataFrame '[Column \"id\" Int, Column \"name\" Text] + print df +@ +-} +module DataFrame.Typed.Lazy ( + -- * Core type + TypedLazyDataFrame, + + -- * Data sources + scanCsv, + scanSeparated, + scanParquet, + fromDataFrame, + fromTypedDataFrame, + + -- * Schema-preserving operations + filter, + take, + + -- * Schema-modifying operations + derive, + select, + + -- * Aggregation + groupBy, + aggregate, + + -- * Joins + join, + + -- * Sort + sortBy, + + -- * Execution + run, + + -- * Re-exports for pipeline construction + module DataFrame.Typed.Expr, + module DataFrame.Typed.Types, + SortOrder (..), +) where + +import Data.Kind (Type) +import Data.Proxy (Proxy (..)) +import qualified Data.Text as T +import GHC.TypeLits (KnownSymbol, Symbol, symbolVal) +import Prelude hiding (filter, take) + +import qualified DataFrame.Internal.Column as C +import qualified DataFrame.Internal.Expression as E +import DataFrame.Internal.Schema (Schema) +import DataFrame.Lazy.Internal.DataFrame (LazyDataFrame) +import qualified DataFrame.Lazy.Internal.DataFrame as L +import DataFrame.Lazy.Internal.LogicalPlan (SortOrder (..)) +import DataFrame.Operations.Join (JoinType) +import DataFrame.Typed.Expr +import DataFrame.Typed.Freeze (unsafeFreeze) +import DataFrame.Typed.Schema +import DataFrame.Typed.Types + +-- | A lazy query with compile-time schema tracking. +newtype TypedLazyDataFrame (cols :: [Type]) = TLD {unTLD :: LazyDataFrame} + +instance Show (TypedLazyDataFrame cols) where + show (TLD ldf) = "TypedLazyDataFrame { " ++ show ldf ++ " }" + +-- | Scan a CSV file with a given schema. +scanCsv :: + Schema -> + T.Text -> + TypedLazyDataFrame cols +scanCsv schema path = TLD (L.scanCsv schema path) + +-- | Scan a character-separated file with a given schema. +scanSeparated :: + Char -> + Schema -> + T.Text -> + TypedLazyDataFrame cols +scanSeparated sep schema path = TLD (L.scanSeparated sep schema path) + +-- | Scan a Parquet file, directory, or glob pattern with a given schema. +scanParquet :: + Schema -> + T.Text -> + TypedLazyDataFrame cols +scanParquet schema path = TLD (L.scanParquet schema path) + +-- | Lift an already-loaded eager 'TypedDataFrame' into a lazy plan. +fromDataFrame :: TypedDataFrame cols -> TypedLazyDataFrame cols +fromDataFrame (TDF df) = TLD (L.fromDataFrame df) + +-- | Synonym for 'fromDataFrame'. +fromTypedDataFrame :: TypedDataFrame cols -> TypedLazyDataFrame cols +fromTypedDataFrame = fromDataFrame + +-- | Keep rows that satisfy the predicate. +filter :: TExpr cols Bool -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +filter (TExpr expr) (TLD ldf) = TLD (L.filter expr ldf) + +-- | Retain at most @n@ rows. +take :: Int -> TypedLazyDataFrame cols -> TypedLazyDataFrame cols +take n (TLD ldf) = TLD (L.take n ldf) + +-- | Add a computed column. +derive :: + forall name a cols. + (KnownSymbol name, C.Columnable a, AssertAbsent name cols) => + TExpr cols a -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame (Snoc cols (Column name a)) +derive (TExpr expr) (TLD ldf) = + TLD (L.derive (T.pack (symbolVal (Proxy @name))) expr ldf) + +-- | Retain only the listed columns. +select :: + forall (names :: [Symbol]) cols. + (AllKnownSymbol names, AssertAllPresent names cols) => + TypedLazyDataFrame cols -> + TypedLazyDataFrame (SubsetSchema names cols) +select (TLD ldf) = TLD (L.select (DataFrame.Typed.Schema.symbolVals @names) ldf) + +-- | A typed lazy grouped query. +newtype TypedLazyGrouped (keys :: [Symbol]) (cols :: [Type]) = TLG + { unTLG :: ([T.Text], LazyDataFrame) + } + +-- | Group by key columns. +groupBy :: + forall (keys :: [Symbol]) cols. + (AllKnownSymbol keys, AssertAllPresent keys cols) => + TypedLazyDataFrame cols -> + TypedLazyGrouped keys cols +groupBy (TLD ldf) = TLG (DataFrame.Typed.Schema.symbolVals @keys, ldf) + +-- | Aggregate a grouped lazy query. +aggregate :: + forall keys cols aggs. + TAgg keys cols aggs -> + TypedLazyGrouped keys cols -> + TypedLazyDataFrame (Append (GroupKeyColumns keys cols) (Reverse aggs)) +aggregate tagg (TLG (keys, ldf)) = + TLD (L.groupBy keys (aggToNamedExprs tagg) ldf) + +-- | Join two lazy queries on a shared key column. +join :: + JoinType -> + T.Text -> + T.Text -> + TypedLazyDataFrame left -> + TypedLazyDataFrame right -> + TypedLazyDataFrame left -- TODO: compute join result schema +join jt leftKey rightKey (TLD left) (TLD right) = + TLD (L.join jt leftKey rightKey left right) + +-- | Sort the result by column name and direction. +sortBy :: + [(T.Text, SortOrder)] -> + TypedLazyDataFrame cols -> + TypedLazyDataFrame cols +sortBy cols (TLD ldf) = TLD (L.sortBy cols ldf) + +-- | Execute the lazy query and return a typed DataFrame. +run :: + forall cols. + (KnownSchema cols) => + TypedLazyDataFrame cols -> + IO (TypedDataFrame cols) +run (TLD ldf) = unsafeFreeze <$> L.runDataFrame ldf + +-- | Convert TAgg to untyped named expressions for the lazy groupBy. +aggToNamedExprs :: TAgg keys cols aggs -> [(T.Text, E.UExpr)] +aggToNamedExprs TAggNil = [] +aggToNamedExprs (TAggCons name (TExpr expr) rest) = + (name, E.UExpr expr) : aggToNamedExprs rest diff --git a/src/DataFrame/Typed/Schema.hs b/src/DataFrame/Typed/Schema.hs --- a/src/DataFrame/Typed/Schema.hs +++ b/src/DataFrame/Typed/Schema.hs @@ -66,7 +66,6 @@ import Type.Reflection (SomeTypeRep, Typeable, someTypeRep) import DataFrame.Internal.Column (Columnable) -import DataFrame.Internal.Types (If) import DataFrame.Typed.Types (Column) -- | Look up the element type of a column by name. @@ -121,11 +120,21 @@ type family ExcludeSchema (names :: [Symbol]) (cols :: [Type]) :: [Type] where ExcludeSchema names '[] = '[] ExcludeSchema names (Column n a ': rest) = - If - (IsElem n names) - (ExcludeSchema names rest) - (Column n a ': ExcludeSchema names rest) + ExcludeSchemaHelper (IsElem n names) n a names rest +type family + ExcludeSchemaHelper + (found :: Bool) + (n :: Symbol) + (a :: Type) + (names :: [Symbol]) + (rest :: [Type]) :: + [Type] + where + ExcludeSchemaHelper 'True n a names rest = ExcludeSchema names rest + ExcludeSchemaHelper 'False n a names rest = + Column n a ': ExcludeSchema names rest + -- | Type-level elem for Symbols type family IsElem (x :: Symbol) (xs :: [Symbol]) :: Bool where IsElem x '[] = 'False @@ -197,14 +206,22 @@ -- | Assert that a column name is present in the schema. type family AssertAllPresent (name :: [Symbol]) (cols :: [Type]) :: Constraint where AssertAllPresent (name ': rest) cols = - If - (HasName name cols) - (AssertAllPresent rest cols) - ( TypeError - ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema") - ) + AssertAllPresentHelper (HasName name cols) name rest cols AssertAllPresent '[] cols = () +type family + AssertAllPresentHelper + (found :: Bool) + (name :: Symbol) + (rest :: [Symbol]) + (cols :: [Type]) :: + Constraint + where + AssertAllPresentHelper 'True name rest cols = AssertAllPresent rest cols + AssertAllPresentHelper 'False name rest cols = + TypeError + ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema") + {- | Strip 'Maybe' from all columns. Used by 'filterAllJust'. @Column "x" (Maybe Double)@ becomes @Column "x" Double@. @@ -232,14 +249,37 @@ type family SharedNames (left :: [Type]) (right :: [Type]) :: [Symbol] where SharedNames '[] right = '[] SharedNames (Column n _ ': rest) right = - If (HasName n right) (n ': SharedNames rest right) (SharedNames rest right) + SharedNamesHelper (HasName n right) n rest right +type family + SharedNamesHelper + (found :: Bool) + (n :: Symbol) + (rest :: [Type]) + (right :: [Type]) :: + [Symbol] + where + SharedNamesHelper 'True n rest right = n ': SharedNames rest right + SharedNamesHelper 'False n rest right = SharedNames rest right + -- | Columns from @left@ whose names do NOT appear in @right@. type family UniqueLeft (left :: [Type]) (rightNames :: [Symbol]) :: [Type] where UniqueLeft '[] _ = '[] UniqueLeft (Column n a ': rest) rn = - If (IsElem n rn) (UniqueLeft rest rn) (Column n a ': UniqueLeft rest rn) + UniqueLeftHelper (IsElem n rn) n a rest rn +type family + UniqueLeftHelper + (found :: Bool) + (n :: Symbol) + (a :: Type) + (rest :: [Type]) + (rn :: [Symbol]) :: + [Type] + where + UniqueLeftHelper 'True n a rest rn = UniqueLeft rest rn + UniqueLeftHelper 'False n a rest rn = Column n a ': UniqueLeft rest rn + -- | Wrap column types in Maybe. type family WrapMaybe (cols :: [Type]) :: [Type] where WrapMaybe '[] = '[] @@ -249,24 +289,58 @@ type family WrapMaybeColumns (names :: [Symbol]) (cols :: [Type]) :: [Type] where WrapMaybeColumns names '[] = '[] WrapMaybeColumns names (Column n a ': rest) = - If - (IsElem n names) - (Column n (Maybe a) ': WrapMaybeColumns names rest) - (Column n a ': WrapMaybeColumns names rest) + WrapMaybeColumnsHelper (IsElem n names) n a names rest +type family + WrapMaybeColumnsHelper + (found :: Bool) + (n :: Symbol) + (a :: Type) + (names :: [Symbol]) + (rest :: [Type]) :: + [Type] + where + WrapMaybeColumnsHelper 'True n a names rest = + Column n (Maybe a) ': WrapMaybeColumns names rest + WrapMaybeColumnsHelper 'False n a names rest = + Column n a ': WrapMaybeColumns names rest + -- | Columns in left whose names collide with right (excluding keys). type family CollidingColumns (left :: [Type]) (right :: [Type]) (keys :: [Symbol]) :: [Type] where CollidingColumns '[] _ _ = '[] CollidingColumns (Column n a ': rest) right keys = - If - (IsElem n keys) - (CollidingColumns rest right keys) - ( If - (HasName n right) - (Column n (These a (Lookup n right)) ': CollidingColumns rest right keys) - (CollidingColumns rest right keys) - ) + CollidingColumnsHelper1 (IsElem n keys) n a rest right keys +type family + CollidingColumnsHelper1 + (isKey :: Bool) + (n :: Symbol) + (a :: Type) + (rest :: [Type]) + (right :: [Type]) + (keys :: [Symbol]) :: + [Type] + where + CollidingColumnsHelper1 'True n a rest right keys = + CollidingColumns rest right keys + CollidingColumnsHelper1 'False n a rest right keys = + CollidingColumnsHelper2 (HasName n right) n a rest right keys + +type family + CollidingColumnsHelper2 + (inRight :: Bool) + (n :: Symbol) + (a :: Type) + (rest :: [Type]) + (right :: [Type]) + (keys :: [Symbol]) :: + [Type] + where + CollidingColumnsHelper2 'True n a rest right keys = + Column n (These a (Lookup n right)) ': CollidingColumns rest right keys + CollidingColumnsHelper2 'False n a rest right keys = + CollidingColumns rest right keys + -- | Inner join result schema. type family InnerJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) :: [Type] where InnerJoinSchema keys left right = @@ -322,18 +396,24 @@ ) ) -------------------------------------------------------------------------------- --- GroupBy helpers -------------------------------------------------------------------------------- - -- | Extract Column entries from a schema whose names appear in @keys@. type family GroupKeyColumns (keys :: [Symbol]) (cols :: [Type]) :: [Type] where GroupKeyColumns keys '[] = '[] GroupKeyColumns keys (Column n a ': rest) = - If - (IsElem n keys) - (Column n a ': GroupKeyColumns keys rest) - (GroupKeyColumns keys rest) + GroupKeyColumnsHelper (IsElem n keys) n a keys rest + +type family + GroupKeyColumnsHelper + (found :: Bool) + (n :: Symbol) + (a :: Type) + (keys :: [Symbol]) + (rest :: [Type]) :: + [Type] + where + GroupKeyColumnsHelper 'True n a keys rest = + Column n a ': GroupKeyColumns keys rest + GroupKeyColumnsHelper 'False n a keys rest = GroupKeyColumns keys rest -- | Provides runtime evidence of a schema: a list of (name, TypeRep) pairs. class KnownSchema (cols :: [Type]) where diff --git a/src/DataFrame/Typed/Types.hs b/src/DataFrame/Typed/Types.hs --- a/src/DataFrame/Typed/Types.hs +++ b/src/DataFrame/Typed/Types.hs @@ -73,8 +73,8 @@ -- | A typed sort order validated against schema @cols@. data TSortOrder (cols :: [Type]) where - Asc :: (Columnable a) => TExpr cols a -> TSortOrder cols - Desc :: (Columnable a) => TExpr cols a -> TSortOrder cols + Asc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols + Desc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols -- | A phantom-typed wrapper over 'GroupedDataFrame'. newtype TypedGrouped (keys :: [Symbol]) (cols :: [Type]) diff --git a/tests/IO/CSV.hs b/tests/IO/CSV.hs new file mode 100644 --- /dev/null +++ b/tests/IO/CSV.hs @@ -0,0 +1,98 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module IO.CSV where + +import qualified Data.ByteString.Lazy as BL +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE +import qualified DataFrame as D +import DataFrame.IO.CSV (fromCsv, fromCsvBytes) +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.DataFrame (DataFrame (..), toCsv) +import qualified DataFrame.Operations.Core as D +import Test.HUnit ( + Test (TestCase, TestLabel), + assertEqual, + assertFailure, + ) + +-- | Happy path: parse a simple CSV string with Int and Text columns. +fromCsvHappyPath :: Test +fromCsvHappyPath = TestLabel "fromCsv_happy_path" $ TestCase $ do + result <- fromCsv "name,age\nAlice,30\nBob,25\nCharlie,35\n" + case result of + Left err -> assertFailure $ "Unexpected Left: " ++ err + Right df -> do + assertEqual "rows" 3 (D.nRows df) + assertEqual "columns" 2 (D.nColumns df) + +-- | Empty input should return a Left. +fromCsvEmpty :: Test +fromCsvEmpty = TestLabel "fromCsv_empty" $ TestCase $ do + result <- fromCsv "" + case result of + Left _ -> return () + Right _ -> assertFailure "Expected Left for empty input" + +-- | fromCsvBytes happy path. +fromCsvBytesHappyPath :: Test +fromCsvBytesHappyPath = TestLabel "fromCsvBytes_happy_path" $ TestCase $ do + let bs = BL.fromStrict (TE.encodeUtf8 "x,y\n1,2\n3,4\n") + df <- fromCsvBytes bs + assertEqual "rows" 2 (D.nRows df) + assertEqual "columns" 2 (D.nColumns df) + +-- | Round trip: toCsv then fromCsv preserves data. +fromCsvRoundTrip :: Test +fromCsvRoundTrip = TestLabel "fromCsv_roundTrip" $ TestCase $ do + let df = + D.fromNamedColumns + [ ("a", DI.fromList @Int [1, 2, 3]) + , ("b", DI.fromList @T.Text ["hello", "world", "test"]) + ] + let csvString = T.unpack (toCsv df) + result <- fromCsv csvString + case result of + Left err -> assertFailure $ "Unexpected Left: " ++ err + Right df' -> do + assertEqual + "round trip dimensions" + (dataframeDimensions df) + (dataframeDimensions df') + assertEqual "round trip data" df df' + +-- | Round trip via fromCsvBytes. +fromCsvBytesRoundTrip :: Test +fromCsvBytesRoundTrip = TestLabel "fromCsvBytes_roundTrip" $ TestCase $ do + let df = + D.fromNamedColumns + [ ("x", DI.fromList @Int [10, 20]) + , ("y", DI.fromList @Double [1.5, 2.5]) + ] + let bs = BL.fromStrict (TE.encodeUtf8 (toCsv df)) + df' <- fromCsvBytes bs + assertEqual + "round trip dimensions" + (dataframeDimensions df) + (dataframeDimensions df') + +-- | Single column CSV. +fromCsvSingleColumn :: Test +fromCsvSingleColumn = TestLabel "fromCsv_single_column" $ TestCase $ do + result <- fromCsv "id\n10\n20\n30\n" + case result of + Left err -> assertFailure $ "Unexpected Left: " ++ err + Right df -> do + assertEqual "rows" 3 (D.nRows df) + assertEqual "columns" 1 (D.nColumns df) + +tests :: [Test] +tests = + [ fromCsvHappyPath + , fromCsvEmpty + , fromCsvBytesHappyPath + , fromCsvRoundTrip + , fromCsvBytesRoundTrip + , fromCsvSingleColumn + ] diff --git a/tests/Main.hs b/tests/Main.hs --- a/tests/Main.hs +++ b/tests/Main.hs @@ -11,6 +11,7 @@ import qualified DecisionTree import qualified Functions +import qualified IO.CSV import qualified IO.JSON import qualified Internal.Parsing import qualified LazyParquet @@ -62,6 +63,7 @@ ++ Operations.Take.tests ++ Operations.Typing.tests ++ Functions.tests + ++ IO.CSV.tests ++ IO.JSON.tests ++ Parquet.tests ++ LazyParquet.tests diff --git a/tests/Parquet.hs b/tests/Parquet.hs --- a/tests/Parquet.hs +++ b/tests/Parquet.hs @@ -4,6 +4,7 @@ module Parquet where import Assertions (assertExpectException) +import Control.Monad (forM_) import qualified DataFrame as D import qualified DataFrame.Functions as F import qualified DataFrame.IO.Parquet as DP @@ -29,7 +30,7 @@ word32ToLittleEndian, word64ToLittleEndian, ) -import DataFrame.Internal.Column (hasMissing) +import DataFrame.Internal.Column (hasElemType, hasMissing) import DataFrame.Internal.DataFrame (unsafeGetColumn) import GHC.IO (unsafePerformIO) import Test.HUnit @@ -41,6 +42,18 @@ , test (DP._readParquetWithOpts (Just True) D.defaultParquetReadOptions) ] +assertColumnNullability :: + String -> [(T.Text, Bool)] -> D.DataFrame -> Assertion +assertColumnNullability label expected df = + forM_ expected $ \(columnName, shouldBeNullable) -> + assertBool + ( label + <> ": expected " + <> T.unpack columnName + <> if shouldBeNullable then " to be nullable" else " to be non-nullable" + ) + (hasMissing (unsafeGetColumn columnName df) == shouldBeNullable) + allTypesPlain :: Test allTypesPlain = testBothReadParquetPaths $ \readParquet -> TestCase @@ -169,6 +182,58 @@ ) ) +safeColumnsWithOpts :: Test +safeColumnsWithOpts = + TestCase $ do + defaultDf <- D.readParquet "./tests/data/alltypes_plain.parquet" + safeDf <- + D.readParquetWithOpts + (D.defaultParquetReadOptions{D.safeColumns = True}) + "./tests/data/alltypes_plain.parquet" + + assertEqual + "safeColumnsWithOpts dimensions" + (D.dimensions defaultDf) + (D.dimensions safeDf) + assertColumnNullability + "default read" + [("id", False), ("bool_col", False)] + defaultDf + assertColumnNullability + "safeColumns read" + [("id", True), ("bool_col", True)] + safeDf + assertBool + "safeColumns id type" + (hasElemType @(Maybe Int32) (unsafeGetColumn "id" safeDf)) + assertBool + "safeColumns bool_col type" + (hasElemType @(Maybe Bool) (unsafeGetColumn "bool_col" safeDf)) + +safeColumnsWithSelectedColumns :: Test +safeColumnsWithSelectedColumns = + TestCase $ do + df <- + D.readParquetWithOpts + ( D.defaultParquetReadOptions + { D.selectedColumns = Just ["id", "bool_col"] + , D.safeColumns = True + } + ) + "./tests/data/alltypes_plain.parquet" + + assertEqual "safeColumnsWithSelectedColumns dimensions" (8, 2) (D.dimensions df) + assertColumnNullability + "safeColumns projected read" + [("id", True), ("bool_col", True)] + df + assertBool + "safeColumns projected id type" + (hasElemType @(Maybe Int32) (unsafeGetColumn "id" df)) + assertBool + "safeColumns projected bool_col type" + (hasElemType @(Maybe Bool) (unsafeGetColumn "bool_col" df)) + predicateWithOptsAcrossFiles :: Test predicateWithOptsAcrossFiles = TestCase @@ -1029,6 +1094,8 @@ , rowRangeWithOpts , predicateWithOpts , predicateUsesNonSelectedColumnWithOpts + , safeColumnsWithOpts + , safeColumnsWithSelectedColumns , predicateWithOptsAcrossFiles , missingSelectedColumnWithOpts , mtCars