dataframe 0.4.1.0 → 0.5.0.0

raw patch · 40 files changed

+2003/−1019 lines, 40 filesdep +thesedep ~dataframedep ~randombinary-addedPVP ok

version bump matches the API change (PVP)

Dependencies added: these

Dependency ranges changed: dataframe, random

API changes (from Hackage documentation)

Files

• CHANGELOG.md +13/−0
• app/Synthesis.hs +1/−2
• benchmark/Main.hs +3/−3
• dataframe.cabal +16/−10
• src/DataFrame.hs +7/−9
• src/DataFrame/DecisionTree.hs +3/−3
• src/DataFrame/Display/Terminal/PrettyPrint.hs +1/−2
• src/DataFrame/Functions.hs +186/−125
• src/DataFrame/IO/CSV.hs +9/−2
• src/DataFrame/IO/Parquet.hs +55/−13
• src/DataFrame/IO/Parquet/Binary.hs +48/−32
• src/DataFrame/IO/Parquet/ColumnStatistics.hs +7/−7
• src/DataFrame/IO/Parquet/Dictionary.hs +50/−41
• src/DataFrame/IO/Parquet/Encoding.hs +30/−18
• src/DataFrame/IO/Parquet/Levels.hs +18/−10
• src/DataFrame/IO/Parquet/Page.hs +291/−262
• src/DataFrame/IO/Parquet/Thrift.hs +45/−26
• src/DataFrame/IO/Parquet/Time.hs +6/−5
• src/DataFrame/IO/Parquet/Types.hs +12/−12
• src/DataFrame/Internal/Column.hs +90/−6
• src/DataFrame/Internal/Expression.hs +190/−168
• src/DataFrame/Internal/Interpreter.hs +81/−188
• src/DataFrame/Operations/Aggregation.hs +7/−1
• src/DataFrame/Operations/Core.hs +3/−2
• src/DataFrame/Operations/Join.hs +24/−3
• src/DataFrame/Operations/Permutation.hs +18/−7
• src/DataFrame/Operations/Statistics.hs +7/−12
• src/DataFrame/Operations/Subset.hs +61/−11
• src/DataFrame/Operations/Transformations.hs +2/−2
• src/DataFrame/Operators.hs +161/−0
• src/DataFrame/Synthesis.hs +1/−1
• tests/GenDataFrame.hs +2/−2
• tests/Main.hs +8/−1
• tests/Monad.hs +27/−0
• tests/Operations/Aggregations.hs +34/−12
• tests/Operations/Join.hs +101/−5
• tests/Operations/Shuffle.hs +84/−0
• tests/Operations/Sort.hs +7/−5
• tests/Parquet.hs +294/−11
• tests/data/transactions.parquet binary

@@ -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).

@@ -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 

@@ -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 ()

@@ -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

@@ -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

@@ -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 ::

@@ -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

@@ -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)

@@ -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

@@ -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)

@@ -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)

@@ -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

@@ -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"

@@ -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)

@@ -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, _) =

@@ -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'

@@ -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

@@ -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 =

@@ -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) 

@@ -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 ::

@@ -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 ++ ")"

@@ -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

@@ -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

@@ -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' ->

@@ -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                         )             ) -- ???

@@ -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)

@@ -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 ::

@@ -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))                             )

@@ -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

@@ -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)

@@ -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) 

@@ -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 ..]

@@ -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

@@ -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]

@@ -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     ]

@@ -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     ]

@@ -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+    ]

@@ -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]

@@ -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+    ]

binary file changed (absent → 1746 bytes)