diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,18 @@ # Revision history for dataframe +## 0.5.0.0 +* SortOrder now takes an expression rather than a string. +* selectBy now respects the original column order. +* Some changes to the internal representation of the expression GADT. +* Added fixity to binary operations for column comparisons. +* Fixed ZSTD decompression for Parquet files. +* Added logical type handling for timestamps and decimals. +* aggregation now handles null dataframes gracefully. +* Drop random dep to 1.2 and use random-shuffle for shuffling instead. +* Join now merges columns to a `These`. +* Added writeTsv function. +* Move expression operators to `Operators` namespace. + ## 0.4.1.0 * Improve signal handling of dataframe repl. * `writeCsv` not correctly writes `Maybe` values (thanks to @mcoady). diff --git a/app/Synthesis.hs b/app/Synthesis.hs --- a/app/Synthesis.hs +++ b/app/Synthesis.hs @@ -8,9 +8,8 @@ import qualified DataFrame.Functions as F import Data.Char -import DataFrame ((|>)) import DataFrame.DecisionTree -import DataFrame.Functions ((.&&), (.=), (.==)) +import DataFrame.Operators hiding (name) import System.Random diff --git a/benchmark/Main.hs b/benchmark/Main.hs --- a/benchmark/Main.hs +++ b/benchmark/Main.hs @@ -6,7 +6,7 @@ import Control.Monad (void) import Criterion.Main -import DataFrame ((|>)) +import DataFrame.Operators import System.Process haskell :: IO () @@ -46,9 +46,9 @@ |> D.groupBy ["ocean_proximity"] |> D.aggregate [ F.minimum (F.col @Double "median_house_value") - `F.as` "minimum_median_house_value" + `as` "minimum_median_house_value" , F.maximum (F.col @Double "median_house_value") - `F.as` "maximum_median_house_value" + `as` "maximum_median_house_value" ] groupByPolars :: IO () diff --git a/dataframe.cabal b/dataframe.cabal --- a/dataframe.cabal +++ b/dataframe.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: dataframe -version: 0.4.1.0 +version: 0.5.0.0 synopsis: A fast, safe, and intuitive DataFrame library. @@ -65,6 +65,7 @@ DataFrame.Operations.Core, DataFrame.Operations.Join, DataFrame.Operations.Merge, + DataFrame.Operators, DataFrame.Operations.Permutation, DataFrame.Operations.Subset, DataFrame.Operations.Statistics, @@ -104,11 +105,13 @@ hashable >= 1.2 && < 2, process ^>= 1.6, snappy-hs ^>= 0.1, - random >= 1.3 && < 2, + random >= 1.2 && < 1.3, + random-shuffle >= 0.0.4 && < 1, regex-tdfa >= 1.3.0 && < 2, scientific >=0.3.1 && <0.4, template-haskell >= 2.0 && < 3, text >= 2.0 && < 3, + these >= 1.1 && < 2, time >= 1.12 && < 2, unordered-containers >= 0.1 && < 1, vector ^>= 0.13, @@ -131,7 +134,7 @@ import: warnings main-is: Benchmark.hs build-depends: base >= 4 && < 5, - dataframe ^>= 0.4, + dataframe >= 0.5 && < 1, random >= 1 && < 2, time >= 1.12 && < 2, vector ^>= 0.13, @@ -143,7 +146,7 @@ import: warnings main-is: Synthesis.hs build-depends: base >= 4 && < 5, - dataframe ^>= 0.4, + dataframe >= 0.5 && < 1, random >= 1 && < 2, text >= 2.0 && < 3 hs-source-dirs: app @@ -157,8 +160,9 @@ directory >= 1.3.0.0 && < 2, filepath >= 1.4 && < 2, process >= 1.6 && < 2, - time >= 1.12 && < 2, - unix >= 2 && < 3 + time >= 1.12 && < 2 + if !os(windows) + build-depends: unix >= 2 && < 3 hs-source-dirs: app default-language: Haskell2010 ghc-options: -rtsopts -threaded -with-rtsopts=-N @@ -171,7 +175,7 @@ build-depends: base >= 4 && < 5, criterion >= 1 && < 2, process >= 1.6 && < 2, - dataframe ^>= 0.4 + dataframe >= 0.5 && < 1 default-language: Haskell2010 ghc-options: -threaded @@ -195,20 +199,22 @@ Operations.Join, Operations.Merge, Operations.ReadCsv, + Operations.Shuffle, Operations.Sort, Operations.Subset, Operations.Statistics, Operations.Take, - Parquet + Parquet, + Monad build-depends: base >= 4 && < 5, - dataframe ^>= 0.4, + dataframe >= 0.5 && < 1, directory >= 1.3.0.0 && < 2, HUnit ^>= 1.6, QuickCheck >= 2 && < 3, - random >= 1 && < 2, random-shuffle >= 0.0.4 && < 1, random >= 1 && < 2, text >= 2.0 && < 3, + these >= 1.1 && < 2, time >= 1.12 && < 2, vector ^>= 0.13, containers >= 0.6.7 && < 0.9 diff --git a/src/DataFrame.hs b/src/DataFrame.hs --- a/src/DataFrame.hs +++ b/src/DataFrame.hs @@ -69,7 +69,7 @@ longitude | 20640 | 0 | 0 | 844 | Double -- 2) Project & filter -ghci> :exposeColumn df +ghci> :declareColumns df ghci> df1 = D.filterWhere (ocean_proximity .== \"ISLAND\") df0 D.|> D.select [F.name median_house_value, F.name median_income, F.name ocean_proximity] -- 3) Add a derived column using the expression DSL @@ -77,10 +77,11 @@ ghci> df2 = D.derive "rooms_per_household" (total_rooms / households) df0 -- 4) Group + aggregate +ghci> import DataFrame.Operators ghci> let grouped = D.groupBy ["ocean_proximity"] df0 ghci> let summary = D.aggregate - [ F.maximum median_house_value \`F.as\` "max_house_value"] + [ F.maximum median_house_value \`as\` "max_house_value"] grouped ghci> D.take 5 summary ---------------------------------- @@ -198,6 +199,9 @@ module Row, module Expression, + -- * Operator symbols. + module Operators, + -- * Display operations module Display, @@ -226,9 +230,6 @@ -- * Plotting module Plot, - - -- * Convenience functions - (|>), ) where @@ -352,7 +353,4 @@ takeLast, ) import DataFrame.Operations.Transformations as Transformations - -import Data.Function ((&)) - -(|>) = (&) +import DataFrame.Operators as Operators diff --git a/src/DataFrame/DecisionTree.hs b/src/DataFrame/DecisionTree.hs --- a/src/DataFrame/DecisionTree.hs +++ b/src/DataFrame/DecisionTree.hs @@ -30,7 +30,7 @@ import qualified Data.Vector.Unboxed as VU import Type.Reflection (typeRep) -import DataFrame.Functions ((./=), (.<), (.<=), (.==), (.>), (.>=)) +import DataFrame.Operators data TreeConfig = TreeConfig { maxTreeDepth :: Int @@ -466,8 +466,8 @@ (If condInner tInner _, _) | cond == condInner -> If cond tInner f (_, If condInner _ fInner) | cond == condInner -> If cond t fInner _ -> If cond t f -pruneExpr (UnaryOp name op e) = UnaryOp name op (pruneExpr e) -pruneExpr (BinaryOp name op l r) = BinaryOp name op (pruneExpr l) (pruneExpr r) +pruneExpr (Unary op e) = Unary op (pruneExpr e) +pruneExpr (Binary op l r) = Binary op (pruneExpr l) (pruneExpr r) pruneExpr e = e buildGreedyTree :: 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 @@ -65,8 +65,7 @@ lines = if properMarkdown then - T.concat [" ", border, " "] - : T.concat ["| ", fillCols colTitleFill consolidatedHeader, " |"] + T.concat ["| ", fillCols colTitleFill consolidatedHeader, " |"] : T.concat ["| ", separator, " |"] : map ((\t -> T.concat ["| ", t, " |"]) . fillCols colValueFill) rows else diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs --- a/src/DataFrame/Functions.hs +++ b/src/DataFrame/Functions.hs @@ -4,6 +4,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} @@ -12,19 +13,16 @@ {-# LANGUAGE TypeApplications #-} {-# LANGUAGE UndecidableInstances #-} -module DataFrame.Functions where +module DataFrame.Functions (module DataFrame.Functions, module DataFrame.Operators) where import DataFrame.Internal.Column import DataFrame.Internal.DataFrame ( DataFrame (..), unsafeGetColumn, ) -import DataFrame.Internal.Expression ( - Expr (..), - NamedExpr, - UExpr (..), - ) +import DataFrame.Internal.Expression hiding (normalize) import DataFrame.Internal.Statistics +import DataFrame.Operations.Core import Control.Applicative import Control.Monad @@ -32,6 +30,7 @@ import qualified Data.Char as Char import Data.Function import Data.Functor +import Data.Int import qualified Data.List as L import qualified Data.Map as M import qualified Data.Maybe as Maybe @@ -39,8 +38,11 @@ import Data.Time import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU +import Data.Word import qualified DataFrame.IO.CSV as CSV import qualified DataFrame.IO.Parquet as Parquet +import DataFrame.IO.Parquet.Thrift +import DataFrame.Operators import Debug.Trace (trace) import Language.Haskell.TH import qualified Language.Haskell.TH.Syntax as TH @@ -48,147 +50,143 @@ import Prelude hiding (maximum, minimum) import Prelude as P -infix 4 .==, .<, .<=, .>=, .>, ./= -infixr 3 .&& -infixr 2 .|| - -name :: (Show a) => Expr a -> T.Text -name (Col n) = n -name other = - error $ - "You must call `name` on a column reference. Not the expression: " ++ show other - -col :: (Columnable a) => T.Text -> Expr a -col = Col - -as :: (Columnable a) => Expr a -> T.Text -> NamedExpr -as expr name = (name, Wrap expr) - -infixr 0 .= -(.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr -(.=) = flip as - -ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a -ifThenElse = If - -lit :: (Columnable a) => a -> Expr a -lit = Lit - lift :: (Columnable a, Columnable b) => (a -> b) -> Expr a -> Expr b -lift = UnaryOp "udf" +lift f = + Unary (MkUnaryOp{unaryFn = f, unaryName = "unaryUdf", unarySymbol = Nothing}) lift2 :: (Columnable c, Columnable b, Columnable a) => (c -> b -> a) -> Expr c -> Expr b -> Expr a -lift2 = BinaryOp "udf" +lift2 f = + Binary + ( MkBinaryOp + { binaryFn = f + , binaryName = "binaryUdf" + , binarySymbol = Nothing + , binaryCommutative = False + , binaryPrecedence = 0 + } + ) +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}) + +lift2Decorated :: + (Columnable c, Columnable b, Columnable a) => + (c -> b -> a) -> + T.Text -> + Maybe T.Text -> + Bool -> + Int -> + Expr c -> + Expr b -> + Expr a +lift2Decorated f name rep comm prec = + Binary + ( MkBinaryOp + { binaryFn = f + , binaryName = name + , binarySymbol = rep + , binaryCommutative = comm + , binaryPrecedence = prec + } + ) + toDouble :: (Columnable a, Real a) => Expr a -> Expr Double -toDouble = UnaryOp "toDouble" realToFrac +toDouble = + Unary + ( MkUnaryOp + { unaryFn = realToFrac + , unaryName = "toDouble" + , unarySymbol = Nothing + } + ) +infix 8 `div` div :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a -div = BinaryOp "div" Prelude.div +div = lift2Decorated Prelude.div "div" (Just "//") False 7 mod :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a -mod = BinaryOp "mod" Prelude.mod - -(.==) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool -(.==) = BinaryOp "eq" (==) - -(./=) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool -(./=) = BinaryOp "neq" (/=) +mod = lift2Decorated Prelude.mod "mod" Nothing False 7 eq :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool -eq = BinaryOp "eq" (==) - -(.<) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool -(.<) = BinaryOp "lt" (<) +eq = (.==) lt :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool -lt = BinaryOp "lt" (<) - --- TODO: Generalize this pattern for other equality functions. -(.>) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool -(.>) = BinaryOp "gt" (>) +lt = (.<) gt :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool gt = (.>) -(.<=) :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool -(.<=) = BinaryOp "leq" (<=) - leq :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool leq = (.<=) -(.>=) :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool -(.>=) = BinaryOp "geq" (>=) - geq :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool -geq = BinaryOp "geq" (>=) +geq = (.>=) and :: Expr Bool -> Expr Bool -> Expr Bool -and = BinaryOp "and" (&&) - -(.&&) :: Expr Bool -> Expr Bool -> Expr Bool -(.&&) = BinaryOp "and" (&&) +and = (.&&) or :: Expr Bool -> Expr Bool -> Expr Bool -or = BinaryOp "or" (||) - -(.||) :: Expr Bool -> Expr Bool -> Expr Bool -(.||) = BinaryOp "or" (||) +or = (.||) not :: Expr Bool -> Expr Bool -not = UnaryOp "not" Prelude.not +not = + Unary + (MkUnaryOp{unaryFn = Prelude.not, unaryName = "not", unarySymbol = Just "~"}) count :: (Columnable a) => Expr a -> Expr Int -count expr = AggFold expr "count" 0 (\acc _ -> acc + 1) +count = Agg (FoldAgg "count" (Just 0) (\acc _ -> acc + 1)) collect :: (Columnable a) => Expr a -> Expr [a] -collect expr = AggFold expr "collect" [] (flip (:)) +collect = Agg (FoldAgg "collect" (Just []) (flip (:))) -mode :: (Columnable a, Eq a) => Expr a -> Expr a -mode expr = - AggVector - expr - "mode" - ( fst - . L.maximumBy (compare `on` snd) - . M.toList - . V.foldl' (\m e -> M.insertWith (+) e 1 m) M.empty +mode :: (Ord a, Columnable a, Eq a) => Expr a -> Expr a +mode = + Agg + ( CollectAgg + "mode" + ( fst + . L.maximumBy (compare `on` snd) + . M.toList + . V.foldl' (\m e -> M.insertWith (+) e 1 m) M.empty + ) ) minimum :: (Columnable a, Ord a) => Expr a -> Expr a -minimum expr = AggReduce expr "minimum" Prelude.min +minimum = Agg (FoldAgg "minimum" Nothing Prelude.min) maximum :: (Columnable a, Ord a) => Expr a -> Expr a -maximum expr = AggReduce expr "maximum" Prelude.max +maximum = Agg (FoldAgg "maximum" Nothing Prelude.max) sum :: forall a. (Columnable a, Num a) => Expr a -> Expr a -sum expr = AggReduce expr "sum" (+) +sum = Agg (FoldAgg "sum" Nothing (+)) {-# SPECIALIZE DataFrame.Functions.sum :: Expr Double -> Expr Double #-} {-# SPECIALIZE DataFrame.Functions.sum :: Expr Int -> Expr Int #-} {-# INLINEABLE DataFrame.Functions.sum #-} sumMaybe :: forall a. (Columnable a, Num a) => Expr (Maybe a) -> Expr a -sumMaybe expr = AggVector expr "sumMaybe" (P.sum . Maybe.catMaybes . V.toList) +sumMaybe = Agg (CollectAgg "sumMaybe" (P.sum . Maybe.catMaybes . V.toList)) mean :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double -mean expr = AggNumericVector expr "mean" mean' +mean = Agg (CollectAgg "mean" mean') {-# SPECIALIZE DataFrame.Functions.mean :: Expr Double -> Expr Double #-} {-# SPECIALIZE DataFrame.Functions.mean :: Expr Int -> Expr Double #-} {-# INLINEABLE DataFrame.Functions.mean #-} meanMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double -meanMaybe expr = AggVector expr "meanMaybe" (mean' . optionalToDoubleVector) +meanMaybe = Agg (CollectAgg "meanMaybe" (mean' . optionalToDoubleVector)) variance :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double -variance expr = AggNumericVector expr "variance" variance' +variance = Agg (CollectAgg "variance" variance') median :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double -median expr = AggNumericVector expr "median" median' +median = Agg (CollectAgg "median" median') medianMaybe :: (Columnable a, Real a) => Expr (Maybe a) -> Expr Double -medianMaybe expr = AggVector expr "meanMaybe" (median' . optionalToDoubleVector) +medianMaybe = Agg (CollectAgg "meanMaybe" (median' . optionalToDoubleVector)) optionalToDoubleVector :: (Real a) => V.Vector (Maybe a) -> VU.Vector Double optionalToDoubleVector = @@ -198,55 +196,53 @@ [] percentile :: Int -> Expr Double -> Expr Double -percentile n expr = - AggNumericVector - expr - (T.pack $ "percentile " ++ show n) - (percentile' n) +percentile n = + Agg + ( CollectAgg + (T.pack $ "percentile " ++ show n) + (percentile' n) + ) stddev :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double -stddev expr = AggNumericVector expr "stddev" (sqrt . variance') +stddev = Agg (CollectAgg "stddev" (sqrt . variance')) stddevMaybe :: forall a. (Columnable a, Real a) => Expr (Maybe a) -> Expr Double -stddevMaybe expr = AggVector expr "stddevMaybe" (sqrt . variance' . optionalToDoubleVector) +stddevMaybe = Agg (CollectAgg "stddevMaybe" (sqrt . variance' . optionalToDoubleVector)) zScore :: Expr Double -> Expr Double zScore c = (c - mean c) / stddev c pow :: (Columnable a, Num a) => Expr a -> Int -> Expr a -pow _ 0 = Lit 1 -pow (Lit n) i = Lit (n ^ i) -pow expr 1 = expr -pow expr i = BinaryOp "pow" (^) expr (lit i) +pow = (.^^) -relu :: (Columnable a, Num a) => Expr a -> Expr a -relu = UnaryOp "relu" (Prelude.max 0) +relu :: (Columnable a, Num a, Ord a) => Expr a -> Expr a +relu = liftDecorated (Prelude.max 0) "relu" Nothing min :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr a -min = BinaryOp "min" Prelude.min +min = lift2Decorated Prelude.min "max" Nothing True 1 max :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr a -max = BinaryOp "max" Prelude.max +max = lift2Decorated Prelude.max "max" Nothing True 1 reduce :: forall a b. (Columnable a, Columnable b) => Expr b -> a -> (a -> b -> a) -> Expr a -reduce expr = AggFold expr "foldUdf" +reduce expr start f = Agg (FoldAgg "foldUdf" (Just start) f) expr toMaybe :: (Columnable a) => Expr a -> Expr (Maybe a) -toMaybe = UnaryOp "toMaybe" Just +toMaybe = liftDecorated Just "toMaybe" Nothing fromMaybe :: (Columnable a) => a -> Expr (Maybe a) -> Expr a -fromMaybe d = UnaryOp ("fromMaybe " <> T.pack (show d)) (Maybe.fromMaybe d) +fromMaybe d = liftDecorated (Maybe.fromMaybe d) "fromMaybe" Nothing isJust :: (Columnable a) => Expr (Maybe a) -> Expr Bool -isJust = UnaryOp "isJust" Maybe.isJust +isJust = liftDecorated Maybe.isJust "isJust" Nothing isNothing :: (Columnable a) => Expr (Maybe a) -> Expr Bool -isNothing = UnaryOp "isNothing" Maybe.isNothing +isNothing = liftDecorated Maybe.isNothing "isNothing" Nothing fromJust :: (Columnable a) => Expr (Maybe a) -> Expr a -fromJust = UnaryOp "fromJust" Maybe.fromJust +fromJust = liftDecorated Maybe.fromJust "fromJust" Nothing whenPresent :: forall a b. @@ -262,7 +258,14 @@ recode :: forall a b. (Columnable a, Columnable b) => [(a, b)] -> Expr a -> Expr (Maybe b) -recode mapping = UnaryOp (T.pack ("recode " ++ show mapping)) (`lookup` mapping) +recode mapping = + Unary + ( MkUnaryOp + { unaryFn = (`lookup` mapping) + , unaryName = "recode " <> T.pack (show mapping) + , unarySymbol = Nothing + } + ) recodeWithCondition :: forall a b. @@ -275,30 +278,54 @@ forall a b. (Columnable a, Columnable b) => b -> [(a, b)] -> Expr a -> Expr b recodeWithDefault d mapping = - UnaryOp - (T.pack ("recodeWithDefault " ++ show d ++ " " ++ show mapping)) - (Maybe.fromMaybe d . (`lookup` mapping)) + Unary + ( MkUnaryOp + { unaryFn = Maybe.fromMaybe d . (`lookup` mapping) + , unaryName = + "recodeWithDefault " <> T.pack (show d) <> " " <> T.pack (show mapping) + , unarySymbol = Nothing + } + ) firstOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a) -firstOrNothing = lift Maybe.listToMaybe +firstOrNothing = liftDecorated Maybe.listToMaybe "firstOrNothing" Nothing lastOrNothing :: (Columnable a) => Expr [a] -> Expr (Maybe a) -lastOrNothing = lift (Maybe.listToMaybe . reverse) +lastOrNothing = liftDecorated (Maybe.listToMaybe . reverse) "lastOrNothing" Nothing splitOn :: T.Text -> Expr T.Text -> Expr [T.Text] -splitOn delim = lift (T.splitOn delim) +splitOn delim = liftDecorated (T.splitOn delim) "splitOn" Nothing match :: T.Text -> Expr T.Text -> Expr (Maybe T.Text) -match regex = lift ((\r -> if T.null r then Nothing else Just r) . (=~ regex)) +match regex = + liftDecorated + ((\r -> if T.null r then Nothing else Just r) . (=~ regex)) + ("match " <> T.pack (show regex)) + Nothing matchAll :: T.Text -> Expr T.Text -> Expr [T.Text] -matchAll regex = lift (getAllTextMatches . (=~ regex)) +matchAll regex = + liftDecorated + (getAllTextMatches . (=~ regex)) + ("matchAll " <> T.pack (show regex)) + Nothing -parseDate :: T.Text -> Expr T.Text -> Expr (Maybe Day) -parseDate format = lift (parseTimeM True defaultTimeLocale (T.unpack format) . T.unpack) +parseDate :: + (ParseTime t, Columnable t) => T.Text -> Expr T.Text -> Expr (Maybe t) +parseDate format = + liftDecorated + (parseTimeM True defaultTimeLocale (T.unpack format) . T.unpack) + ("parseDate " <> format) + Nothing daysBetween :: Expr Day -> Expr Day -> Expr Int -daysBetween d1 d2 = lift fromIntegral (lift2 diffDays d1 d2) +daysBetween = + lift2Decorated + (\d1 d2 -> fromIntegral (diffDays d1 d2)) + "daysBetween" + Nothing + True + 2 bind :: forall a b m. @@ -413,23 +440,57 @@ (CSV.readSeparated (CSV.defaultReadOptions{CSV.numColumns = Just 100}) path) declareColumns df --- TODO: We don't have to read the whole file, we can just read the schema. declareColumnsFromParquetFile :: String -> DecsQ declareColumnsFromParquetFile path = do - df <- liftIO (Parquet.readParquet path) + metadata <- liftIO (Parquet.readMetadataFromPath path) + let df = schemaToEmptyDataFrame (schema metadata) declareColumns df +schemaToEmptyDataFrame :: [SchemaElement] -> DataFrame +schemaToEmptyDataFrame elems = + let leafElems = filter (\e -> numChildren e == 0) elems + in fromNamedColumns (map schemaElemToColumn leafElems) + +schemaElemToColumn :: SchemaElement -> (T.Text, Column) +schemaElemToColumn elem = + let name = elementName elem + in (name, emptyColumnForType (elementType elem)) + +emptyColumnForType :: TType -> Column +emptyColumnForType = \case + BOOL -> fromList @Bool [] + BYTE -> fromList @Word8 [] + I16 -> fromList @Int16 [] + I32 -> fromList @Int32 [] + I64 -> fromList @Int64 [] + I96 -> fromList @Int64 [] + FLOAT -> fromList @Float [] + DOUBLE -> fromList @Double [] + STRING -> fromList @T.Text [] + other -> error $ "Unsupported parquet type for column: " <> show other + declareColumnsFromCsvWithOpts :: CSV.ReadOptions -> String -> DecsQ declareColumnsFromCsvWithOpts opts path = do df <- liftIO (CSV.readSeparated opts path) declareColumns df declareColumns :: DataFrame -> DecsQ -declareColumns df = +declareColumns = declareColumnsWithPrefix' Nothing + +declareColumnsWithPrefix :: T.Text -> DataFrame -> DecsQ +declareColumnsWithPrefix prefix = declareColumnsWithPrefix' (Just prefix) + +declareColumnsWithPrefix' :: Maybe T.Text -> DataFrame -> DecsQ +declareColumnsWithPrefix' prefix df = let names = (map fst . L.sortBy (compare `on` snd) . M.toList . columnIndices) df types = map (columnTypeString . (`unsafeGetColumn` df)) names - specs = zipWith (\name type_ -> (name, sanitize name, type_)) names types + specs = + zipWith + ( \name type_ -> (name, maybe "" (sanitize . (<> "_")) prefix <> sanitize name, type_) + ) + names + types in fmap concat $ forM specs $ \(raw, nm, tyStr) -> do ty <- typeFromString (words tyStr) 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 @@ -237,10 +237,14 @@ @ -} readSeparated :: ReadOptions -> FilePath -> IO DataFrame -readSeparated !opts !path = do - let sep = columnSeparator opts +readSeparated opts !path = do let stripUtf8Bom bs = fromMaybe bs (BL.stripPrefix "\xEF\xBB\xBF" bs) csvData <- stripUtf8Bom <$> BL.readFile path + decodeSeparated opts csvData + +decodeSeparated :: ReadOptions -> BL.ByteString -> IO DataFrame +decodeSeparated !opts csvData = do + let sep = columnSeparator opts let decodeOpts = Csv.defaultDecodeOptions{Csv.decDelimiter = fromIntegral (ord sep)} let stream = CsvStream.decodeWith decodeOpts Csv.NoHeader csvData @@ -387,6 +391,9 @@ writeCsv :: FilePath -> DataFrame -> IO () writeCsv = writeSeparated ',' + +writeTsv :: FilePath -> DataFrame -> IO () +writeTsv = writeSeparated '\t' writeSeparated :: -- | Separator 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 @@ -1,22 +1,23 @@ +{-# LANGUAGE NumericUnderscores #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.IO.Parquet ( - readParquet, - readParquetFiles, -) where +module DataFrame.IO.Parquet where import Control.Monad import Data.Bits import qualified Data.ByteString as BSO -import Data.Char import Data.Either import Data.IORef import Data.Int import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T +import Data.Text.Encoding +import Data.Time +import Data.Time.Clock.POSIX (posixSecondsToUTCTime) import Data.Word import qualified DataFrame.Internal.Column as DI import DataFrame.Internal.DataFrame (DataFrame) @@ -31,6 +32,7 @@ import DataFrame.IO.Parquet.Types import System.Directory (doesDirectoryExist) +import qualified Data.Vector.Unboxed as VU import System.FilePath (()) {- | Read a parquet file from path and load it into a dataframe. @@ -47,6 +49,7 @@ let columnNames = map fst columnPaths colMap <- newIORef (M.empty :: M.Map T.Text DI.Column) + lTypeMap <- newIORef (M.empty :: M.Map T.Text LogicalType) contents <- BSO.readFile path @@ -82,8 +85,7 @@ else colDataPageOffset let colLength = columnTotalCompressedSize metadata - let columnBytes = - map (BSO.index contents . fromIntegral) [colStart .. (colStart + colLength - 1)] + let columnBytes = BSO.take (fromIntegral colLength) (BSO.drop (fromIntegral colStart) contents) pages <- readAllPages (columnCodec metadata) columnBytes @@ -97,6 +99,7 @@ let schemaTail = drop 1 (schema fileMetadata) let colPath = columnPathInSchema (columnMetaData colChunk) let (maxDef, maxRep) = levelsForPath schemaTail colPath + let lType = logicalType (schemaTail !! colIdx) column <- processColumnPages (maxDef, maxRep) @@ -104,13 +107,17 @@ (columnType metadata) primaryEncoding maybeTypeLength + lType modifyIORef colMap (M.insertWith DI.concatColumnsEither colName column) + modifyIORef lTypeMap (M.insert colName lType) finalColMap <- readIORef colMap + finalLTypeMap <- readIORef lTypeMap let orderedColumns = map - (\name -> (name, finalColMap M.! name)) + ( \name -> (name, applyLogicalType (finalLTypeMap M.! name) $ finalColMap M.! name) + ) (filter (`M.member` finalColMap) columnNames) pure $ DI.fromNamedColumns orderedColumns @@ -176,8 +183,9 @@ ParquetType -> ParquetEncoding -> Maybe Int32 -> + LogicalType -> IO DI.Column -processColumnPages (maxDef, maxRep) pages pType _ maybeTypeLength = do +processColumnPages (maxDef, maxRep) pages pType _ maybeTypeLength lType = do let dictPages = filter isDictionaryPage pages let dataPages = filter isDataPage pages @@ -225,13 +233,13 @@ in pure (toMaybeDouble maxDef defLvls vals) PBYTE_ARRAY -> let (raws, _) = readNByteArrays nPresent afterLvls - texts = map (T.pack . map (chr . fromIntegral)) raws + texts = map decodeUtf8 raws in pure (toMaybeText maxDef defLvls texts) PFIXED_LEN_BYTE_ARRAY -> case maybeTypeLength of Just len -> let (raws, _) = splitFixed nPresent (fromIntegral len) afterLvls - texts = map (T.pack . map (chr . fromIntegral)) raws + texts = map decodeUtf8 raws in pure (toMaybeText maxDef defLvls texts) Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type length" PARQUET_TYPE_UNKNOWN -> error "Cannot read unknown Parquet type" @@ -277,13 +285,13 @@ in pure (toMaybeDouble maxDef defLvls vals) PBYTE_ARRAY -> let (raws, _) = readNByteArrays nPresent afterLvls - texts = map (T.pack . map (chr . fromIntegral)) raws + texts = map decodeUtf8 raws in pure (toMaybeText maxDef defLvls texts) PFIXED_LEN_BYTE_ARRAY -> case maybeTypeLength of Just len -> let (raws, _) = splitFixed nPresent (fromIntegral len) afterLvls - texts = map (T.pack . map (chr . fromIntegral)) raws + texts = map decodeUtf8 raws in pure (toMaybeText maxDef defLvls texts) Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type length" PARQUET_TYPE_UNKNOWN -> error "Cannot read unknown Parquet type" @@ -300,3 +308,37 @@ (c : cs) -> pure $ L.foldl' (\l r -> fromRight (error "concat failed") (DI.concatColumns l r)) c cs + +applyLogicalType :: LogicalType -> DI.Column -> DI.Column +applyLogicalType (TimestampType isUTC unit) col = + fromRight col $ + DI.mapColumn + (microsecondsToUTCTime . (* (1_000_000 `div` unitDivisor unit))) + col +applyLogicalType (DecimalType precision scale) col + | precision <= 9 = case DI.toVector @Int32 @VU.Vector col of + Right xs -> + DI.fromUnboxedVector $ + VU.map (\raw -> fromIntegral @Int32 @Double raw / 10 ^ scale) xs + Left _ -> col + | precision <= 18 = case DI.toVector @Int64 @VU.Vector col of + Right xs -> + DI.fromUnboxedVector $ + VU.map (\raw -> fromIntegral @Int64 @Double raw / 10 ^ scale) xs + Left _ -> col + | otherwise = col +applyLogicalType _ col = col + +microsecondsToUTCTime :: Int64 -> UTCTime +microsecondsToUTCTime us = + posixSecondsToUTCTime (fromIntegral us / 1_000_000) + +unitDivisor :: TimeUnit -> Int64 +unitDivisor MILLISECONDS = 1_000 +unitDivisor MICROSECONDS = 1_000_000 +unitDivisor NANOSECONDS = 1_000_000_000 +unitDivisor TIME_UNIT_UNKNOWN = 1 + +applyScale :: Int32 -> Int32 -> Double +applyScale scale rawValue = + fromIntegral rawValue / (10 ^ scale) 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 @@ -10,32 +10,45 @@ import Data.Int import Data.Word -littleEndianWord32 :: [Word8] -> Word32 +littleEndianWord32 :: BS.ByteString -> Word32 littleEndianWord32 bytes - | length bytes >= 4 = + | BS.length bytes >= 4 = foldr (.|.) 0 - (zipWith (\b i -> fromIntegral b `shiftL` i) (take 4 bytes) [0, 8, 16, 24]) - | otherwise = littleEndianWord32 (take 4 $ bytes ++ repeat 0) + ( zipWith + (\b i -> fromIntegral b `shiftL` i) + (BS.unpack $ BS.take 4 bytes) + [0, 8, 16, 24] + ) + | otherwise = + littleEndianWord32 (BS.take 4 $ bytes `BS.append` BS.pack [0, 0, 0, 0]) -littleEndianWord64 :: [Word8] -> Word64 +littleEndianWord64 :: BS.ByteString -> Word64 littleEndianWord64 bytes = foldr (.|.) 0 - (zipWith (\b i -> fromIntegral b `shiftL` i) (take 8 bytes) [0, 8 ..]) + ( zipWith + (\b i -> fromIntegral b `shiftL` i) + (BS.unpack $ BS.take 8 bytes) + [0, 8 ..] + ) -littleEndianInt32 :: [Word8] -> Int32 +littleEndianInt32 :: BS.ByteString -> Int32 littleEndianInt32 = fromIntegral . littleEndianWord32 -word64ToLittleEndian :: Word64 -> [Word8] -word64ToLittleEndian w = map (\i -> fromIntegral (w `shiftR` i)) [0, 8, 16, 24, 32, 40, 48, 56] +word64ToLittleEndian :: Word64 -> BS.ByteString +word64ToLittleEndian w = + BS.map + (\i -> fromIntegral (w `shiftR` fromIntegral i)) + (BS.pack [0, 8, 16, 24, 32, 40, 48, 56]) -word32ToLittleEndian :: Word32 -> [Word8] -word32ToLittleEndian w = map (\i -> fromIntegral (w `shiftR` i)) [0, 8, 16, 24] +word32ToLittleEndian :: Word32 -> BS.ByteString +word32ToLittleEndian w = + BS.map (\i -> fromIntegral (w `shiftR` fromIntegral i)) (BS.pack [0, 8, 16, 24]) -readUVarInt :: [Word8] -> (Word64, [Word8]) +readUVarInt :: BS.ByteString -> (Word64, BS.ByteString) readUVarInt xs = loop xs 0 0 0 where {- @@ -44,31 +57,34 @@ - The high bit (0x80) is the continuation flag: 1 = more bytes follow, 0 = last byte Why the magic number 10: For a 64‑bit integer we need at most ceil(64 / 7) = 10 bytes -} - loop :: [Word8] -> Word64 -> Int -> Int -> (Word64, [Word8]) + loop :: BS.ByteString -> Word64 -> Int -> Int -> (Word64, BS.ByteString) loop bs result _ 10 = (result, bs) - loop (b : bs) result shift i - | b < 0x80 = (result .|. (fromIntegral b `shiftL` shift), bs) - | otherwise = - let payloadBits = fromIntegral (b .&. 0x7f) :: Word64 - in loop bs (result .|. (payloadBits `shiftL` shift)) (shift + 7) (i + 1) - loop [] _ _ _ = error "readUVarInt: not enough input bytes" + loop xs result shift i = case BS.uncons xs of + Nothing -> error "readUVarInt: not enough input bytes" + Just (b, bs) -> + if b < 0x80 + then (result .|. (fromIntegral b `shiftL` shift), bs) + else + let payloadBits = fromIntegral (b .&. 0x7f) :: Word64 + in loop bs (result .|. (payloadBits `shiftL` shift)) (shift + 7) (i + 1) -readVarIntFromBytes :: (Integral a) => [Word8] -> (a, [Word8]) +readVarIntFromBytes :: (Integral a) => BS.ByteString -> (a, BS.ByteString) readVarIntFromBytes bs = (fromIntegral n, rem) where (n, rem) = loop 0 0 bs - loop _ result [] = (result, []) - loop shift result (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 + loop shift 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 -readIntFromBytes :: (Integral a) => [Word8] -> (a, [Word8]) +readIntFromBytes :: (Integral a) => BS.ByteString -> (a, BS.ByteString) readIntFromBytes bs = let (n, rem) = readVarIntFromBytes bs u = fromIntegral n :: Word32 in (fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), rem) -readInt32FromBytes :: [Word8] -> (Int32, [Word8]) +readInt32FromBytes :: BS.ByteString -> (Int32, BS.ByteString) readInt32FromBytes bs = let (n', rem) = readVarIntFromBytes @Int64 bs n = fromIntegral n' :: Int32 @@ -110,20 +126,20 @@ nameSize <- readVarIntFromBuffer @Int buf pos map (chr . fromIntegral) <$> replicateM nameSize (readAndAdvance pos buf) -readByteStringFromBytes :: [Word8] -> ([Word8], [Word8]) +readByteStringFromBytes :: BS.ByteString -> (BS.ByteString, BS.ByteString) readByteStringFromBytes xs = let (size, rem) = readVarIntFromBytes @Int xs in - splitAt size rem + BS.splitAt size rem -readByteString :: BS.ByteString -> IORef Int -> IO [Word8] +readByteString :: BS.ByteString -> IORef Int -> IO BS.ByteString readByteString buf pos = do size <- readVarIntFromBuffer @Int buf pos - replicateM size (readAndAdvance pos buf) + BS.pack <$> replicateM size (readAndAdvance pos buf) -readByteString' :: BS.ByteString -> Int64 -> IO [Word8] -readByteString' buf size = mapM (`readSingleByte` buf) [0 .. (size - 1)] +readByteString' :: BS.ByteString -> Int64 -> IO BS.ByteString +readByteString' buf size = BS.pack <$> mapM (`readSingleByte` buf) [0 .. (size - 1)] readSingleByte :: Int64 -> BS.ByteString -> IO Word8 readSingleByte pos buffer = return $ BS.index buffer (fromIntegral pos) diff --git a/src/DataFrame/IO/Parquet/ColumnStatistics.hs b/src/DataFrame/IO/Parquet/ColumnStatistics.hs --- a/src/DataFrame/IO/Parquet/ColumnStatistics.hs +++ b/src/DataFrame/IO/Parquet/ColumnStatistics.hs @@ -1,19 +1,19 @@ module DataFrame.IO.Parquet.ColumnStatistics where -import Data.Int -import Data.Word +import qualified Data.ByteString as BS +import Data.Int (Int64) data ColumnStatistics = ColumnStatistics - { columnMin :: [Word8] - , columnMax :: [Word8] + { columnMin :: BS.ByteString + , columnMax :: BS.ByteString , columnNullCount :: Int64 , columnDistictCount :: Int64 - , columnMinValue :: [Word8] - , columnMaxValue :: [Word8] + , columnMinValue :: BS.ByteString + , columnMaxValue :: BS.ByteString , isColumnMaxValueExact :: Bool , isColumnMinValueExact :: Bool } deriving (Show, Eq) emptyColumnStatistics :: ColumnStatistics -emptyColumnStatistics = ColumnStatistics [] [] 0 0 [] [] False False +emptyColumnStatistics = ColumnStatistics BS.empty BS.empty 0 0 BS.empty BS.empty False False diff --git a/src/DataFrame/IO/Parquet/Dictionary.hs b/src/DataFrame/IO/Parquet/Dictionary.hs --- a/src/DataFrame/IO/Parquet/Dictionary.hs +++ b/src/DataFrame/IO/Parquet/Dictionary.hs @@ -4,12 +4,13 @@ import Control.Monad import Data.Bits +import qualified Data.ByteString as BS import Data.Char import Data.Int import Data.Maybe import qualified Data.Text as T +import Data.Text.Encoding import Data.Time -import Data.Word import DataFrame.IO.Parquet.Binary import DataFrame.IO.Parquet.Encoding import DataFrame.IO.Parquet.Levels @@ -27,7 +28,7 @@ dictCardinality (DDouble ds) = length ds dictCardinality (DText ds) = length ds -readDictVals :: ParquetType -> [Word8] -> Maybe Int32 -> DictVals +readDictVals :: ParquetType -> BS.ByteString -> Maybe Int32 -> DictVals readDictVals PBOOLEAN bs (Just count) = DBool (take (fromIntegral count) $ readPageBool bs) readDictVals PINT32 bs _ = DInt32 (readPageInt32 bs) readDictVals PINT64 bs _ = DInt64 (readPageInt64 bs) @@ -38,60 +39,68 @@ readDictVals PFIXED_LEN_BYTE_ARRAY bs (Just len) = DText (readPageFixedBytes bs (fromIntegral len)) readDictVals t _ _ = error $ "Unsupported dictionary type: " ++ show t -readPageInt32 :: [Word8] -> [Int32] -readPageInt32 [] = [] -readPageInt32 xs = littleEndianInt32 (take 4 xs) : readPageInt32 (drop 4 xs) +readPageInt32 :: BS.ByteString -> [Int32] +readPageInt32 xs + | BS.null xs = [] + | otherwise = littleEndianInt32 (BS.take 4 xs) : readPageInt32 (BS.drop 4 xs) -readPageWord64 :: [Word8] -> [Double] -readPageWord64 [] = [] -readPageWord64 xs = - castWord64ToDouble (littleEndianWord64 (take 8 xs)) : readPageWord64 (drop 8 xs) +readPageWord64 :: BS.ByteString -> [Double] +readPageWord64 xs + | BS.null xs = [] + | otherwise = + castWord64ToDouble (littleEndianWord64 (BS.take 8 xs)) + : readPageWord64 (BS.drop 8 xs) -readPageBytes :: [Word8] -> [T.Text] -readPageBytes [] = [] -readPageBytes xs = - let lenBytes = fromIntegral (littleEndianInt32 $ take 4 xs) - totalBytesRead = lenBytes + 4 - in T.pack (map (chr . fromIntegral) $ take lenBytes (drop 4 xs)) - : readPageBytes (drop totalBytesRead xs) +readPageBytes :: BS.ByteString -> [T.Text] +readPageBytes xs + | BS.null xs = [] + | otherwise = + let lenBytes = fromIntegral (littleEndianInt32 $ BS.take 4 xs) + totalBytesRead = lenBytes + 4 + in T.pack (map (chr . fromIntegral) $ take lenBytes (BS.unpack (BS.drop 4 xs))) + : readPageBytes (BS.drop totalBytesRead xs) -readPageBool :: [Word8] -> [Bool] -readPageBool [] = [] +readPageBool :: BS.ByteString -> [Bool] readPageBool bs = - concatMap (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7]) bs + concatMap (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7]) (BS.unpack bs) -readPageInt64 :: [Word8] -> [Int64] -readPageInt64 [] = [] -readPageInt64 xs = fromIntegral (littleEndianWord64 (take 8 xs)) : readPageInt64 (drop 8 xs) +readPageInt64 :: BS.ByteString -> [Int64] +readPageInt64 xs + | BS.null xs = [] + | otherwise = + fromIntegral (littleEndianWord64 (BS.take 8 xs)) : readPageInt64 (BS.drop 8 xs) -readPageFloat :: [Word8] -> [Float] -readPageFloat [] = [] -readPageFloat xs = - castWord32ToFloat (littleEndianWord32 (take 4 xs)) : readPageFloat (drop 4 xs) +readPageFloat :: BS.ByteString -> [Float] +readPageFloat xs + | BS.null xs = [] + | otherwise = + castWord32ToFloat (littleEndianWord32 (BS.take 4 xs)) + : readPageFloat (BS.drop 4 xs) -readNInt96Times :: Int -> [Word8] -> ([UTCTime], [Word8]) +readNInt96Times :: Int -> BS.ByteString -> ([UTCTime], BS.ByteString) readNInt96Times 0 bs = ([], bs) readNInt96Times k bs = - let timestamp96 = take 12 bs + let timestamp96 = BS.take 12 bs utcTime = int96ToUTCTime timestamp96 - bs' = drop 12 bs + bs' = BS.drop 12 bs (times, rest) = readNInt96Times (k - 1) bs' in (utcTime : times, rest) -readPageInt96Times :: [Word8] -> [UTCTime] -readPageInt96Times [] = [] -readPageInt96Times bs = - let (times, _) = readNInt96Times (length bs `div` 12) bs - in times +readPageInt96Times :: BS.ByteString -> [UTCTime] +readPageInt96Times bs + | BS.null bs = [] + | otherwise = + let (times, _) = readNInt96Times (BS.length bs `div` 12) bs + in times -readPageFixedBytes :: [Word8] -> Int -> [T.Text] -readPageFixedBytes [] _ = [] -readPageFixedBytes xs len = - let chunk = take len xs - text = T.pack (map (chr . fromIntegral) chunk) - in text : readPageFixedBytes (drop len xs) len +readPageFixedBytes :: BS.ByteString -> Int -> [T.Text] +readPageFixedBytes xs len + | BS.null xs = [] + | otherwise = + decodeUtf8 (BS.take len xs) : readPageFixedBytes (BS.drop len xs) len -decodeDictV1 :: Maybe DictVals -> Int -> [Int] -> Int -> [Word8] -> IO DI.Column +decodeDictV1 :: + Maybe DictVals -> Int -> [Int] -> Int -> BS.ByteString -> IO DI.Column decodeDictV1 dictValsM maxDef defLvls nPresent bytes = case dictValsM of Nothing -> error "Dictionary-encoded page but dictionary is missing" diff --git a/src/DataFrame/IO/Parquet/Encoding.hs b/src/DataFrame/IO/Parquet/Encoding.hs --- a/src/DataFrame/IO/Parquet/Encoding.hs +++ b/src/DataFrame/IO/Parquet/Encoding.hs @@ -1,6 +1,7 @@ module DataFrame.IO.Parquet.Encoding where import Data.Bits +import qualified Data.ByteString as BS import Data.List (foldl') import Data.Word import DataFrame.IO.Parquet.Binary @@ -16,39 +17,50 @@ bytesForBW :: Int -> Int bytesForBW bw = (bw + 7) `div` 8 -unpackBitPacked :: Int -> Int -> [Word8] -> ([Word32], [Word8]) +unpackBitPacked :: Int -> Int -> BS.ByteString -> ([Word32], BS.ByteString) unpackBitPacked bw count bs | count <= 0 = ([], bs) - | null bs = ([], bs) + | BS.null bs = ([], bs) | otherwise = let totalBits = bw * count totalBytes = (totalBits + 7) `div` 8 - chunk = take totalBytes bs - rest = drop totalBytes bs - bits = concatMap (\b -> map (\i -> (b `shiftR` i) .&. 1) [0 .. 7]) chunk - toN xs = foldl' (\a (i, b) -> a .|. (b `shiftL` i)) 0 (zip [0 ..] xs) + chunk = BS.take totalBytes bs + rest = BS.drop totalBytes bs + bits = + BS.concatMap + (\b -> BS.map (\i -> (b `shiftR` fromIntegral i) .&. 1) (BS.pack [0 .. 7])) + chunk + toN :: BS.ByteString -> Word32 + toN = + fst + . BS.foldl' + (\(a, i) b -> (a .|. (fromIntegral b `shiftL` i), i + 1)) + (0 :: Word32, 0 :: Int) - extractValues _ [] = [] + extractValues :: Int -> BS.ByteString -> [Word32] extractValues n bitsLeft + | BS.null bitsLeft = [] | n <= 0 = [] - | length bitsLeft < bw = [] + | BS.length bitsLeft < bw = [] | otherwise = - let (this, bitsLeft') = splitAt bw bitsLeft + let (this, bitsLeft') = BS.splitAt bw bitsLeft in toN this : extractValues (n - 1) bitsLeft' vals = extractValues count bits - in (map fromIntegral vals, rest) + in (vals, rest) -decodeRLEBitPackedHybrid :: Int -> Int -> [Word8] -> ([Word32], [Word8]) +decodeRLEBitPackedHybrid :: + Int -> Int -> BS.ByteString -> ([Word32], BS.ByteString) decodeRLEBitPackedHybrid bw need bs | bw == 0 = (replicate need 0, bs) | otherwise = go need bs [] where mask :: Word32 mask = if bw == 32 then maxBound else (1 `shiftL` bw) - 1 + go :: Int -> BS.ByteString -> [Word32] -> ([Word32], BS.ByteString) go 0 rest acc = (reverse acc, rest) go n rest acc - | null rest = (reverse acc, rest) + | BS.null rest = (reverse acc, rest) | otherwise = let (hdr64, afterHdr) = readUVarInt rest isPacked = (hdr64 .&. 1) == 1 @@ -63,17 +75,17 @@ else let runLen = fromIntegral (hdr64 `shiftR` 1) :: Int nbytes = bytesForBW bw - word32 = littleEndianWord32 (take 4 afterHdr) - afterV = drop nbytes afterHdr + word32 = littleEndianWord32 (BS.take 4 afterHdr) + afterV = BS.drop nbytes afterHdr val = word32 .&. mask takeN = min n runLen in go (n - takeN) afterV (replicate takeN val ++ acc) -decodeDictIndicesV1 :: Int -> Int -> [Word8] -> ([Int], [Word8]) +decodeDictIndicesV1 :: Int -> Int -> BS.ByteString -> ([Int], BS.ByteString) decodeDictIndicesV1 need dictCard bs = - case bs of - [] -> error "empty dictionary index stream" - (w0 : rest0) -> + case BS.uncons bs of + Nothing -> error "empty dictionary index stream" + Just (w0, rest0) -> let bw = fromIntegral w0 :: Int (u32s, rest1) = decodeRLEBitPackedHybrid bw need rest0 in (map fromIntegral u32s, rest1) diff --git a/src/DataFrame/IO/Parquet/Levels.hs b/src/DataFrame/IO/Parquet/Levels.hs --- a/src/DataFrame/IO/Parquet/Levels.hs +++ b/src/DataFrame/IO/Parquet/Levels.hs @@ -1,5 +1,6 @@ module DataFrame.IO.Parquet.Levels where +import qualified Data.ByteString as BS import Data.Int import Data.List import qualified Data.Text as T @@ -10,7 +11,8 @@ import DataFrame.IO.Parquet.Thrift import DataFrame.IO.Parquet.Types -readLevelsV1 :: Int -> Int -> Int -> [Word8] -> ([Int], [Int], [Word8]) +readLevelsV1 :: + Int -> Int -> Int -> BS.ByteString -> ([Int], [Int], BS.ByteString) readLevelsV1 n maxDef maxRep bs = let bwDef = bitWidthForMaxLevel maxDef bwRep = bitWidthForMaxLevel maxRep @@ -19,9 +21,9 @@ if bwRep == 0 then (replicate n 0, bs) else - let repLength = littleEndianWord32 (take 4 bs) - repData = take (fromIntegral repLength) (drop 4 bs) - afterRepData = drop (4 + fromIntegral repLength) bs + let repLength = littleEndianWord32 (BS.take 4 bs) + repData = BS.take (fromIntegral repLength) (BS.drop 4 bs) + afterRepData = BS.drop (4 + fromIntegral repLength) bs (repVals, _) = decodeRLEBitPackedHybrid bwRep n repData in (repVals, afterRepData) @@ -29,18 +31,24 @@ if bwDef == 0 then (replicate n 0, afterRep) else - let defLength = littleEndianWord32 (take 4 afterRep) - defData = take (fromIntegral defLength) (drop 4 afterRep) - afterDefData = drop (4 + fromIntegral defLength) afterRep + let defLength = littleEndianWord32 (BS.take 4 afterRep) + defData = BS.take (fromIntegral defLength) (BS.drop 4 afterRep) + afterDefData = BS.drop (4 + fromIntegral defLength) afterRep (defVals, _) = decodeRLEBitPackedHybrid bwDef n defData in (defVals, afterDefData) in (map fromIntegral defLvlsU32, map fromIntegral repLvlsU32, afterDef) readLevelsV2 :: - Int -> Int -> Int -> Int32 -> Int32 -> [Word8] -> ([Int], [Int], [Word8]) + Int -> + Int -> + Int -> + Int32 -> + Int32 -> + BS.ByteString -> + ([Int], [Int], BS.ByteString) readLevelsV2 n maxDef maxRep defLen repLen bs = - let (repBytes, afterRepBytes) = splitAt (fromIntegral repLen) bs - (defBytes, afterDefBytes) = splitAt (fromIntegral defLen) afterRepBytes + let (repBytes, afterRepBytes) = BS.splitAt (fromIntegral repLen) bs + (defBytes, afterDefBytes) = BS.splitAt (fromIntegral defLen) afterRepBytes bwDef = bitWidthForMaxLevel maxDef bwRep = bitWidthForMaxLevel maxRep (repLvlsU32, _) = 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 @@ -4,13 +4,12 @@ module DataFrame.IO.Parquet.Page where import qualified Codec.Compression.GZip as GZip -import Codec.Compression.Zstd.Streaming +import qualified Codec.Compression.Zstd.Streaming as Zstd import Data.Bits -import qualified Data.ByteString as BSO +import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as LB import Data.Int -import Data.Maybe (fromMaybe, listToMaybe) -import Data.Word +import Data.Maybe (fromMaybe) import DataFrame.IO.Parquet.Binary import DataFrame.IO.Parquet.Thrift import DataFrame.IO.Parquet.Types @@ -28,315 +27,345 @@ DictionaryPageHeader{..} -> True _ -> False -readPage :: CompressionCodec -> [Word8] -> IO (Maybe Page, [Word8]) -readPage c [] = pure (Nothing, []) -readPage c columnBytes = do - let (hdr, rem) = readPageHeader emptyPageHeader columnBytes 0 - let compressed = take (fromIntegral $ compressedPageSize hdr) rem +readPage :: CompressionCodec -> BS.ByteString -> IO (Maybe Page, BS.ByteString) +readPage c columnBytes = + if BS.null columnBytes + then pure (Nothing, BS.empty) + else do + let (hdr, rem) = readPageHeader emptyPageHeader columnBytes 0 - fullData <- case c of - ZSTD -> do - Consume dFunc <- decompress - Consume dFunc' <- dFunc (BSO.pack compressed) - Done res <- dFunc' BSO.empty - pure res - SNAPPY -> case Snappy.decompress (BSO.pack compressed) of - Left e -> error (show e) - Right res -> pure res - UNCOMPRESSED -> pure (BSO.pack compressed) - GZIP -> pure (LB.toStrict (GZip.decompress (LB.pack compressed))) - other -> error ("Unsupported compression type: " ++ show other) - pure - ( Just $ Page hdr (BSO.unpack fullData) - , drop (fromIntegral $ compressedPageSize hdr) rem - ) + let compressed = BS.take (fromIntegral $ compressedPageSize hdr) rem -readPageHeader :: PageHeader -> [Word8] -> Int16 -> (PageHeader, [Word8]) -readPageHeader hdr [] _ = (hdr, []) + fullData <- case c of + ZSTD -> do + result <- Zstd.decompress + drainZstd result compressed [] + where + drainZstd (Zstd.Consume f) input acc = do + result <- f input + drainZstd result BS.empty acc + drainZstd (Zstd.Produce chunk next) _ acc = do + result <- next + drainZstd result BS.empty (acc <> [chunk]) + drainZstd (Zstd.Done final) _ acc = + pure $ BS.concat (acc <> [final]) + drainZstd (Zstd.Error msg msg2) _ _ = + error ("ZSTD error: " ++ msg ++ " " ++ msg2) + SNAPPY -> case Snappy.decompress compressed of + Left e -> error (show e) + Right res -> pure res + UNCOMPRESSED -> pure compressed + GZIP -> pure (LB.toStrict (GZip.decompress (BS.fromStrict compressed))) + other -> error ("Unsupported compression type: " ++ show other) + pure + ( Just $ Page hdr fullData + , BS.drop (fromIntegral $ compressedPageSize hdr) rem + ) + +readPageHeader :: + PageHeader -> BS.ByteString -> Int16 -> (PageHeader, BS.ByteString) readPageHeader hdr xs lastFieldId = - let - fieldContents = readField' xs lastFieldId - in - case fieldContents of - Nothing -> (hdr, drop 1 xs) - Just (rem, elemType, identifier) -> case identifier of - 1 -> - let - (pType, rem') = readInt32FromBytes rem - in - readPageHeader (hdr{pageHeaderPageType = pageTypeFromInt pType}) rem' identifier - 2 -> - let - (uncompressedPageSize, rem') = readInt32FromBytes rem - in - readPageHeader - (hdr{uncompressedPageSize = uncompressedPageSize}) - rem' - identifier - 3 -> - let - (compressedPageSize, rem') = readInt32FromBytes rem - in - readPageHeader (hdr{compressedPageSize = compressedPageSize}) rem' identifier - 4 -> - let - (crc, rem') = readInt32FromBytes rem - in - readPageHeader (hdr{pageHeaderCrcChecksum = crc}) rem' identifier - 5 -> - let - (dataPageHeader, rem') = readPageTypeHeader emptyDataPageHeader rem 0 - in - readPageHeader (hdr{pageTypeHeader = dataPageHeader}) rem' identifier - 6 -> error "Index page header not supported" - 7 -> - let - (dictionaryPageHeader, rem') = readPageTypeHeader emptyDictionaryPageHeader rem 0 - in - readPageHeader (hdr{pageTypeHeader = dictionaryPageHeader}) rem' identifier - 8 -> - let - (dataPageHeaderV2, rem') = readPageTypeHeader emptyDataPageHeaderV2 rem 0 - in - readPageHeader (hdr{pageTypeHeader = dataPageHeaderV2}) rem' identifier - n -> error $ "Unknown page header field " ++ show n + if BS.null xs + then (hdr, BS.empty) + else + let + fieldContents = readField' xs lastFieldId + in + case fieldContents of + Nothing -> (hdr, BS.drop 1 xs) + Just (rem, elemType, identifier) -> case identifier of + 1 -> + let + (pType, rem') = readInt32FromBytes rem + in + readPageHeader (hdr{pageHeaderPageType = pageTypeFromInt pType}) rem' identifier + 2 -> + let + (uncompressedPageSize, rem') = readInt32FromBytes rem + in + readPageHeader + (hdr{uncompressedPageSize = uncompressedPageSize}) + rem' + identifier + 3 -> + let + (compressedPageSize, rem') = readInt32FromBytes rem + in + readPageHeader (hdr{compressedPageSize = compressedPageSize}) rem' identifier + 4 -> + let + (crc, rem') = readInt32FromBytes rem + in + readPageHeader (hdr{pageHeaderCrcChecksum = crc}) rem' identifier + 5 -> + let + (dataPageHeader, rem') = readPageTypeHeader emptyDataPageHeader rem 0 + in + readPageHeader (hdr{pageTypeHeader = dataPageHeader}) rem' identifier + 6 -> error "Index page header not supported" + 7 -> + let + (dictionaryPageHeader, rem') = readPageTypeHeader emptyDictionaryPageHeader rem 0 + in + readPageHeader (hdr{pageTypeHeader = dictionaryPageHeader}) rem' identifier + 8 -> + let + (dataPageHeaderV2, rem') = readPageTypeHeader emptyDataPageHeaderV2 rem 0 + in + readPageHeader (hdr{pageTypeHeader = dataPageHeaderV2}) rem' identifier + n -> error $ "Unknown page header field " ++ show n readPageTypeHeader :: - PageTypeHeader -> [Word8] -> Int16 -> (PageTypeHeader, [Word8]) -readPageTypeHeader hdr [] _ = (hdr, []) + PageTypeHeader -> BS.ByteString -> Int16 -> (PageTypeHeader, BS.ByteString) readPageTypeHeader INDEX_PAGE_HEADER _ _ = error "readPageTypeHeader: unsupported INDEX_PAGE_HEADER" readPageTypeHeader PAGE_TYPE_HEADER_UNKNOWN _ _ = error "readPageTypeHeader: unsupported PAGE_TYPE_HEADER_UNKNOWN" readPageTypeHeader hdr@(DictionaryPageHeader{..}) xs lastFieldId = - let - fieldContents = readField' xs lastFieldId - in - case fieldContents of - Nothing -> (hdr, drop 1 xs) - Just (rem, elemType, identifier) -> case identifier of - 1 -> - let - (numValues, rem') = readInt32FromBytes rem - in - readPageTypeHeader - (hdr{dictionaryPageHeaderNumValues = numValues}) - rem' - identifier - 2 -> - let - (enc, rem') = readInt32FromBytes rem - in - readPageTypeHeader - (hdr{dictionaryPageHeaderEncoding = parquetEncodingFromInt enc}) - rem' - identifier - 3 -> - let - isSorted = fromMaybe (error "readPageTypeHeader: not enough bytes") (listToMaybe rem) - in - readPageTypeHeader - (hdr{dictionaryPageIsSorted = isSorted == compactBooleanTrue}) - -- TODO(mchavinda): The bool logic here is a little tricky. - -- If the field is a bool then you can get the value - -- from the byte (and you don't have to drop a field). - -- 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 - identifier - n -> - error $ "readPageTypeHeader: unsupported identifier " ++ show n + if BS.null xs + then (hdr, BS.empty) + else + let + fieldContents = readField' xs lastFieldId + in + case fieldContents of + Nothing -> (hdr, BS.drop 1 xs) + Just (rem, elemType, identifier) -> case identifier of + 1 -> + let + (numValues, rem') = readInt32FromBytes rem + in + readPageTypeHeader + (hdr{dictionaryPageHeaderNumValues = numValues}) + rem' + identifier + 2 -> + let + (enc, rem') = readInt32FromBytes rem + in + readPageTypeHeader + (hdr{dictionaryPageHeaderEncoding = parquetEncodingFromInt enc}) + rem' + identifier + 3 -> + let + isSorted = fromMaybe (error "readPageTypeHeader: not enough bytes") (rem BS.!? 0) + in + readPageTypeHeader + (hdr{dictionaryPageIsSorted = isSorted == compactBooleanTrue}) + -- TODO(mchavinda): The bool logic here is a little tricky. + -- If the field is a bool then you can get the value + -- from the byte (and you don't have to drop a field). + -- 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 + identifier + n -> + error $ "readPageTypeHeader: unsupported identifier " ++ show n readPageTypeHeader hdr@(DataPageHeader{..}) xs lastFieldId = - let - fieldContents = readField' xs lastFieldId - in - case fieldContents of - Nothing -> (hdr, drop 1 xs) - Just (rem, elemType, identifier) -> case identifier of - 1 -> - let - (numValues, rem') = readInt32FromBytes rem - in - readPageTypeHeader (hdr{dataPageHeaderNumValues = numValues}) rem' identifier - 2 -> - let - (enc, rem') = readInt32FromBytes rem - in - readPageTypeHeader - (hdr{dataPageHeaderEncoding = parquetEncodingFromInt enc}) - rem' - identifier - 3 -> - let - (enc, rem') = readInt32FromBytes rem - in - readPageTypeHeader - (hdr{definitionLevelEncoding = parquetEncodingFromInt enc}) - rem' - identifier - 4 -> - let - (enc, rem') = readInt32FromBytes rem - in - readPageTypeHeader - (hdr{repetitionLevelEncoding = parquetEncodingFromInt enc}) - rem' - identifier - 5 -> - let - (stats, rem') = readStatisticsFromBytes emptyColumnStatistics rem 0 - in - readPageTypeHeader (hdr{dataPageHeaderStatistics = stats}) rem' identifier - n -> error $ show n + if BS.null xs + then (hdr, BS.empty) + else + let + fieldContents = readField' xs lastFieldId + in + case fieldContents of + Nothing -> (hdr, BS.drop 1 xs) + Just (rem, elemType, identifier) -> case identifier of + 1 -> + let + (numValues, rem') = readInt32FromBytes rem + in + readPageTypeHeader (hdr{dataPageHeaderNumValues = numValues}) rem' identifier + 2 -> + let + (enc, rem') = readInt32FromBytes rem + in + readPageTypeHeader + (hdr{dataPageHeaderEncoding = parquetEncodingFromInt enc}) + rem' + identifier + 3 -> + let + (enc, rem') = readInt32FromBytes rem + in + readPageTypeHeader + (hdr{definitionLevelEncoding = parquetEncodingFromInt enc}) + rem' + identifier + 4 -> + let + (enc, rem') = readInt32FromBytes rem + in + readPageTypeHeader + (hdr{repetitionLevelEncoding = parquetEncodingFromInt enc}) + rem' + identifier + 5 -> + let + (stats, rem') = readStatisticsFromBytes emptyColumnStatistics rem 0 + in + readPageTypeHeader (hdr{dataPageHeaderStatistics = stats}) rem' identifier + n -> error $ show n readPageTypeHeader hdr@(DataPageHeaderV2{..}) xs lastFieldId = - let - fieldContents = readField' xs lastFieldId - in - case fieldContents of - Nothing -> (hdr, drop 1 xs) - Just (rem, elemType, identifier) -> case identifier of - 1 -> - let - (numValues, rem') = readInt32FromBytes rem - in - readPageTypeHeader (hdr{dataPageHeaderV2NumValues = numValues}) rem' identifier - 2 -> - let - (numNulls, rem') = readInt32FromBytes rem - in - readPageTypeHeader (hdr{dataPageHeaderV2NumNulls = numNulls}) rem' identifier - 3 -> - let - (numRows, rem') = readInt32FromBytes rem - in - readPageTypeHeader (hdr{dataPageHeaderV2NumRows = numRows}) rem' identifier - 4 -> - let - (enc, rem') = readInt32FromBytes rem - in - readPageTypeHeader - (hdr{dataPageHeaderV2Encoding = parquetEncodingFromInt enc}) - rem' - identifier - 5 -> - let - (n, rem') = readInt32FromBytes rem - in - readPageTypeHeader (hdr{definitionLevelByteLength = n}) rem' identifier - 6 -> - let - (n, rem') = readInt32FromBytes rem - in - readPageTypeHeader (hdr{repetitionLevelByteLength = n}) rem' identifier - 7 -> - let - (isCompressed, rem') = case rem of - b : bytes -> ((b .&. 0x0f) == compactBooleanTrue, bytes) - [] -> (True, []) - in - readPageTypeHeader - (hdr{dataPageHeaderV2IsCompressed = isCompressed}) - rem' - identifier - 8 -> - let - (stats, rem') = readStatisticsFromBytes emptyColumnStatistics rem 0 - in - readPageTypeHeader - (hdr{dataPageHeaderV2Statistics = stats}) - rem' - identifier - n -> error $ show n + if BS.null xs + then (hdr, BS.empty) + else + let + fieldContents = readField' xs lastFieldId + in + case fieldContents of + Nothing -> (hdr, BS.drop 1 xs) + Just (rem, elemType, identifier) -> case identifier of + 1 -> + let + (numValues, rem') = readInt32FromBytes rem + in + readPageTypeHeader (hdr{dataPageHeaderV2NumValues = numValues}) rem' identifier + 2 -> + let + (numNulls, rem') = readInt32FromBytes rem + in + readPageTypeHeader (hdr{dataPageHeaderV2NumNulls = numNulls}) rem' identifier + 3 -> + let + (numRows, rem') = readInt32FromBytes rem + in + readPageTypeHeader (hdr{dataPageHeaderV2NumRows = numRows}) rem' identifier + 4 -> + let + (enc, rem') = readInt32FromBytes rem + in + readPageTypeHeader + (hdr{dataPageHeaderV2Encoding = parquetEncodingFromInt enc}) + rem' + identifier + 5 -> + let + (n, rem') = readInt32FromBytes rem + in + readPageTypeHeader (hdr{definitionLevelByteLength = n}) rem' identifier + 6 -> + let + (n, rem') = readInt32FromBytes rem + in + readPageTypeHeader (hdr{repetitionLevelByteLength = n}) rem' identifier + 7 -> + let + (isCompressed, rem') = case BS.uncons rem of + Just (b, bytes) -> ((b .&. 0x0f) == compactBooleanTrue, bytes) + Nothing -> (True, BS.empty) + in + readPageTypeHeader + (hdr{dataPageHeaderV2IsCompressed = isCompressed}) + rem' + identifier + 8 -> + let + (stats, rem') = readStatisticsFromBytes emptyColumnStatistics rem 0 + in + readPageTypeHeader + (hdr{dataPageHeaderV2Statistics = stats}) + rem' + identifier + n -> error $ show n -readField' :: [Word8] -> Int16 -> Maybe ([Word8], TType, Int16) -readField' [] _ = Nothing -readField' (x : xs) lastFieldId - | x .&. 0x0f == 0 = Nothing - | otherwise = - let modifier = fromIntegral ((x .&. 0xf0) `shiftR` 4) :: Int16 - (identifier, rem) = - if modifier == 0 - then readIntFromBytes @Int16 xs - else (lastFieldId + modifier, xs) - elemType = toTType (x .&. 0x0f) - in Just (rem, elemType, identifier) +readField' :: BS.ByteString -> Int16 -> Maybe (BS.ByteString, TType, Int16) +readField' bs lastFieldId = case BS.uncons bs of + Nothing -> Nothing + Just (x, xs) -> + if x .&. 0x0f == 0 + then Nothing + else + let modifier = fromIntegral ((x .&. 0xf0) `shiftR` 4) :: Int16 + (identifier, rem) = + if modifier == 0 + then readIntFromBytes @Int16 xs + else (lastFieldId + modifier, xs) + elemType = toTType (x .&. 0x0f) + in Just (rem, elemType, identifier) -readAllPages :: CompressionCodec -> [Word8] -> IO [Page] +readAllPages :: CompressionCodec -> BS.ByteString -> IO [Page] readAllPages codec bytes = go bytes [] where - go [] acc = return (reverse acc) - go bs acc = do - (maybePage, remaining) <- readPage codec bs - case maybePage of - Nothing -> return (reverse acc) - Just page -> go remaining (page : acc) + go bs acc = + if BS.null bs + then return (reverse acc) + else do + (maybePage, remaining) <- readPage codec bs + case maybePage of + Nothing -> return (reverse acc) + Just page -> go remaining (page : acc) -readNInt32 :: Int -> [Word8] -> ([Int32], [Word8]) +readNInt32 :: Int -> BS.ByteString -> ([Int32], BS.ByteString) readNInt32 0 bs = ([], bs) readNInt32 k bs = - let x = littleEndianInt32 (take 4 bs) - bs' = drop 4 bs + let x = littleEndianInt32 (BS.take 4 bs) + bs' = BS.drop 4 bs (xs, rest) = readNInt32 (k - 1) bs' in (x : xs, rest) -readNDouble :: Int -> [Word8] -> ([Double], [Word8]) +readNDouble :: Int -> BS.ByteString -> ([Double], BS.ByteString) readNDouble 0 bs = ([], bs) readNDouble k bs = - let x = castWord64ToDouble (littleEndianWord64 (take 8 bs)) - bs' = drop 8 bs + let x = castWord64ToDouble (littleEndianWord64 (BS.take 8 bs)) + bs' = BS.drop 8 bs (xs, rest) = readNDouble (k - 1) bs' in (x : xs, rest) -readNByteArrays :: Int -> [Word8] -> ([[Word8]], [Word8]) +readNByteArrays :: Int -> BS.ByteString -> ([BS.ByteString], BS.ByteString) readNByteArrays 0 bs = ([], bs) readNByteArrays k bs = - let len = fromIntegral (littleEndianInt32 (take 4 bs)) :: Int - body = take len (drop 4 bs) - bs' = drop (4 + len) bs + let len = fromIntegral (littleEndianInt32 (BS.take 4 bs)) :: Int + body = BS.take len (BS.drop 4 bs) + bs' = BS.drop (4 + len) bs (xs, rest) = readNByteArrays (k - 1) bs' in (body : xs, rest) -readNBool :: Int -> [Word8] -> ([Bool], [Word8]) +readNBool :: Int -> BS.ByteString -> ([Bool], BS.ByteString) readNBool 0 bs = ([], bs) readNBool count bs = let totalBytes = (count + 7) `div` 8 - chunk = take totalBytes bs - rest = drop totalBytes bs - bits = concatMap (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7]) chunk + chunk = BS.take totalBytes bs + rest = BS.drop totalBytes bs + bits = + concatMap + (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7]) + (BS.unpack chunk) bools = take count bits in (bools, rest) -readNInt64 :: Int -> [Word8] -> ([Int64], [Word8]) +readNInt64 :: Int -> BS.ByteString -> ([Int64], BS.ByteString) readNInt64 0 bs = ([], bs) readNInt64 k bs = - let x = fromIntegral (littleEndianWord64 (take 8 bs)) - bs' = drop 8 bs + let x = fromIntegral (littleEndianWord64 (BS.take 8 bs)) + bs' = BS.drop 8 bs (xs, rest) = readNInt64 (k - 1) bs' in (x : xs, rest) -readNFloat :: Int -> [Word8] -> ([Float], [Word8]) +readNFloat :: Int -> BS.ByteString -> ([Float], BS.ByteString) readNFloat 0 bs = ([], bs) readNFloat k bs = - let x = castWord32ToFloat (littleEndianWord32 (take 4 bs)) - bs' = drop 4 bs + let x = castWord32ToFloat (littleEndianWord32 (BS.take 4 bs)) + bs' = BS.drop 4 bs (xs, rest) = readNFloat (k - 1) bs' in (x : xs, rest) -splitFixed :: Int -> Int -> [Word8] -> ([[Word8]], [Word8]) +splitFixed :: Int -> Int -> BS.ByteString -> ([BS.ByteString], BS.ByteString) splitFixed 0 _ bs = ([], bs) splitFixed k len bs = - let body = take len bs - bs' = drop len bs + let body = BS.take len bs + bs' = BS.drop len bs (xs, rest) = splitFixed (k - 1) len bs' in (body : xs, rest) readStatisticsFromBytes :: - ColumnStatistics -> [Word8] -> Int16 -> (ColumnStatistics, [Word8]) + ColumnStatistics -> BS.ByteString -> Int16 -> (ColumnStatistics, BS.ByteString) readStatisticsFromBytes cs xs lastFieldId = let fieldContents = readField' xs lastFieldId in case fieldContents of - Nothing -> (cs, drop 1 xs) + Nothing -> (cs, BS.drop 1 xs) Just (rem, elemType, identifier) -> case identifier of 1 -> let @@ -369,19 +398,19 @@ in readStatisticsFromBytes (cs{columnMinValue = minInBytes}) rem' identifier 7 -> - case rem of - [] -> + case BS.uncons rem of + Nothing -> error "readStatisticsFromBytes: not enough bytes" - (isMaxValueExact : rem') -> + Just (isMaxValueExact, rem') -> readStatisticsFromBytes (cs{isColumnMaxValueExact = isMaxValueExact == compactBooleanTrue}) rem' identifier 8 -> - case rem of - [] -> + case BS.uncons rem of + Nothing -> error "readStatisticsFromBytes: not enough bytes" - (isMinValueExact : rem') -> + Just (isMinValueExact, rem') -> readStatisticsFromBytes (cs{isColumnMinValueExact = isMinValueExact == compactBooleanTrue}) rem' 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 @@ -93,7 +93,7 @@ , createdBy :: Maybe String , columnOrders :: [ColumnOrder] , encryptionAlgorithm :: EncryptionAlgorithm - , footerSigningKeyMetadata :: [Word8] + , footerSigningKeyMetadata :: BS.ByteString } deriving (Show, Eq) @@ -149,7 +149,7 @@ , createdBy = Nothing , columnOrders = [] , encryptionAlgorithm = ENCRYPTION_ALGORITHM_UNKNOWN - , footerSigningKeyMetadata = [] + , footerSigningKeyMetadata = BS.empty } data ColumnMetaData = ColumnMetaData @@ -182,7 +182,7 @@ , columnChunkColumnIndexOffset :: Int64 , columnChunkColumnIndexLength :: Int32 , cryptoMetadata :: ColumnCryptoMetadata - , encryptedColumnMetadata :: [Word8] + , encryptedColumnMetadata :: BS.ByteString } deriving (Show, Eq) @@ -234,7 +234,16 @@ emptyColumnChunk :: ColumnChunk emptyColumnChunk = - ColumnChunk "" 0 emptyColumnMetadata 0 0 0 0 COLUMN_CRYPTO_METADATA_UNKNOWN [] + ColumnChunk + "" + 0 + emptyColumnMetadata + 0 + 0 + 0 + 0 + COLUMN_CRYPTO_METADATA_UNKNOWN + BS.empty emptyKeyValue :: KeyValue emptyKeyValue = KeyValue{key = "", value = ""} @@ -754,13 +763,23 @@ Just (elemType, identifier) -> case identifier of 1 -> do readAesGcmV1 - (AesGcmV1{aadPrefix = [], aadFileUnique = [], supplyAadPrefix = False}) + ( AesGcmV1 + { aadPrefix = BS.empty + , aadFileUnique = BS.empty + , supplyAadPrefix = False + } + ) buf pos 0 2 -> do readAesGcmCtrV1 - (AesGcmCtrV1{aadPrefix = [], aadFileUnique = [], supplyAadPrefix = False}) + ( AesGcmCtrV1 + { aadPrefix = BS.empty + , aadFileUnique = BS.empty + , supplyAadPrefix = False + } + ) buf pos 0 @@ -792,21 +811,21 @@ Just (elemType, identifier) -> case identifier of 1 -> do aadPrefix <- readByteString buf pos - readAesGcmCtrV1 (v{aadPrefix = aadPrefix}) buf pos lastFieldId + readAesGcmCtrV1 (v{aadPrefix = aadPrefix}) buf pos identifier 2 -> do aadFileUnique <- readByteString buf pos readAesGcmCtrV1 (v{aadFileUnique = aadFileUnique}) buf pos - lastFieldId + identifier 3 -> do supplyAadPrefix <- readAndAdvance pos buf readAesGcmCtrV1 (v{supplyAadPrefix = supplyAadPrefix == compactBooleanTrue}) buf pos - lastFieldId + identifier _ -> return ENCRYPTION_ALGORITHM_UNKNOWN readAesGcmCtrV1 _ _ _ _ = error "readAesGcmCtrV1 called with non AesGcmCtrV1" @@ -824,17 +843,17 @@ Just (elemType, identifier) -> case identifier of 1 -> do aadPrefix <- readByteString buf pos - readAesGcmV1 (v{aadPrefix = aadPrefix}) buf pos lastFieldId + readAesGcmV1 (v{aadPrefix = aadPrefix}) buf pos identifier 2 -> do aadFileUnique <- readByteString buf pos - readAesGcmV1 (v{aadFileUnique = aadFileUnique}) buf pos lastFieldId + readAesGcmV1 (v{aadFileUnique = aadFileUnique}) buf pos identifier 3 -> do supplyAadPrefix <- readAndAdvance pos buf readAesGcmV1 (v{supplyAadPrefix = supplyAadPrefix == compactBooleanTrue}) buf pos - lastFieldId + identifier _ -> return ENCRYPTION_ALGORITHM_UNKNOWN readAesGcmV1 _ _ _ _ = error "readAesGcmV1 called with non AesGcmV1" @@ -1101,10 +1120,10 @@ Just (elemType, identifier) -> case identifier of 1 -> do scale' <- readInt32FromBuffer buf pos - readDecimalType precision scale' buf pos lastFieldId + readDecimalType precision scale' buf pos identifier 2 -> do precision' <- readInt32FromBuffer buf pos - readDecimalType precision' scale buf pos lastFieldId + readDecimalType precision' scale buf pos identifier _ -> error $ "UNKNOWN field ID for DecimalType" ++ show identifier readTimeType :: @@ -1117,15 +1136,14 @@ readTimeType isAdjustedToUTC unit buf pos lastFieldId = do fieldContents <- readField buf pos lastFieldId case fieldContents of - Nothing -> return (TimeType isAdjustedToUTC unit) + Nothing -> return (TimeType{isAdjustedToUTC = isAdjustedToUTC, unit = unit}) Just (elemType, identifier) -> case identifier of 1 -> do - -- TODO: Check for empty - isAdjustedToUTC' <- (== compactBooleanTrue) <$> readAndAdvance pos buf - readTimeType isAdjustedToUTC' unit buf pos lastFieldId + let isAdjustedToUTC' = elemType == toTType compactBooleanTrue + readTimeType isAdjustedToUTC' unit buf pos identifier 2 -> do unit' <- readUnit TIME_UNIT_UNKNOWN buf pos 0 - readTimeType isAdjustedToUTC unit' buf pos lastFieldId + readTimeType isAdjustedToUTC unit' buf pos identifier _ -> error $ "UNKNOWN field ID for TimeType" ++ show identifier readTimestampType :: @@ -1138,17 +1156,15 @@ readTimestampType isAdjustedToUTC unit buf pos lastFieldId = do fieldContents <- readField buf pos lastFieldId case fieldContents of - Nothing -> return (TimestampType isAdjustedToUTC unit) + Nothing -> return (TimestampType{isAdjustedToUTC = isAdjustedToUTC, unit = unit}) Just (elemType, identifier) -> case identifier of 1 -> do - -- TODO: Check for empty - isAdjustedToUTC' <- (== compactBooleanTrue) <$> readNoAdvance pos buf - readTimestampType False unit buf pos lastFieldId + let isAdjustedToUTC' = elemType == toTType compactBooleanTrue + readTimestampType isAdjustedToUTC' unit buf pos identifier 2 -> do - _ <- readField buf pos identifier unit' <- readUnit TIME_UNIT_UNKNOWN buf pos 0 - readTimestampType isAdjustedToUTC unit' buf pos lastFieldId - _ -> error $ "UNKNOWN field ID for TimestampType" ++ show identifier + readTimestampType isAdjustedToUTC 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 @@ -1157,9 +1173,12 @@ Nothing -> return unit Just (elemType, identifier) -> case identifier of 1 -> do + _ <- readField buf pos 0 readUnit MILLISECONDS buf pos identifier 2 -> do + _ <- readField buf pos 0 readUnit MICROSECONDS buf pos identifier 3 -> do + _ <- readField buf pos 0 readUnit NANOSECONDS buf pos identifier n -> error $ "Unknown time unit: " ++ show n diff --git a/src/DataFrame/IO/Parquet/Time.hs b/src/DataFrame/IO/Parquet/Time.hs --- a/src/DataFrame/IO/Parquet/Time.hs +++ b/src/DataFrame/IO/Parquet/Time.hs @@ -2,16 +2,17 @@ module DataFrame.IO.Parquet.Time where +import qualified Data.ByteString as BS import Data.Time import Data.Word import DataFrame.IO.Parquet.Binary -int96ToUTCTime :: [Word8] -> UTCTime +int96ToUTCTime :: BS.ByteString -> UTCTime int96ToUTCTime bytes - | length bytes /= 12 = error "INT96 must be exactly 12 bytes" + | BS.length bytes /= 12 = error "INT96 must be exactly 12 bytes" | otherwise = - let (nanosBytes, julianBytes) = splitAt 8 bytes + let (nanosBytes, julianBytes) = BS.splitAt 8 bytes nanosSinceMidnight = littleEndianWord64 nanosBytes julianDay = littleEndianWord32 julianBytes in julianDayAndNanosToUTCTime (fromIntegral julianDay) nanosSinceMidnight @@ -38,13 +39,13 @@ -- I include this here even though it's unused because we'll likely use -- it for the writer. Since int96 is deprecated this is only included for completeness anyway. -utcTimeToInt96 :: UTCTime -> [Word8] +utcTimeToInt96 :: UTCTime -> BS.ByteString utcTimeToInt96 (UTCTime day diffTime) = let julianDay = dayToJulianDay day nanosSinceMidnight = floor (realToFrac diffTime * 1_000_000_000) nanosBytes = word64ToLittleEndian nanosSinceMidnight julianBytes = word32ToLittleEndian (fromIntegral julianDay) - in nanosBytes ++ julianBytes + in nanosBytes `BS.append` julianBytes dayToJulianDay :: Day -> Integer dayToJulianDay day = diff --git a/src/DataFrame/IO/Parquet/Types.hs b/src/DataFrame/IO/Parquet/Types.hs --- a/src/DataFrame/IO/Parquet/Types.hs +++ b/src/DataFrame/IO/Parquet/Types.hs @@ -1,9 +1,9 @@ module DataFrame.IO.Parquet.Types where +import qualified Data.ByteString as BS import Data.Int import qualified Data.Text as T import Data.Time -import Data.Word data ParquetType = PBOOLEAN @@ -125,26 +125,26 @@ emptyGeospatialStatistics = GeospatialStatistics emptyBoundingBox [] data ColumnStatistics = ColumnStatistics - { columnMin :: [Word8] - , columnMax :: [Word8] + { columnMin :: BS.ByteString + , columnMax :: BS.ByteString , columnNullCount :: Int64 , columnDistictCount :: Int64 - , columnMinValue :: [Word8] - , columnMaxValue :: [Word8] + , columnMinValue :: BS.ByteString + , columnMaxValue :: BS.ByteString , isColumnMaxValueExact :: Bool , isColumnMinValueExact :: Bool } deriving (Show, Eq) emptyColumnStatistics :: ColumnStatistics -emptyColumnStatistics = ColumnStatistics [] [] 0 0 [] [] False False +emptyColumnStatistics = ColumnStatistics BS.empty BS.empty 0 0 BS.empty BS.empty False False data ColumnCryptoMetadata = COLUMN_CRYPTO_METADATA_UNKNOWN | ENCRYPTION_WITH_FOOTER_KEY | EncryptionWithColumnKey { columnCryptPathInSchema :: [String] - , columnKeyMetadata :: [Word8] + , columnKeyMetadata :: BS.ByteString } deriving (Show, Eq) @@ -166,13 +166,13 @@ data EncryptionAlgorithm = ENCRYPTION_ALGORITHM_UNKNOWN | AesGcmV1 - { aadPrefix :: [Word8] - , aadFileUnique :: [Word8] + { aadPrefix :: BS.ByteString + , aadFileUnique :: BS.ByteString , supplyAadPrefix :: Bool } | AesGcmCtrV1 - { aadPrefix :: [Word8] - , aadFileUnique :: [Word8] + { aadPrefix :: BS.ByteString + , aadFileUnique :: BS.ByteString , supplyAadPrefix :: Bool } deriving (Show, Eq) @@ -189,7 +189,7 @@ data Page = Page { pageHeader :: PageHeader - , pageBytes :: [Word8] + , pageBytes :: BS.ByteString } deriving (Show, Eq) 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 @@ -30,6 +30,7 @@ import Control.Monad.ST (runST) import Data.Kind (Type) import Data.Maybe +import Data.These import Data.Type.Equality (TestEquality (..)) import DataFrame.Errors import DataFrame.Internal.Parsing @@ -57,15 +58,21 @@ TColumn :: (Columnable a) => Column -> TypedColumn a instance (Eq a) => Eq (TypedColumn a) where + (==) :: (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 +-- | Gets the underlying vector from a TypedColumn. +vectorFromTypedColumn :: TypedColumn a -> VB.Vector a +vectorFromTypedColumn (TColumn value) = either throw id (toVector value) + -- | Checks if a column contains missing values. hasMissing :: Column -> Bool hasMissing (OptionalColumn column) = True @@ -351,12 +358,8 @@ {-# INLINE atIndicesStable #-} atIndicesWithNulls :: VB.Vector (Maybe Int) -> Column -> Column -atIndicesWithNulls indices column - | VB.all isJust indices = - atIndicesStable - (VU.fromList (catMaybes (VB.toList indices))) - column - | otherwise = case column of +atIndicesWithNulls indices column = + case column of BoxedColumn col -> OptionalColumn $ VB.map (fmap (col VB.!)) indices UnboxedColumn col -> @@ -614,6 +617,87 @@ zipColumns (OptionalColumn optcolumn) (UnboxedColumn column) = BoxedColumn (VG.zip optcolumn (VB.convert column)) zipColumns (OptionalColumn optcolumn) (OptionalColumn optother) = BoxedColumn (VG.zip optcolumn optother) {-# INLINE zipColumns #-} + +-- | Merge two columns using `These`. +mergeColumns :: Column -> Column -> Column +mergeColumns (BoxedColumn column) (BoxedColumn other) = BoxedColumn (VG.zipWith These column other) +mergeColumns (BoxedColumn column) (UnboxedColumn other) = + BoxedColumn + ( VB.generate + (min (VG.length column) (VG.length other)) + (\i -> These (column VG.! i) (other VG.! i)) + ) +mergeColumns (BoxedColumn column) (OptionalColumn other) = + BoxedColumn + ( VB.generate + (min (VG.length column) (VG.length other)) + ( \i -> + if isNothing (other VG.! i) + then This (column VG.! i) + else These (column VG.! i) (fromJust $ other VG.! i) + ) + ) +mergeColumns (UnboxedColumn column) (BoxedColumn other) = + BoxedColumn + ( VB.generate + (min (VG.length column) (VG.length other)) + (\i -> These (column VG.! i) (other VG.! i)) + ) +mergeColumns (UnboxedColumn column) (UnboxedColumn other) = + BoxedColumn + ( VB.generate + (min (VG.length column) (VG.length other)) + (\i -> These (column VG.! i) (other VG.! i)) + ) +mergeColumns (UnboxedColumn column) (OptionalColumn other) = + BoxedColumn + ( VB.generate + (min (VG.length column) (VG.length other)) + ( \i -> + if isNothing (other VG.! i) + then This (column VG.! i) + else These (column VG.! i) (fromJust $ other VG.! i) + ) + ) +mergeColumns (OptionalColumn column) (BoxedColumn other) = + BoxedColumn + ( VB.generate + (min (VG.length column) (VG.length other)) + ( \i -> + if isNothing (column VG.! i) + then That (other VG.! i) + else These (fromJust $ column VG.! i) (other VG.! i) + ) + ) +mergeColumns (OptionalColumn column) (UnboxedColumn other) = + BoxedColumn + ( VB.generate + (min (VG.length column) (VG.length other)) + ( \i -> + if isNothing (column VG.! i) + then That (other VG.! i) + else These (fromJust $ column VG.! i) (other VG.! i) + ) + ) +mergeColumns (OptionalColumn column) (OptionalColumn other) = + OptionalColumn + ( VB.generate + (min (VG.length column) (VG.length other)) + ( \i -> + if isNothing (column VG.! i) && isNothing (other VG.! i) + then Nothing + else + ( if isNothing (column VG.! i) + then Just (That (fromJust $ other VG.! i)) + else + ( if isNothing (other VG.! i) + then Just (This (fromJust $ column VG.! i)) + else Just (These (fromJust $ column VG.! i) (fromJust $ other VG.! i)) + ) + ) + ) + ) +{-# INLINE mergeColumns #-} -- | An internal, column version of zipWith. zipWithColumns :: 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 @@ -17,75 +17,100 @@ import qualified Data.Text as T import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) import qualified Data.Vector.Generic as VG -import qualified Data.Vector.Unboxed as VU import DataFrame.Internal.Column import Type.Reflection (Typeable, typeOf, typeRep) +data UnaryOp a b = MkUnaryOp + { unaryFn :: a -> b + , unaryName :: T.Text + , unarySymbol :: Maybe T.Text + } + +data BinaryOp a b c = MkBinaryOp + { binaryFn :: a -> b -> c + , binaryName :: T.Text + , binarySymbol :: Maybe T.Text + , binaryCommutative :: Bool + , binaryPrecedence :: Int + } + +data AggStrategy a b where + CollectAgg :: + (VG.Vector v b, Typeable v) => T.Text -> (v b -> a) -> AggStrategy a b + FoldAgg :: T.Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b + data Expr a where Col :: (Columnable a) => T.Text -> Expr a Lit :: (Columnable a) => a -> Expr a - UnaryOp :: - (Columnable a, Columnable b) => T.Text -> (b -> a) -> Expr b -> Expr a - BinaryOp :: + Unary :: + (Columnable a, Columnable b) => UnaryOp b a -> Expr b -> Expr a + Binary :: (Columnable c, Columnable b, Columnable a) => - T.Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a + BinaryOp c b a -> Expr c -> Expr b -> Expr a If :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a - AggVector :: - (VG.Vector v b, Typeable v, Columnable a, Columnable b) => - Expr b -> T.Text -> (v b -> a) -> Expr a - AggReduce :: (Columnable a) => Expr a -> T.Text -> (a -> a -> a) -> Expr a - -- TODO(mchav): Numeric reduce might be superfluous since expressions are already type checked. - AggNumericVector :: - (Columnable a, Columnable b, VU.Unbox a, VU.Unbox b, Num a, Num b) => - Expr b -> T.Text -> (VU.Vector b -> a) -> Expr a - AggFold :: - forall a b. - (Columnable a, Columnable b) => Expr b -> T.Text -> a -> (a -> b -> a) -> Expr a + Agg :: (Columnable a, Columnable b) => AggStrategy a b -> Expr b -> Expr a data UExpr where - Wrap :: (Columnable a) => Expr a -> UExpr + UExpr :: (Columnable a) => Expr a -> UExpr instance Show UExpr where show :: UExpr -> String - show (Wrap expr) = show expr + show (UExpr expr) = show expr type NamedExpr = (T.Text, UExpr) instance (Num a, Columnable a) => Num (Expr a) where (+) :: Expr a -> Expr a -> Expr a - (+) (Lit x) (Lit y) = Lit (x + y) - (+) e1 e2 - | e1 == e2 = BinaryOp "mult" (*) e1 (Lit 2) - | otherwise = BinaryOp "add" (+) e1 e2 + (+) = + Binary + ( MkBinaryOp + { binaryFn = (+) + , binaryName = "add" + , binarySymbol = Just "+" + , binaryCommutative = True + , binaryPrecedence = 6 + } + ) (-) :: Expr a -> Expr a -> Expr a - (-) (Lit x) (Lit y) = Lit (x - y) - (-) e1 e2 = BinaryOp "sub" (-) e1 e2 + (-) = + Binary + ( MkBinaryOp + { binaryFn = (-) + , binaryName = "sub" + , binarySymbol = Just "-" + , binaryCommutative = False + , binaryPrecedence = 6 + } + ) (*) :: Expr a -> Expr a -> Expr a - (*) (Lit 0) _ = Lit 0 - (*) _ (Lit 0) = Lit 0 - (*) (Lit 1) e = e - (*) e (Lit 1) = e - (*) (Lit x) (Lit y) = Lit (x * y) - (*) e1 e2 - | e1 == e2 = BinaryOp "pow" (^) e1 (Lit @Int 2) - | otherwise = BinaryOp "mult" (*) e1 e2 + (*) = + Binary + ( MkBinaryOp + { binaryFn = (*) + , binaryName = "mult" + , binarySymbol = Just "*" + , binaryCommutative = True + , binaryPrecedence = 7 + } + ) fromInteger :: Integer -> Expr a fromInteger = Lit . fromInteger negate :: Expr a -> Expr a - negate (Lit n) = Lit (negate n) - negate expr = UnaryOp "negate" negate expr + negate = + Unary + (MkUnaryOp{unaryFn = negate, unaryName = "negate", unarySymbol = Nothing}) abs :: (Num a) => Expr a -> Expr a - abs (Lit n) = Lit (abs n) - abs expr = UnaryOp "abs" abs expr + abs = Unary (MkUnaryOp{unaryFn = abs, unaryName = "abs", unarySymbol = Nothing}) signum :: (Num a) => Expr a -> Expr a - signum (Lit n) = Lit (signum n) - signum expr = UnaryOp "signum" signum expr + signum = + Unary + (MkUnaryOp{unaryFn = signum, unaryName = "signum", unarySymbol = Nothing}) add :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a add = (+) @@ -101,8 +126,16 @@ fromRational = Lit . fromRational (/) :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a - (/) (Lit l1) (Lit l2) = Lit (l1 / l2) - (/) e1 e2 = BinaryOp "divide" (/) e1 e2 + (/) = + Binary + ( MkBinaryOp + { binaryFn = (/) + , binaryName = "divide" + , binarySymbol = Just "/" + , binaryCommutative = True + , binaryPrecedence = 7 + } + ) divide :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a divide = (/) @@ -115,75 +148,96 @@ pi :: (Floating a, Columnable a) => Expr a pi = Lit pi exp :: (Floating a, Columnable a) => Expr a -> Expr a - exp = UnaryOp "exp" exp + exp = Unary (MkUnaryOp{unaryFn = exp, unaryName = "exp", unarySymbol = Nothing}) + sqrt :: (Floating a, Columnable a) => Expr a -> Expr a + sqrt = + Unary (MkUnaryOp{unaryFn = sqrt, unaryName = "sqrt", unarySymbol = Nothing}) + (**) :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a + (**) = + Binary + ( MkBinaryOp + { binaryFn = (**) + , binaryName = "exponentiate" + , binarySymbol = Just "**" + , binaryCommutative = False + , binaryPrecedence = 8 + } + ) log :: (Floating a, Columnable a) => Expr a -> Expr a - log = UnaryOp "log" log + log = Unary (MkUnaryOp{unaryFn = log, unaryName = "log", unarySymbol = Nothing}) + logBase :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a + logBase = + Binary + ( MkBinaryOp + { binaryFn = logBase + , binaryName = "logBase" + , binarySymbol = Nothing + , binaryCommutative = False + , binaryPrecedence = 1 + } + ) sin :: (Floating a, Columnable a) => Expr a -> Expr a - sin = UnaryOp "sin" sin + sin = Unary (MkUnaryOp{unaryFn = sin, unaryName = "sin", unarySymbol = Nothing}) cos :: (Floating a, Columnable a) => Expr a -> Expr a - cos = UnaryOp "cos" cos + cos = Unary (MkUnaryOp{unaryFn = cos, unaryName = "cos", unarySymbol = Nothing}) + tan :: (Floating a, Columnable a) => Expr a -> Expr a + tan = Unary (MkUnaryOp{unaryFn = tan, unaryName = "tan", unarySymbol = Nothing}) asin :: (Floating a, Columnable a) => Expr a -> Expr a - asin = UnaryOp "asin" asin + asin = + Unary (MkUnaryOp{unaryFn = asin, unaryName = "asin", unarySymbol = Nothing}) acos :: (Floating a, Columnable a) => Expr a -> Expr a - acos = UnaryOp "acos" acos + acos = + Unary (MkUnaryOp{unaryFn = acos, unaryName = "acos", unarySymbol = Nothing}) atan :: (Floating a, Columnable a) => Expr a -> Expr a - atan = UnaryOp "atan" atan + atan = + Unary (MkUnaryOp{unaryFn = atan, unaryName = "atan", unarySymbol = Nothing}) sinh :: (Floating a, Columnable a) => Expr a -> Expr a - sinh = UnaryOp "sinh" sinh + sinh = + Unary (MkUnaryOp{unaryFn = sinh, unaryName = "sinh", unarySymbol = Nothing}) cosh :: (Floating a, Columnable a) => Expr a -> Expr a - cosh = UnaryOp "cosh" cosh + cosh = + Unary (MkUnaryOp{unaryFn = cosh, unaryName = "cosh", unarySymbol = Nothing}) asinh :: (Floating a, Columnable a) => Expr a -> Expr a - asinh = UnaryOp "asinh" sinh + asinh = + Unary + (MkUnaryOp{unaryFn = asinh, unaryName = "asinh", unarySymbol = Nothing}) acosh :: (Floating a, Columnable a) => Expr a -> Expr a - acosh = UnaryOp "acosh" acosh + acosh = + Unary + (MkUnaryOp{unaryFn = acosh, unaryName = "acosh", unarySymbol = Nothing}) atanh :: (Floating a, Columnable a) => Expr a -> Expr a - atanh = UnaryOp "atanh" atanh + atanh = + Unary + (MkUnaryOp{unaryFn = atanh, unaryName = "atanh", unarySymbol = Nothing}) instance (Show a) => Show (Expr a) where show :: forall a. (Show a) => Expr a -> String show (Col name) = "(col @" ++ show (typeRep @a) ++ " " ++ show name ++ ")" show (Lit value) = "(lit (" ++ show value ++ "))" show (If cond l r) = "(ifThenElse " ++ show cond ++ " " ++ show l ++ " " ++ show r ++ ")" - show (UnaryOp name f value) = "(" ++ T.unpack name ++ " " ++ show value ++ ")" - show (BinaryOp name f a b) = "(" ++ T.unpack name ++ " " ++ show a ++ " " ++ show b ++ ")" - show (AggNumericVector expr op _) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")" - show (AggVector expr op _) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")" - show (AggReduce expr op _) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")" - show (AggFold expr op _ _) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")" + show (Unary op value) = "(" ++ T.unpack (unaryName op) ++ " " ++ show value ++ ")" + show (Binary op a b) = "(" ++ T.unpack (binaryName op) ++ " " ++ show a ++ " " ++ show b ++ ")" + show (Agg (CollectAgg op _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")" + show (Agg (FoldAgg op _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")" normalize :: (Eq a, Ord a, Show a, Typeable a) => Expr a -> Expr a normalize expr = case expr of Col name -> Col name Lit val -> Lit val If cond th el -> If (normalize cond) (normalize th) (normalize el) - UnaryOp name f e -> UnaryOp name f (normalize e) - BinaryOp name f e1 e2 - | isCommutative name -> + Unary op e -> Unary op (normalize e) + Binary op e1 e2 + | binaryCommutative op -> let n1 = normalize e1 n2 = normalize e2 in case testEquality (typeOf n1) (typeOf n2) of Nothing -> expr Just Refl -> if compareExpr n1 n2 == GT - then BinaryOp name f n2 n1 -- Swap to canonical order - else BinaryOp name f n1 n2 - | otherwise -> BinaryOp name f (normalize e1) (normalize e2) - AggVector e name f -> AggVector (normalize e) name f - AggReduce e name f -> AggReduce (normalize e) name f - AggNumericVector e name f -> AggNumericVector (normalize e) name f - AggFold e name init f -> AggFold (normalize e) name init f - -isCommutative :: T.Text -> Bool -isCommutative name = - name - `elem` [ "add" - , "mult" - , "min" - , "max" - , "eq" - , "and" - , "or" - ] + then Binary op n2 n1 -- Swap to canonical order + else Binary op n1 n2 + | otherwise -> Binary op (normalize e1) (normalize e2) + Agg strat e -> Agg strat (normalize e) -- Compare expressions for ordering (used in normalization) compareExpr :: Expr a -> Expr a -> Ordering @@ -193,12 +247,10 @@ exprKey (Col name) = "0:" ++ T.unpack name exprKey (Lit val) = "1:" ++ show val exprKey (If c t e) = "2:" ++ exprKey c ++ exprKey t ++ exprKey e - exprKey (UnaryOp name _ e) = "3:" ++ T.unpack name ++ exprKey e - exprKey (BinaryOp name _ e1 e2) = "4:" ++ T.unpack name ++ exprKey e1 ++ exprKey e2 - exprKey (AggVector e name _) = "5:" ++ T.unpack name ++ exprKey e - exprKey (AggReduce e name _) = "6:" ++ T.unpack name ++ exprKey e - exprKey (AggNumericVector e name _) = "7:" ++ T.unpack name ++ exprKey e - exprKey (AggFold e name _ _) = "8:" ++ T.unpack name ++ 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 (Agg (CollectAgg name _) e) = "5:" ++ T.unpack name ++ exprKey e + exprKey (Agg (FoldAgg name _ _) e) = "5:" ++ T.unpack name ++ exprKey e instance (Eq a, Columnable a) => Eq (Expr a) where (==) l r = eqNormalized (normalize l) (normalize r) @@ -212,18 +264,12 @@ eqNormalized (Lit v1) (Lit v2) = v1 == v2 eqNormalized (If c1 t1 e1) (If c2 t2 e2) = c1 == c2 && t1 `exprEq` t2 && e1 `exprEq` e2 - eqNormalized (UnaryOp n1 _ e1) (UnaryOp n2 _ e2) = - n1 == n2 && e1 `exprEq` e2 - eqNormalized (BinaryOp n1 _ e1a e1b) (BinaryOp n2 _ e2a e2b) = - n1 == n2 && e1a `exprEq` e2a && e1b `exprEq` e2b - eqNormalized (AggVector e1 n1 _) (AggVector e2 n2 _) = - n1 == n2 && e1 `exprEq` e2 - eqNormalized (AggReduce e1 n1 _) (AggReduce e2 n2 _) = + eqNormalized (Unary op1 e1) (Unary op2 e2) = unaryName op1 == unaryName op2 && e1 `exprEq` e2 + eqNormalized (Binary op1 e1a e1b) (Binary op2 e2a e2b) = binaryName op1 == binaryName op2 && e1a `exprEq` e2a && e1b `exprEq` e2b + eqNormalized (Agg (CollectAgg n1 _) e1) (Agg (CollectAgg n2 _) e2) = n1 == n2 && e1 `exprEq` e2 - eqNormalized (AggNumericVector e1 n1 _) (AggNumericVector e2 n2 _) = + eqNormalized (Agg (FoldAgg n1 _ _) e1) (Agg (FoldAgg n2 _ _) e2) = n1 == n2 && e1 `exprEq` e2 - eqNormalized (AggFold e1 n1 i1 _) (AggFold e2 n2 i2 _) = - n1 == n2 && e1 `exprEq` e2 && i1 == i2 eqNormalized _ _ = False instance (Ord a, Columnable a) => Ord (Expr a) where @@ -233,35 +279,23 @@ (Lit v1, Lit v2) -> compare v1 v2 (If c1 t1 e1', If c2 t2 e2') -> compare c1 c2 <> exprComp t1 t2 <> exprComp e1' e2' - (UnaryOp n1 _ e1', UnaryOp n2 _ e2') -> - compare n1 n2 <> exprComp e1' e2' - (BinaryOp n1 _ a1 b1, BinaryOp n2 _ a2 b2) -> - compare n1 n2 <> exprComp a1 a2 <> exprComp b1 b2 - (AggVector e1' n1 _, AggVector e2' n2 _) -> - compare n1 n2 <> exprComp e1' e2' - (AggReduce e1' n1 _, AggReduce e2' n2 _) -> - compare n1 n2 <> exprComp e1' e2' - (AggNumericVector e1' n1 _, AggNumericVector e2' n2 _) -> - compare n1 n2 <> exprComp e1' e2' - (AggFold e1' n1 i1 _, AggFold e2' n2 i2 _) -> - compare n1 n2 <> exprComp e1' e2' <> compare i1 i2 + (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' -- Different constructors - compare by priority (Col _, _) -> LT (_, Col _) -> GT (Lit _, _) -> LT (_, Lit _) -> GT - (UnaryOp{}, _) -> LT - (_, UnaryOp{}) -> GT - (BinaryOp{}, _) -> LT - (_, BinaryOp{}) -> GT + (Unary{}, _) -> LT + (_, Unary{}) -> GT + (Binary{}, _) -> LT + (_, Binary{}) -> GT (If{}, _) -> LT (_, If{}) -> GT - (AggVector{}, _) -> LT - (_, AggVector{}) -> GT - (AggReduce{}, _) -> LT - (_, AggReduce{}) -> GT - (AggNumericVector{}, _) -> LT - (_, AggNumericVector{}) -> GT + (Agg{}, _) -> LT exprComp :: (Columnable b, Columnable c) => Expr b -> Expr c -> Ordering exprComp e1 e2 = case testEquality (typeOf e1) (typeOf e2) of @@ -283,70 +317,58 @@ (Lit _) -> expr (If cond l r) -> If (replaceExpr new old cond) (replaceExpr new old l) (replaceExpr new old r) - (UnaryOp name f value) -> UnaryOp name f (replaceExpr new old value) - (BinaryOp name f l r) -> BinaryOp name f (replaceExpr new old l) (replaceExpr new old r) - (AggNumericVector expr op f) -> AggNumericVector (replaceExpr new old expr) op f - (AggVector expr op f) -> AggVector (replaceExpr new old expr) op f - (AggReduce expr op f) -> AggReduce (replaceExpr new old expr) op f - (AggFold expr op acc f) -> AggFold (replaceExpr new old expr) op acc f + (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) eSize :: Expr a -> Int eSize (Col _) = 1 eSize (Lit _) = 1 eSize (If c l r) = 1 + eSize c + eSize l + eSize r -eSize (UnaryOp _ _ e) = 1 + eSize e -eSize (BinaryOp _ _ l r) = 1 + eSize l + eSize r -eSize (AggNumericVector expr op _) = eSize expr + 1 -eSize (AggVector expr op _) = eSize expr + 1 -eSize (AggReduce expr op _) = eSize expr + 1 -eSize (AggFold expr op _ _) = eSize expr + 1 +eSize (Unary _ e) = 1 + eSize e +eSize (Binary _ l r) = 1 + eSize l + eSize r +eSize (Agg strategy expr) = eSize expr + 1 getColumns :: Expr a -> [T.Text] getColumns (Col cName) = [cName] getColumns expr@(Lit _) = [] getColumns (If cond l r) = getColumns cond <> getColumns l <> getColumns r -getColumns (UnaryOp name f value) = getColumns value -getColumns (BinaryOp name f l r) = getColumns l <> getColumns r -getColumns (AggNumericVector expr op f) = getColumns expr -getColumns (AggVector expr op f) = getColumns expr -getColumns (AggReduce expr op f) = getColumns expr -getColumns (AggFold expr op acc f) = getColumns expr +getColumns (Unary op value) = getColumns value +getColumns (Binary op l r) = getColumns l <> getColumns r +getColumns (Agg strategy expr) = getColumns expr prettyPrint :: Expr a -> String -prettyPrint = go 0 +prettyPrint = go 0 0 where - go :: Int -> Expr a -> String - go prec expr = case expr of + indent :: Int -> String + indent n = replicate (n * 2) ' ' + + go :: Int -> Int -> Expr a -> String + go depth prec expr = case expr of Col name -> T.unpack name Lit value -> show value If cond t e -> - "if (" ++ go 0 cond ++ ") then (" ++ go 0 t ++ ") else (" ++ go 0 e ++ ")" - UnaryOp name _ arg -> T.unpack name ++ "(" ++ go 0 arg ++ ")" - BinaryOp name _ l r -> - let p = opPrec name - inner = go p l ++ " " ++ opSym name ++ " " ++ go p r + let inner = + "if " + ++ go (depth + 1) 0 cond + ++ "\n" + ++ indent (depth + 1) + ++ "then " + ++ go (depth + 1) 0 t + ++ "\n" + ++ indent (depth + 1) + ++ "else " + ++ go (depth + 1) 0 e + in if prec > 0 then "(" ++ inner ++ ")" else inner + Unary op arg -> case unarySymbol op of + Nothing -> T.unpack (unaryName op) ++ "(" ++ go depth 0 arg ++ ")" + Just sym -> T.unpack sym ++ "(" ++ go depth 0 arg ++ ")" + Binary op l r -> + let p = binaryPrecedence op + inner = case binarySymbol op of + Just name -> go depth p l ++ " " ++ T.unpack name ++ " " ++ go depth p r + Nothing -> + T.unpack (binaryName op) ++ "(" ++ go depth p l ++ ", " ++ go depth p r ++ ")" in if prec > p then "(" ++ inner ++ ")" else inner - AggVector arg op _ -> T.unpack op ++ "(" ++ go 0 arg ++ ")" - AggReduce arg op _ -> T.unpack op ++ "(" ++ go 0 arg ++ ")" - AggNumericVector arg op _ -> T.unpack op ++ "(" ++ go 0 arg ++ ")" - AggFold arg op _ _ -> T.unpack op ++ "(" ++ go 0 arg ++ ")" - - opPrec :: T.Text -> Int - opPrec "add" = 1 - opPrec "sub" = 1 - opPrec "mult" = 2 - opPrec "divide" = 2 - opPrec "pow" = 3 - opPrec _ = 0 - - opSym :: T.Text -> String - opSym "add" = "+" - opSym "sub" = "-" - opSym "mult" = "*" - opSym "divide" = "/" - opSym "pow" = "^" - opSym "lt" = "<" - opSym "leq" = "<=" - opSym "gt" = ">" - opSym "geq" = ">=" - opSym other = T.unpack other + Agg (CollectAgg op _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")" + Agg (FoldAgg op _ _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")" 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 @@ -17,7 +17,6 @@ import Control.Monad.ST (runST) import Data.Bifunctor import qualified Data.Map as M -import Data.Maybe (fromMaybe, isJust) import qualified Data.Text as T import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) import qualified Data.Vector as V @@ -43,36 +42,36 @@ where columnNotFound = Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) -- Unary operations. -interpret df expr@(UnaryOp _ (f :: c -> d) value) = first (handleInterpretException (show expr)) $ do - (TColumn value') <- interpret @c df value - fmap TColumn (mapColumn f value') +interpret df expr@(Unary (op :: UnaryOp b c) value) = first (handleInterpretException (show expr)) $ do + (TColumn value') <- interpret df value + fmap TColumn (mapColumn (unaryFn op) value') -- Variations of binary operations. -interpret df expr@(BinaryOp _ (f :: c -> d -> e) left right) = first (handleInterpretException (show expr)) $ case (left, right) of - (Lit left, Lit right) -> interpret df (Lit (f left right)) +interpret df expr@(Binary (op :: BinaryOp c b a) left right) = first (handleInterpretException (show expr)) $ case (left, right) of + (Lit left, Lit right) -> interpret df (Lit (binaryFn op left right)) (Lit left, right) -> do -- If we have a literal then we don't have to materialise -- the column. - (TColumn value') <- interpret @d df right - fmap TColumn (mapColumn (f left) value') + (TColumn value') <- interpret df right + fmap TColumn (mapColumn (binaryFn op left) value') (left, Lit right) -> do -- Same as the above except the right side is the -- literl. - (TColumn value') <- interpret @c df left - fmap TColumn (mapColumn (`f` right) value') + (TColumn value') <- interpret df left + fmap TColumn (mapColumn (flip (binaryFn op) right) value') (_, _) -> do -- In the general case we interpret and zip. - (TColumn left') <- interpret @c df left - (TColumn right') <- interpret @d df right - fmap TColumn (zipWithColumns f left' right') + (TColumn left') <- interpret df left + (TColumn right') <- interpret df right + fmap TColumn (zipWithColumns (binaryFn op) left' right') -- Conditionals interpret df expr@(If cond l r) = first (handleInterpretException (show expr)) $ do - (TColumn conditions) <- interpret @Bool df cond - (TColumn left) <- interpret @a df l - (TColumn right) <- interpret @a df r + (TColumn conditions) <- interpret df cond + (TColumn left) <- interpret df l + (TColumn right) <- interpret df r let branch (c :: Bool) (l' :: a, r' :: a) = if c then l' else r' fmap TColumn (zipWithColumns branch conditions (zipColumns left right)) -interpret df expression@(AggVector expr op (f :: v b -> c)) = do - (TColumn column) <- interpret @b df expr +interpret df expression@(Agg (CollectAgg op (f :: v b -> c)) expr) = do + (TColumn column) <- interpret df expr -- Helper for errors. Should probably find a way of throwing this -- without leaking the fact that we use `Vector` to users. let aggTypeError expected = @@ -93,30 +92,17 @@ (BoxedColumn col) -> processColumn col (OptionalColumn col) -> processColumn col (UnboxedColumn col) -> processColumn col -interpret df expression@(AggReduce expr op (f :: a -> a -> a)) = first (handleInterpretException (show expr)) $ do - (TColumn column) <- interpret @a df expr - value <- foldl1Column f column - pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value -interpret df expression@(AggNumericVector expr op (f :: VU.Vector b -> c)) = first (handleInterpretException (show expression)) $ do - (TColumn column) <- interpret @b df expr - case column of - (UnboxedColumn (v :: VU.Vector d)) -> case testEquality (typeRep @d) (typeRep @b) of - Just Refl -> - pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) (f v) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - (Right (typeRep @b)) - (Right (typeRep @d)) - (Just "interpret") - (Just (show expression)) - ) - _ -> error "Trying to apply numeric computation to non-numeric column" -interpret df expression@(AggFold expr op start (f :: (a -> b -> a))) = first (handleInterpretException (show expression)) $ do - (TColumn column) <- interpret @b df expr - value <- foldlColumn f start column +interpret df expression@(Agg (FoldAgg op (Just v) f) expr) = first (handleInterpretException (show expr)) $ do + (TColumn column) <- interpret df expr + value <- foldlColumn f v column pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value +interpret df expression@(Agg (FoldAgg op Nothing (f :: a -> b -> a)) expr) = first (handleInterpretException (show expr)) $ do + (TColumn column) <- interpret df expr + case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> do + value <- foldl1Column f column + pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value + Nothing -> error "Type error" data AggregationResult a = UnAggregated Column @@ -138,8 +124,8 @@ Just (OptionalColumn col) -> Right $ UnAggregated $ fromVector $ mkUnaggregatedColumnBoxed col os indices Just (UnboxedColumn col) -> Right $ UnAggregated $ fromVector $ mkUnaggregatedColumnUnboxed col os indices -interpretAggregation gdf expression@(UnaryOp _ (f :: c -> d) expr) = - case interpretAggregation @c gdf expr of +interpretAggregation gdf expression@(Unary (op :: UnaryOp b a) expr) = + case interpretAggregation @b gdf expr of Left (TypeMismatchException context) -> Left $ TypeMismatchException @@ -150,105 +136,95 @@ ) Left e -> Left e Right (UnAggregated unaggregated) -> case unaggregated of - BoxedColumn (col :: V.Vector b) -> case testEquality (typeRep @b) (typeRep @(V.Vector c)) of - Just Refl -> case sUnbox @d of - SFalse -> Right $ UnAggregated $ fromVector $ V.map (V.map f) col + BoxedColumn (col :: V.Vector c) -> case testEquality (typeRep @c) (typeRep @(V.Vector b)) of + Just Refl -> case sUnbox @a of + SFalse -> Right $ UnAggregated $ fromVector $ V.map (V.map (unaryFn op)) col STrue -> Right $ UnAggregated $ fromVector $ - V.map (V.convert @V.Vector @d @VU.Vector . V.map f) col - Nothing -> case testEquality (typeRep @b) (typeRep @(VU.Vector c)) of - Nothing -> Left $ nestedTypeException @b @c (show expression) - Just Refl -> case (sUnbox @c, sUnbox @a) of + V.map (V.convert @V.Vector @a @VU.Vector . V.map (unaryFn op)) col + Nothing -> case testEquality (typeRep @c) (typeRep @(VU.Vector b)) of + Nothing -> Left $ nestedTypeException @c @a (show expression) + Just Refl -> case (sUnbox @b, sUnbox @a) of (SFalse, _) -> Left $ InternalException "Boxed type inside an unboxed column" - (STrue, STrue) -> Right $ UnAggregated $ fromVector $ V.map (VU.map f) col - (STrue, _) -> Right $ UnAggregated $ fromVector $ V.map (V.map f . VU.convert) col + (STrue, STrue) -> Right $ UnAggregated $ fromVector $ V.map (VU.map (unaryFn op)) col + (STrue, _) -> + Right $ UnAggregated $ fromVector $ V.map (V.map (unaryFn op) . VU.convert) col _ -> Left $ InternalException "Aggregated into a non-boxed column" - Right (Aggregated (TColumn aggregated)) -> case mapColumn f aggregated of + Right (Aggregated (TColumn aggregated)) -> case mapColumn (unaryFn op) aggregated of Left e -> Left e Right col -> Right $ Aggregated $ TColumn col -interpretAggregation gdf expression@(BinaryOp name (f :: c -> d -> e) left (Lit (right :: g))) = case testEquality (typeRep @g) (typeRep @d) of - Just Refl -> interpretAggregation gdf (UnaryOp name (`f` right) left) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @g) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) -interpretAggregation gdf expression@(BinaryOp name (f :: c -> d -> e) (Lit (left :: g)) right) = case testEquality (typeRep @g) (typeRep @c) of - Just Refl -> interpretAggregation gdf (UnaryOp name (f left) right) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @g) - , expectedType = Right (typeRep @c) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) -interpretAggregation gdf expression@(BinaryOp _ (f :: c -> d -> e) left right) = - case (interpretAggregation @c gdf left, interpretAggregation @d gdf right) of - (Right (Aggregated (TColumn left')), Right (Aggregated (TColumn right'))) -> case zipWithColumns f left' right' of +interpretAggregation gdf expression@(Binary (op :: t c b a) left (Lit (right :: b))) = + interpretAggregation + gdf + (Unary (MkUnaryOp (flip (binaryFn op) right) "udf" Nothing) left) +interpretAggregation gdf expression@(Binary (op :: BinaryOp c b a) (Lit (left :: c)) right) = + interpretAggregation + gdf + (Unary (MkUnaryOp (binaryFn op left) "udf" Nothing) right) +interpretAggregation gdf expression@(Binary (op :: BinaryOp c b a) left right) = + case (interpretAggregation gdf left, interpretAggregation gdf right) of + (Right (Aggregated (TColumn left')), Right (Aggregated (TColumn right'))) -> case zipWithColumns (binaryFn op) left' right' of Left e -> Left e Right col -> Right $ Aggregated $ TColumn col (Right (UnAggregated left'), Right (UnAggregated right')) -> case (left', right') of (BoxedColumn (l :: V.Vector m), BoxedColumn (r :: V.Vector n)) -> case testEquality (typeRep @m) (typeRep @(VU.Vector c)) of - Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector d)) of - Just Refl -> case (sUnbox @c, sUnbox @d, sUnbox @e) of + Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector b)) of + Just Refl -> case (sUnbox @c, sUnbox @b, sUnbox @a) of (STrue, STrue, STrue) -> - Right $ UnAggregated $ fromVector $ V.zipWith (VU.zipWith f) l r + Right $ UnAggregated $ fromVector $ V.zipWith (VU.zipWith (binaryFn op)) l r (STrue, STrue, SFalse) -> Right $ UnAggregated $ fromVector $ - V.zipWith (\l' r' -> V.zipWith f (V.convert l') (V.convert r')) l r + V.zipWith (\l' r' -> V.zipWith (binaryFn op) (V.convert l') (V.convert r')) l r (_, _, _) -> Left $ InternalException "Boxed vectors contain unboxed types" - Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector d)) of + Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector b)) of Just Refl -> case sUnbox @c of STrue -> Right $ UnAggregated $ fromVector $ - V.zipWith (V.zipWith f . V.convert) l r + V.zipWith (V.zipWith (binaryFn op) . V.convert) l r SFalse -> Left $ InternalException "Unboxed vectors contain boxed types" - Nothing -> Left $ nestedTypeException @n @d (show right) + Nothing -> Left $ nestedTypeException @n @b (show right) Nothing -> case testEquality (typeRep @m) (typeRep @(V.Vector c)) of Nothing -> Left $ nestedTypeException @m @c (show left) - Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector d)) of - Just Refl -> case (sUnbox @d, sUnbox @e) of + Just Refl -> case testEquality (typeRep @n) (typeRep @(VU.Vector b)) of + Just Refl -> case (sUnbox @b, sUnbox @a) of (STrue, STrue) -> Right $ UnAggregated $ fromVector $ V.zipWith - (\l' r' -> V.convert @V.Vector @e @VU.Vector $ V.zipWith f l' (V.convert r')) + ( \l' r' -> + V.convert @V.Vector @a @VU.Vector $ V.zipWith (binaryFn op) l' (V.convert r') + ) l r (STrue, SFalse) -> Right $ UnAggregated $ fromVector $ - V.zipWith (\l' r' -> V.zipWith f l' (V.convert r')) l r + V.zipWith (\l' r' -> V.zipWith (binaryFn op) l' (V.convert r')) l r (_, _) -> Left $ InternalException "Unboxed vectors contain boxed types" - Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector d)) of - Just Refl -> case sUnbox @e of + Nothing -> case testEquality (typeRep @n) (typeRep @(V.Vector b)) of + Just Refl -> case sUnbox @a of SFalse -> Right $ UnAggregated $ fromVector $ - V.zipWith (V.zipWith f . V.convert) l r + V.zipWith (V.zipWith (binaryFn op) . V.convert) l r STrue -> Right $ UnAggregated $ fromVector $ - V.zipWith (\l' r' -> V.convert @V.Vector @e @VU.Vector $ V.zipWith f l' r') l r - Nothing -> Left $ nestedTypeException @n @d (show right) + V.zipWith + (\l' r' -> V.convert @V.Vector @a @VU.Vector $ V.zipWith (binaryFn op) l' r') + l + r + Nothing -> Left $ nestedTypeException @n @b (show right) _ -> Left $ InternalException "Aggregated into a non-boxed column" (Right _, Right _) -> Left $ @@ -571,7 +547,7 @@ } ) (_, _, Left e) -> Left e -interpretAggregation gdf@(Grouped df names indices os) expression@(AggVector expr op (f :: v b -> c)) = +interpretAggregation gdf@(Grouped df names indices os) expression@(Agg (CollectAgg op (f :: v b -> c)) expr) = case interpretAggregation @b gdf expr of Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(v b)) (typeRep @d) of Nothing -> Left $ nestedTypeException @d @b (show expr) @@ -630,93 +606,9 @@ } ) (Left e) -> Left e -interpretAggregation gdf@(Grouped df names indices os) expression@(AggNumericVector (Col name) op (f :: VU.Vector b -> c)) = - case getColumn name df of - -- TODO(mchavinda): Fix the compedium of type errors here - -- This is mostly done help with the benchmarking. - Nothing -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) - Just (BoxedColumn col) -> error "Type mismatch." - Just (OptionalColumn col) -> error "Type mismatch." - Just (UnboxedColumn (col :: VU.Vector d)) -> case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> case testEquality (typeRep @c) (typeRep @a) of - Just Refl -> - Right $ - Aggregated $ - TColumn $ - fromUnboxedVector $ - mkAggregatedColumnUnboxed col os indices f - Nothing -> error "Type mismatch" - Nothing -> error "Type mismatch" -interpretAggregation gdf@(Grouped df names indices os) expression@(AggNumericVector expr op (f :: VU.Vector b -> c)) = - case interpretAggregation @b gdf expr of - (Left (TypeMismatchException context)) -> - Left $ - TypeMismatchException - ( context - { callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expression) - } - ) - (Left e) -> Left e - Right (UnAggregated (BoxedColumn (col :: V.Vector d))) -> case testEquality (typeRep @(VU.Vector b)) (typeRep @d) of - Nothing -> case testEquality (typeRep @(VU.Vector Int)) (typeRep @d) of - Nothing -> case testEquality (typeRep @(V.Vector Integer)) (typeRep @d) of - Nothing -> Left $ nestedTypeException @d @b (show expr) - Just Refl -> - Right $ - Aggregated $ - TColumn $ - fromVector $ - V.map (f . VU.convert . V.map fromIntegral) col - Just Refl -> - Right $ - Aggregated $ - TColumn $ - fromVector $ - V.map f (V.map (VU.map fromIntegral) col) - Just Refl -> Right $ Aggregated $ TColumn $ fromVector $ V.map f col - Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column" - Right (Aggregated (TColumn (BoxedColumn (col :: V.Vector d)))) -> case testEquality (typeRep @Integer) (typeRep @d) of - Just Refl -> interpretAggregation @c gdf (Lit ((f . V.convert . V.map fromIntegral) col)) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expr) - } - ) - Right (Aggregated (TColumn (UnboxedColumn (col :: VU.Vector d)))) -> case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> interpretAggregation @c gdf (Lit (f col)) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expr) - } - ) - Right (Aggregated (TColumn (OptionalColumn (col :: V.Vector (Maybe d))))) -> case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> - interpretAggregation @c - gdf - (Lit ((f . V.convert . V.map (fromMaybe 0) . V.filter isJust) col)) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "interpretAggregation" - , errorColumnName = Just (show expr) - } - ) -interpretAggregation gdf@(Grouped df names indices os) expression@(AggReduce (Col name) op (f :: a -> a -> a)) = - case getColumn name df of +interpretAggregation gdf@(Grouped df names indices os) expression@(Agg (FoldAgg op Nothing (f :: a -> b -> a)) (Col name)) = case testEquality (typeRep @a) (typeRep @b) of + Nothing -> error "Type mismatch" + Just Refl -> case getColumn name df of Nothing -> Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) Just (BoxedColumn (col :: V.Vector d)) -> case testEquality (typeRep @a) (typeRep @d) of Nothing -> error "Type mismatch" @@ -742,8 +634,9 @@ fromUnboxedVector $ mkReducedColumnUnboxed col os indices f Nothing -> error "Type mismatch" -interpretAggregation gdf@(Grouped df names indices os) expression@(AggReduce expr op (f :: a -> a -> a)) = - case interpretAggregation @a gdf expr of +interpretAggregation gdf@(Grouped df names indices os) expression@(Agg (FoldAgg op Nothing (f :: a -> b -> a)) expr) = case testEquality (typeRep @a) (typeRep @b) of + Nothing -> error "Type mismatch" + Just Refl -> case interpretAggregation @a gdf expr of (Left (TypeMismatchException context)) -> Left $ TypeMismatchException @@ -774,8 +667,8 @@ Right (Aggregated (TColumn column)) -> case foldl1Column f column of Left e -> Left e Right value -> interpretAggregation @a gdf (Lit value) -interpretAggregation gdf@(Grouped df names indices os) expression@(AggFold expr op s (f :: (a -> b -> a))) = - case interpretAggregation @b gdf expr of +interpretAggregation gdf@(Grouped df names indices os) expression@(Agg (FoldAgg op (Just s) (f :: a -> b -> a)) expr) = + case interpretAggregation gdf expr of (Left (TypeMismatchException context)) -> Left $ TypeMismatchException 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 @@ -53,6 +53,12 @@ (T.pack $ show $ names L.\\ columnNames df) "groupBy" (columnNames df) + | nRows df == 0 = + Grouped + df + names + VU.empty + (VU.fromList [0]) | otherwise = Grouped df @@ -261,7 +267,7 @@ (VU.map (valueIndices VU.!) (VU.init offsets)) (select groupingColumns df) - f (name, Wrap (expr :: Expr a)) d = + f (name, UExpr (expr :: Expr a)) d = let value = case interpretAggregation @a gdf expr of Left e -> throw e diff --git a/src/DataFrame/Operations/Core.hs b/src/DataFrame/Operations/Core.hs --- a/src/DataFrame/Operations/Core.hs +++ b/src/DataFrame/Operations/Core.hs @@ -680,7 +680,8 @@ @ -} -valueCounts :: forall a. (Columnable a) => Expr a -> DataFrame -> [(a, Int)] +valueCounts :: + forall a. (Ord a, Columnable a) => Expr a -> DataFrame -> [(a, Int)] valueCounts expr df = case columnAsVector expr df of Left e -> throw e Right column' -> @@ -702,7 +703,7 @@ @ -} valueProportions :: - forall a. (Columnable a) => Expr a -> DataFrame -> [(a, Double)] + forall a. (Ord a, Columnable a) => Expr a -> DataFrame -> [(a, Double)] valueProportions expr df = case columnAsVector expr df of Left e -> throw e Right column' -> 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 @@ -137,7 +137,14 @@ then df else ( if name `elem` leftColumns - then insertIfPresent ("Right_" <> name) (D.getColumn name expandedRight) df + then + insertIfPresent + name + ( D.mergeColumns + <$> D.getColumn name expandedLeft + <*> D.getColumn name expandedRight + ) + df else insertIfPresent name (D.getColumn name expandedRight) df ) ) @@ -215,7 +222,14 @@ then df else ( if name `elem` leftColumns - then insertIfPresent ("Right_" <> name) (D.getColumn name expandedRight) df + then + insertIfPresent + name + ( D.mergeColumns + <$> D.getColumn name expandedLeft + <*> D.getColumn name expandedRight + ) + df else insertIfPresent name (D.getColumn name expandedRight) df ) ) @@ -322,7 +336,14 @@ _ -> error "Join should have optional keys." else ( if name `elem` leftColumns - then insertIfPresent ("Right_" <> name) (D.getColumn name expandedRight) df + then + insertIfPresent + name + ( D.mergeColumns + <$> D.getColumn name expandedLeft + <*> D.getColumn name expandedRight + ) + df else insertIfPresent name (D.getColumn name expandedRight) df ) ) -- ??? 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 @@ -1,3 +1,6 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE InstanceSigs #-} {-# LANGUAGE OverloadedStrings #-} module DataFrame.Operations.Permutation where @@ -11,19 +14,27 @@ import DataFrame.Errors (DataFrameException (..)) import DataFrame.Internal.Column import DataFrame.Internal.DataFrame (DataFrame (..)) +import DataFrame.Internal.Expression import DataFrame.Internal.Row import DataFrame.Operations.Core import System.Random +import System.Random.Shuffle (shuffle') -- | Sort order taken as a parameter by the 'sortBy' function. -data SortOrder - = Asc T.Text - | Desc T.Text - deriving (Eq) +data SortOrder where + Asc :: (Columnable a) => Expr a -> SortOrder + Desc :: (Columnable a) => Expr a -> SortOrder +instance Eq SortOrder where + (==) :: SortOrder -> SortOrder -> Bool + (==) (Asc _) (Asc _) = True + (==) (Desc _) (Desc _) = True + (==) _ _ = False + getSortColumnName :: SortOrder -> T.Text -getSortColumnName (Asc n) = n -getSortColumnName (Desc n) = n +getSortColumnName (Asc (Col n)) = n +getSortColumnName (Desc (Col n)) = n +getSortColumnName _ = error "Sorting on compound column" mustFlipCompare :: SortOrder -> Bool mustFlipCompare (Asc _) = True @@ -65,4 +76,4 @@ df{columns = V.map (atIndicesStable indexes) (columns df)} shuffledIndices :: (RandomGen g) => g -> Int -> VU.Vector Int -shuffledIndices pureGen k = VU.fromList (fst (uniformShuffleList [0 .. (k - 1)] pureGen)) +shuffledIndices pureGen k = VU.fromList (shuffle' [0 .. (k - 1)] k pureGen) 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 @@ -57,33 +57,28 @@ Percentage (%) | 44.26% | 31.74% | 0.02% | 11.09% | 12.88% @ -} -frequencies :: T.Text -> DataFrame -> DataFrame -frequencies name df = +frequencies :: forall a. (Columnable a) => Expr a -> DataFrame -> DataFrame +frequencies expr df = let - counts :: forall a. (Columnable a) => [(a, Int)] - counts = valueCounts (Col @a name) df + counts = valueCounts expr df calculatePercentage cs k = toAny $ toPct2dp (fromIntegral k / fromIntegral (P.sum $ map snd cs)) initDf = empty & insertVector "Statistic" (V.fromList ["Count" :: T.Text, "Percentage (%)"]) - freqs :: forall v a. (VG.Vector v a, Columnable a) => v a -> DataFrame freqs col = L.foldl' ( \d (col, k) -> insertVector (showValue @a col) - (V.fromList [toAny k, calculatePercentage (counts @a) k]) + (V.fromList [toAny k, calculatePercentage counts k]) d ) initDf counts in - case getColumn name df of - Nothing -> - throw $ ColumnNotFoundException name "frequencies" (M.keys $ columnIndices df) - Just ((BoxedColumn (column :: V.Vector a))) -> freqs column - Just ((OptionalColumn (column :: V.Vector a))) -> freqs column - Just ((UnboxedColumn (column :: VU.Vector a))) -> freqs column + case columnAsVector expr df of + Left err -> throw err + Right column -> freqs column -- | Calculates the mean of a given column as a standalone value. mean :: 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 @@ -10,6 +10,7 @@ import qualified Data.List as L import qualified Data.Map as M +import qualified Data.Set as S import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Generic as VG @@ -295,9 +296,10 @@ -- | O(n) select columns by column predicate name. selectBy :: [SelectionCriteria] -> DataFrame -> DataFrame -selectBy xs df = select columnsWithProperties df +selectBy xs df = select finalSelection df where - columnsWithProperties = L.foldl' columnWithProperty [] xs + finalSelection = Prelude.filter (`S.member` columnsWithProperties) (columnNames df) + columnsWithProperties = S.fromList (L.foldl' columnWithProperty [] xs) columnWithProperty acc (ColumnName name) = acc ++ [name] columnWithProperty acc (ColumnNameProperty f) = acc ++ L.filter f (columnNames df) columnWithProperty acc (ColumnTextRange (from, to)) = @@ -339,7 +341,19 @@ in df & insertUnboxedVector "__rand__" rand - & filterWhere (BinaryOp "geq" (>=) (Col @Double "__rand__") (Lit (1 - p))) + & filterWhere + ( Binary + ( MkBinaryOp + { binaryFn = (>=) + , binaryName = "geq" + , binarySymbol = Just ">=" + , binaryCommutative = False + , binaryPrecedence = 1 + } + ) + (Col @Double "__rand__") + (Lit (1 - p)) + ) & exclude ["__rand__"] {- | Split a dataset into two. The first in the tuple gets a sample of p (0 <= p <= 1) and the second gets (1 - p). This is useful for creating test and train splits. @@ -359,10 +373,34 @@ withRand = df & insertUnboxedVector "__rand__" rand in ( withRand - & filterWhere (BinaryOp "leq" (<=) (Col @Double "__rand__") (Lit p)) + & filterWhere + ( Binary + ( MkBinaryOp + { binaryFn = (<=) + , binaryName = "leq" + , binarySymbol = Just "<=" + , binaryCommutative = False + , binaryPrecedence = 1 + } + ) + (Col @Double "__rand__") + (Lit p) + ) & exclude ["__rand__"] , withRand - & filterWhere (BinaryOp "gt" (>) (Col @Double "__rand__") (Lit p)) + & filterWhere + ( Binary + ( MkBinaryOp + { binaryFn = (>) + , binaryName = "gt" + , binarySymbol = Just ">" + , binaryCommutative = False + , binaryPrecedence = 1 + } + ) + (Col @Double "__rand__") + (Lit p) + ) & exclude ["__rand__"] ) @@ -384,9 +422,15 @@ singleFold n d = d & filterWhere - ( BinaryOp - "geq" - (>=) + ( Binary + ( MkBinaryOp + { binaryFn = (>=) + , binaryName = "geq" + , binarySymbol = Just ">=" + , binaryCommutative = False + , binaryPrecedence = 1 + } + ) (Col @Double "__rand__") (Lit (fromIntegral n * partitionSize)) ) @@ -397,9 +441,15 @@ d'' = d & filterWhere - ( BinaryOp - "lt" - (<) + ( Binary + ( MkBinaryOp + { binaryFn = (<) + , binaryName = "lt" + , binarySymbol = Just "<" + , binaryCommutative = False + , binaryPrecedence = 1 + } + ) (Col @Double "__rand__") (Lit (fromIntegral n * partitionSize)) ) diff --git a/src/DataFrame/Operations/Transformations.hs b/src/DataFrame/Operations/Transformations.hs --- a/src/DataFrame/Operations/Transformations.hs +++ b/src/DataFrame/Operations/Transformations.hs @@ -71,7 +71,7 @@ Left e -> throw e Right (TColumn value) -> (insertColumn name value df) - { derivingExpressions = M.insert name (Wrap expr) (derivingExpressions df) + { derivingExpressions = M.insert name (UExpr expr) (derivingExpressions df) } {- | O(k) Apply a function to an expression in a dataframe and @@ -91,7 +91,7 @@ deriveMany :: [NamedExpr] -> DataFrame -> DataFrame deriveMany exprs df = let - f (name, Wrap (expr :: Expr a)) d = + f (name, UExpr (expr :: Expr a)) d = case interpret @a df expr of Left e -> throw e Right (TColumn value) -> insertColumn name value d diff --git a/src/DataFrame/Operators.hs b/src/DataFrame/Operators.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operators.hs @@ -0,0 +1,161 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} + +module DataFrame.Operators where + +import Data.Function ((&)) +import qualified Data.Text as T +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Expression ( + BinaryOp ( + MkBinaryOp, + binaryCommutative, + binaryFn, + binaryName, + binaryPrecedence, + binarySymbol + ), + Expr (Binary, Col, If, Lit), + NamedExpr, + UExpr (UExpr), + ) + +infix 8 .^^ +infix 4 .==, .<, .<=, .>=, .>, ./= +infixr 3 .&& +infixr 2 .|| +infixr 0 .= + +(|>) :: a -> (a -> b) -> b +(|>) = (&) + +as :: (Columnable a) => Expr a -> T.Text -> NamedExpr +as expr name = (name, UExpr expr) + +name :: (Show a) => Expr a -> T.Text +name (Col n) = n +name other = + error $ + "You must call `name` on a column reference. Not the expression: " ++ show other + +col :: (Columnable a) => T.Text -> Expr a +col = Col + +ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a +ifThenElse = If + +lit :: (Columnable a) => a -> Expr a +lit = Lit + +(.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr +(.=) = flip as + +(.==) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool +(.==) = + Binary + ( MkBinaryOp + { binaryFn = (==) + , binaryName = "eq" + , binarySymbol = Just "==" + , binaryCommutative = True + , binaryPrecedence = 4 + } + ) + +(./=) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool +(./=) = + Binary + ( MkBinaryOp + { binaryFn = (/=) + , binaryName = "neq" + , binarySymbol = Just "/=" + , binaryCommutative = True + , binaryPrecedence = 4 + } + ) + +(.<) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool +(.<) = + Binary + ( MkBinaryOp + { binaryFn = (<) + , binaryName = "lt" + , binarySymbol = Just "<" + , binaryCommutative = False + , binaryPrecedence = 4 + } + ) + +(.>) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool +(.>) = + Binary + ( MkBinaryOp + { binaryFn = (>) + , binaryName = "gt" + , binarySymbol = Just ">" + , binaryCommutative = False + , binaryPrecedence = 4 + } + ) + +(.<=) :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool +(.<=) = + Binary + ( MkBinaryOp + { binaryFn = (<=) + , binaryName = "leq" + , binarySymbol = Just "<=" + , binaryCommutative = False + , binaryPrecedence = 4 + } + ) + +(.>=) :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool +(.>=) = + Binary + ( MkBinaryOp + { binaryFn = (>=) + , binaryName = "geq" + , binarySymbol = Just ">=" + , binaryCommutative = False + , binaryPrecedence = 4 + } + ) + +(.&&) :: Expr Bool -> Expr Bool -> Expr Bool +(.&&) = + Binary + ( MkBinaryOp + { binaryFn = (&&) + , binaryName = "and" + , binarySymbol = Just "&&" + , binaryCommutative = True + , binaryPrecedence = 3 + } + ) + +(.||) :: Expr Bool -> Expr Bool -> Expr Bool +(.||) = + Binary + ( MkBinaryOp + { binaryFn = (||) + , binaryName = "or" + , binarySymbol = Just "||" + , binaryCommutative = True + , binaryPrecedence = 2 + } + ) + +(.^^) :: (Columnable a, Num a) => Expr a -> Int -> Expr a +(.^^) expr i = + Binary + ( MkBinaryOp + { binaryFn = (^) + , binaryName = "pow" + , binarySymbol = Just "^" + , binaryCommutative = False + , binaryPrecedence = 8 + } + ) + expr + (Lit i) diff --git a/src/DataFrame/Synthesis.hs b/src/DataFrame/Synthesis.hs --- a/src/DataFrame/Synthesis.hs +++ b/src/DataFrame/Synthesis.hs @@ -37,8 +37,8 @@ import qualified Data.Text as T import Data.Type.Equality import qualified Data.Vector.Unboxed as VU -import DataFrame.Functions ((.&&), (.<=), (.>), (.||)) import qualified DataFrame.Operations.Core as D +import DataFrame.Operators import Debug.Trace (trace) import Type.Reflection (typeRep) diff --git a/tests/GenDataFrame.hs b/tests/GenDataFrame.hs --- a/tests/GenDataFrame.hs +++ b/tests/GenDataFrame.hs @@ -21,8 +21,8 @@ genDataFrame :: Gen DataFrame genDataFrame = do - numRows <- choose (0, 100) - numCols <- choose (0, 10) + numRows <- choose (100, 1000) + numCols <- choose (1, 10) colNames <- V.fromList <$> vectorOf numCols genUniqueColName cols <- V.fromList <$> vectorOf numCols (genColumn numRows) let indices = M.fromList $ zip (V.toList colNames) [0 ..] diff --git a/tests/Main.hs b/tests/Main.hs --- a/tests/Main.hs +++ b/tests/Main.hs @@ -18,6 +18,7 @@ import Test.QuickCheck import qualified Functions +import qualified Monad import qualified Operations.Aggregations import qualified Operations.Apply import qualified Operations.Core @@ -28,6 +29,7 @@ import qualified Operations.Join import qualified Operations.Merge import qualified Operations.ReadCsv +import qualified Operations.Shuffle import qualified Operations.Sort import qualified Operations.Statistics import qualified Operations.Subset @@ -5119,6 +5121,7 @@ ++ Operations.Join.tests ++ Operations.Merge.tests ++ Operations.ReadCsv.tests + ++ Operations.Shuffle.tests ++ Operations.Sort.tests ++ Operations.Statistics.tests ++ Operations.Take.tests @@ -5138,6 +5141,10 @@ mapM (quickCheckWithResult stdArgs) Operations.Subset.tests - if failures result > 0 || errors result > 0 || not (all isSuccessful propRes) + monadRes <- mapM (quickCheckWithResult stdArgs) Monad.tests + if failures result > 0 + || errors result > 0 + || not (all isSuccessful propRes) + || not (all isSuccessful monadRes) then Exit.exitFailure else Exit.exitSuccess diff --git a/tests/Monad.hs b/tests/Monad.hs new file mode 100644 --- /dev/null +++ b/tests/Monad.hs @@ -0,0 +1,27 @@ +module Monad where + +import qualified DataFrame as D +import DataFrame.Internal.DataFrame +import DataFrame.Monad +import GenDataFrame () +import System.Random +import Test.QuickCheck +import Test.QuickCheck.Monadic + +roundToTwoPlaces x = fromIntegral (round (x * 100)) / 100.0 + +prop_sampleM :: DataFrame -> Gen (Gen Property) +prop_sampleM df = monadic' $ do + p <- run $ choose (0.0 :: Double, 1.0 :: Double) + let expectedRate = roundToTwoPlaces p + seed <- run $ choose (0, 1000) + let rowCount = D.nRows df + let colCount = D.nColumns df + pre (colCount > 1 && rowCount > 100) + let finalDf = execFrameM df (sampleM (mkStdGen seed) expectedRate) + let finalRowCount = D.nRows finalDf + let realRate = roundToTwoPlaces $ fromIntegral finalRowCount / fromIntegral rowCount + let diff = abs $ expectedRate - realRate + assert (diff <= 0.11) + +tests = [prop_sampleM] diff --git a/tests/Operations/Aggregations.hs b/tests/Operations/Aggregations.hs --- a/tests/Operations/Aggregations.hs +++ b/tests/Operations/Aggregations.hs @@ -9,6 +9,7 @@ import qualified DataFrame.Internal.Column as DI import Data.Function +import DataFrame.Operators import Test.HUnit values :: [(T.Text, DI.Column)] @@ -36,8 +37,8 @@ ) ( testData & D.groupBy ["test1"] - & D.aggregate [F.count (F.col @Int "test2") `F.as` "test2"] - & D.sortBy [D.Asc "test1"] + & D.aggregate [F.count (F.col @Int "test2") `as` "test2"] + & D.sortBy [D.Asc (F.col @Int "test1")] ) ) @@ -53,8 +54,8 @@ ) ( testData & D.groupBy ["test1"] - & D.aggregate [F.mean (F.col @Int "test2") `F.as` "test2"] - & D.sortBy [D.Asc "test1"] + & D.aggregate [F.mean (F.col @Int "test2") `as` "test2"] + & D.sortBy [D.Asc (F.col @Int "test1")] ) ) @@ -71,9 +72,9 @@ ( testData & D.groupBy ["test1"] & D.aggregate - [ F.mean (F.lift (fromIntegral @Int @Double) (F.col @Int "test2")) `F.as` "test2" + [ F.mean (F.lift (fromIntegral @Int @Double) (F.col @Int "test2")) `as` "test2" ] - & D.sortBy [D.Asc "test1"] + & D.sortBy [D.Asc (F.col @Int "test1")] ) ) @@ -89,8 +90,8 @@ ) ( testData & D.groupBy ["test1"] - & D.aggregate [F.mean (F.col @Int "test2" + F.col @Int "test2") `F.as` "test2"] - & D.sortBy [D.Asc "test1"] + & D.aggregate [F.mean (F.col @Int "test2" + F.col @Int "test2") `as` "test2"] + & D.sortBy [D.Asc (F.col @Int "test1")] ) ) @@ -108,9 +109,9 @@ & D.groupBy ["test1"] & D.aggregate [ F.maximum (F.lift (fromIntegral @Int @Double) (F.col @Int "test2")) - `F.as` "test2" + `as` "test2" ] - & D.sortBy [D.Asc "test1"] + & D.sortBy [D.Asc (F.col @Int "test1")] ) ) @@ -127,11 +128,29 @@ ( testData & D.groupBy ["test1"] & D.aggregate - [F.maximum (F.col @Int "test2" + F.col @Int "test2") `F.as` "test2"] - & D.sortBy [D.Asc "test1"] + [F.maximum (F.col @Int "test2" + F.col @Int "test2") `as` "test2"] + & D.sortBy [D.Asc (F.col @Int "test1")] ) ) +aggregationOnNoRows :: Test +aggregationOnNoRows = + TestCase + ( assertEqual + "Aggregation on DataFrame with no rows" + ( D.fromNamedColumns + [ ("test1", DI.fromList ([] :: [Int])) + , ("sum(test2)", DI.fromList ([] :: [Int])) + ] + ) + ( testData + & D.drop 12 + & D.groupBy ["test1"] + & D.aggregate + [F.sum (F.col @Int "test2") `as` "sum(test2)"] + ) + ) + tests :: [Test] tests = [ TestLabel "foldAggregation" foldAggregation @@ -148,4 +167,7 @@ , TestLabel "reduceAggregationOfUnaggregatedBinaryOp" reduceAggregationOfUnaggregatedBinaryOp + , TestLabel + "aggregationOnNoRows" + aggregationOnNoRows ] diff --git a/tests/Operations/Join.hs b/tests/Operations/Join.hs --- a/tests/Operations/Join.hs +++ b/tests/Operations/Join.hs @@ -4,7 +4,9 @@ module Operations.Join where import Data.Text (Text) +import Data.These import qualified DataFrame as D +import qualified DataFrame.Functions as F import DataFrame.Operations.Join import Test.HUnit @@ -33,7 +35,7 @@ , ("B", D.fromList ["B0" :: Text, "B1", "B2"]) ] ) - (D.sortBy [D.Asc "key"] (innerJoin ["key"] df1 df2)) + (D.sortBy [D.Asc (F.col @Text "key")] (innerJoin ["key"] df1 df2)) ) testLeftJoin :: Test @@ -47,7 +49,7 @@ , ("B", D.fromList [Just "B0", Just "B1" :: Maybe Text, Just "B2"]) ] ) - (D.sortBy [D.Asc "key"] (leftJoin ["key"] df2 df1)) + (D.sortBy [D.Asc (F.col @Text "key")] (leftJoin ["key"] df2 df1)) ) testRightJoin :: Test @@ -57,11 +59,11 @@ "Test right join with single key" ( D.fromNamedColumns [ ("key", D.fromList ["K0" :: Text, "K1", "K2"]) - , ("A", D.fromList ["A0" :: Text, "A1", "A2"]) + , ("A", D.fromList [Just "A0" :: Maybe Text, Just "A1", Just "A2"]) , ("B", D.fromList ["B0" :: Text, "B1", "B2"]) ] ) - (D.sortBy [D.Asc "key"] (rightJoin ["key"] df2 df1)) + (D.sortBy [D.Asc (F.col @Text "key")] (rightJoin ["key"] df2 df1)) ) staffDf :: D.DataFrame @@ -108,13 +110,107 @@ ) ] ) - (D.sortBy [D.Asc "Name"] (fullOuterJoin ["Name"] studentDf staffDf)) + (D.sortBy [D.Asc (F.col @Text "Name")] (fullOuterJoin ["Name"] studentDf staffDf)) ) +dfL :: D.DataFrame +dfL = + D.fromNamedColumns + [ ("key", D.fromList ["K0" :: Text, "K1", "K2"]) + , ("X", D.fromList ["LX0" :: Text, "LX1", "LX2"]) + , ("Lonly", D.fromList ["L0" :: Text, "L1", "L2"]) + ] + +dfR :: D.DataFrame +dfR = + D.fromNamedColumns + [ ("key", D.fromList ["K0" :: Text, "K1", "K3"]) + , ("X", D.fromList ["RX0" :: Text, "RX1", "RX3"]) + , ("Ronly", D.fromList [10 :: Int, 11, 13]) + ] + +testInnerJoinWithCollisions :: Test +testInnerJoinWithCollisions = + TestCase + ( assertEqual + "Test inner join with single key" + ( D.fromNamedColumns + [ ("key", D.fromList ["K0" :: Text, "K1"]) + , ("X", D.fromList [These "LX0" "RX0" :: These Text Text, These "LX1" "RX1"]) + , ("Lonly", D.fromList ["L0" :: Text, "L1"]) + , ("Ronly", D.fromList [10 :: Int, 11]) + ] + ) + (D.sortBy [D.Asc (F.col @Text "key")] (innerJoin ["key"] dfR dfL)) + ) + +testLeftJoinWithCollisions :: Test +testLeftJoinWithCollisions = + TestCase + ( assertEqual + "Test left join with single key" + ( D.fromNamedColumns + [ ("key", D.fromList ["K0" :: Text, "K1", "K2"]) + , + ( "X" + , D.fromList [These "LX0" "RX0" :: These Text Text, These "LX1" "RX1", This "LX2"] + ) + , ("Lonly", D.fromList ["L0" :: Text, "L1", "L2"]) + , ("Ronly", D.fromList [Just 10 :: Maybe Int, Just 11, Nothing]) + ] + ) + (D.sortBy [D.Asc (F.col @Text "key")] (leftJoin ["key"] dfR dfL)) + ) + +testRightJoinWithCollisions :: Test +testRightJoinWithCollisions = + TestCase + ( assertEqual + "Test right join with single key" + ( D.fromNamedColumns + [ ("key", D.fromList ["K0" :: Text, "K1", "K3"]) + , + ( "X" + , D.fromList [These "RX0" "LX0" :: These Text Text, These "RX1" "LX1", This "RX3"] + ) + , ("Ronly", D.fromList [10 :: Int, 11, 13]) + , ("Lonly", D.fromList [Just "L0" :: Maybe Text, Just "L1", Nothing]) + ] + ) + (D.sortBy [D.Asc (F.col @Text "key")] (rightJoin ["key"] dfR dfL)) + ) + +testOuterJoinWithCollisions :: Test +testOuterJoinWithCollisions = + TestCase + ( assertEqual + "Test right join with single key" + ( D.fromNamedColumns + [ ("key", D.fromList ["K0" :: Text, "K1", "K2", "K3"]) + , + ( "X" + , D.fromList + [ Just (These "LX0" "RX0") :: Maybe (These Text Text) + , Just (These "LX1" "RX1") + , Just (This "LX2") + , Just (That "RX3") + ] + ) + , ("Lonly", D.fromList [Just "L0" :: Maybe Text, Just "L1", Just "L2", Nothing]) + , ("Ronly", D.fromList [Just 10 :: Maybe Int, Just 11, Nothing, Just 13]) + ] + ) + (D.sortBy [D.Asc (F.col @Text "key")] (fullOuterJoin ["key"] dfR dfL)) + ) + tests :: [Test] tests = [ TestLabel "innerJoin" testInnerJoin , TestLabel "leftJoin" testLeftJoin , TestLabel "rightJoin" testRightJoin , TestLabel "fullOuterJoin" testFullOuterJoin + , TestLabel "innerJoinWithCollisions" testInnerJoinWithCollisions + , TestLabel "leftJoinWithCollisions" testLeftJoinWithCollisions + , TestLabel "rightJoinWithCollisions" testRightJoinWithCollisions + , TestLabel "outerJoinWithCollisions" testOuterJoinWithCollisions ] diff --git a/tests/Operations/Shuffle.hs b/tests/Operations/Shuffle.hs new file mode 100644 --- /dev/null +++ b/tests/Operations/Shuffle.hs @@ -0,0 +1,84 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Operations.Shuffle where + +import qualified DataFrame as D + +import DataFrame.Operations.Permutation (shuffle) +import System.Random (mkStdGen) +import Test.HUnit (Test (..), assertEqual) + +testDataFrame :: D.DataFrame +testDataFrame = + D.fromNamedColumns + [ ("numbers", D.fromList @Int [1 .. 26]) + ] + +-- Test that shuffling does anything at all +shuffleShuffles :: Test +shuffleShuffles = + let gen = mkStdGen 1234 + shuffled = shuffle gen testDataFrame + initialNumbers = D.extractNumericColumn "numbers" testDataFrame + shuffledNumbers = D.extractNumericColumn "numbers" shuffled + in TestCase + ( assertEqual + "Shuffled column unequal to initial column" + False + (initialNumbers == shuffledNumbers) + ) + +shufflePreservesColumnNames :: Test +shufflePreservesColumnNames = + let gen = mkStdGen 837 + shuffled = shuffle gen testDataFrame + in TestCase + ( assertEqual + "Column names are unchanged" + (D.columnNames shuffled) + (D.columnNames testDataFrame) + ) + +-- Test that un-shuffling restores the original dataframe +-- which is known to be sorted in this case +shufflePreservesData :: Test +shufflePreservesData = + let gen = mkStdGen 1234 + shuffled = shuffle gen testDataFrame + sortedShuffled = D.sortBy [D.Asc (D.col @Int "numbers")] shuffled + in TestCase + (assertEqual "sort recovers initial numbers" testDataFrame sortedShuffled) + +-- Test that shuffling isn't doing anything sneaky with summoning +-- random numbers somehow +shuffleSameSeedIsSameShuffle :: Test +shuffleSameSeedIsSameShuffle = + let gen = mkStdGen 1234 + shuffled1 = shuffle gen testDataFrame + shuffled2 = shuffle gen testDataFrame + in TestCase + (assertEqual "shuffle with same seed gives same result" shuffled1 shuffled2) + +-- Test that different seeds give different results +shuffleDifferentSeedIsDifferent :: Test +shuffleDifferentSeedIsDifferent = + let gen1 = mkStdGen 1234 + gen2 = mkStdGen 4321 + shuffled1 = shuffle gen1 testDataFrame + shuffled2 = shuffle gen2 testDataFrame + in TestCase + ( assertEqual + "shuffle with different seeds gives different results" + False + (shuffled1 == shuffled2) + ) + +tests :: [Test] +tests = + [ TestLabel "shuffleShuffles" shuffleShuffles + , TestLabel "shufflePreservesData" shufflePreservesData + , TestLabel "shufflePreservesColumnNames" shufflePreservesColumnNames + , TestLabel "shuffleSameSeedIsSameShuffle" shuffleSameSeedIsSameShuffle + , TestLabel "shuffleDifferentSeedIsDifferent" shuffleDifferentSeedIsDifferent + ] diff --git a/tests/Operations/Sort.hs b/tests/Operations/Sort.hs --- a/tests/Operations/Sort.hs +++ b/tests/Operations/Sort.hs @@ -1,4 +1,5 @@ {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} module Operations.Sort where @@ -6,6 +7,7 @@ import Data.Char import qualified Data.Text as T import qualified DataFrame as D +import qualified DataFrame.Functions as F import qualified DataFrame.Internal.Column as DI import System.Random import System.Random.Shuffle (shuffle') @@ -40,7 +42,7 @@ , ("test2", DI.fromList ['a' .. 'z']) ] ) - (D.sortBy [D.Asc "test1"] testData) + (D.sortBy [D.Asc (F.col @Int "test1")] testData) ) sortByDescendingWAI :: Test @@ -53,7 +55,7 @@ , ("test2", DI.fromList $ reverse ['a' .. 'z']) ] ) - (D.sortBy [D.Desc "test1"] testData) + (D.sortBy [D.Desc (F.col @Int "test1")] testData) ) sortByTwoColumns :: Test @@ -62,7 +64,7 @@ ( assertEqual "Sorting moreTestData (which is already sorted) is idempotent." moreTestData - (D.sortBy [D.Asc "test1", D.Asc "test2"] moreTestData) + (D.sortBy [D.Asc (F.col @Int "test1"), D.Asc (F.col @Int "test2")] moreTestData) ) sortByOneColumnAscOneColumnDesc :: Test @@ -75,7 +77,7 @@ , ("test2", DI.fromList $ [10 :: Int, 9 .. 1] ++ [10, 9 .. 1]) ] ) - (D.sortBy [D.Asc "test1", D.Desc "test2"] moreTestData) + (D.sortBy [D.Asc (F.col @Int "test1"), D.Desc (F.col @Int "test2")] moreTestData) ) sortByColumnDoesNotExist :: Test @@ -84,7 +86,7 @@ ( assertExpectException "[Error Case]" (D.columnNotFound "[\"test0\"]" "sortBy" (D.columnNames testData)) - (print $ D.sortBy [D.Asc "test0"] testData) + (print $ D.sortBy [D.Asc (F.col @Int "test0")] testData) ) tests :: [Test] diff --git a/tests/Parquet.hs b/tests/Parquet.hs --- a/tests/Parquet.hs +++ b/tests/Parquet.hs @@ -40,14 +40,14 @@ , ( "timestamp_col" , D.fromList - [ UTCTime (fromGregorian 2009 3 1) (secondsToDiffTime 0) - , UTCTime (fromGregorian 2009 3 1) (secondsToDiffTime 60) - , UTCTime (fromGregorian 2009 4 1) (secondsToDiffTime 0) - , UTCTime (fromGregorian 2009 4 1) (secondsToDiffTime 60) - , UTCTime (fromGregorian 2009 2 1) (secondsToDiffTime 0) - , UTCTime (fromGregorian 2009 2 1) (secondsToDiffTime 60) - , UTCTime (fromGregorian 2009 1 1) (secondsToDiffTime 0) - , UTCTime (fromGregorian 2009 1 1) (secondsToDiffTime 60) + [ UTCTime{utctDay = fromGregorian 2009 3 1, utctDayTime = secondsToDiffTime 0} + , UTCTime{utctDay = fromGregorian 2009 3 1, utctDayTime = secondsToDiffTime 60} + , UTCTime{utctDay = fromGregorian 2009 4 1, utctDayTime = secondsToDiffTime 0} + , UTCTime{utctDay = fromGregorian 2009 4 1, utctDayTime = secondsToDiffTime 60} + , UTCTime{utctDay = fromGregorian 2009 2 1, utctDayTime = secondsToDiffTime 0} + , UTCTime{utctDay = fromGregorian 2009 2 1, utctDayTime = secondsToDiffTime 60} + , UTCTime{utctDay = fromGregorian 2009 1 1, utctDayTime = secondsToDiffTime 0} + , UTCTime{utctDay = fromGregorian 2009 1 1, utctDayTime = secondsToDiffTime 60} ] ) ] @@ -61,6 +61,113 @@ (unsafePerformIO (D.readParquet "./tests/data/alltypes_plain.parquet")) ) +allTypesTinyPagesDimensions :: Test +allTypesTinyPagesDimensions = + TestCase + ( assertEqual + "allTypesTinyPages last few" + (7300, 13) + ( unsafePerformIO + (fmap D.dimensions (D.readParquet "./tests/data/alltypes_tiny_pages.parquet")) + ) + ) + +tinyPagesLast10 :: D.DataFrame +tinyPagesLast10 = + D.fromNamedColumns + [ ("id", D.fromList @Int32 (reverse [6174 .. 6183])) + , ("bool_col", D.fromList @Bool (Prelude.take 10 (cycle [False, True]))) + , ("tinyint_col", D.fromList @Int32 [3, 2, 1, 0, 9, 8, 7, 6, 5, 4]) + , ("smallint_col", D.fromList @Int32 [3, 2, 1, 0, 9, 8, 7, 6, 5, 4]) + , ("int_col", D.fromList @Int32 [3, 2, 1, 0, 9, 8, 7, 6, 5, 4]) + , ("bigint_col", D.fromList @Int64 [30, 20, 10, 0, 90, 80, 70, 60, 50, 40]) + , + ( "float_col" + , D.fromList @Float [3.3, 2.2, 1.1, 0, 9.9, 8.8, 7.7, 6.6, 5.5, 4.4] + ) + , + ( "date_string_col" + , D.fromList @Text + [ "09/11/10" + , "09/11/10" + , "09/11/10" + , "09/11/10" + , "09/10/10" + , "09/10/10" + , "09/10/10" + , "09/10/10" + , "09/10/10" + , "09/10/10" + ] + ) + , + ( "string_col" + , D.fromList @Text ["3", "2", "1", "0", "9", "8", "7", "6", "5", "4"] + ) + , + ( "timestamp_col" + , D.fromList @UTCTime + [ UTCTime + { utctDay = fromGregorian 2010 9 10 + , utctDayTime = secondsToDiffTime 85384 + } + , UTCTime + { utctDay = fromGregorian 2010 9 10 + , utctDayTime = secondsToDiffTime 85324 + } + , UTCTime + { utctDay = fromGregorian 2010 9 10 + , utctDayTime = secondsToDiffTime 85264 + } + , UTCTime + { utctDay = fromGregorian 2010 9 10 + , utctDayTime = secondsToDiffTime 85204 + } + , UTCTime + { utctDay = fromGregorian 2010 9 9 + , utctDayTime = secondsToDiffTime 85144 + } + , UTCTime + { utctDay = fromGregorian 2010 9 9 + , utctDayTime = secondsToDiffTime 85084 + } + , UTCTime + { utctDay = fromGregorian 2010 9 9 + , utctDayTime = secondsToDiffTime 85024 + } + , UTCTime + { utctDay = fromGregorian 2010 9 9 + , utctDayTime = secondsToDiffTime 84964 + } + , UTCTime + { utctDay = fromGregorian 2010 9 9 + , utctDayTime = secondsToDiffTime 84904 + } + , UTCTime + { utctDay = fromGregorian 2010 9 9 + , utctDayTime = secondsToDiffTime 84844 + } + ] + ) + , ("year", D.fromList @Int32 (replicate 10 2010)) + , ("month", D.fromList @Int32 (replicate 10 9)) + ] + +allTypesTinyPagesLastFew :: Test +allTypesTinyPagesLastFew = + TestCase + ( assertEqual + "allTypesTinyPages dimensions" + tinyPagesLast10 + ( unsafePerformIO + -- Excluding doubles because they are weird to compare. + ( fmap + (D.takeLast 10 . D.exclude ["double_col"]) + (D.readParquet "./tests/data/alltypes_tiny_pages.parquet") + ) + ) + ) + allTypesPlainSnappy :: Test allTypesPlainSnappy = TestCase @@ -79,6 +186,176 @@ (unsafePerformIO (D.readParquet "./tests/data/alltypes_dictionary.parquet")) ) +transactions :: D.DataFrame +transactions = + D.fromNamedColumns + [ ("transaction_id", D.fromList [1 :: Int32, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]) + , + ( "event_time" + , D.fromList + [ UTCTime + { utctDay = fromGregorian 2024 1 3 + , utctDayTime = secondsToDiffTime 29564 + picosecondsToDiffTime 2311000000 + } + , UTCTime + { utctDay = fromGregorian 2024 1 3 + , utctDayTime = secondsToDiffTime 35101 + picosecondsToDiffTime 118900000000 + } + , UTCTime + { utctDay = fromGregorian 2024 1 4 + , utctDayTime = secondsToDiffTime 39802 + picosecondsToDiffTime 774512000000 + } + , UTCTime + { utctDay = fromGregorian 2024 1 5 + , utctDayTime = secondsToDiffTime 53739 + picosecondsToDiffTime 1000000 + } + , UTCTime + { utctDay = fromGregorian 2024 1 6 + , utctDayTime = secondsToDiffTime 8278 + picosecondsToDiffTime 543210000000 + } + , UTCTime + { utctDay = fromGregorian 2024 1 6 + , utctDayTime = secondsToDiffTime 8284 + picosecondsToDiffTime 211000000000 + } + , UTCTime + { utctDay = fromGregorian 2024 1 7 + , utctDayTime = secondsToDiffTime 63000 + } + , UTCTime + { utctDay = fromGregorian 2024 1 8 + , utctDayTime = secondsToDiffTime 24259 + picosecondsToDiffTime 390000000000 + } + , UTCTime + { utctDay = fromGregorian 2024 1 9 + , utctDayTime = secondsToDiffTime 48067 + picosecondsToDiffTime 812345000000 + } + , UTCTime + { utctDay = fromGregorian 2024 1 10 + , utctDayTime = secondsToDiffTime 82799 + picosecondsToDiffTime 999999000000 + } + , UTCTime + { utctDay = fromGregorian 2024 1 11 + , utctDayTime = secondsToDiffTime 36000 + picosecondsToDiffTime 100000000000 + } + , UTCTime + { utctDay = fromGregorian 2024 1 12 + , utctDayTime = secondsToDiffTime 56028 + picosecondsToDiffTime 667891000000 + } + ] + ) + , + ( "user_email" + , D.fromList + [ "alice@example.com" :: Text + , "bob@example.com" + , "carol@example.com" + , "alice@example.com" + , "dave@example.com" + , "dave@example.com" + , "eve@example.com" + , "frank@example.com" + , "grace@example.com" + , "dave@example.com" + , "alice@example.com" + , "heidi@example.com" + ] + ) + , + ( "transaction_type" + , D.fromList + [ "purchase" :: Text + , "purchase" + , "refund" + , "purchase" + , "purchase" + , "purchase" + , "purchase" + , "withdrawal" + , "purchase" + , "purchase" + , "purchase" + , "refund" + ] + ) + , + ( "amount" + , D.fromList + [ 142.50 :: Double + , 29.99 + , 89.00 + , 2399.00 + , 15.00 + , 15.00 + , 450.75 + , 200.00 + , 55.20 + , 3200.00 + , 74.99 + , 120.00 + ] + ) + , + ( "currency" + , D.fromList + [ "USD" :: Text + , "USD" + , "EUR" + , "USD" + , "GBP" + , "GBP" + , "USD" + , "EUR" + , "CAD" + , "USD" + , "USD" + , "GBP" + ] + ) + , + ( "status" + , D.fromList + [ "approved" :: Text + , "approved" + , "approved" + , "declined" + , "approved" + , "declined" + , "approved" + , "approved" + , "approved" + , "flagged" + , "approved" + , "approved" + ] + ) + , + ( "location" + , D.fromList + [ "New York, US" :: Text + , "London, GB" + , "Berlin, DE" + , "New York, US" + , "Manchester, GB" + , "Lagos, NG" + , "San Francisco, US" + , "Paris, FR" + , "Toronto, CA" + , "New York, US" + , "New York, US" + , "Edinburgh, GB" + ] + ) + ] + +transactionsTest :: Test +transactionsTest = + TestCase + ( assertEqual + "transactions" + transactions + (unsafePerformIO (D.readParquet "./tests/data/transactions.parquet")) + ) + mtCarsDataset :: D.DataFrame mtCarsDataset = D.fromNamedColumns @@ -537,7 +814,13 @@ (unsafePerformIO (D.readParquet "./tests/data/mtcars.parquet")) ) --- Uncomment to run parquet tests. --- Currently commented because they don't run with github CI tests :: [Test] -tests = [allTypesPlain, allTypesPlainSnappy, allTypesDictionary, mtCars] +tests = + [ allTypesPlain + , allTypesPlainSnappy + , allTypesDictionary + , mtCars + , allTypesTinyPagesLastFew + , allTypesTinyPagesDimensions + , transactionsTest + ] diff --git a/tests/data/transactions.parquet b/tests/data/transactions.parquet new file mode 100644 Binary files /dev/null and b/tests/data/transactions.parquet differ