diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Revision history for dataframe +## 0.3.0.3 +* Improved parquet reader. The reader now supports most parquet files found in the wild. +* Improve CSV parsing: Parse bytestring and convert to text only at the end. Remove some redundancies in parsing with suggestions from @Jhingon. +* Faster correlation computation. +* Update version of granite that ships with dataframe and add new scatterBy plot. + ## 0.3.0.2 * Re-enable Parquet. * Change columnInfo to describeColumns diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -117,7 +117,26 @@ ## Supported input formats * CSV -* Apache Parquet (still buggy and experimental) +* Apache Parquet + * Supports all primitive parquet types plain and uncompressed. + * Can decode both v1 and v2 data pages. + * Supports Snappy and ZSTD compression. + * Supports RLE/bitpacking encoding for primitive types + * Backward compatible with INT96 type. + * From the parquet-testing repo we can successfully read the following: + * alltypes_dictionary.parquet + * alltypes_plain.parquet + * alltypes_plain.snappy.parquet + * alltypes_tiny_pages_plain.parquet + * binary_truncated_min_max.parquet + * datapage_v1-corrupt-checksum.parquet + * datapage_v1-snappy-compressed-checksum.parquet + * datapage_v1-uncompressed-checksum.parquet + * We will continue adding functionality/coverage as needed. + + +## Supported output formats +* CSV ## Future work * Apache arrow compatability diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -1,17 +1,18 @@ - --- Useful Haskell extensions. -{-# LANGUAGE OverloadedStrings #-} -- Allow string literal to be interpreted as any other string type. -{-# LANGUAGE TypeApplications #-} -- Convenience syntax for specifiying the type `sum a b :: Int` vs `sum @Int a b'. {-# LANGUAGE NumericUnderscores #-} +-- Useful Haskell extensions. +-- Allow string literal to be interpreted as any other string type. +{-# LANGUAGE OverloadedStrings #-} +-- Convenience syntax for specifiying the type `sum a b :: Int` vs `sum @Int a b'. +{-# LANGUAGE TypeApplications #-} import qualified DataFrame as D -- import for general functionality. import qualified DataFrame.Functions as F -- import for column expressions. import DataFrame ((|>)) -- import chaining operator with unqualified. -import qualified Data.Vector.Unboxed as VU import Control.Monad (replicateM) import Data.Time +import qualified Data.Vector.Unboxed as VU import System.Random.Stateful main :: IO () diff --git a/benchmark/Main.hs b/benchmark/Main.hs --- a/benchmark/Main.hs +++ b/benchmark/Main.hs @@ -2,59 +2,65 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeApplications #-} -import qualified DataFrame as D -import qualified DataFrame.Functions as F import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM +import qualified DataFrame as D +import qualified DataFrame.Functions as F import Control.Monad (replicateM) import Criterion.Main -import DataFrame ((|>)) import Data.Time +import DataFrame ((|>)) import System.Process import System.Random.Stateful haskell :: IO () haskell = do - output <- readProcess "cabal" ["run", "dataframe"] "" - putStrLn output + output <- readProcess "cabal" ["run", "dataframe", "-O2"] "" + putStrLn output polars :: IO () polars = do - output <- readProcess "./benchmark/dataframe_benchmark/bin/python3" ["./benchmark/polars/polars_benchmark.py"] "" - putStrLn output + output <- readProcess "./benchmark/dataframe_benchmark/bin/python3" ["./benchmark/polars/polars_benchmark.py"] "" + putStrLn output pandas :: IO () pandas = do - output <- readProcess "./benchmark/dataframe_benchmark/bin/python3" ["./benchmark/pandas/pandas_benchmark.py"] "" - putStrLn output + output <- readProcess "./benchmark/dataframe_benchmark/bin/python3" ["./benchmark/pandas/pandas_benchmark.py"] "" + putStrLn output groupByHaskell :: IO () groupByHaskell = do - df <- D.readCsv "./data/housing.csv" - print $ df |> D.groupBy ["ocean_proximity"] - |> D.aggregate [ (F.minimum (F.col @Double "median_house_value")) `F.as` "minimum_median_house_value" - , (F.maximum (F.col @Double "median_house_value")) `F.as` "maximum_median_house_value"] + df <- D.readCsv "./data/housing.csv" + print $ + df + |> D.groupBy ["ocean_proximity"] + |> D.aggregate + [ (F.minimum (F.col @Double "median_house_value")) `F.as` "minimum_median_house_value" + , (F.maximum (F.col @Double "median_house_value")) `F.as` "maximum_median_house_value" + ] groupByPolars :: IO () groupByPolars = do - output <- readProcess "./benchmark/dataframe_benchmark/bin/python3" ["./benchmark/polars/group_by.py"] "" - putStrLn output + output <- readProcess "./benchmark/dataframe_benchmark/bin/python3" ["./benchmark/polars/group_by.py"] "" + putStrLn output groupByPandas :: IO () groupByPandas = do - output <- readProcess "./benchmark/dataframe_benchmark/bin/python3" ["./benchmark/pandas/group_by.py"] "" - putStrLn output + output <- readProcess "./benchmark/dataframe_benchmark/bin/python3" ["./benchmark/pandas/group_by.py"] "" + putStrLn output main = do - output <- readProcess "cabal" ["build", "-O2"] "" - putStrLn output - defaultMain [ - bgroup "stats" [ bench "simpleStatsHaskell" $ nfIO haskell - , bench "simpleStatsPandas" $ nfIO pandas - , bench "simpleStatsPolars" $ nfIO polars - , bench "groupByHaskell" $ nfIO groupByHaskell - , bench "groupByPolars" $ nfIO groupByPolars - , bench "groupByPandas" $ nfIO groupByPandas - ] - ] + output <- readProcess "cabal" ["build", "-O2"] "" + putStrLn output + defaultMain + [ bgroup + "stats" + [ bench "simpleStatsHaskell" $ nfIO haskell + , bench "simpleStatsPandas" $ nfIO pandas + , bench "simpleStatsPolars" $ nfIO polars + , bench "groupByHaskell" $ nfIO groupByHaskell + , bench "groupByPolars" $ nfIO groupByPolars + , bench "groupByPandas" $ nfIO groupByPandas + ] + ] diff --git a/dataframe.cabal b/dataframe.cabal --- a/dataframe.cabal +++ b/dataframe.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: dataframe -version: 0.3.0.2 +version: 0.3.0.3 synopsis: A fast, safe, and intuitive DataFrame library. @@ -22,7 +22,7 @@ location: https://github.com/mchav/dataframe library - -- default-extensions: StrictData + default-extensions: Strict exposed-modules: DataFrame, DataFrame.Lazy, DataFrame.Functions @@ -48,10 +48,15 @@ DataFrame.Display.Terminal.Plot, DataFrame.IO.CSV, DataFrame.IO.Parquet, + DataFrame.IO.Parquet.Binary + DataFrame.IO.Parquet.Dictionary + DataFrame.IO.Parquet.Levels + DataFrame.IO.Parquet.Thrift DataFrame.IO.Parquet.ColumnStatistics, DataFrame.IO.Parquet.Compression, DataFrame.IO.Parquet.Encoding, DataFrame.IO.Parquet.Page, + DataFrame.IO.Parquet.Time, DataFrame.IO.Parquet.Types, DataFrame.Lazy.IO.CSV, DataFrame.Lazy.Internal.DataFrame @@ -59,10 +64,11 @@ array ^>= 0.5, attoparsec >= 0.12 && <= 0.14.4, bytestring >= 0.11 && <= 0.12.2.0, + bytestring-lexing >= 0.5 && < 0.6, containers >= 0.6.7 && < 0.8, directory >= 1.3.0.0 && <= 1.3.9.0, filepath >= 1.0.0.0 && <= 1.5.4.0, - granite ^>= 0.1.0.2, + granite ^>= 0.2, hashable >= 1.2 && <= 1.5.0.0, snappy-hs ^>= 0.1, statistics >= 0.16.2.1 && <= 0.16.3.0, @@ -97,10 +103,15 @@ DataFrame.Display.Terminal.Plot, DataFrame.IO.CSV, DataFrame.IO.Parquet, + DataFrame.IO.Parquet.Binary + DataFrame.IO.Parquet.Dictionary + DataFrame.IO.Parquet.Levels + DataFrame.IO.Parquet.Thrift DataFrame.IO.Parquet.ColumnStatistics, DataFrame.IO.Parquet.Compression, DataFrame.IO.Parquet.Encoding, DataFrame.IO.Parquet.Page, + DataFrame.IO.Parquet.Time, DataFrame.IO.Parquet.Types, DataFrame.Operations.Join, DataFrame.Operations.Merge, @@ -110,9 +121,10 @@ array ^>= 0.5, attoparsec >= 0.12 && <= 0.14.4, bytestring >= 0.11 && <= 0.12.2.0, + bytestring-lexing >= 0.5 && < 0.6, containers >= 0.6.7 && < 0.8, directory >= 1.3.0.0 && <= 1.3.9.0, - granite ^>= 0.1.0.2, + granite ^>= 0.2, hashable >= 1.2 && <= 1.5.0.0, random >= 1 && <= 1.3.1, snappy-hs ^>= 0.1, @@ -152,16 +164,67 @@ Operations.Filter, Operations.GroupBy, Operations.InsertColumn, - Operations.Sort, + Operations.Sort, + Operations.Statistics, Operations.Take, - Parquet - build-depends: base >= 4.17.2.0 && < 4.22, - HUnit ^>= 1.6, - random >= 1, - random-shuffle >= 0.0.4, - text >= 2.0, - time >= 1.12, - vector ^>= 0.13, - dataframe - hs-source-dirs: tests + Parquet, + + -- Dataframe modules. + DataFrame, + DataFrame.Internal.Types, + DataFrame.Internal.Expression, + DataFrame.Internal.Parsing, + DataFrame.Internal.Column, + DataFrame.Display.Terminal.PrettyPrint, + DataFrame.Display.Terminal.Colours, + DataFrame.Internal.DataFrame, + DataFrame.Internal.Row, + DataFrame.Errors, + DataFrame.Operations.Core, + DataFrame.Operations.Subset, + DataFrame.Operations.Sorting, + DataFrame.Operations.Statistics, + DataFrame.Operations.Transformations, + DataFrame.Operations.Typing, + DataFrame.Operations.Aggregation, + DataFrame.Display.Terminal.Plot, + DataFrame.IO.CSV, + DataFrame.IO.Parquet, + DataFrame.IO.Parquet.Binary + DataFrame.IO.Parquet.Dictionary + DataFrame.IO.Parquet.Levels + DataFrame.IO.Parquet.Thrift + DataFrame.IO.Parquet.ColumnStatistics, + DataFrame.IO.Parquet.Compression, + DataFrame.IO.Parquet.Encoding, + DataFrame.IO.Parquet.Page, + DataFrame.IO.Parquet.Time, + DataFrame.IO.Parquet.Types, + DataFrame.Operations.Join, + DataFrame.Operations.Merge, + DataFrame.Lazy.IO.CSV, + DataFrame.Functions + build-depends: base >= 4.17.2.0 && < 4.22, + array ^>= 0.5, + attoparsec >= 0.12 && <= 0.14.4, + bytestring >= 0.11 && <= 0.12.2.0, + bytestring-lexing >= 0.5 && < 0.6, + containers >= 0.6.7 && < 0.8, + directory >= 1.3.0.0 && <= 1.3.9.0, + granite ^>= 0.2, + hashable >= 1.2 && <= 1.5.0.0, + HUnit ^>= 1.6, + random >= 1, + random-shuffle >= 0.0.4, + random >= 1 && <= 1.3.1, + snappy-hs ^>= 0.1, + statistics >= 0.16.2.1 && <= 0.16.3.0, + template-haskell >= 2.0 && <= 2.30, + text >= 2.0 && <= 2.1.2, + time >= 1.12 && <= 1.14, + vector ^>= 0.13, + vector-algorithms ^>= 0.9, + zstd >= 0.1.2.0 && <= 0.1.3.0 + hs-source-dirs: tests, + src default-language: Haskell2010 diff --git a/src/DataFrame.hs b/src/DataFrame.hs --- a/src/DataFrame.hs +++ b/src/DataFrame.hs @@ -1,30 +1,261 @@ {-# LANGUAGE OverloadedStrings #-} -module DataFrame - ( module D, + +{- | +Module : DataFrame +Copyright : (c) 2025 +License : GPL-3.0 +Maintainer : mschavinda@gmail.com +Stability : experimental +Portability : POSIX + +Batteries-included entry point for the DataFrame library. + +This module re-exports the most commonly used pieces of the @dataframe@ library so you +can get productive fast in GHCi, IHaskell, or scripts. + +__Naming convention__ +* Use the @D.@ (\"DataFrame\") prefix for core table operations. +* Use the @F.@ (\"Functions\") prefix for the expression DSL (columns, math, aggregations). + +Example session: + +@ +-- GHCi quality-of-life: +ghci> :set -XOverloadedStrings -XTypeApplications +ghci> :module + DataFrame as D, DataFrame.Functions as F, Data.Text (Text) +@ + += Quick start +Load a CSV, select a few columns, filter, derive a column, then group + aggregate: + +@ +-- 1) Load data +ghci> df0 <- D.readCsv "data/housing.csv" +ghci> D.describeColumns df0 +-------------------------------------------------------------------------------------------------------------------- +index | Column Name | # Non-null Values | # Null Values | # Partially parsed | # Unique Values | Type +------|--------------------|-------------------|---------------|--------------------|-----------------|------------- + Int | Text | Int | Int | Int | Int | Text +------|--------------------|-------------------|---------------|--------------------|-----------------|------------- +0 | ocean_proximity | 20640 | 0 | 0 | 5 | Text +1 | median_house_value | 20640 | 0 | 0 | 3842 | Double +2 | median_income | 20640 | 0 | 0 | 12928 | Double +3 | households | 20640 | 0 | 0 | 1815 | Double +4 | population | 20640 | 0 | 0 | 3888 | Double +5 | total_bedrooms | 20640 | 0 | 0 | 1924 | Maybe Double +6 | total_rooms | 20640 | 0 | 0 | 5926 | Double +7 | housing_median_age | 20640 | 0 | 0 | 52 | Double +8 | latitude | 20640 | 0 | 0 | 862 | Double +9 | longitude | 20640 | 0 | 0 | 844 | Double + +-- 2) Project & filter +ghci> let df1 = df1 = D.filter @Text "ocean_proximity" (== "ISLAND") df0 D.|> D.select ["median_house_value", "median_income", "ocean_proximity"] + +-- 3) Add a derived column using the expression DSL +-- (col types are explicit via TypeApplications) +ghci> df2 = D.derive "rooms_per_household" (F.col @Double "total_rooms" / F.col @Double "households") df0 + +-- 4) Group + aggregate +ghci> let grouped = D.groupBy ["ocean_proximity"] df0 +ghci> let summary = + D.aggregate + [ F.maximum (F.col @Double "median_house_value") `F.as` "max_house_value"] + grouped +ghci> D.take 5 summary +----------------------------------------- +index | ocean_proximity | max_house_value +------|-----------------|---------------- + Int | Text | Double +------|-----------------|---------------- +0 | <1H OCEAN | 500001.0 +1 | INLAND | 500001.0 +2 | ISLAND | 450000.0 +3 | NEAR BAY | 500001.0 +4 | NEAR OCEAN | 500001.0 +@ + +== Simple operations (cheat sheet) +Most users only need a handful of verbs: + +__I/O__ + + * @D.readCsv :: FilePath -> IO DataFrame@ + * @D.writeCsv :: FilePath -> DataFrame -> IO ()@ + * @D.readParquet :: FilePath -> IO DataFrame@ + +__Exploration__ + + * @D.take :: Int -> DataFrame -> DataFrame@ + * @D.takeLast :: Int -> DataFrame -> DataFrame@ + * @D.describeColumns :: DataFrame -> DataFrame@ + * @D.summarize :: DataFrame -> DataFrame@ + +__Row ops__ + + * @D.filter :: Columnable a => Text -> (a -> Bool) -> DataFrame -> DataFrame@ + * @D.sortBy :: SortOrder -> [Text] -> DataFrame -> DataFrame@ + +__Column ops__ + + * @D.select :: [Text] -> DataFrame -> DataFrame@ + * @D.exclude :: [Text] -> DataFrame -> DataFrame@ + * @D.rename :: [(Text,Text)] -> DataFrame -> DataFrame@ + * @D.derive :: Text -> D.Expr a -> DataFrame -> DataFrame@ + +__Group & aggregate__ + + * @D.groupBy :: [Text] -> DataFrame -> GroupedDataFrame@ + * @D.aggregate :: [(Text, F.UExpr)] -> GroupedDataFrame -> DataFrame@ + +__Joins__ + + * @D.innerJoin / D.leftJoin / D.rightJoin / D.fullJoin@ + +== Expression DSL (F.*) at a glance +Columns (typed): + +@ +F.col @Text "ocean_proximity" +F.col @Double "total_rooms" +F.lit @Double 1.0 +@ + +Math & comparisons (overloaded by type): + +@ +(+), (-), (*), (/), abs, log, exp, round +(F.eq), (F.gt), (F.geq), (F.lt), (F.leq) +@ + +Aggregations (for 'D.aggregate'): + +@ +F.count @a (F.col @a "c") +F.sum @Double (F.col @Double "x") +F.mean @Double (F.col @Double "x") +F.min @t (F.col @t "x") +F.max @t (F.col @t "x") +@ + +== REPL power-tool: ':exposeColumns' + +Use @:exposeColumns @ in GHCi/IHaskell to turn each column of a bound 'DataFrame' +into a local binding with the same (mangled if needed) name and the column's concrete +vector type. This is great for quick ad-hoc analysis, plotting, or hand-rolled checks. + +@ +-- Suppose df has columns: "passengers" :: Int, "fare" :: Double, "payment" :: Text +ghci> :set -XTemplateHaskell +ghci> :exposeColumns df + +-- Now you have in scope: +ghci> :type passengers +passengers :: Expr Int + +ghci> :type fare +fare :: Expr Double + +ghci> :type payment +payment :: Expr Text + +-- You can use them directly: +ghci> D.derive "fare_with_tip" (fare * F.lit 1.2) +@ + +Notes: + +* Name mangling: spaces and non-identifier characters are replaced (e.g. @"trip id"@ -> @trip_id@). +* Optional/nullable columns are exposed as @Expr (Maybe a)@. +-} +module DataFrame ( + -- * Core data structures + module Dataframe, + module Column, + module Expression, + + -- * Core dataframe operations + module Core, + + -- * I/O + module CSV, + module Parquet, + + -- * Operations + module Subset, + module Transformations, + module Aggregation, + module Sorting, + module Merge, + module Join, + module Statistics, + + -- * Errors + module Errors, + + -- * Plotting + module Plot, + + -- * Convenience functions (|>), - printSessionSchema - ) +) where -import DataFrame.Internal.Types as D -import DataFrame.Internal.Expression as D -import DataFrame.Internal.Parsing as D -import DataFrame.Internal.Column as D -import DataFrame.Internal.DataFrame as D hiding (columnIndices, columns) -import DataFrame.Internal.Row as D hiding (mkRowRep) -import DataFrame.Errors as D -import DataFrame.Operations.Core as D -import DataFrame.Operations.Merge as D -import DataFrame.Operations.Join as D -import DataFrame.Operations.Subset as D -import DataFrame.Operations.Sorting as D -import DataFrame.Operations.Statistics as D -import DataFrame.Operations.Transformations as D -import DataFrame.Operations.Typing as D -import DataFrame.Operations.Aggregation as D -import DataFrame.Display.Terminal.Plot as D -import DataFrame.IO.CSV as D -import DataFrame.IO.Parquet as D +import DataFrame.Display.Terminal.Plot as Plot +import DataFrame.Errors as Errors +import DataFrame.IO.CSV as CSV (ReadOptions (..), defaultOptions, readCsv, readSeparated, readTsv) +import DataFrame.IO.Parquet as Parquet (readParquet) +import DataFrame.Internal.Column as Column ( + Column, + fromList, + fromUnboxedVector, + fromVector, + toList, + toVector, + ) +import DataFrame.Internal.DataFrame as Dataframe ( + DataFrame, + GroupedDataFrame, + columnAsVector, + empty, + toMatrix, + ) +import DataFrame.Internal.Expression as Expression (Expr) +import DataFrame.Operations.Aggregation as Aggregation (aggregate, distinct, groupBy) +import DataFrame.Operations.Core as Core hiding (ColumnInfo (..), nulls, partiallyParsed, renameSafe) +import DataFrame.Operations.Join as Join +import DataFrame.Operations.Merge as Merge +import DataFrame.Operations.Sorting as Sorting +import DataFrame.Operations.Statistics as Statistics ( + correlation, + frequencies, + interQuartileRange, + mean, + median, + skewness, + standardDeviation, + sum, + summarize, + variance, + ) +import DataFrame.Operations.Subset as Subset ( + cube, + drop, + dropLast, + exclude, + filter, + filterAllJust, + filterBy, + filterJust, + filterWhere, + range, + select, + selectBy, + selectIntRange, + selectRange, + take, + takeLast, + ) +import DataFrame.Operations.Transformations as Transformations import Data.Function import Data.List @@ -32,16 +263,3 @@ import qualified Data.Text.IO as T (|>) = (&) - -printSessionSchema :: D.DataFrame -> IO () -printSessionSchema df = T.putStrLn $ let - (lhs, rhs) = foldr go ([], []) (D.columnNames df) - columnRep name = let - colType = T.pack (D.columnTypeString (D.unsafeGetColumn name df)) - in "F.col @(" <> colType <> ") \"" <> name <> "\"" - go name (l, r) = (T.toLower name:l, columnRep name:r) - in T.unlines [":{", "{-# LANGUAGE TypeApplications #-}", - "import qualified DataFrame.Functions as F", - "import Data.Text (Text)", - "(" <> T.intercalate "," lhs <> ")" <> " = " <> "(" <> T.intercalate "," rhs <> ")", - ":}"] diff --git a/src/DataFrame/Display/Terminal/Plot.hs b/src/DataFrame/Display/Terminal/Plot.hs --- a/src/DataFrame/Display/Terminal/Plot.hs +++ b/src/DataFrame/Display/Terminal/Plot.hs @@ -1,563 +1,640 @@ {-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE ExplicitNamespaces #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE GADTs #-} - -module DataFrame.Display.Terminal.Plot where - -import Control.Monad -import qualified Data.List as L -import qualified Data.Map as M -import qualified Data.Text as T -import qualified Data.Vector as V -import qualified Data.Vector.Generic as VG -import qualified Data.Vector.Unboxed as VU -import Data.Maybe (fromMaybe, catMaybes) -import Data.Typeable (Typeable) -import Data.Type.Equality (type (:~:)(Refl), TestEquality(testEquality)) -import Type.Reflection (typeRep) -import GHC.Stack (HasCallStack) - -import DataFrame.Internal.Column (Column(..), Columnable) -import DataFrame.Internal.DataFrame (DataFrame(..)) -import DataFrame.Operations.Core -import Granite - -data PlotConfig = PlotConfig - { plotType :: PlotType - , plotTitle :: String - , plotSettings :: Plot - } - -data PlotType - = Histogram' - | Scatter' - | Line' - | Bar' - | BoxPlot' - | Pie' - | StackedBar' - | Heatmap' - deriving (Eq, Show) - -defaultPlotConfig :: PlotType -> PlotConfig -defaultPlotConfig ptype = PlotConfig - { plotType = ptype - , plotTitle = "" - , plotSettings = defPlot - } - -plotHistogram :: HasCallStack => T.Text -> DataFrame -> IO () -plotHistogram colName df = plotHistogramWith colName (defaultPlotConfig Histogram') df - -plotHistogramWith :: HasCallStack => T.Text -> PlotConfig -> DataFrame -> IO () -plotHistogramWith colName config df = do - let values = extractNumericColumn colName df - (minVal, maxVal) = if null values then (0, 1) else (minimum values, maximum values) - putStrLn $ histogram (plotTitle config) (bins 30 minVal maxVal) values (plotSettings config) - -plotScatter :: HasCallStack => T.Text -> T.Text -> DataFrame -> IO () -plotScatter xCol yCol df = plotScatterWith xCol yCol (defaultPlotConfig Scatter') df - -plotScatterWith :: HasCallStack => T.Text -> T.Text -> PlotConfig -> DataFrame -> IO () -plotScatterWith xCol yCol config df = do - let xVals = extractNumericColumn xCol df - yVals = extractNumericColumn yCol df - points = zip xVals yVals - putStrLn $ scatter (plotTitle config) [(T.unpack xCol ++ " vs " ++ T.unpack yCol, points)] (plotSettings config) - -plotLines :: HasCallStack => [T.Text] -> DataFrame -> IO () -plotLines colNames df = plotLinesWith colNames (defaultPlotConfig Line') df - -plotLinesWith :: HasCallStack => [T.Text] -> PlotConfig -> DataFrame -> IO () -plotLinesWith colNames config df = do - seriesData <- forM colNames $ \col -> do - let values = extractNumericColumn col df - indices = map fromIntegral [0..length values - 1] - return (T.unpack col, zip indices values) - putStrLn $ lineGraph (plotTitle config) seriesData (plotSettings config) - -plotBoxPlots :: HasCallStack => [T.Text] -> DataFrame -> IO () -plotBoxPlots colNames df = plotBoxPlotsWith colNames (defaultPlotConfig BoxPlot') df - -plotBoxPlotsWith :: HasCallStack => [T.Text] -> PlotConfig -> DataFrame -> IO () -plotBoxPlotsWith colNames config df = do - boxData <- forM colNames $ \col -> do - let values = extractNumericColumn col df - return (T.unpack col, values) - putStrLn $ boxPlot (plotTitle config) boxData (plotSettings config) - -plotStackedBars :: HasCallStack => T.Text -> [T.Text] -> DataFrame -> IO () -plotStackedBars categoryCol valueColumns df = - plotStackedBarsWith categoryCol valueColumns (defaultPlotConfig StackedBar') df - -plotStackedBarsWith :: HasCallStack => T.Text -> [T.Text] -> PlotConfig -> DataFrame -> IO () -plotStackedBarsWith categoryCol valueColumns config df = do - let categories = extractStringColumn categoryCol df - uniqueCategories = L.nub categories - - stackData <- forM uniqueCategories $ \cat -> do - let indices = [i | (i, c) <- zip [0..] categories, c == cat] - seriesData <- forM valueColumns $ \col -> do - let allValues = extractNumericColumn col df - values = [allValues !! i | i <- indices, i < length allValues] - return (T.unpack col, sum values) - return (cat, seriesData) - - putStrLn $ stackedBars (plotTitle config) stackData (plotSettings config) - -plotHeatmap :: HasCallStack => DataFrame -> IO () -plotHeatmap df = plotHeatmapWith (defaultPlotConfig Heatmap') df - -plotHeatmapWith :: HasCallStack => PlotConfig -> DataFrame -> IO () -plotHeatmapWith config df = do - let numericCols = filter (isNumericColumn df) (columnNames df) - matrix = map (\col -> extractNumericColumn col df) numericCols - putStrLn $ heatmap (plotTitle config) matrix (plotSettings config) - -isNumericColumn :: DataFrame -> T.Text -> Bool -isNumericColumn df colName = - case M.lookup colName (columnIndices df) of - Nothing -> False - Just idx -> - let col = columns df V.! idx - in case col of - BoxedColumn (vec :: V.Vector a) -> - case testEquality (typeRep @a) (typeRep @Double) of - Just _ -> True - Nothing -> case testEquality (typeRep @a) (typeRep @Int) of - Just _ -> True - Nothing -> case testEquality (typeRep @a) (typeRep @Float) of - Just _ -> True - Nothing -> False - UnboxedColumn (vec :: VU.Vector a) -> - case testEquality (typeRep @a) (typeRep @Double) of - Just _ -> True - Nothing -> case testEquality (typeRep @a) (typeRep @Int) of - Just _ -> True - Nothing -> case testEquality (typeRep @a) (typeRep @Float) of - Just _ -> True - Nothing -> False - -- Haven't dealt with optionals yet. - _ -> False - -plotAllHistograms :: HasCallStack => DataFrame -> IO () -plotAllHistograms df = do - let numericCols = filter (isNumericColumn df) (columnNames df) - forM_ numericCols $ \col -> do - plotHistogram col df - -plotCorrelationMatrix :: HasCallStack => DataFrame -> IO () -plotCorrelationMatrix df = do - let numericCols = filter (isNumericColumn df) (columnNames df) - let correlations = map (\col1 -> - map (\col2 -> let - vals1 = extractNumericColumn col1 df - vals2 = extractNumericColumn col2 df - in correlation vals1 vals2) numericCols) numericCols - putStrLn $ heatmap "Correlation Matrix" correlations defPlot - where - correlation xs ys = - let n = fromIntegral $ length xs - meanX = sum xs / n - meanY = sum ys / n - covXY = sum [(x - meanX) * (y - meanY) | (x,y) <- zip xs ys] / n - stdX = sqrt $ sum [(x - meanX)^2 | x <- xs] / n - stdY = sqrt $ sum [(y - meanY)^2 | y <- ys] / n - in covXY / (stdX * stdY) - -quickPlot :: HasCallStack => [T.Text] -> DataFrame -> IO () -quickPlot [] df = plotAllHistograms df >> putStrLn "Plotted all numeric columns" -quickPlot [col] df = plotHistogram col df -quickPlot [col1, col2] df = plotScatter col1 col2 df -quickPlot cols df = plotLines cols df - -plotBars :: HasCallStack => T.Text -> DataFrame -> IO () -plotBars colName df = plotBarsWith colName Nothing (defaultPlotConfig Bar') df - -plotBarsWith :: HasCallStack => T.Text -> Maybe T.Text -> PlotConfig -> DataFrame -> IO () -plotBarsWith colName groupByCol config df = - case groupByCol of - Nothing -> plotSingleBars colName config df - Just grpCol -> plotGroupedBarsWith grpCol colName config df - -plotSingleBars :: HasCallStack => T.Text -> PlotConfig -> DataFrame -> IO () -plotSingleBars colName config df = do - let barData = getCategoricalCounts colName df - case barData of - Just counts -> do - let grouped = groupWithOther 10 counts - putStrLn $ bars (plotTitle config) grouped (plotSettings config) - Nothing -> do - let values = extractNumericColumn colName df - if length values > 20 - then do - let labels = map (\i -> "Item " ++ show i) [1..length values] - paired = zip labels values - grouped = groupWithOther 10 paired - putStrLn $ bars (plotTitle config) grouped (plotSettings config) - else do - let labels = map (\i -> "Item " ++ show i) [1..length values] - putStrLn $ bars (plotTitle config) (zip labels values) (plotSettings config) - -plotBarsTopN :: HasCallStack => Int -> T.Text -> DataFrame -> IO () -plotBarsTopN n colName df = plotBarsTopNWith n colName (defaultPlotConfig Bar') df - -plotBarsTopNWith :: HasCallStack => Int -> T.Text -> PlotConfig -> DataFrame -> IO () -plotBarsTopNWith n colName config df = do - let barData = getCategoricalCounts colName df - case barData of - Just counts -> do - let grouped = groupWithOther n counts - putStrLn $ bars (plotTitle config) grouped (plotSettings config) - Nothing -> do - let values = extractNumericColumn colName df - labels = map (\i -> "Item " ++ show i) [1..length values] - paired = zip labels values - grouped = groupWithOther n paired - putStrLn $ bars (plotTitle config) grouped (plotSettings config) - -plotGroupedBarsWith :: HasCallStack => T.Text -> T.Text -> PlotConfig -> DataFrame -> IO () -plotGroupedBarsWith groupCol valCol config df = do - let isNumeric = isNumericColumnCheck valCol df - - if isNumeric - then do - let groups = extractStringColumn groupCol df - values = extractNumericColumn valCol df - grouped = M.toList $ M.fromListWith (+) (zip groups values) - finalGroups = groupWithOther 10 grouped - putStrLn $ bars (plotTitle config) finalGroups (plotSettings config) - else do - let groups = extractStringColumn groupCol df - vals = extractStringColumn valCol df - pairs = zip groups vals - counts = M.toList $ M.fromListWith (+) - [(g ++ " - " ++ v, 1) | (g, v) <- pairs] - finalCounts = groupWithOther 15 [(k, fromIntegral v) | (k, v) <- counts] - putStrLn $ bars (plotTitle config) finalCounts (plotSettings config) - -plotValueCounts :: HasCallStack => T.Text -> DataFrame -> IO () -plotValueCounts colName df = plotValueCountsWith colName 10 (defaultPlotConfig Bar') df - -plotValueCountsWith :: HasCallStack => T.Text -> Int -> PlotConfig -> DataFrame -> IO () -plotValueCountsWith colName maxBars config df = do - let counts = getCategoricalCounts colName df - case counts of - Just c -> do - let grouped = groupWithOther maxBars c - config' = config { plotTitle = if null (plotTitle config) - then "Value counts for " ++ T.unpack colName - else plotTitle config } - putStrLn $ bars (T.unpack colName) grouped (plotSettings config') - Nothing -> error $ "Could not get value counts for column " ++ T.unpack colName - -plotBarsWithPercentages :: HasCallStack => T.Text -> DataFrame -> IO () -plotBarsWithPercentages colName df = do - let counts = getCategoricalCounts colName df - case counts of - Just c -> do - let total = sum (map snd c) - percentages = [(label ++ " (" ++ show (round (100 * val / total) :: Int) ++ "%)", val) - | (label, val) <- c] - grouped = groupWithOther 10 percentages - putStrLn $ bars ("Distribution of " ++ T.unpack colName) grouped defPlot - Nothing -> error $ "Could not get value counts for column " ++ T.unpack colName - -smartPlotBars :: HasCallStack => T.Text -> DataFrame -> IO () -smartPlotBars colName df = do - let counts = getCategoricalCounts colName df - case counts of - Just c -> do - let numUnique = length c - config = (defaultPlotConfig Bar') { - plotTitle = T.unpack colName ++ " (" ++ show numUnique ++ " unique values)" - } - if numUnique <= 12 - then putStrLn $ bars (plotTitle config) c (plotSettings config) - else if numUnique <= 20 - then do - let grouped = groupWithOther 12 c - putStrLn $ bars (plotTitle config ++ " - Top 12 + Other") grouped (plotSettings config) - else do - let grouped = groupWithOther 10 c - otherCount = numUnique - 10 - putStrLn $ bars (plotTitle config ++ " - Top 10 + Other (" ++ show otherCount ++ " items)") - grouped (plotSettings config) - Nothing -> plotBars colName df - -plotCategoricalSummary :: HasCallStack => DataFrame -> IO () -plotCategoricalSummary df = do - let cols = columnNames df - forM_ cols $ \col -> do - let counts = getCategoricalCounts col df - case counts of - Just c -> when (length c > 1) $ do - let numUnique = length c - putStrLn $ "\n=== " ++ T.unpack col ++ " (" ++ show numUnique ++ " unique values) ===" - if numUnique > 15 then plotBarsTopN 10 col df else plotBars col df - Nothing -> return () - -getCategoricalCounts :: HasCallStack => T.Text -> DataFrame -> Maybe [(String, Double)] -getCategoricalCounts colName df = - case M.lookup colName (columnIndices df) of - Nothing -> error $ "Column " ++ T.unpack colName ++ " not found" - Just idx -> - let col = columns df V.! idx - in case col of - BoxedColumn vec -> let counts = countValues vec - in Just [(show k, fromIntegral v) | (k, v) <- counts] - UnboxedColumn vec -> let counts = countValuesUnboxed vec - in Just [(show k, fromIntegral v) | (k, v) <- counts] - where - countValues :: (Ord a, Show a) => V.Vector a -> [(a, Int)] - countValues vec = M.toList $ V.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec - - countValuesUnboxed :: (Ord a, Show a, VU.Unbox a) => VU.Vector a -> [(a, Int)] - countValuesUnboxed vec = M.toList $ VU.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec - -isNumericColumnCheck :: T.Text -> DataFrame -> Bool -isNumericColumnCheck colName df = - case M.lookup colName (columnIndices df) of - Nothing -> False - Just idx -> - let col = columns df V.! idx - in case col of - BoxedColumn (vec :: V.Vector a) -> isNumericType @a - UnboxedColumn (vec :: VU.Vector a) -> isNumericType @a - -isNumericType :: forall a. Typeable a => Bool -isNumericType = - case testEquality (typeRep @a) (typeRep @Double) of - Just _ -> True - Nothing -> case testEquality (typeRep @a) (typeRep @Int) of - Just _ -> True - Nothing -> case testEquality (typeRep @a) (typeRep @Float) of - Just _ -> True - Nothing -> case testEquality (typeRep @a) (typeRep @Integer) of - Just _ -> True - Nothing -> False - -extractStringColumn :: HasCallStack => T.Text -> DataFrame -> [String] -extractStringColumn colName df = - case M.lookup colName (columnIndices df) of - Nothing -> error $ "Column " ++ T.unpack colName ++ " not found" - Just idx -> - let col = columns df V.! idx - in case col of - BoxedColumn vec -> V.toList $ V.map show vec - UnboxedColumn vec -> V.toList $ VG.map show (VG.convert vec) - -extractNumericColumn :: HasCallStack => T.Text -> DataFrame -> [Double] -extractNumericColumn colName df = - case M.lookup colName (columnIndices df) of - Nothing -> error $ "Column " ++ T.unpack colName ++ " not found" - Just idx -> - let col = columns df V.! idx - in case col of - BoxedColumn vec -> vectorToDoubles vec - UnboxedColumn vec -> unboxedVectorToDoubles vec - -vectorToDoubles :: forall a. (Typeable a, Show a) => V.Vector a -> [Double] -vectorToDoubles vec = - case testEquality (typeRep @a) (typeRep @Double) of - Just Refl -> V.toList vec - Nothing -> case testEquality (typeRep @a) (typeRep @Int) of - Just Refl -> V.toList $ V.map fromIntegral vec - Nothing -> case testEquality (typeRep @a) (typeRep @Integer) of - Just Refl -> V.toList $ V.map fromIntegral vec - Nothing -> case testEquality (typeRep @a) (typeRep @Float) of - Just Refl -> V.toList $ V.map realToFrac vec - Nothing -> error $ "Column is not numeric (type: " ++ show (typeRep @a) ++ ")" - -unboxedVectorToDoubles :: forall a. (Typeable a, VU.Unbox a, Show a) => VU.Vector a -> [Double] -unboxedVectorToDoubles vec = - case testEquality (typeRep @a) (typeRep @Double) of - Just Refl -> VU.toList vec - Nothing -> case testEquality (typeRep @a) (typeRep @Int) of - Just Refl -> VU.toList $ VU.map fromIntegral vec - Nothing -> case testEquality (typeRep @a) (typeRep @Float) of - Just Refl -> VU.toList $ VU.map realToFrac vec - Nothing -> error $ "Column is not numeric (type: " ++ show (typeRep @a) ++ ")" - -groupWithOther :: Int -> [(String, Double)] -> [(String, Double)] -groupWithOther n items = - let sorted = L.sortOn (negate . snd) items - (topN, rest) = splitAt n sorted - otherSum = sum (map snd rest) - result = if null rest || otherSum == 0 - then topN - else topN ++ [("Other (" ++ show (length rest) ++ " items)", otherSum)] - in result - -plotPie :: HasCallStack => T.Text -> Maybe T.Text -> DataFrame -> IO () -plotPie valCol labelCol df = plotPieWith valCol labelCol (defaultPlotConfig Pie') df - -plotPieWith :: HasCallStack => T.Text -> Maybe T.Text -> PlotConfig -> DataFrame -> IO () -plotPieWith valCol labelCol config df = do - let categoricalData = getCategoricalCounts valCol df - case categoricalData of - Just counts -> do - let grouped = groupWithOtherForPie 8 counts - putStrLn $ pie (plotTitle config) grouped (plotSettings config) - Nothing -> do - let values = extractNumericColumn valCol df - labels = case labelCol of - Nothing -> map (\i -> "Item " ++ show i) [1..length values] - Just lCol -> extractStringColumn lCol df - let pieData = zip labels values - grouped = if length pieData > 10 - then groupWithOtherForPie 8 pieData - else pieData - putStrLn $ pie (plotTitle config) grouped (plotSettings config) - -groupWithOtherForPie :: Int -> [(String, Double)] -> [(String, Double)] -groupWithOtherForPie n items = - let total = sum (map snd items) - sorted = L.sortOn (negate . snd) items - (topN, rest) = splitAt n sorted - otherSum = sum (map snd rest) - otherPct = round (100 * otherSum / total) :: Int - result = if null rest || otherSum == 0 - then topN - else topN ++ [("Other (" ++ show (length rest) ++ " items, " ++ - show otherPct ++ "%)", otherSum)] - in result - -plotPieWithPercentages :: HasCallStack => T.Text -> DataFrame -> IO () -plotPieWithPercentages colName df = plotPieWithPercentagesConfig colName (defaultPlotConfig Pie') df - -plotPieWithPercentagesConfig :: HasCallStack => T.Text -> PlotConfig -> DataFrame -> IO () -plotPieWithPercentagesConfig colName config df = do - let counts = getCategoricalCounts colName df - case counts of - Just c -> do - let total = sum (map snd c) - withPct = [(label ++ " (" ++ show (round (100 * val / total) :: Int) ++ "%)", val) - | (label, val) <- c] - grouped = groupWithOtherForPie 8 withPct - putStrLn $ pie (plotTitle config) grouped (plotSettings config) - Nothing -> do - let values = extractNumericColumn colName df - total = sum values - labels = map (\i -> "Item " ++ show i) [1..length values] - withPct = [(label ++ " (" ++ show (round (100 * val / total) :: Int) ++ "%)", val) - | (label, val) <- zip labels values] - grouped = groupWithOtherForPie 8 withPct - putStrLn $ pie (plotTitle config) grouped (plotSettings config) - -plotPieTopN :: HasCallStack => Int -> T.Text -> DataFrame -> IO () -plotPieTopN n colName df = plotPieTopNWith n colName (defaultPlotConfig Pie') df - -plotPieTopNWith :: HasCallStack => Int -> T.Text -> PlotConfig -> DataFrame -> IO () -plotPieTopNWith n colName config df = do - let counts = getCategoricalCounts colName df - case counts of - Just c -> do - let grouped = groupWithOtherForPie n c - putStrLn $ pie (plotTitle config) grouped (plotSettings config) - Nothing -> do - let values = extractNumericColumn colName df - labels = map (\i -> "Item " ++ show i) [1..length values] - paired = zip labels values - grouped = groupWithOtherForPie n paired - putStrLn $ pie (plotTitle config) grouped (plotSettings config) - -smartPlotPie :: HasCallStack => T.Text -> DataFrame -> IO () -smartPlotPie colName df = do - let counts = getCategoricalCounts colName df - case counts of - Just c -> do - let numUnique = length c - total = sum (map snd c) - significant = filter (\(_, v) -> v / total >= 0.01) c - config = (defaultPlotConfig Pie') { - plotTitle = T.unpack colName ++ " Distribution" - } - if length significant <= 6 - then putStrLn $ pie (plotTitle config) significant (plotSettings config) - else if length significant <= 10 - then do - let grouped = groupWithOtherForPie 8 c - putStrLn $ pie (plotTitle config) grouped (plotSettings config) - else do - let grouped = groupWithOtherForPie 6 c - putStrLn $ pie (plotTitle config) grouped (plotSettings config) - Nothing -> plotPie colName Nothing df - -plotPieGrouped :: HasCallStack => T.Text -> T.Text -> DataFrame -> IO () -plotPieGrouped groupCol valCol df = plotPieGroupedWith groupCol valCol (defaultPlotConfig Pie') df - -plotPieGroupedWith :: HasCallStack => T.Text -> T.Text -> PlotConfig -> DataFrame -> IO () -plotPieGroupedWith groupCol valCol config df = do - let isNumeric = isNumericColumnCheck valCol df - - if isNumeric - then do - let groups = extractStringColumn groupCol df - values = extractNumericColumn valCol df - grouped = M.toList $ M.fromListWith (+) (zip groups values) - finalGroups = groupWithOtherForPie 8 grouped - putStrLn $ pie (plotTitle config) finalGroups (plotSettings config) - else do - let groups = extractStringColumn groupCol df - vals = extractStringColumn valCol df - combined = zipWith (\g v -> g ++ " - " ++ v) groups vals - counts = M.toList $ M.fromListWith (+) [(c, 1) | c <- combined] - finalCounts = groupWithOtherForPie 10 [(k, fromIntegral v) | (k, v) <- counts] - putStrLn $ pie (plotTitle config) finalCounts (plotSettings config) - -plotPieComparison :: HasCallStack => [T.Text] -> DataFrame -> IO () -plotPieComparison cols df = do - forM_ cols $ \col -> do - let counts = getCategoricalCounts col df - case counts of - Just c -> when (length c > 1 && length c <= 20) $ do - putStrLn $ "\n=== " ++ T.unpack col ++ " Distribution ===" - smartPlotPie col df - Nothing -> return () - -plotBinaryPie :: HasCallStack => T.Text -> DataFrame -> IO () -plotBinaryPie colName df = do - let counts = getCategoricalCounts colName df - case counts of - Just c -> - if length c == 2 - then do - let total = sum (map snd c) - withPct = [(label ++ " (" ++ show (round (100 * val / total) :: Int) ++ "%)", val) - | (label, val) <- c] - putStrLn $ pie (T.unpack colName ++ " Proportion") withPct defPlot - else error $ "Column " ++ T.unpack colName ++ " is not binary (has " ++ - show (length c) ++ " unique values)" - Nothing -> error $ "Column " ++ T.unpack colName ++ " is not categorical" - -plotMarketShare :: HasCallStack => T.Text -> DataFrame -> IO () -plotMarketShare colName df = plotMarketShareWith colName (defaultPlotConfig Pie') df - -plotMarketShareWith :: HasCallStack => T.Text -> PlotConfig -> DataFrame -> IO () -plotMarketShareWith colName config df = do - let counts = getCategoricalCounts colName df - case counts of - Just c -> do - let total = sum (map snd c) - sorted = L.sortOn (negate . snd) c - significantShares = takeWhile (\(_, v) -> v / total >= 0.02) sorted - otherSum = sum [v | (_, v) <- c, v `notElem` map snd significantShares] - - formatShare (label, val) = - let pct = round (100 * val / total) :: Int - in (label ++ " (" ++ show pct ++ "%)", val) - - shares = map formatShare significantShares - finalShares = if otherSum > 0 && otherSum / total >= 0.01 - then shares ++ [("Others (<2% each)", otherSum)] - else shares - - let config' = config { plotTitle = if null (plotTitle config) - then T.unpack colName ++ " Market Share" - else plotTitle config } - putStrLn $ pie (plotTitle config') finalShares (plotSettings config') - Nothing -> error $ "Column " ++ T.unpack colName ++ " is not categorical" +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.Display.Terminal.Plot where + +import Control.Monad +import qualified Data.List as L +import qualified Data.Map as M +import Data.Maybe (catMaybes, fromMaybe) +import qualified Data.Text as T +import qualified Data.Text.IO as T +import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) +import Data.Typeable (Typeable) +import qualified Data.Vector as V +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Unboxed as VU +import GHC.Stack (HasCallStack) +import Type.Reflection (typeRep) + +import DataFrame.Internal.Column (Column (..), Columnable) +import qualified DataFrame.Internal.Column as D +import DataFrame.Internal.DataFrame (DataFrame (..)) +import DataFrame.Operations.Core +import qualified DataFrame.Operations.Subset as D +import Granite + +data PlotConfig = PlotConfig + { plotType :: PlotType + , plotSettings :: Plot + } + +data PlotType + = Histogram + | Scatter + | Line + | Bar + | BoxPlot + | Pie + | StackedBar + | Heatmap + deriving (Eq, Show) + +defaultPlotConfig :: PlotType -> PlotConfig +defaultPlotConfig ptype = + PlotConfig + { plotType = ptype + , plotSettings = defPlot + } + +plotHistogram :: (HasCallStack) => T.Text -> DataFrame -> IO () +plotHistogram colName df = plotHistogramWith colName (defaultPlotConfig Histogram) df + +plotHistogramWith :: (HasCallStack) => T.Text -> PlotConfig -> DataFrame -> IO () +plotHistogramWith colName config df = do + let values = extractNumericColumn colName df + (minVal, maxVal) = if null values then (0, 1) else (minimum values, maximum values) + T.putStrLn $ histogram (bins 30 minVal maxVal) values (plotSettings config) + +plotScatter :: (HasCallStack) => T.Text -> T.Text -> DataFrame -> IO () +plotScatter xCol yCol df = plotScatterWith xCol yCol (defaultPlotConfig Scatter) df + +plotScatterWith :: (HasCallStack) => T.Text -> T.Text -> PlotConfig -> DataFrame -> IO () +plotScatterWith xCol yCol config df = do + let xVals = extractNumericColumn xCol df + yVals = extractNumericColumn yCol df + points = zip xVals yVals + T.putStrLn $ scatter [(xCol <> " vs " <> yCol, points)] (plotSettings config) + +plotScatterBy :: (HasCallStack) => T.Text -> T.Text -> T.Text -> DataFrame -> IO () +plotScatterBy xCol yCol grouping df = plotScatterByWith xCol yCol grouping (defaultPlotConfig Scatter) df + +plotScatterByWith :: (HasCallStack) => T.Text -> T.Text -> T.Text -> PlotConfig -> DataFrame -> IO () +plotScatterByWith xCol yCol grouping config df = do + let vals = extractStringColumn grouping df + let df' = insertColumn grouping (D.fromList vals) df + xs <- forM (L.nub vals) $ \col -> do + let filtered = D.filter grouping (== col) df' + xVals = extractNumericColumn xCol filtered + yVals = extractNumericColumn yCol filtered + points = zip xVals yVals + pure (col, points) + T.putStrLn $ scatter xs (plotSettings config) + +plotLines :: (HasCallStack) => T.Text -> [T.Text] -> DataFrame -> IO () +plotLines xAxis colNames df = plotLinesWith xAxis colNames (defaultPlotConfig Line) df + +plotLinesWith :: (HasCallStack) => T.Text -> [T.Text] -> PlotConfig -> DataFrame -> IO () +plotLinesWith xAxis colNames config df = do + seriesData <- forM colNames $ \col -> do + let values = extractNumericColumn col df + indices = extractNumericColumn xAxis df + return (col, zip indices values) + T.putStrLn $ lineGraph seriesData (plotSettings config) + +plotBoxPlots :: (HasCallStack) => [T.Text] -> DataFrame -> IO () +plotBoxPlots colNames df = plotBoxPlotsWith colNames (defaultPlotConfig BoxPlot) df + +plotBoxPlotsWith :: (HasCallStack) => [T.Text] -> PlotConfig -> DataFrame -> IO () +plotBoxPlotsWith colNames config df = do + boxData <- forM colNames $ \col -> do + let values = extractNumericColumn col df + return (col, values) + T.putStrLn $ boxPlot boxData (plotSettings config) + +plotStackedBars :: (HasCallStack) => T.Text -> [T.Text] -> DataFrame -> IO () +plotStackedBars categoryCol valueColumns df = + plotStackedBarsWith categoryCol valueColumns (defaultPlotConfig StackedBar) df + +plotStackedBarsWith :: (HasCallStack) => T.Text -> [T.Text] -> PlotConfig -> DataFrame -> IO () +plotStackedBarsWith categoryCol valueColumns config df = do + let categories = extractStringColumn categoryCol df + uniqueCategories = L.nub categories + + stackData <- forM uniqueCategories $ \cat -> do + let indices = [i | (i, c) <- zip [0 ..] categories, c == cat] + seriesData <- forM valueColumns $ \col -> do + let allValues = extractNumericColumn col df + values = [allValues !! i | i <- indices, i < length allValues] + return (col, sum values) + return (cat, seriesData) + + T.putStrLn $ stackedBars stackData (plotSettings config) + +plotHeatmap :: (HasCallStack) => DataFrame -> IO () +plotHeatmap df = plotHeatmapWith (defaultPlotConfig Heatmap) df + +plotHeatmapWith :: (HasCallStack) => PlotConfig -> DataFrame -> IO () +plotHeatmapWith config df = do + let numericCols = filter (isNumericColumn df) (columnNames df) + matrix = map (\col -> extractNumericColumn col df) numericCols + T.putStrLn $ heatmap matrix (plotSettings config) + +isNumericColumn :: DataFrame -> T.Text -> Bool +isNumericColumn df colName = + case M.lookup colName (columnIndices df) of + Nothing -> False + Just idx -> + let col = columns df V.! idx + in case col of + BoxedColumn (vec :: V.Vector a) -> + case testEquality (typeRep @a) (typeRep @Double) of + Just _ -> True + Nothing -> case testEquality (typeRep @a) (typeRep @Int) of + Just _ -> True + Nothing -> case testEquality (typeRep @a) (typeRep @Float) of + Just _ -> True + Nothing -> False + UnboxedColumn (vec :: VU.Vector a) -> + case testEquality (typeRep @a) (typeRep @Double) of + Just _ -> True + Nothing -> case testEquality (typeRep @a) (typeRep @Int) of + Just _ -> True + Nothing -> case testEquality (typeRep @a) (typeRep @Float) of + Just _ -> True + Nothing -> False + -- Haven't dealt with optionals yet. + _ -> False + +plotAllHistograms :: (HasCallStack) => DataFrame -> IO () +plotAllHistograms df = do + let numericCols = filter (isNumericColumn df) (columnNames df) + forM_ numericCols $ \col -> do + T.putStrLn col + plotHistogram col df + +plotCorrelationMatrix :: (HasCallStack) => DataFrame -> IO () +plotCorrelationMatrix df = do + let numericCols = filter (isNumericColumn df) (columnNames df) + let correlations = + map + ( \col1 -> + map + ( \col2 -> + let + vals1 = extractNumericColumn col1 df + vals2 = extractNumericColumn col2 df + in + correlation vals1 vals2 + ) + numericCols + ) + numericCols + print (zip [0 ..] numericCols) + T.putStrLn $ heatmap correlations (defPlot{plotTitle = "Correlation Matrix"}) + where + correlation xs ys = + let n = fromIntegral $ length xs + meanX = sum xs / n + meanY = sum ys / n + covXY = sum [(x - meanX) * (y - meanY) | (x, y) <- zip xs ys] / n + stdX = sqrt $ sum [(x - meanX) ^ 2 | x <- xs] / n + stdY = sqrt $ sum [(y - meanY) ^ 2 | y <- ys] / n + in covXY / (stdX * stdY) + +plotBars :: (HasCallStack) => T.Text -> DataFrame -> IO () +plotBars colName df = plotBarsWith colName Nothing (defaultPlotConfig Bar) df + +plotBarsWith :: (HasCallStack) => T.Text -> Maybe T.Text -> PlotConfig -> DataFrame -> IO () +plotBarsWith colName groupByCol config df = + case groupByCol of + Nothing -> plotSingleBars colName config df + Just grpCol -> plotGroupedBarsWith grpCol colName config df + +plotSingleBars :: (HasCallStack) => T.Text -> PlotConfig -> DataFrame -> IO () +plotSingleBars colName config df = do + let barData = getCategoricalCounts colName df + case barData of + Just counts -> do + let grouped = groupWithOther 10 counts + T.putStrLn $ bars grouped (plotSettings config) + Nothing -> do + let values = extractNumericColumn colName df + if length values > 20 + then do + let labels = map (\i -> "Item " <> T.pack (show i)) [1 .. length values] + paired = zip labels values + grouped = groupWithOther 10 paired + T.putStrLn $ bars grouped (plotSettings config) + else do + let labels = map (\i -> "Item " <> T.pack (show i)) [1 .. length values] + T.putStrLn $ bars (zip labels values) (plotSettings config) + +plotBarsTopN :: (HasCallStack) => Int -> T.Text -> DataFrame -> IO () +plotBarsTopN n colName df = plotBarsTopNWith n colName (defaultPlotConfig Bar) df + +plotBarsTopNWith :: (HasCallStack) => Int -> T.Text -> PlotConfig -> DataFrame -> IO () +plotBarsTopNWith n colName config df = do + let barData = getCategoricalCounts colName df + case barData of + Just counts -> do + let grouped = groupWithOther n counts + T.putStrLn $ bars grouped (plotSettings config) + Nothing -> do + let values = extractNumericColumn colName df + labels = map (\i -> "Item " <> T.pack (show i)) [1 .. length values] + paired = zip labels values + grouped = groupWithOther n paired + T.putStrLn $ bars grouped (plotSettings config) + +plotGroupedBarsWith :: (HasCallStack) => T.Text -> T.Text -> PlotConfig -> DataFrame -> IO () +plotGroupedBarsWith = plotGroupedBarsWithN 10 + +plotGroupedBarsWithN :: (HasCallStack) => Int -> T.Text -> T.Text -> PlotConfig -> DataFrame -> IO () +plotGroupedBarsWithN n groupCol valCol config df = do + let isNumeric = isNumericColumnCheck valCol df + + if isNumeric + then do + let groups = extractStringColumn groupCol df + values = extractNumericColumn valCol df + grouped = M.toList $ M.fromListWith (+) (zip groups values) + finalGroups = groupWithOther n grouped + T.putStrLn $ bars finalGroups (plotSettings config) + else do + let groups = extractStringColumn groupCol df + vals = extractStringColumn valCol df + pairs = zip groups vals + counts = + M.toList $ + M.fromListWith + (+) + [(g <> " - " <> v, 1) | (g, v) <- pairs] + finalCounts = groupWithOther n [(k, fromIntegral v) | (k, v) <- counts] + T.putStrLn $ bars finalCounts (plotSettings config) + +plotValueCounts :: (HasCallStack) => T.Text -> DataFrame -> IO () +plotValueCounts colName df = plotValueCountsWith colName 10 (defaultPlotConfig Bar) df + +plotValueCountsWith :: (HasCallStack) => T.Text -> Int -> PlotConfig -> DataFrame -> IO () +plotValueCountsWith colName maxBars config df = do + let counts = getCategoricalCounts colName df + case counts of + Just c -> do + let grouped = groupWithOther maxBars c + config' = + config + { plotSettings = + (plotSettings config) + { plotTitle = + if T.null (plotTitle (plotSettings config)) + then "Value counts for " <> colName + else plotTitle (plotSettings config) + } + } + T.putStrLn $ bars grouped (plotSettings config') + Nothing -> error $ "Could not get value counts for column " ++ T.unpack colName + +plotBarsWithPercentages :: (HasCallStack) => T.Text -> DataFrame -> IO () +plotBarsWithPercentages colName df = do + let counts = getCategoricalCounts colName df + case counts of + Just c -> do + let total = sum (map snd c) + percentages = + [ (label <> " (" <> T.pack (show (round (100 * val / total) :: Int)) <> "%)", val) + | (label, val) <- c + ] + grouped = groupWithOther 10 percentages + T.putStrLn $ bars grouped (defPlot{plotTitle = "Distribution of " <> colName}) + Nothing -> error $ "Could not get value counts for column " ++ T.unpack colName + +smartPlotBars :: (HasCallStack) => T.Text -> DataFrame -> IO () +smartPlotBars colName df = do + let counts = getCategoricalCounts colName df + case counts of + Just c -> do + let numUnique = length c + config = + (defaultPlotConfig Bar) + { plotSettings = (plotSettings (defaultPlotConfig Bar)){plotTitle = colName <> " (" <> T.pack (show numUnique) <> " unique values)"} + } + if numUnique <= 12 + then T.putStrLn $ bars c (plotSettings config) + else + if numUnique <= 20 + then do + let grouped = groupWithOther 12 c + T.putStrLn $ bars grouped (plotSettings config) + else do + let grouped = groupWithOther 10 c + otherCount = numUnique - 10 + T.putStrLn $ + bars + grouped + (plotSettings config) + Nothing -> plotBars colName df + +plotCategoricalSummary :: (HasCallStack) => DataFrame -> IO () +plotCategoricalSummary df = do + let cols = columnNames df + forM_ cols $ \col -> do + let counts = getCategoricalCounts col df + case counts of + Just c -> when (length c > 1) $ do + let numUnique = length c + putStrLn $ "\n=== " ++ T.unpack col ++ " (" ++ show numUnique ++ " unique values) ===" + if numUnique > 15 then plotBarsTopN 10 col df else plotBars col df + Nothing -> return () + +getCategoricalCounts :: (HasCallStack) => T.Text -> DataFrame -> Maybe [(T.Text, Double)] +getCategoricalCounts colName df = + case M.lookup colName (columnIndices df) of + Nothing -> error $ "Column " ++ T.unpack colName ++ " not found" + Just idx -> + let col = columns df V.! idx + in case col of + BoxedColumn vec -> + let counts = countValues vec + in Just [(T.pack (show k), fromIntegral v) | (k, v) <- counts] + UnboxedColumn vec -> + let counts = countValuesUnboxed vec + in Just [(T.pack (show k), fromIntegral v) | (k, v) <- counts] + _ -> Nothing + where + countValues :: (Ord a, Show a) => V.Vector a -> [(a, Int)] + countValues vec = M.toList $ V.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec + + countValuesUnboxed :: (Ord a, Show a, VU.Unbox a) => VU.Vector a -> [(a, Int)] + countValuesUnboxed vec = M.toList $ VU.foldr' (\x acc -> M.insertWith (+) x 1 acc) M.empty vec + +isNumericColumnCheck :: T.Text -> DataFrame -> Bool +isNumericColumnCheck colName df = + case M.lookup colName (columnIndices df) of + Nothing -> False + Just idx -> + let col = columns df V.! idx + in case col of + BoxedColumn (vec :: V.Vector a) -> isNumericType @a + UnboxedColumn (vec :: VU.Vector a) -> isNumericType @a + _ -> False + +isNumericType :: forall a. (Typeable a) => Bool +isNumericType = + case testEquality (typeRep @a) (typeRep @Double) of + Just _ -> True + Nothing -> case testEquality (typeRep @a) (typeRep @Int) of + Just _ -> True + Nothing -> case testEquality (typeRep @a) (typeRep @Float) of + Just _ -> True + Nothing -> case testEquality (typeRep @a) (typeRep @Integer) of + Just _ -> True + Nothing -> False + +extractStringColumn :: (HasCallStack) => T.Text -> DataFrame -> [T.Text] +extractStringColumn colName df = + case M.lookup colName (columnIndices df) of + Nothing -> error $ "Column " ++ T.unpack colName ++ " not found" + Just idx -> + let col = columns df V.! idx + in case col of + BoxedColumn vec -> V.toList $ V.map (T.pack . show) vec + UnboxedColumn vec -> V.toList $ VG.map (T.pack . show) (VG.convert vec) + OptionalColumn vec -> V.toList $ V.map (T.pack . show) vec + +extractNumericColumn :: (HasCallStack) => T.Text -> DataFrame -> [Double] +extractNumericColumn colName df = + case M.lookup colName (columnIndices df) of + Nothing -> error $ "Column " ++ T.unpack colName ++ " not found" + Just idx -> + let col = columns df V.! idx + in case col of + BoxedColumn vec -> vectorToDoubles vec + UnboxedColumn vec -> unboxedVectorToDoubles vec + _ -> [] + +vectorToDoubles :: forall a. (Typeable a, Show a) => V.Vector a -> [Double] +vectorToDoubles vec = + case testEquality (typeRep @a) (typeRep @Double) of + Just Refl -> V.toList vec + Nothing -> case testEquality (typeRep @a) (typeRep @Int) of + Just Refl -> V.toList $ V.map fromIntegral vec + Nothing -> case testEquality (typeRep @a) (typeRep @Integer) of + Just Refl -> V.toList $ V.map fromIntegral vec + Nothing -> case testEquality (typeRep @a) (typeRep @Float) of + Just Refl -> V.toList $ V.map realToFrac vec + Nothing -> error $ "Column is not numeric (type: " ++ show (typeRep @a) ++ ")" + +unboxedVectorToDoubles :: forall a. (Typeable a, VU.Unbox a, Show a) => VU.Vector a -> [Double] +unboxedVectorToDoubles vec = + case testEquality (typeRep @a) (typeRep @Double) of + Just Refl -> VU.toList vec + Nothing -> case testEquality (typeRep @a) (typeRep @Int) of + Just Refl -> VU.toList $ VU.map fromIntegral vec + Nothing -> case testEquality (typeRep @a) (typeRep @Float) of + Just Refl -> VU.toList $ VU.map realToFrac vec + Nothing -> error $ "Column is not numeric (type: " ++ show (typeRep @a) ++ ")" + +groupWithOther :: Int -> [(T.Text, Double)] -> [(T.Text, Double)] +groupWithOther n items = + let sorted = L.sortOn (negate . snd) items + (topN, rest) = splitAt n sorted + otherSum = sum (map snd rest) + result = + if null rest || otherSum == 0 + then topN + else topN ++ [("Other (" <> T.pack (show (length rest)) <> " items)", otherSum)] + in result + +plotPie :: (HasCallStack) => T.Text -> Maybe T.Text -> DataFrame -> IO () +plotPie valCol labelCol df = plotPieWith valCol labelCol (defaultPlotConfig Pie) df + +plotPieWith :: (HasCallStack) => T.Text -> Maybe T.Text -> PlotConfig -> DataFrame -> IO () +plotPieWith valCol labelCol config df = do + let categoricalData = getCategoricalCounts valCol df + case categoricalData of + Just counts -> do + let grouped = groupWithOtherForPie 8 counts + T.putStrLn $ pie grouped (plotSettings config) + Nothing -> do + let values = extractNumericColumn valCol df + labels = case labelCol of + Nothing -> map (\i -> "Item " <> T.pack (show i)) [1 .. length values] + Just lCol -> extractStringColumn lCol df + let pieData = zip labels values + grouped = + if length pieData > 10 + then groupWithOtherForPie 8 pieData + else pieData + T.putStrLn $ pie grouped (plotSettings config) + +groupWithOtherForPie :: Int -> [(T.Text, Double)] -> [(T.Text, Double)] +groupWithOtherForPie n items = + let total = sum (map snd items) + sorted = L.sortOn (negate . snd) items + (topN, rest) = splitAt n sorted + otherSum = sum (map snd rest) + otherPct = round (100 * otherSum / total) :: Int + result = + if null rest || otherSum == 0 + then topN + else + topN + ++ [ + ( "Other (" + <> T.pack (show (length rest)) + <> " items, " + <> T.pack (show otherPct) + <> "%)" + , otherSum + ) + ] + in result + +plotPieWithPercentages :: (HasCallStack) => T.Text -> DataFrame -> IO () +plotPieWithPercentages colName df = plotPieWithPercentagesConfig colName (defaultPlotConfig Pie) df + +plotPieWithPercentagesConfig :: (HasCallStack) => T.Text -> PlotConfig -> DataFrame -> IO () +plotPieWithPercentagesConfig colName config df = do + let counts = getCategoricalCounts colName df + case counts of + Just c -> do + let total = sum (map snd c) + withPct = + [ (label <> " (" <> T.pack (show (round (100 * val / total) :: Int)) <> "%)", val) + | (label, val) <- c + ] + grouped = groupWithOtherForPie 8 withPct + T.putStrLn $ pie grouped (plotSettings config) + Nothing -> do + let values = extractNumericColumn colName df + total = sum values + labels = map (\i -> "Item " <> T.pack (show i)) [1 .. length values] + withPct = + [ (label <> " (" <> T.pack (show (round (100 * val / total) :: Int)) <> "%)", val) + | (label, val) <- zip labels values + ] + grouped = groupWithOtherForPie 8 withPct + T.putStrLn $ pie grouped (plotSettings config) + +plotPieTopN :: (HasCallStack) => Int -> T.Text -> DataFrame -> IO () +plotPieTopN n colName df = plotPieTopNWith n colName (defaultPlotConfig Pie) df + +plotPieTopNWith :: (HasCallStack) => Int -> T.Text -> PlotConfig -> DataFrame -> IO () +plotPieTopNWith n colName config df = do + let counts = getCategoricalCounts colName df + case counts of + Just c -> do + let grouped = groupWithOtherForPie n c + T.putStrLn $ pie grouped (plotSettings config) + Nothing -> do + let values = extractNumericColumn colName df + labels = map (\i -> "Item " <> T.pack (show i)) [1 .. length values] + paired = zip labels values + grouped = groupWithOtherForPie n paired + T.putStrLn $ pie grouped (plotSettings config) + +smartPlotPie :: (HasCallStack) => T.Text -> DataFrame -> IO () +smartPlotPie colName df = do + let counts = getCategoricalCounts colName df + case counts of + Just c -> do + let numUnique = length c + total = sum (map snd c) + significant = filter (\(_, v) -> v / total >= 0.01) c + config = + (defaultPlotConfig Pie) + { plotSettings = (plotSettings (defaultPlotConfig Pie)){plotTitle = colName <> " Distribution"} + } + if length significant <= 6 + then T.putStrLn $ pie significant (plotSettings config) + else + if length significant <= 10 + then do + let grouped = groupWithOtherForPie 8 c + T.putStrLn $ pie grouped (plotSettings config) + else do + let grouped = groupWithOtherForPie 6 c + T.putStrLn $ pie grouped (plotSettings config) + Nothing -> plotPie colName Nothing df + +plotPieGrouped :: (HasCallStack) => T.Text -> T.Text -> DataFrame -> IO () +plotPieGrouped groupCol valCol df = plotPieGroupedWith groupCol valCol (defaultPlotConfig Pie) df + +plotPieGroupedWith :: (HasCallStack) => T.Text -> T.Text -> PlotConfig -> DataFrame -> IO () +plotPieGroupedWith groupCol valCol config df = do + let isNumeric = isNumericColumnCheck valCol df + + if isNumeric + then do + let groups = extractStringColumn groupCol df + values = extractNumericColumn valCol df + grouped = M.toList $ M.fromListWith (+) (zip groups values) + finalGroups = groupWithOtherForPie 8 grouped + T.putStrLn $ pie finalGroups (plotSettings config) + else do + let groups = extractStringColumn groupCol df + vals = extractStringColumn valCol df + combined = zipWith (\g v -> g <> " - " <> v) groups vals + counts = M.toList $ M.fromListWith (+) [(c, 1) | c <- combined] + finalCounts = groupWithOtherForPie 10 [(k, fromIntegral v) | (k, v) <- counts] + T.putStrLn $ pie finalCounts (plotSettings config) + +plotPieComparison :: (HasCallStack) => [T.Text] -> DataFrame -> IO () +plotPieComparison cols df = do + forM_ cols $ \col -> do + let counts = getCategoricalCounts col df + case counts of + Just c -> when (length c > 1 && length c <= 20) $ do + putStrLn $ "\n=== " ++ T.unpack col ++ " Distribution ===" + smartPlotPie col df + Nothing -> return () + +plotBinaryPie :: (HasCallStack) => T.Text -> DataFrame -> IO () +plotBinaryPie colName df = do + let counts = getCategoricalCounts colName df + case counts of + Just c -> + if length c == 2 + then do + let total = sum (map snd c) + withPct = + [ (label <> " (" <> T.pack (show (round (100 * val / total) :: Int)) <> "%)", val) + | (label, val) <- c + ] + T.putStrLn $ pie withPct defPlot + else + error $ + "Column " + ++ T.unpack colName + ++ " is not binary (has " + ++ show (length c) + ++ " unique values)" + Nothing -> error $ "Column " ++ T.unpack colName ++ " is not categorical" + +plotMarketShare :: (HasCallStack) => T.Text -> DataFrame -> IO () +plotMarketShare colName df = plotMarketShareWith colName (defaultPlotConfig Pie) df + +plotMarketShareWith :: (HasCallStack) => T.Text -> PlotConfig -> DataFrame -> IO () +plotMarketShareWith colName config df = do + let counts = getCategoricalCounts colName df + case counts of + Just c -> do + let total = sum (map snd c) + sorted = L.sortOn (negate . snd) c + significantShares = takeWhile (\(_, v) -> v / total >= 0.02) sorted + otherSum = sum [v | (_, v) <- c, v `notElem` map snd significantShares] + + formatShare (label, val) = + let pct = round (100 * val / total) :: Int + in (label <> " (" <> T.pack (show pct) <> "%)", val) + + shares = map formatShare significantShares + finalShares = + if otherSum > 0 && otherSum / total >= 0.01 + then shares <> [("Others (<2% each)", otherSum)] + else shares + + let config' = + config + -- { plotSettings = (plotSettings config) { + -- plotTitle = colName <> ": market share" + -- } + -- } + T.putStrLn $ pie finalShares (plotSettings config') + Nothing -> error $ "Column " ++ T.unpack colName ++ " is not categorical" diff --git a/src/DataFrame/Display/Terminal/PrettyPrint.hs b/src/DataFrame/Display/Terminal/PrettyPrint.hs --- a/src/DataFrame/Display/Terminal/PrettyPrint.hs +++ b/src/DataFrame/Display/Terminal/PrettyPrint.hs @@ -1,4 +1,5 @@ {-# LANGUAGE OverloadedStrings #-} + module DataFrame.Display.Terminal.PrettyPrint where import qualified Data.Text as T @@ -13,10 +14,10 @@ -- a type for describing table columns data ColDesc t = ColDesc - { colTitleFill :: Filler, - colTitle :: T.Text, - colValueFill :: Filler - } + { colTitleFill :: Filler + , colTitle :: T.Text + , colValueFill :: Filler + } -- functions that fill a string (s) to a given width (n) by adding pad -- character (c) to align left, right, or center @@ -45,13 +46,14 @@ showTable :: Bool -> [T.Text] -> [T.Text] -> [[T.Text]] -> T.Text showTable properMarkdown header types rows = - let consolidatedHeader = if properMarkdown then zipWith (\h t -> h <> "
" <> t) header types else header - cs = map (\h -> ColDesc center h left) consolidatedHeader - widths = [maximum $ map T.length col | col <- transpose $ consolidatedHeader : types : rows] - border = T.intercalate "---" [T.replicate width (T.singleton '-') | width <- widths] - separator = T.intercalate "-|-" [T.replicate width (T.singleton '-') | width <- widths] - fillCols fill cols = T.intercalate " | " [fill c width col | (c, width, col) <- zip3 cs widths cols] - lines = if properMarkdown - then border : fillCols colTitleFill consolidatedHeader : separator : map (fillCols colValueFill) rows - else border : fillCols colTitleFill consolidatedHeader : separator : fillCols colTitleFill types : separator : map (fillCols colValueFill) rows - in T.unlines lines + let consolidatedHeader = if properMarkdown then zipWith (\h t -> h <> "
" <> t) header types else header + cs = map (\h -> ColDesc center h left) consolidatedHeader + widths = [maximum $ map T.length col | col <- transpose $ consolidatedHeader : types : rows] + border = T.intercalate "---" [T.replicate width (T.singleton '-') | width <- widths] + separator = T.intercalate "-|-" [T.replicate width (T.singleton '-') | width <- widths] + fillCols fill cols = T.intercalate " | " [fill c width col | (c, width, col) <- zip3 cs widths cols] + lines = + if properMarkdown + then border : fillCols colTitleFill consolidatedHeader : separator : map (fillCols colValueFill) rows + else border : fillCols colTitleFill consolidatedHeader : separator : fillCols colTitleFill types : separator : map (fillCols colValueFill) rows + in T.unlines lines diff --git a/src/DataFrame/Errors.hs b/src/DataFrame/Errors.hs --- a/src/DataFrame/Errors.hs +++ b/src/DataFrame/Errors.hs @@ -1,8 +1,8 @@ {-# LANGUAGE DeriveAnyClass #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE RankNTypes #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} module DataFrame.Errors where @@ -11,87 +11,99 @@ import Control.Exception import Data.Array import Data.Either -import DataFrame.Display.Terminal.Colours import Data.Typeable (Typeable) +import DataFrame.Display.Terminal.Colours import Type.Reflection (TypeRep) data TypeErrorContext a b = MkTypeErrorContext - { userType :: Either String (TypeRep a) - , expectedType :: Either String (TypeRep b) - , errorColumnName :: Maybe String - , callingFunctionName :: Maybe String - } + { userType :: Either String (TypeRep a) + , expectedType :: Either String (TypeRep b) + , errorColumnName :: Maybe String + , callingFunctionName :: Maybe String + } data DataFrameException where - TypeMismatchException :: forall a b. (Typeable a, Typeable b) - => TypeErrorContext a b - -> DataFrameException + TypeMismatchException :: + forall a b. + (Typeable a, Typeable b) => + TypeErrorContext a b -> + DataFrameException ColumnNotFoundException :: T.Text -> T.Text -> [T.Text] -> DataFrameException + EmptyDataSetException :: T.Text -> DataFrameException deriving (Exception) instance Show DataFrameException where show :: DataFrameException -> String - show (TypeMismatchException context) = let - errorString = typeMismatchError (either id show (userType context)) (either id show (expectedType context)) - in addCallPointInfo (errorColumnName context) (callingFunctionName context) errorString + show (TypeMismatchException context) = + let + errorString = typeMismatchError (either id show (userType context)) (either id show (expectedType context)) + in + addCallPointInfo (errorColumnName context) (callingFunctionName context) errorString show (ColumnNotFoundException columnName callPoint availableColumns) = columnNotFound columnName callPoint availableColumns + show (EmptyDataSetException callPoint) = emptyDataSetError callPoint columnNotFound :: T.Text -> T.Text -> [T.Text] -> String columnNotFound name callPoint columns = - red "\n\n[ERROR] " - ++ "Column not found: " - ++ T.unpack name - ++ " for operation " - ++ T.unpack callPoint - ++ "\n\tDid you mean " - ++ T.unpack (guessColumnName name columns) - ++ "?\n\n" + red "\n\n[ERROR] " + ++ "Column not found: " + ++ T.unpack name + ++ " for operation " + ++ T.unpack callPoint + ++ "\n\tDid you mean " + ++ T.unpack (guessColumnName name columns) + ++ "?\n\n" typeMismatchError :: String -> String -> String typeMismatchError givenType expectedType = - red $ - red "\n\n[Error]: Type Mismatch" - ++ "\n\tWhile running your code I tried to " - ++ "get a column of type: " - ++ red (show givenType) - ++ " but the column in the dataframe was actually of type: " - ++ green (show expectedType) + red $ + red "\n\n[Error]: Type Mismatch" + ++ "\n\tWhile running your code I tried to " + ++ "get a column of type: " + ++ red (show givenType) + ++ " but the column in the dataframe was actually of type: " + ++ green (show expectedType) +emptyDataSetError :: T.Text -> String +emptyDataSetError callPoint = + red "\n\n[ERROR] " + ++ T.unpack callPoint + ++ " cannot be called on empty data sets" + addCallPointInfo :: Maybe String -> Maybe String -> String -> String addCallPointInfo (Just name) (Just cp) err = - err - ++ ( "\n\tThis happened when calling function " - ++ brightGreen cp - ++ " on the column " - ++ brightGreen name - ++ "\n\n" - ++ typeAnnotationSuggestion cp - ) + err + ++ ( "\n\tThis happened when calling function " + ++ brightGreen cp + ++ " on the column " + ++ brightGreen name + ++ "\n\n" + ++ typeAnnotationSuggestion cp + ) addCallPointInfo Nothing (Just cp) err = err ++ "\n" ++ typeAnnotationSuggestion cp addCallPointInfo (Just name) Nothing err = - err - ++ ( "\n\tOn the column " - ++ name - ++ "\n\n" - ) + err + ++ ( "\n\tOn the column " + ++ name + ++ "\n\n" + ) addCallPointInfo Nothing Nothing err = err typeAnnotationSuggestion :: String -> String typeAnnotationSuggestion cp = - "\n\n\tTry adding a type at the end of the function e.g " - ++ "change\n\t\t" - ++ red (cp ++ " ...") - ++ " to \n\t\t" - ++ green ("(" ++ cp ++ " ... :: )") - ++ "\n\tor add " - ++ "{-# LANGUAGE TypeApplications #-} to the top of your " - ++ "file then change the call to \n\t\t" - ++ brightGreen (cp ++ " @ ....") + "\n\n\tTry adding a type at the end of the function e.g " + ++ "change\n\t\t" + ++ red (cp ++ " ...") + ++ " to \n\t\t" + ++ green ("(" ++ cp ++ " ... :: )") + ++ "\n\tor add " + ++ "{-# LANGUAGE TypeApplications #-} to the top of your " + ++ "file then change the call to \n\t\t" + ++ brightGreen (cp ++ " @ ....") guessColumnName :: T.Text -> [T.Text] -> T.Text guessColumnName userInput columns = case map (\k -> (editDistance userInput k, k)) columns of - [] -> "" - res -> (snd . minimum) res + [] -> "" + res -> (snd . minimum) res editDistance :: T.Text -> T.Text -> Int editDistance xs ys = table ! (m, n) @@ -107,8 +119,8 @@ dist (0, j) = j dist (i, 0) = i dist (i, j) = - minimum - [ table ! (i - 1, j) + 1, - table ! (i, j - 1) + 1, - if x ! i == y ! j then table ! (i - 1, j - 1) else 1 + table ! (i - 1, j - 1) - ] + minimum + [ table ! (i - 1, j) + 1 + , table ! (i, j - 1) + 1 + , if x ! i == y ! j then table ! (i - 1, j - 1) else 1 + table ! (i - 1, j - 1) + ] diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs --- a/src/DataFrame/Functions.hs +++ b/src/DataFrame/Functions.hs @@ -1,49 +1,50 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE TemplateHaskell #-} + module DataFrame.Functions where import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame (DataFrame(..), unsafeGetColumn) -import DataFrame.Internal.Expression (Expr(..), UExpr(..)) +import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn) +import DataFrame.Internal.Expression (Expr (..), UExpr (..)) -import Control.Monad -import Data.Function +import Control.Monad +import qualified Data.Char as Char +import Data.Function import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T +import qualified Data.Vector as VB import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector as VB -import Language.Haskell.TH -import qualified Language.Haskell.TH.Syntax as TH -import qualified Data.Char as Char import Debug.Trace (traceShow) +import Language.Haskell.TH +import qualified Language.Haskell.TH.Syntax as TH import Type.Reflection (typeRep) -col :: Columnable a => T.Text -> Expr a +col :: (Columnable a) => T.Text -> Expr a col = Col -as :: Columnable a => Expr a -> T.Text -> (T.Text, UExpr) +as :: (Columnable a) => Expr a -> T.Text -> (T.Text, UExpr) as expr name = (name, Wrap expr) -lit :: Columnable a => a -> Expr a +lit :: (Columnable a) => a -> Expr a lit = Lit lift :: (Columnable a, Columnable b) => (a -> b) -> Expr a -> Expr b lift = Apply "udf" -lift2 :: (Columnable c, Columnable b, Columnable a) => (c -> b -> a) -> Expr c -> Expr b -> Expr a +lift2 :: (Columnable c, Columnable b, Columnable a) => (c -> b -> a) -> Expr c -> Expr b -> Expr a lift2 = BinOp "udf" eq :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool @@ -61,120 +62,126 @@ geq :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool geq = BinOp "geq" (>=) -count :: Columnable a => Expr a -> Expr Int +count :: (Columnable a) => Expr a -> Expr Int count (Col name) = GeneralAggregate name "count" VG.length count _ = error "Argument can only be a column reference not an unevaluated expression" -minimum :: Columnable a => Expr a -> Expr a +minimum :: (Columnable a) => Expr a -> Expr a minimum (Col name) = ReductionAggregate name "minimum" min -maximum :: Columnable a => Expr a -> Expr a +maximum :: (Columnable a) => Expr a -> Expr a maximum (Col name) = ReductionAggregate name "maximum" max -sum :: forall a . (Columnable a, Num a, VU.Unbox a) => Expr a -> Expr a +sum :: forall a. (Columnable a, Num a, VU.Unbox a) => Expr a -> Expr a sum (Col name) = NumericAggregate name "sum" VG.sum mean :: (Columnable a, Num a, VU.Unbox a) => Expr a -> Expr Double -mean (Col name) = let - mean' samp = let - (!total, !n) = VG.foldl' (\(!total, !n) v -> (total + v, n + 1)) (0 :: Double, 0 :: Int) samp - in total / fromIntegral n - in NumericAggregate name "mean" mean' +mean (Col name) = + let + mean' samp = + let + (!total, !n) = VG.foldl' (\(!total, !n) v -> (total + v, n + 1)) (0 :: Double, 0 :: Int) samp + in + total / fromIntegral n + in + NumericAggregate name "mean" mean' -- See Section 2.4 of the Haskell Report https://www.haskell.org/definition/haskell2010.pdf isReservedId :: T.Text -> Bool isReservedId t = case t of - "case" -> True - "class" -> True - "data" -> True - "default" -> True - "deriving" -> True - "do" -> True - "else" -> True - "foreign" -> True - "if" -> True - "import" -> True - "in" -> True - "infix" -> True - "infixl" -> True - "infixr" -> True - "instance" -> True - "let" -> True - "module" -> True - "newtype" -> True - "of" -> True - "then" -> True - "type" -> True - "where" -> True - _ -> False + "case" -> True + "class" -> True + "data" -> True + "default" -> True + "deriving" -> True + "do" -> True + "else" -> True + "foreign" -> True + "if" -> True + "import" -> True + "in" -> True + "infix" -> True + "infixl" -> True + "infixr" -> True + "instance" -> True + "let" -> True + "module" -> True + "newtype" -> True + "of" -> True + "then" -> True + "type" -> True + "where" -> True + _ -> False isVarId :: T.Text -> Bool isVarId t = case T.uncons t of --- We might want to check c == '_' || Char.isLower c --- since the haskell report considers '_' a lowercase character --- However, to prevent an edge case where a user may have a --- "Name" and an "_Name_" in the same scope, wherein we'd end up --- with duplicate "_Name_"s, we eschew the check for '_' here. - Just (c, _) -> Char.isLower c && Char.isAlpha c - Nothing -> False + -- We might want to check c == '_' || Char.isLower c + -- since the haskell report considers '_' a lowercase character + -- However, to prevent an edge case where a user may have a + -- "Name" and an "_Name_" in the same scope, wherein we'd end up + -- with duplicate "_Name_"s, we eschew the check for '_' here. + Just (c, _) -> Char.isLower c && Char.isAlpha c + Nothing -> False isHaskellIdentifier :: T.Text -> Bool -isHaskellIdentifier t = not (isVarId t) || isReservedId t +isHaskellIdentifier t = not (isVarId t) || isReservedId t sanitize :: T.Text -> T.Text sanitize t - | isValid = t - | isHaskellIdentifier t' = "_" <> t' <> "_" - | otherwise = t' + | isValid = t + | isHaskellIdentifier t' = "_" <> t' <> "_" + | otherwise = t' where - isValid - = not (isHaskellIdentifier t) - && isVarId t - && T.all Char.isAlphaNum t + isValid = + not (isHaskellIdentifier t) + && isVarId t + && T.all Char.isAlphaNum t t' = T.map replaceInvalidCharacters . T.filter (not . parentheses) $ t replaceInvalidCharacters c - | Char.isUpper c = Char.toLower c - | Char.isSpace c = '_' - | Char.isPunctuation c = '_' -- '-' will also become a '_' - | Char.isSymbol c = '_' - | Char.isAlphaNum c = c -- Blanket condition - | otherwise = '_' -- If we're unsure we'll default to an underscore + | Char.isUpper c = Char.toLower c + | Char.isSpace c = '_' + | Char.isPunctuation c = '_' -- '-' will also become a '_' + | Char.isSymbol c = '_' + | Char.isAlphaNum c = c -- Blanket condition + | otherwise = '_' -- If we're unsure we'll default to an underscore parentheses c = case c of - '(' -> True - ')' -> True - '{' -> True - '}' -> True - '[' -> True - ']' -> True - _ -> False + '(' -> True + ')' -> True + '{' -> True + '}' -> True + '[' -> True + ']' -> True + _ -> False typeFromString :: [String] -> Q Type -typeFromString [] = fail "No type specified" +typeFromString [] = fail "No type specified" typeFromString [t] = do - maybeType <- lookupTypeName t - case maybeType of - Just name -> return (ConT name) - Nothing -> fail $ "Unsupported type: " ++ t -typeFromString [tycon,t1] = do - outer <- typeFromString [tycon] - inner <- typeFromString [t1] - return (AppT outer inner) -typeFromString [tycon,t1,t2] = do - outer <- typeFromString [tycon] - lhs <- typeFromString [t1] - rhs <- typeFromString [t2] - return (AppT (AppT outer lhs) rhs) + maybeType <- lookupTypeName t + case maybeType of + Just name -> return (ConT name) + Nothing -> fail $ "Unsupported type: " ++ t +typeFromString [tycon, t1] = do + outer <- typeFromString [tycon] + inner <- typeFromString [t1] + return (AppT outer inner) +typeFromString [tycon, t1, t2] = do + outer <- typeFromString [tycon] + lhs <- typeFromString [t1] + rhs <- typeFromString [t2] + return (AppT (AppT outer lhs) rhs) typeFromString s = fail $ "Unsupported type: " ++ (unwords s) declareColumns :: DataFrame -> DecsQ -declareColumns df = let - names = (map fst . L.sortBy (compare `on` snd). M.toList . columnIndices) df +declareColumns df = + let + names = (map fst . L.sortBy (compare `on` snd) . M.toList . columnIndices) df types = map (columnTypeString . (`unsafeGetColumn` df)) names specs = zipWith (\name type_ -> (sanitize name, type_)) names types - in fmap concat $ forM specs $ \(nm, tyStr) -> do - traceShow nm (pure ()) - ty <- typeFromString (words tyStr) - let n = mkName (T.unpack nm) - sig <- sigD n [t| Expr $(pure ty) |] - val <- valD (varP n) (normalB [| col $(TH.lift nm) |]) [] - pure [sig, val] + in + fmap concat $ forM specs $ \(nm, tyStr) -> do + traceShow nm (pure ()) + ty <- typeFromString (words tyStr) + let n = mkName (T.unpack nm) + sig <- sigD n [t|Expr $(pure ty)|] + val <- valD (varP n) (normalB [|col $(TH.lift nm)|]) [] + pure [sig, val] diff --git a/src/DataFrame/IO/CSV.hs b/src/DataFrame/IO/CSV.hs --- a/src/DataFrame/IO/CSV.hs +++ b/src/DataFrame/IO/CSV.hs @@ -1,245 +1,376 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE GADTs #-} {-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NumericUnderscores #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE RankNTypes #-} + module DataFrame.IO.CSV where -import qualified Data.ByteString.Char8 as C +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as Builder +import qualified Data.ByteString.Char8 as C8 import qualified Data.List as L -import qualified Data.Map as M -import qualified Data.Set as S +import qualified Data.Map.Strict as M import qualified Data.Text as T -import qualified Data.Text.Lazy as TL -import qualified Data.Text.Lazy.IO as TLIO +import qualified Data.Text.Encoding as TE import qualified Data.Text.IO as TIO import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Mutable as VM +import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM -import Control.Applicative ((<$>), (<|>), (<*>), (<*), (*>), many) -import Control.Monad (forM_, zipWithM_, unless, void) -import Data.Attoparsec.Text +import Control.Applicative (many, (*>), (<$>), (<*), (<*>), (<|>)) +import Control.Monad (forM_, unless, void, when, zipWithM_) +import Control.Monad.ST (runST) +import Data.Attoparsec.ByteString.Char8 hiding (endOfLine) +import Data.Bits (shiftL) import Data.Char -import DataFrame.Internal.Column (Column(..), MutableColumn(..), freezeColumn', writeColumn, columnLength) -import DataFrame.Internal.DataFrame (DataFrame(..)) -import DataFrame.Internal.Parsing -import DataFrame.Operations.Typing -import Data.Foldable (fold) +import Data.Either import Data.Function (on) import Data.IORef import Data.Maybe -import Data.Text.Encoding (decodeUtf8Lenient) -import Data.Type.Equality - ( TestEquality (testEquality), - type (:~:) (Refl) - ) +import Data.Type.Equality ( + TestEquality (testEquality), + type (:~:) (Refl), + ) +import DataFrame.Internal.Column (Column (..), MutableColumn (..), columnLength, freezeColumn', writeColumn) +import DataFrame.Internal.DataFrame (DataFrame (..)) +import DataFrame.Internal.Parsing +import DataFrame.Operations.Typing import GHC.IO.Handle (Handle) -import Prelude hiding (concat, takeWhile) import System.IO import Type.Reflection +import Prelude hiding (concat, takeWhile) --- | Record for CSV read options. -data ReadOptions = ReadOptions { - hasHeader :: Bool, - inferTypes :: Bool, - safeRead :: Bool -} +data GrowingVector a = GrowingVector + { gvData :: !(IORef (VM.IOVector a)) + , gvSize :: !(IORef Int) + , gvCapacity :: !(IORef Int) + } --- | By default we assume the file has a header, we infer the types on read --- and we convert any rows with nullish objects into Maybe (safeRead). +data GrowingUnboxedVector a = GrowingUnboxedVector + { guvData :: !(IORef (VUM.IOVector a)) + , guvSize :: !(IORef Int) + , guvCapacity :: !(IORef Int) + } + +data GrowingColumn + = GrowingInt !(GrowingUnboxedVector Int) !(IORef [Int]) + | GrowingDouble !(GrowingUnboxedVector Double) !(IORef [Int]) + | GrowingText !(GrowingVector T.Text) !(IORef [Int]) + +-- | CSV read parameters. +data ReadOptions = ReadOptions + { hasHeader :: Bool + -- ^ Whether or not the CSV file has a header. (default: True) + , inferTypes :: Bool + -- ^ Whether to try and infer types. (default: True) + , safeRead :: Bool + -- ^ Whether to partially parse values into `Maybe`/Either`. (default: True) + , chunkSize :: Int + -- ^ Default chunk size (in bytes) for csv reader. (default: 512'000) + } + defaultOptions :: ReadOptions -defaultOptions = ReadOptions { hasHeader = True, inferTypes = True, safeRead = True } +defaultOptions = + ReadOptions + { hasHeader = True + , inferTypes = True + , safeRead = True + , chunkSize = 512_000 + } --- | Reads a CSV file from the given path. --- Note this file stores intermediate temporary files --- while converting the CSV from a row to a columnar format. +newGrowingVector :: Int -> IO (GrowingVector a) +newGrowingVector !initCap = do + vec <- VM.unsafeNew initCap + GrowingVector <$> newIORef vec <*> newIORef 0 <*> newIORef initCap + +newGrowingUnboxedVector :: (VUM.Unbox a) => Int -> IO (GrowingUnboxedVector a) +newGrowingUnboxedVector !initCap = do + vec <- VUM.unsafeNew initCap + GrowingUnboxedVector <$> newIORef vec <*> newIORef 0 <*> newIORef initCap + +appendGrowingVector :: GrowingVector a -> a -> IO () +appendGrowingVector (GrowingVector vecRef sizeRef capRef) !val = do + size <- readIORef sizeRef + cap <- readIORef capRef + vec <- readIORef vecRef + + vec' <- + if size >= cap + then do + let !newCap = cap `shiftL` 1 + newVec <- VM.unsafeGrow vec newCap + writeIORef vecRef newVec + writeIORef capRef newCap + return newVec + else return vec + + VM.unsafeWrite vec' size val + writeIORef sizeRef $! size + 1 + +appendGrowingUnboxedVector :: (VUM.Unbox a) => GrowingUnboxedVector a -> a -> IO () +appendGrowingUnboxedVector (GrowingUnboxedVector vecRef sizeRef capRef) !val = do + size <- readIORef sizeRef + cap <- readIORef capRef + vec <- readIORef vecRef + + vec' <- + if size >= cap + then do + let !newCap = cap `shiftL` 1 + newVec <- VUM.unsafeGrow vec newCap + writeIORef vecRef newVec + writeIORef capRef newCap + return newVec + else return vec + + VUM.unsafeWrite vec' size val + writeIORef sizeRef $! size + 1 + +freezeGrowingVector :: GrowingVector a -> IO (V.Vector a) +freezeGrowingVector (GrowingVector vecRef sizeRef _) = do + vec <- readIORef vecRef + size <- readIORef sizeRef + V.freeze (VM.slice 0 size vec) + +freezeGrowingUnboxedVector :: (VUM.Unbox a) => GrowingUnboxedVector a -> IO (VU.Vector a) +freezeGrowingUnboxedVector (GrowingUnboxedVector vecRef sizeRef _) = do + vec <- readIORef vecRef + size <- readIORef sizeRef + VU.freeze (VUM.slice 0 size vec) + +{- | Read CSV file from path and load it into a dataframe. + +==== __Example__ +@ +ghci> D.readCsv "./data/taxi.csv" df + +@ +-} readCsv :: String -> IO DataFrame readCsv = readSeparated ',' defaultOptions --- | Reads a tab separated file from the given path. --- Note this file stores intermediate temporary files --- while converting the CSV from a row to a columnar format. +{- | Read TSV (tab separated) file from path and load it into a dataframe. + +==== __Example__ +@ +ghci> D.readTsv "./data/taxi.tsv" df + +@ +-} readTsv :: String -> IO DataFrame readTsv = readSeparated '\t' defaultOptions --- | Reads a character separated file into a dataframe using mutable vectors. +{- | Read text file with specified delimiter into a dataframe. + +==== __Example__ +@ +ghci> D.readSeparated ';' D.defaultOptions "./data/taxi.txt" df + +@ +-} readSeparated :: Char -> ReadOptions -> String -> IO DataFrame -readSeparated c opts path = do - totalRows <- countRows c path +readSeparated !sep !opts !path = do withFile path ReadMode $ \handle -> do - firstRow <- map T.strip . parseSep c <$> TIO.hGetLine handle - let columnNames = if hasHeader opts - then map (T.filter (/= '\"')) firstRow - else map (T.singleton . intToDigit) [0..(length firstRow - 1)] - -- If there was no header rewind the file cursor. + hSetBuffering handle (BlockBuffering (Just (chunkSize opts))) + + firstLine <- C8.hGetLine handle + let firstRow = parseLine sep firstLine + columnNames = + if hasHeader opts + then map (T.filter (/= '\"') . TE.decodeUtf8Lenient) firstRow + else map (T.singleton . intToDigit) [0 .. length firstRow - 1] + unless (hasHeader opts) $ hSeek handle AbsoluteSeek 0 - -- Initialize mutable vectors for each column - let numColumns = length columnNames - let numRows = if hasHeader opts then totalRows - 1 else totalRows - -- Use this row to infer the types of the rest of the column. - -- TODO: this isn't robust but in so far as this is a guess anyway - -- it's probably fine. But we should probably sample n rows and pick - -- the most likely type from the sample. - dataRow <- map T.strip . parseSep c <$> TIO.hGetLine handle + dataLine <- C8.hGetLine handle + let dataRow = parseLine sep dataLine + growingCols <- initializeColumns dataRow opts - -- This array will track the indices of all null values for each column. - -- If any exist then the column will be an optional type. - nullIndices <- VM.unsafeNew numColumns - VM.set nullIndices [] - mutableCols <- VM.unsafeNew numColumns - getInitialDataVectors numRows mutableCols dataRow + processRow 0 dataRow growingCols - -- Read rows into the mutable vectors - fillColumns numRows c mutableCols nullIndices handle + processFile handle sep growingCols (chunkSize opts) 1 - -- Freeze the mutable vectors into immutable ones - nulls' <- V.unsafeFreeze nullIndices - cols <- V.mapM (freezeColumn mutableCols nulls' opts) (V.generate numColumns id) - return $ DataFrame { - columns = cols, - columnIndices = M.fromList (zip columnNames [0..]), - dataframeDimensions = (maybe 0 columnLength (cols V.!? 0), V.length cols) - } -{-# INLINE readSeparated #-} + frozenCols <- V.fromList <$> mapM freezeGrowingColumn growingCols -getInitialDataVectors :: Int -> VM.IOVector MutableColumn -> [T.Text] -> IO () -getInitialDataVectors n mCol xs = do - forM_ (zip [0..] xs) $ \(i, x) -> do - col <- case inferValueType x of - "Int" -> MUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c) - "Double" -> MUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c) - _ -> MBoxedColumn <$> ((VM.unsafeNew n :: IO (VM.IOVector T.Text)) >>= \c -> VM.unsafeWrite c 0 x >> return c) - VM.unsafeWrite mCol i col -{-# INLINE getInitialDataVectors #-} + let numRows = maybe 0 columnLength (frozenCols V.!? 0) -inferValueType :: T.Text -> T.Text -inferValueType s = let - example = s - in case readInt example of - Just _ -> "Int" - Nothing -> case readDouble example of - Just _ -> "Double" - Nothing -> "Other" -{-# INLINE inferValueType #-} + return $ + DataFrame + { columns = frozenCols + , columnIndices = M.fromList (zip columnNames [0 ..]) + , dataframeDimensions = (numRows, V.length frozenCols) + } --- | Reads rows from the handle and stores values in mutable vectors. -fillColumns :: Int -> Char -> VM.IOVector MutableColumn -> VM.IOVector [(Int, T.Text)] -> Handle -> IO () -fillColumns n c mutableCols nullIndices handle = do - input <- newIORef (mempty :: T.Text) - forM_ [1..n] $ \i -> do - isEOF <- hIsEOF handle - input' <- readIORef input - unless (isEOF && input' == mempty) $ do - parseWith (TIO.hGetChunk handle) (parseRow c) input' >>= \case - Fail unconsumed ctx er -> do - erpos <- hTell handle - fail $ "Failed to parse CSV file around " <> show erpos <> " byte; due: " - <> show er <> "; context: " <> show ctx - Partial c -> do - fail "Partial handler is called" - Done (unconsumed :: T.Text) (row :: [T.Text]) -> do - writeIORef input unconsumed - zipWithM_ (writeValue mutableCols nullIndices i) [0..] row -{-# INLINE fillColumns #-} +initializeColumns :: [BS.ByteString] -> ReadOptions -> IO [GrowingColumn] +initializeColumns row opts = mapM initColumn row + where + initColumn :: BS.ByteString -> IO GrowingColumn + initColumn bs = do + nullsRef <- newIORef [] + let val = TE.decodeUtf8Lenient bs + if inferTypes opts + then case inferType val of + IntType -> GrowingInt <$> newGrowingUnboxedVector 1_024 <*> pure nullsRef + DoubleType -> GrowingDouble <$> newGrowingUnboxedVector 1_024 <*> pure nullsRef + TextType -> GrowingText <$> newGrowingVector 1_024 <*> pure nullsRef + else GrowingText <$> newGrowingVector 1_024 <*> pure nullsRef --- | Writes a value into the appropriate column, resizing the vector if necessary. -writeValue :: VM.IOVector MutableColumn -> VM.IOVector [(Int, T.Text)] -> Int -> Int -> T.Text -> IO () -writeValue mutableCols nullIndices count colIndex value = do - col <- VM.unsafeRead mutableCols colIndex - res <- writeColumn count value col - let modify value = VM.unsafeModify nullIndices ((count, value) :) colIndex - either modify (const (return ())) res -{-# INLINE writeValue #-} +data InferredType = IntType | DoubleType | TextType --- | Freezes a mutable vector into an immutable one, trimming it to the actual row count. -freezeColumn :: VM.IOVector MutableColumn -> V.Vector [(Int, T.Text)] -> ReadOptions -> Int -> IO Column -freezeColumn mutableCols nulls opts colIndex = do - col <- VM.unsafeRead mutableCols colIndex - freezeColumn' (nulls V.! colIndex) col -{-# INLINE freezeColumn #-} +inferType :: T.Text -> InferredType +inferType !t + | T.null t = TextType + | isJust (readInt t) = IntType + | isJust (readDouble t) = DoubleType + | otherwise = TextType -parseSep :: Char -> T.Text -> [T.Text] -parseSep c s = either error id (parseOnly (record c) s) -{-# INLINE parseSep #-} +processRow :: Int -> [BS.ByteString] -> [GrowingColumn] -> IO () +processRow !rowIdx !vals !cols = zipWithM_ (processValue rowIdx) vals cols + where + processValue :: Int -> BS.ByteString -> GrowingColumn -> IO () + processValue !idx !bs !col = do + let !val = TE.decodeUtf8Lenient bs + case col of + GrowingInt gv nulls -> + case readByteStringInt bs of + Just !i -> appendGrowingUnboxedVector gv i + Nothing -> do + appendGrowingUnboxedVector gv 0 + modifyIORef' nulls (idx :) + GrowingDouble gv nulls -> + case readByteStringDouble bs of + Just !d -> appendGrowingUnboxedVector gv d + Nothing -> do + appendGrowingUnboxedVector gv 0.0 + modifyIORef' nulls (idx :) + GrowingText gv nulls -> + if isNull val + then do + appendGrowingVector gv T.empty + modifyIORef' nulls (idx :) + else appendGrowingVector gv val -record :: Char -> Parser [T.Text] -record c = - field c `sepBy1` char c - "record" -{-# INLINE record #-} +isNull :: T.Text -> Bool +isNull t = T.null t || t == "NA" || t == "NULL" || t == "null" -parseRow :: Char -> Parser [T.Text] -parseRow c = (record c <* lineEnd) "record-new-line" +processFile :: Handle -> Char -> [GrowingColumn] -> Int -> Int -> IO () +processFile !handle !sep !cols !chunk r = do + let go remain !rowIdx = do + parseWith (C8.hGetNonBlocking handle chunk) (parseRow sep) remain >>= \case + Fail unconsumed ctx er -> do + erpos <- hTell handle + fail $ + "Failed to parse CSV file around " + <> show erpos + <> " byte; due: " + <> show er + <> "; context: " + <> show ctx + Partial c -> do + fail "Partial handler is called" + Done (unconsumed :: C8.ByteString) (row :: [C8.ByteString]) -> do + processRow rowIdx row cols + unless (row == [] || unconsumed == mempty) $ go unconsumed $! rowIdx + 1 + return () + go "" r -field :: Char -> Parser T.Text -field c = - quotedField <|> unquotedField c - "field" -{-# INLINE field #-} +parseLine :: Char -> BS.ByteString -> [BS.ByteString] +parseLine !sep = fromRight [] . parseOnly (record sep) -unquotedTerminators :: Char -> S.Set Char -unquotedTerminators sep = S.fromList [sep, '\n', '\r', '"'] +parseRow :: Char -> Parser [C8.ByteString] +parseRow sep = record sep <* (endOfLine <|> endOfInput) "CSV row" -unquotedField :: Char -> Parser T.Text -unquotedField sep = - takeWhile (not . (`S.member` terminators)) "unquoted field" - where terminators = unquotedTerminators sep -{-# INLINE unquotedField #-} +record :: Char -> Parser [BS.ByteString] +record sep = field sep `sepBy` char sep "CSV record" -quotedField :: Parser T.Text -quotedField = char '"' *> contents <* char '"' "quoted field" - where - contents = fold <$> many (unquote <|> unescape) - where - unquote = takeWhile1 (notInClass "\"\\") - unescape = char '\\' *> do - T.singleton <$> do - char '\\' <|> char '"' -{-# INLINE quotedField #-} +field :: Char -> Parser BS.ByteString +field sep = quotedField <|> unquotedField sep "CSV field" -lineEnd :: Parser () -lineEnd = - (endOfLine <|> endOfInput) - "end of line" -{-# INLINE lineEnd #-} +unquotedField :: Char -> Parser BS.ByteString +unquotedField sep = takeWhile (\c -> c /= sep && c /= '\n' && c /= '\r') "unquoted field" --- | First pass to count rows for exact allocation -countRows :: Char -> FilePath -> IO Int -countRows c path = withFile path ReadMode $! go 0 "" - where - go n input h = do - isEOF <- hIsEOF h - if isEOF && input == mempty - then pure n - else - parseWith (TIO.hGetChunk h) (parseRow c) input >>= \case - Fail unconsumed ctx er -> do - erpos <- hTell h - fail $ "Failed to parse CSV file around " <> show erpos <> " byte; due: " - <> show er <> "; context: " <> show ctx <> " " <> show unconsumed - Partial c -> do - fail $ "Partial handler is called; n = " <> show n - Done (unconsumed :: T.Text) _ -> - go (n + 1) unconsumed h -{-# INLINE countRows #-} +quotedField :: Parser BS.ByteString +quotedField = + do + char '"' + contents <- Builder.toLazyByteString <$> parseQuotedContents + char '"' + return $ BS.toStrict contents + "quoted field" + where + parseQuotedContents = mconcat <$> many quotedChar + quotedChar = + (Builder.byteString <$> takeWhile1 (/= '"')) + <|> (char '"' *> char '"' *> pure (Builder.char8 '"')) + "quoted field content" +endOfLine :: Parser () +endOfLine = + (void (string "\r\n") <|> void (char '\n') <|> void (char '\r')) + "line ending" + +freezeGrowingColumn :: GrowingColumn -> IO Column +freezeGrowingColumn (GrowingInt gv nullsRef) = do + vec <- freezeGrowingUnboxedVector gv + nulls <- readIORef nullsRef + if null nulls + then return $ UnboxedColumn vec + else do + let size = VU.length vec + mvec <- VM.new size + forM_ [0 .. size - 1] $ \i -> do + if i `elem` nulls + then VM.write mvec i Nothing + else VM.write mvec i (Just (vec VU.! i)) + BoxedColumn <$> V.freeze mvec +freezeGrowingColumn (GrowingDouble gv nullsRef) = do + vec <- freezeGrowingUnboxedVector gv + nulls <- readIORef nullsRef + if null nulls + then return $ UnboxedColumn vec + else do + let size = VU.length vec + mvec <- VM.new size + forM_ [0 .. size - 1] $ \i -> do + if i `elem` nulls + then VM.write mvec i Nothing + else VM.write mvec i (Just (vec VU.! i)) + BoxedColumn <$> V.freeze mvec +freezeGrowingColumn (GrowingText gv nullsRef) = do + vec <- freezeGrowingVector gv + nulls <- readIORef nullsRef + if null nulls + then return $ BoxedColumn vec + else do + let size = V.length vec + mvec <- VM.new size + forM_ [0 .. size - 1] $ \i -> do + if i `elem` nulls + then VM.write mvec i Nothing + else VM.write mvec i (Just (vec V.! i)) + BoxedColumn <$> V.freeze mvec + writeCsv :: String -> DataFrame -> IO () writeCsv = writeSeparated ',' -writeSeparated :: Char -- ^ Separator - -> String -- ^ Path to write to - -> DataFrame - -> IO () -writeSeparated c filepath df = withFile filepath WriteMode $ \handle ->do +writeSeparated :: + -- | Separator + Char -> + -- | Path to write to + String -> + DataFrame -> + IO () +writeSeparated c filepath df = withFile filepath WriteMode $ \handle -> do let (rows, columns) = dataframeDimensions df let headers = map fst (L.sortBy (compare `on` snd) (M.toList (columnIndices df))) TIO.hPutStrLn handle (T.intercalate ", " headers) - forM_ [0..(rows - 1)] $ \i -> do + forM_ [0 .. (rows - 1)] $ \i -> do let row = getRowAsText df i TIO.hPutStrLn handle (T.intercalate ", " row) @@ -249,43 +380,46 @@ indexMap = M.fromList (map (\(a, b) -> (b, a)) $ M.toList (columnIndices df)) go k (BoxedColumn (c :: V.Vector a)) acc = case c V.!? i of Just e -> textRep : acc - where textRep = case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> e - Nothing -> case typeRep @a of - App t1 t2 -> case eqTypeRep t1 (typeRep @Maybe) of - Just HRefl -> case testEquality t2 (typeRep @T.Text) of - Just Refl -> fromMaybe "null" e - Nothing -> (fromOptional . (T.pack . show)) e - where fromOptional s - | T.isPrefixOf "Just " s = T.drop (T.length "Just ") s - | otherwise = "null" - Nothing -> (T.pack . show) e - _ -> (T.pack . show) e + where + textRep = case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> e + Nothing -> case typeRep @a of + App t1 t2 -> case eqTypeRep t1 (typeRep @Maybe) of + Just HRefl -> case testEquality t2 (typeRep @T.Text) of + Just Refl -> fromMaybe "null" e + Nothing -> (fromOptional . (T.pack . show)) e + where + fromOptional s + | T.isPrefixOf "Just " s = T.drop (T.length "Just ") s + | otherwise = "null" + Nothing -> (T.pack . show) e + _ -> (T.pack . show) e Nothing -> error $ "Column " - ++ T.unpack (indexMap M.! k) - ++ " has less items than " - ++ "the other columns at index " - ++ show i + ++ T.unpack (indexMap M.! k) + ++ " has less items than " + ++ "the other columns at index " + ++ show i go k (UnboxedColumn c) acc = case c VU.!? i of Just e -> T.pack (show e) : acc Nothing -> error $ "Column " - ++ T.unpack (indexMap M.! k) - ++ " has less items than " - ++ "the other columns at index " - ++ show i + ++ T.unpack (indexMap M.! k) + ++ " has less items than " + ++ "the other columns at index " + ++ show i go k (OptionalColumn (c :: V.Vector (Maybe a))) acc = case c V.!? i of Just e -> textRep : acc - where textRep = case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> fromMaybe "Nothing" e - Nothing -> (T.pack . show) e + where + textRep = case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> fromMaybe "Nothing" e + Nothing -> (T.pack . show) e Nothing -> error $ "Column " - ++ T.unpack (indexMap M.! k) - ++ " has less items than " - ++ "the other columns at index " - ++ show i + ++ T.unpack (indexMap M.! k) + ++ " has less items than " + ++ "the other columns at index " + ++ show i diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs --- a/src/DataFrame/IO/Parquet.hs +++ b/src/DataFrame/IO/Parquet.hs @@ -1,1560 +1,250 @@ -{-# LANGUAGE ExplicitNamespaces #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE RecordWildCards #-} - -module DataFrame.IO.Parquet ( - readParquet - ) where - -import Codec.Compression.Zstd.Streaming -import Control.Monad -import qualified Data.ByteString as BSO -import qualified Data.ByteString.Char8 as BS -import Data.Char -import Data.Foldable -import qualified Data.Vector.Unboxed as VU -import Data.IORef -import qualified Data.Map as M -import Data.Maybe -import qualified Data.Text as T -import DataFrame.Internal.DataFrame (DataFrame) -import qualified DataFrame.Internal.DataFrame as DI -import qualified DataFrame.Internal.Column as DI -import qualified DataFrame.Operations.Core as DI -import qualified Snappy as Snappy -import Foreign -import GHC.Float -import GHC.IO (unsafePerformIO) -import System.IO - -import DataFrame.IO.Parquet.ColumnStatistics -import DataFrame.IO.Parquet.Compression -import DataFrame.IO.Parquet.Encoding -import DataFrame.IO.Parquet.Page -import DataFrame.IO.Parquet.Types - -footerSize :: Integer -footerSize = 8 - -data PageEncodingStats = PageEncodingStats - { pageEncodingPageType :: PageType, - pageEncoding :: ParquetEncoding, - pagesWithEncoding :: Int32 - } - deriving (Show, Eq) - -emptyPageEncodingStats :: PageEncodingStats -emptyPageEncodingStats = PageEncodingStats PAGE_TYPE_UNKNOWN PARQUET_ENCODING_UNKNOWN 0 - -data SizeStatistics = SizeStatisics - { unencodedByteArrayDataTypes :: Int64, - repetitionLevelHistogram :: [Int64], - definitionLevelHistogram :: [Int64] - } - deriving (Show, Eq) - -emptySizeStatistics :: SizeStatistics -emptySizeStatistics = SizeStatisics 0 [] [] - -data BoundingBox = BoundingBox - { xmin :: Double, - xmax :: Double, - ymin :: Double, - ymax :: Double, - zmin :: Double, - zmax :: Double, - mmin :: Double, - mmax :: Double - } - deriving (Show, Eq) - -emptyBoundingBox :: BoundingBox -emptyBoundingBox = BoundingBox 0 0 0 0 0 0 0 0 - -data GeospatialStatistics = GeospatialStatistics - { bbox :: BoundingBox, - geospatialTypes :: [Int32] - } - deriving (Show, Eq) - -emptyGeospatialStatistics :: GeospatialStatistics -emptyGeospatialStatistics = GeospatialStatistics emptyBoundingBox [] - -emptyKeyValue :: KeyValue -emptyKeyValue = KeyValue {key = "", value = ""} - -data ColumnMetaData = ColumnMetaData - { columnType :: ParquetType, - columnEncodings :: [ParquetEncoding], - columnPathInSchema :: [String], - columnCodec :: CompressionCodec, - columnNumValues :: Int64, - columnTotalUncompressedSize :: Int64, - columnTotalCompressedSize :: Int64, - columnKeyValueMetadata :: [KeyValue], - columnDataPageOffset :: Int64, - columnIndexPageOffset :: Int64, - columnDictionaryPageOffset :: Int64, - columnStatistics :: ColumnStatistics, - columnEncodingStats :: [PageEncodingStats], - bloomFilterOffset :: Int64, - bloomFilterLength :: Int32, - columnSizeStatistics :: SizeStatistics, - columnGeospatialStatistics :: GeospatialStatistics - } - deriving (Show, Eq) - -emptyColumnMetadata :: ColumnMetaData -emptyColumnMetadata = ColumnMetaData PARQUET_TYPE_UNKNOWN [] [] COMPRESSION_CODEC_UNKNOWN 0 0 0 [] 0 0 0 emptyColumnStatistics [] 0 0 emptySizeStatistics emptyGeospatialStatistics - -data ColumnCryptoMetadata - = COLUMN_CRYPTO_METADATA_UNKNOWN - | ENCRYPTION_WITH_FOOTER_KEY - | EncryptionWithColumnKey - { columnCryptPathInSchema :: [String], - columnKeyMetadata :: [Word8] - } - deriving (Show, Eq) - -data ColumnChunk = ColumnChunk - { columnChunkFilePath :: String, - columnChunkMetadataFileOffset :: Int64, - columnMetaData :: ColumnMetaData, - columnChunkOffsetIndexOffset :: Int64, - columnChunkOffsetIndexLength :: Int32, - columnChunkColumnIndexOffset :: Int64, - columnChunkColumnIndexLength :: Int32, - cryptoMetadata :: ColumnCryptoMetadata, - encryptedColumnMetadata :: [Word8] - } - deriving (Show, Eq) - -emptyColumnChunk :: ColumnChunk -emptyColumnChunk = ColumnChunk "" 0 emptyColumnMetadata 0 0 0 0 COLUMN_CRYPTO_METADATA_UNKNOWN [] - -data SortingColumn = SortingColumn - { columnIndex :: Int32, - columnOrderDescending :: Bool, - nullFirst :: Bool - } - deriving (Show, Eq) - -emptySortingColumn :: SortingColumn -emptySortingColumn = SortingColumn 0 False False - -data RowGroup = RowGroup - { rowGroupColumns :: [ColumnChunk], - totalByteSize :: Int64, - rowGroupNumRows :: Int64, - rowGroupSortingColumns :: [SortingColumn], - fileOffset :: Int64, - totalCompressedSize :: Int64, - ordinal :: Int16 - } - deriving (Show, Eq) - -emptyRowGroup :: RowGroup -emptyRowGroup = RowGroup [] 0 0 [] 0 0 0 - -data ColumnOrder - = TYPE_ORDER - | COLUMN_ORDER_UNKNOWN - deriving (Show, Eq) - -data EncryptionAlgorithm - = ENCRYPTION_ALGORITHM_UNKNOWN - | AesGcmV1 - { aadPrefix :: [Word8], - aadFileUnique :: [Word8], - supplyAadPrefix :: Bool - } - | AesGcmCtrV1 - { aadPrefix :: [Word8], - aadFileUnique :: [Word8], - supplyAadPrefix :: Bool - } - deriving (Show, Eq) - -data KeyValue = KeyValue - { key :: String, - value :: String - } - deriving (Show, Eq) - -data FileMetadata = FileMetaData - { version :: Int32, - schema :: [SchemaElement], - numRows :: Integer, - rowGroups :: [RowGroup], - keyValueMetadata :: [KeyValue], - createdBy :: Maybe String, - columnOrders :: [ColumnOrder], - encryptionAlgorithm :: EncryptionAlgorithm, - footerSigningKeyMetadata :: [Word8] - } - deriving (Show, Eq) - -defaultMetadata :: FileMetadata -defaultMetadata = - FileMetaData - { version = 0, - schema = [], - numRows = 0, - rowGroups = [], - keyValueMetadata = [], - createdBy = Nothing, - columnOrders = [], - encryptionAlgorithm = ENCRYPTION_ALGORITHM_UNKNOWN, - footerSigningKeyMetadata = [] - } - -readParquet :: String -> IO DataFrame -readParquet path = withBinaryFile path ReadMode $ \handle -> do - (size, magicString) <- readMetadataSizeFromFooter handle - when (magicString /= "PAR1") $ error "Invalid Parquet file" - - colMap <- newIORef (M.empty :: (M.Map T.Text DI.Column)) - colNames <- newIORef ([] :: [T.Text]) - - fileMetadata <- readMetadata handle size - forM_ (rowGroups fileMetadata) $ \r -> do - forM_ (rowGroupColumns r) $ \c -> do - let metadata = columnMetaData c - let colDataPageOffset = columnDataPageOffset metadata - let colDictionaryPageOffset = columnDictionaryPageOffset metadata - let colStart = if colDictionaryPageOffset > 0 && colDataPageOffset > colDictionaryPageOffset - then colDictionaryPageOffset - else colDataPageOffset - let colLength = columnTotalCompressedSize metadata - columnBytes <-readBytes handle colStart colLength - (maybePage, res) <- readPage (columnCodec metadata) columnBytes - case maybePage of - Just p -> if isDictionaryPage p - then do - (maybePage', res') <- readPage (columnCodec metadata) res - let p' = fromMaybe (error "Empty page") maybePage' - let schemaElem = filter (\se -> (elementName se) == (T.pack $ head (columnPathInSchema metadata))) (schema fileMetadata) - let rep = if null schemaElem then UNKNOWN_REPETITION_TYPE else ((repetitionType . head) schemaElem) - when (rep == REPEATED || rep == UNKNOWN_REPETITION_TYPE) (error $ "REPETITION TYPE NOT SUPPORTED: " ++ show rep) - - case ((definitionLevelEncoding . pageTypeHeader . pageHeader ) p') of - ERLE -> do - let rleColumn = case columnType metadata of - PBYTE_ARRAY -> readByteArrayColumn (pageBytes p) - PDOUBLE -> readDoubleColumn (pageBytes p) - PINT32 -> readInt32Column (pageBytes p) - t -> error $ "UNKNOWN TYPE: " ++ (show t) - let nbytes = littleEndianInt32 (take 4 (pageBytes p')) - let rleDecoder = MkRleDecoder (drop 4 (pageBytes p')) 1 0 0 - - let bytesAfterLevels = - case rep of - REQUIRED -> pageBytes p' -- no def-levels to skip - OPTIONAL -> -- skip def-level RLE block (len + payload) - let nbytes = littleEndianInt32 (take 4 (pageBytes p')) - in drop (fromIntegral nbytes + 4) (pageBytes p') - REPEATED -> error "REPEATED not supported" - let bitWidth = head bytesAfterLevels - let indexDecoder = MkRleDecoder (tail bytesAfterLevels) (fromIntegral bitWidth) 0 0 - - let finalCol = DI.takeColumn ((fromIntegral . dataPageHeaderNumValues . pageTypeHeader . pageHeader) p') (decodeDictionary rleColumn rleDecoder indexDecoder) - let colName = T.pack $ head (columnPathInSchema metadata) - - modifyIORef' colNames (++[colName]) - modifyIORef' colMap (\m -> M.insertWith (\l r -> fromMaybe (error "UNEXPECTED") (DI.concatColumns l r)) colName finalCol m) - other -> error $ "UNSUPPORTED ENCODING: " ++ (show other) - else do - -- p is the first page we read; since it’s not a dictionary page, it must be a Data Page - let schemaElem = filter (\se -> (elementName se) == (T.pack $ head (columnPathInSchema metadata))) (schema fileMetadata) - let rep = if null schemaElem then UNKNOWN_REPETITION_TYPE else ((repetitionType . head) schemaElem) - when (rep == REPEATED || rep == UNKNOWN_REPETITION_TYPE) (error $ "REPETITION TYPE NOT SUPPORTED: " ++ show rep) - - -- Determine how many values to read from this page - let nVals = case pageTypeHeader (pageHeader p) of - DataPageHeader{..} -> fromIntegral dataPageHeaderNumValues - DataPageHeaderV2{..} -> fromIntegral dataPageHeaderV2NumValues - _ -> error "Unexpected page header for data page" - - -- For REQUIRED fields there are no level streams to skip in Data Page v1. - -- For Data Page v2, def/rep sections are at the front; REQUIRED => both lengths are 0. - let bytesAfterLevels = case pageTypeHeader (pageHeader p) of - DataPageHeader{..} -> pageBytes p -- v1, REQUIRED => no level bytes - DataPageHeaderV2{..} -> drop (fromIntegral (definitionLevelByteLength + repetitionLevelByteLength)) (pageBytes p) - _ -> pageBytes p - - -- Check the value encoding; support PLAIN here - let enc = case pageTypeHeader (pageHeader p) of - DataPageHeader{..} -> dataPageHeaderEncoding - DataPageHeaderV2{..} -> dataPageHeaderV2Encoding - _ -> PARQUET_ENCODING_UNKNOWN - when (enc /= EPLAIN) (error $ "Unsupported non-dictionary encoding: " ++ show enc) - - -- Read values in PLAIN format - let ptype = columnType metadata - finalCol <- - case ptype of - PINT32 -> let (col, _) = readPlainColumn PINT32 nVals bytesAfterLevels in pure col - PDOUBLE -> let (col, _) = readPlainColumn PDOUBLE nVals bytesAfterLevels in pure col - PBYTE_ARRAY -> - let (col, _) = readPlainColumn PBYTE_ARRAY nVals bytesAfterLevels in pure col - PFIXED_LEN_BYTE_ARRAY -> - let len = fromIntegral (typeLength (head schemaElem)) - (col, _) = readPlainFixedLenColumn len nVals bytesAfterLevels - in pure col - other -> error $ "PLAIN not implemented for: " ++ show other - - let colName = T.pack $ head (columnPathInSchema metadata) - modifyIORef' colNames (++[colName]) - modifyIORef' colMap (\m -> M.insertWith (\l r -> fromMaybe (error "UNEXPECTED") (DI.concatColumns l r)) colName finalCol m) - Nothing -> pure () - - c' <- readIORef colMap - colNames' <- readIORef colNames - let asscList = map (\name -> (name, c' M.! name)) colNames' - pure $ DI.fromNamedColumns asscList - --- Reads exactly `n` values in PLAIN encoding and returns (column, remainingBytes) -readPlainColumn :: ParquetType -> Int -> [Word8] -> (DI.Column, [Word8]) -readPlainColumn PINT32 n bs = - let (vals, rest) = readNInt32 n bs - in (DI.fromList vals, rest) - -readPlainColumn PDOUBLE n bs = - let (vals, rest) = readNDouble n bs - in (DI.fromList vals, rest) - --- PLAIN BYTE_ARRAY: 4-byte little-endian length followed by bytes -readPlainColumn PBYTE_ARRAY n bs = - let (vals, rest) = readNByteArrays n bs - in (DI.fromList (map (T.pack . map (chr . fromIntegral)) vals), rest) - --- PLAIN FIXED_LEN_BYTE_ARRAY: each value is `len` bytes (from schema element) -readPlainColumn PFIXED_LEN_BYTE_ARRAY n bs = - error "FIXED_LEN_BYTE_ARRAY requires typeLength from schema; use readPlainFixedLenColumn len" - -readPlainFixedLenColumn :: Int -> Int -> [Word8] -> (DI.Column, [Word8]) -readPlainFixedLenColumn len n bs = - let (vals, rest) = splitFixed n len bs - in (DI.fromList (map (T.pack . map (chr . fromIntegral)) vals), rest) - --- Helpers -readNInt32 :: Int -> [Word8] -> ([Int32],[Word8]) -readNInt32 0 bs = ([], bs) -readNInt32 k bs = - let x = littleEndianInt32 (take 4 bs) - bs' = drop 4 bs - (xs, rest) = readNInt32 (k-1) bs' - in (x:xs, rest) - -readNDouble :: Int -> [Word8] -> ([Double],[Word8]) -readNDouble 0 bs = ([], bs) -readNDouble k bs = - let x = castWord64ToDouble (littleEndianWord64 (take 8 bs)) - bs' = drop 8 bs - (xs, rest) = readNDouble (k-1) bs' - in (x:xs, rest) - -readNByteArrays :: Int -> [Word8] -> ([[Word8]],[Word8]) -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 - (xs, rest) = readNByteArrays (k-1) bs' - in (body:xs, rest) - -splitFixed :: Int -> Int -> [Word8] -> ([[Word8]],[Word8]) -splitFixed 0 _ bs = ([], bs) -splitFixed k len bs = - let body = take len bs - bs' = drop len bs - (xs, rest) = splitFixed (k-1) len bs' - in (body:xs, rest) - -decodeDictionary :: DI.Column -> RleDecoder -> RleDecoder -> DI.Column -decodeDictionary col rleDecoder indexDecoder - | repCount indexDecoder > 0 = error "UNIMPLEMENTED: Repetition not supported" - | litCount indexDecoder > 0 = decodeDictionary (DI.atIndicesStable (VU.map fromIntegral (getIndices indexDecoder)) col) rleDecoder (indexDecoder { litCount = 0 }) - | otherwise = let - (finished, indexDecoder') = advance indexDecoder - in if finished then col else decodeDictionary col rleDecoder indexDecoder' - -advance :: RleDecoder -> (Bool, RleDecoder) -advance indexDecoder - | (rleDecoderData indexDecoder) == [] = (True, indexDecoder) - | otherwise = let - (indicator, remaining) = readUVarInt (rleDecoderData indexDecoder) - isLiteral = (indicator .&. 1) /= 0 - countValues = (fromIntegral (indicator `shiftR` 1) :: Int32) - litCount = if isLiteral then (countValues * 8) else 0 - in if isLiteral then (False, indexDecoder { rleDecoderData = remaining, litCount = litCount }) else (True, indexDecoder) -- (error "NON-LITERAL TYPES NOT YET SUPPORTED") - -getIndices :: RleDecoder -> VU.Vector Word32 -getIndices indexDecoder - | rleBitWidth indexDecoder == 5 = unpackWidth5 (rleDecoderData indexDecoder) - | rleBitWidth indexDecoder == 1 = unpackWidth1 (rleDecoderData indexDecoder) - | rleBitWidth indexDecoder == 2 = unpackWidth2 (rleDecoderData indexDecoder) - | rleBitWidth indexDecoder == 3 = unpackWidth3 (rleDecoderData indexDecoder) - | otherwise = error $ "Unsupported bit width: " ++ (show (rleBitWidth indexDecoder)) - -unpackWidth5 :: [Word8] -> VU.Vector Word32 -unpackWidth5 [] = VU.empty -unpackWidth5 bytes = let - n0 = littleEndianWord32 $ take 4 bytes - n1 = littleEndianWord32 $ take 4 $ drop 4 bytes - n2 = littleEndianWord32 $ take 4 $ drop 8 bytes - n3 = littleEndianWord32 $ take 4 $ drop 12 bytes - n4 = littleEndianWord32 $ take 4 $ drop 16 bytes - out0 = (n0 .>>. 0) `mod` (1 .<<. 5) - out1 = (n0 .>>. 5) `mod` (1 .<<. 5) - out2 = (n0 .>>. 10) `mod` (1 .<<. 5) - out3 = (n0 .>>. 15) `mod` (1 .<<. 5) - out4 = (n0 .>>. 20) `mod` (1 .<<. 5) - out5 = (n0 .>>. 25) `mod` (1 .<<. 5) - out6 = (n0 .>>. 30) .|. ((n1 `mod` (1 .<<. 3)) .<<. (5 - 3)) - out7 = (n1 .>>. 3) `mod` (1 .<<. 5) - out8 = (n1 .>>. 8) `mod` (1 .<<. 5) - out9 = (n1 .>>. 13) `mod` (1 .<<. 5) - out10 = (n1 .>>. 18) `mod` (1 .<<. 5) - out11 = (n1 .>>. 23) `mod` (1 .<<. 5) - out12 = (n1 .>>. 28) .|. (n2 `mod` (1 .<<. 1)) .<<. (5 - 1) - out13 = (n2 .>>. 1) `mod` (1 .<<. 5) - out14 = (n2 .>>. 6) `mod` (1 .<<. 5) - out15 = (n2 .>>. 11) `mod` (1 .<<. 5) - out16 = (n2 .>>. 16) `mod` (1 .<<. 5) - out17 = (n2 .>>. 21) `mod` (1 .<<. 5) - out18 = (n2 .>>. 26) `mod` (1 .<<. 5) - out19 = (n2 .>>. 31) .|. (n3 `mod` (1 .<<. 4)) .<<. (5 - 4) - out20 = (n3 .>>. 4) `mod` (1 .<<. 5) - out21 = (n3 .>>. 9) `mod` (1 .<<. 5) - out22 = (n3 .>>. 14) `mod` (1 .<<. 5) - out23 = (n3 .>>. 19) `mod` (1 .<<. 5) - out24 = (n3 .>>. 24) `mod` (1 .<<. 5) - out25 = (n3 .>>. 29) .|. (n4 `mod` (1 .<<. 2)) .<<. (5 - 2) - out26 = (n4 .>>. 2) `mod` (1 .<<. 5) - out27 = (n4 .>>. 7) `mod` (1 .<<. 5) - out28 = (n4 .>>. 12) `mod` (1 .<<. 5) - out29 = (n4 .>>. 17) `mod` (1 .<<. 5) - out30 = (n4 .>>. 22) `mod` (1 .<<. 5) - out31 = (n4 .>>. 27) - in (VU.fromList [out0,out1,out2,out3,out4,out5,out6,out7,out8,out9,out10,out11,out12,out13,out14,out15,out16,out17,out18,out19,out20,out21,out22,out23,out24,out25,out26,out27,out28,out29,out30,out31]) VU.++ (unpackWidth5 (drop 20 bytes)) - -unpackWidth2, unpackWidth1, unpackWidth3 :: [Word8] -> VU.Vector Word32 -unpackWidth1 [] = VU.empty -unpackWidth1 bytes = let - n = littleEndianWord32 $ take 4 bytes - in VU.fromList (map (\i -> (n .>>. i) .&. 1) [0..31]) VU.++ (unpackWidth1 (drop 4 bytes)) -unpackWidth2 [] = VU.empty -unpackWidth2 bytes = let - n = littleEndianWord32 $ take 4 bytes - in VU.fromList (map (\i -> (n .>>. (i * 2)) `mod` (1 .<<. 2)) [0..14] ++ [n .>>. 30]) VU.++ (unpackWidth2 (drop 4 bytes)) -unpackWidth3 [] = VU.empty -unpackWidth3 bytes = let - n0 = littleEndianWord32 $ take 4 bytes - n1 = littleEndianWord32 $ take 4 $ drop 4 bytes - n2 = littleEndianWord32 $ take 4 $ drop 8 bytes - out0 = (n0 .>>. 0) `mod` (1 .<<. 3) - out1 = (n0 .>>. 3) `mod` (1 .<<. 3) - out2 = (n0 .>>. 6) `mod` (1 .<<. 3) - out3 = (n0 .>>. 9) `mod` (1 .<<. 3) - out4 = (n0 .>>. 12) `mod` (1 .<<. 3) - out5 = (n0 .>>. 15) `mod` (1 .<<. 3) - out6 = (n0 .>>. 18) `mod` (1 .<<. 3) - out7 = (n0 .>>. 21) `mod` (1 .<<. 3) - out8 = (n0 .>>. 24) `mod` (1 .<<. 3) - out9 = (n0 .>>. 27) `mod` (1 .<<. 3) - out10 = (n0 .>>. 30) .|. (n1 `mod` (1 .<<. 1)) .<<. (3 - 1) - out11 = (n1 .>>. 1) `mod` (1 .<<. 3) - out12 = (n1 .>>. 4) `mod` (1 .<<. 3) - out13 = (n1 .>>. 7) `mod` (1 .<<. 3) - out14 = (n1 .>>. 10) `mod` (1 .<<. 3) - out15 = (n1 .>>. 13) `mod` (1 .<<. 3) - out16 = (n1 .>>. 16) `mod` (1 .<<. 3) - out17 = (n1 .>>. 19) `mod` (1 .<<. 3) - out18 = (n1 .>>. 22) `mod` (1 .<<. 3) - out19 = (n1 .>>. 25) `mod` (1 .<<. 3) - out20 = (n1 .>>. 28) `mod` (1 .<<. 3) - out21 = ((n1 .>>. 31) `mod` (1 .<<. 3)) .|. (n2 `mod` (1 .<<. 2)) .<<. (3 - 2) - out22 = (n2 .>>. 2) `mod` (1 .<<. 3) - out23 = (n2 .>>. 5) `mod` (1 .<<. 3) - out24 = (n2 .>>. 8) `mod` (1 .<<. 3) - out25 = (n2 .>>. 11) `mod` (1 .<<. 3) - out26 = (n2 .>>. 14) `mod` (1 .<<. 3) - out27 = (n2 .>>. 17) `mod` (1 .<<. 3) - out28 = (n2 .>>. 20) `mod` (1 .<<. 3) - out29 = (n2 .>>. 23) `mod` (1 .<<. 3) - out30 = (n2 .>>. 26) `mod` (1 .<<. 3) - out31 = (n2 .>>. 29) - in (VU.fromList [out0,out1,out2,out3,out4,out5,out6,out7,out8,out9,out10,out11,out12,out13,out14,out15,out16,out17,out18,out19,out20,out21,out22,out23,out24,out25,out26,out27,out28,out29,out30,out31]) VU.++ (unpackWidth3 (drop 12 bytes)) - -data RleDecoder = MkRleDecoder { rleDecoderData :: [Word8] - , rleBitWidth :: Int32 - , repCount :: Int32 - , litCount :: Int32 - } deriving (Show, Eq) - -expandDictionary :: [Word8] -> [Word8] -expandDictionary (bitWidth:rest) = rest - -readInt32Column :: [Word8] -> DI.Column -readInt32Column = DI.fromList . readPageInt32 - -readDoubleColumn :: [Word8] -> DI.Column -readDoubleColumn = DI.fromList . readPageWord64 - -readByteArrayColumn :: [Word8] -> DI.Column -readByteArrayColumn = DI.fromList .readPageBytes - -readPageInt32 :: [Word8] -> [Int32] -readPageInt32 [] = [] -readPageInt32 xs = (fromIntegral (littleEndianInt32 (take 4 xs))) : readPageInt32 (drop 4 xs) - -readPageWord64 :: [Word8] -> [Double] -readPageWord64 [] = [] -readPageWord64 xs = (castWord64ToDouble (littleEndianWord64 (take 8 xs))) : readPageWord64 (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) - -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 - - -- Weird round about way to uncompress zstd files compressed using the - -- streaming API - 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) - comp -> error ("UNSUPPORTED_COMPRESSION TYPE: " ++ (show comp)) - pure $ (Just $ Page hdr (BSO.unpack fullData), drop (fromIntegral $ compressedPageSize hdr) rem) - -data Page = Page { pageHeader :: PageHeader - , pageBytes :: [Word8] } deriving (Show, Eq) - -data PageHeader = PageHeader { pageHeaderPageType :: PageType - , uncompressedPageSize :: Int32 - , compressedPageSize ::Int32 - , pageHeaderCrcChecksum :: Int32 - , pageTypeHeader :: PageTypeHeader - } deriving (Show, Eq) - - -emptyPageHeader = PageHeader PAGE_TYPE_UNKNOWN 0 0 0 PAGE_TYPE_HEADER_UNKNOWN - -isDataPage :: Page -> Bool -isDataPage page = case pageTypeHeader (pageHeader page) of - DataPageHeader {..} -> True - DataPageHeaderV2 {..} -> True - _ -> False - -isDictionaryPage :: Page -> Bool -isDictionaryPage page = case pageTypeHeader (pageHeader page) of - DictionaryPageHeader {..} -> True - _ -> False - -data PageTypeHeader = DataPageHeader { dataPageHeaderNumValues :: Int32 - , dataPageHeaderEncoding :: ParquetEncoding - , definitionLevelEncoding :: ParquetEncoding - , repetitionLevelEncoding :: ParquetEncoding - , dataPageHeaderStatistics :: ColumnStatistics - } - | DataPageHeaderV2 { dataPageHeaderV2NumValues :: Int32 - , dataPageHeaderV2NumNulls :: Int32 - , dataPageHeaderV2NumRows :: Int32 - , dataPageHeaderV2Encoding :: ParquetEncoding - , definitionLevelByteLength :: Int32 - , repetitionLevelByteLength :: Int32 - , dataPageHeaderV2IsCompressed :: Bool - , dataPageHeaderV2Statistics :: ColumnStatistics - } - | DictionaryPageHeader { dictionaryPageHeaderNumValues :: Int32 - , dictionaryPageHeaderEncoding :: ParquetEncoding - , dictionaryPageIsSorted :: Bool - } - | INDEX_PAGE_HEADER - | PAGE_TYPE_HEADER_UNKNOWN deriving (Show, Eq) - -emptyDictionaryPageHeader = DictionaryPageHeader 0 PARQUET_ENCODING_UNKNOWN False -emptyDataPageHeader = DataPageHeader 0 PARQUET_ENCODING_UNKNOWN PARQUET_ENCODING_UNKNOWN PARQUET_ENCODING_UNKNOWN emptyColumnStatistics -emptyDataPageHeaderV2 = DataPageHeaderV2 0 0 0 PARQUET_ENCODING_UNKNOWN 0 0 False emptyColumnStatistics - -readPageHeader :: PageHeader -> [Word8] -> Int16 -> (PageHeader, [Word8]) -readPageHeader hdr [] _ = (hdr, []) -readPageHeader hdr xs lastFieldId = let - fieldContents = readField' xs lastFieldId - in case fieldContents of - Nothing -> (hdr, tail 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 - 5 -> let - (dataPageHeader, rem') = readPageTypeHeader emptyDataPageHeader rem 0 - in readPageHeader (hdr {pageTypeHeader = dataPageHeader}) rem' identifier - 7 -> let - (dictionaryPageHeader, rem') = readPageTypeHeader emptyDictionaryPageHeader rem 0 - in readPageHeader (hdr {pageTypeHeader = dictionaryPageHeader}) rem' identifier - n -> error $ show n - -readPageTypeHeader :: PageTypeHeader -> [Word8] -> Int16 -> (PageTypeHeader, [Word8]) -readPageTypeHeader hdr [] _ = (hdr, []) -readPageTypeHeader hdr@(DictionaryPageHeader {..}) xs lastFieldId = let - fieldContents = readField' xs lastFieldId - in case fieldContents of - Nothing -> (hdr, tail 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: rem') = rem - in readPageTypeHeader (hdr {dictionaryPageIsSorted = isSorted == compactBooleanTrue}) rem' identifier - n -> error $ show n -readPageTypeHeader hdr@(DataPageHeader {..}) xs lastFieldId = let - fieldContents = readField' xs lastFieldId - in case fieldContents of - Nothing -> (hdr, tail 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 - -readStatisticsFromBytes :: ColumnStatistics -> [Word8] -> Int16 -> (ColumnStatistics, [Word8]) -readStatisticsFromBytes cs xs lastFieldId = let - fieldContents = readField' xs lastFieldId - in case fieldContents of - Nothing -> (cs, tail xs) - Just (rem, elemType, identifier) -> case identifier of - 1 -> let - (maxInBytes, rem') = readByteStringFromBytes rem - in readStatisticsFromBytes (cs {columnMax = maxInBytes}) rem' identifier - 2 -> let - (minInBytes, rem') = readByteStringFromBytes rem - in readStatisticsFromBytes (cs {columnMin = minInBytes}) rem' identifier - 3 -> let - (nullCount, rem') = readIntFromBytes @Int64 rem - in readStatisticsFromBytes (cs {columnNullCount = nullCount}) rem' identifier - 4 -> let - (distinctCount, rem') = readIntFromBytes @Int64 rem - in readStatisticsFromBytes (cs {columnDistictCount = distinctCount}) rem' identifier - 5 -> let - (maxInBytes, rem') = readByteStringFromBytes rem - in readStatisticsFromBytes (cs {columnMaxValue = maxInBytes}) rem' identifier - 6 -> let - (minInBytes, rem') = readByteStringFromBytes rem - in readStatisticsFromBytes (cs {columnMinValue = minInBytes}) rem' identifier - 7 -> let - (isMaxValueExact: rem') = rem - in readStatisticsFromBytes (cs {isColumnMaxValueExact = isMaxValueExact == compactBooleanTrue}) rem' identifier - 8 -> let - (isMinValueExact: rem') = rem - in readStatisticsFromBytes (cs {isColumnMinValueExact = isMinValueExact == compactBooleanTrue}) rem' identifier - n -> error $ show n - -readBytes :: Handle -> Int64 -> Int64 -> IO [Word8] -readBytes handle colStart colLen = do - buf <- mallocBytes (fromIntegral colLen) :: IO (Ptr Word8) - hSeek handle AbsoluteSeek (fromIntegral colStart) - _ <- hGetBuf handle buf (fromIntegral colLen) - columnBytes <- readByteString' buf colLen - free buf - pure columnBytes - -numBytesInFile :: Handle -> IO Integer -numBytesInFile handle = do - hSeek handle SeekFromEnd 0 - hTell handle - -readMetadataSizeFromFooter :: Handle -> IO (Integer, BS.ByteString) -readMetadataSizeFromFooter handle = do - footerOffSet <- numBytesInFile handle - - buf <- mallocBytes (fromIntegral footerSize) :: IO (Ptr Word8) - - hSeek handle AbsoluteSeek (fromIntegral $! footerOffSet - footerSize) - n <- hGetBuf handle buf (fromIntegral footerSize) - - -- The bytes that store the metadata size. - sizeBytes <- mapM (\i -> fromIntegral <$> (peekElemOff buf i :: IO Word8) :: IO Int32) [0 .. 3] - let size = fromIntegral $ foldl' (.|.) 0 $! zipWith shift sizeBytes [0, 8, 16, 24] - - magicStringBytes <- mapM (\i -> peekElemOff buf i :: IO Word8) [4 .. 7] - let magicString = BSO.pack magicStringBytes - free buf - return (size, magicString) - -readMetadata :: Handle -> Integer -> IO FileMetadata -readMetadata handle size = do - metaDataBuf <- mallocBytes (fromIntegral size) :: IO (Ptr Word8) - footerOffSet <- numBytesInFile handle - - hSeek handle AbsoluteSeek (fromIntegral $! footerOffSet - footerSize - size) - - metadataBytesRead <- hGetBuf handle metaDataBuf (fromIntegral size) - let lastFieldId = 0 - let fieldStack = [] - bufferPos <- newIORef (0 :: Int) - metadata <- readFileMetaData defaultMetadata metaDataBuf bufferPos lastFieldId fieldStack - free metaDataBuf - return metadata - -readFileMetaData :: FileMetadata -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO FileMetadata -readFileMetaData metadata metaDataBuf bufferPos lastFieldId fieldStack = do - fieldContents <- readField metaDataBuf bufferPos lastFieldId fieldStack - case fieldContents of - Nothing -> return metadata - Just (elemType, identifier) -> case identifier of - 1 -> do - version <- readIntFromBuffer @Int32 metaDataBuf bufferPos - readFileMetaData (metadata {version = version}) metaDataBuf bufferPos identifier fieldStack - 2 -> do - -- We can do some type checking/exception handling here. - -- Check elemType == List - sizeAndType <- readAndAdvance bufferPos metaDataBuf - let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int - -- type of the contents of the list. - let elemType = toTType sizeAndType - schemaElements <- replicateM sizeOnly (readSchemaElement defaultSchemaElement metaDataBuf bufferPos 0 []) - readFileMetaData (metadata {schema = schemaElements}) metaDataBuf bufferPos identifier fieldStack - 3 -> do - numRows <- readIntFromBuffer @Int64 metaDataBuf bufferPos - readFileMetaData (metadata {numRows = fromIntegral numRows}) metaDataBuf bufferPos identifier fieldStack - 4 -> do - -- We can do some type checking/exception handling here. - -- Check elemType == List - sizeAndType <- readAndAdvance bufferPos metaDataBuf - let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int - -- type of the contents of the list. - let elemType = toTType sizeAndType - rowGroups <- replicateM sizeOnly (readRowGroup emptyRowGroup metaDataBuf bufferPos 0 []) - readFileMetaData (metadata {rowGroups = rowGroups}) metaDataBuf bufferPos identifier fieldStack - 5 -> do - -- We can do some type checking/exception handling here. - -- Check elemType == List - sizeAndType <- readAndAdvance bufferPos metaDataBuf - let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int - -- type of the contents of the list. - let elemType = toTType sizeAndType - keyValueMetadata <- replicateM sizeOnly (readKeyValue emptyKeyValue metaDataBuf bufferPos 0 []) - readFileMetaData (metadata {keyValueMetadata = keyValueMetadata}) metaDataBuf bufferPos identifier fieldStack - 6 -> do - createdBy <- readString metaDataBuf bufferPos - readFileMetaData (metadata {createdBy = Just createdBy}) metaDataBuf bufferPos identifier fieldStack - 7 -> do - sizeAndType <- readAndAdvance bufferPos metaDataBuf - let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int - let elemType = toTType sizeAndType - columnOrders <- replicateM sizeOnly (readColumnOrder metaDataBuf bufferPos 0 []) - readFileMetaData (metadata {columnOrders = columnOrders}) metaDataBuf bufferPos identifier fieldStack - 8 -> do - encryptionAlgorithm <- readEncryptionAlgorithm metaDataBuf bufferPos 0 [] - readFileMetaData (metadata {encryptionAlgorithm = encryptionAlgorithm}) metaDataBuf bufferPos identifier fieldStack - 9 -> do - footerSigningKeyMetadata <- readByteString metaDataBuf bufferPos - readFileMetaData (metadata {footerSigningKeyMetadata = footerSigningKeyMetadata}) metaDataBuf bufferPos identifier fieldStack - n -> return $ error $ "UNIMPLEMENTED " ++ show n - -readEncryptionAlgorithm :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO EncryptionAlgorithm -readEncryptionAlgorithm buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return ENCRYPTION_ALGORITHM_UNKNOWN - Just (elemType, identifier) -> case identifier of - 1 -> do - readAesGcmV1 (AesGcmV1 {aadPrefix = [], aadFileUnique = [], supplyAadPrefix = False}) buf pos 0 [] - 2 -> do - readAesGcmCtrV1 (AesGcmCtrV1 {aadPrefix = [], aadFileUnique = [], supplyAadPrefix = False}) buf pos 0 [] - n -> return ENCRYPTION_ALGORITHM_UNKNOWN - -readAesGcmV1 :: EncryptionAlgorithm -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO EncryptionAlgorithm -readAesGcmV1 v@(AesGcmV1 aadPrefix aadFileUnique supplyAadPrefix) buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return v - Just (elemType, identifier) -> case identifier of - 1 -> do - aadPrefix <- readByteString buf pos - readAesGcmV1 (v {aadPrefix = aadPrefix}) buf pos lastFieldId fieldStack - 2 -> do - aadFileUnique <- readByteString buf pos - readAesGcmV1 (v {aadFileUnique = aadFileUnique}) buf pos lastFieldId fieldStack - 3 -> do - supplyAadPrefix <- readAndAdvance pos buf - readAesGcmV1 (v {supplyAadPrefix = supplyAadPrefix == compactBooleanTrue}) buf pos lastFieldId fieldStack - _ -> return ENCRYPTION_ALGORITHM_UNKNOWN - -readAesGcmCtrV1 :: EncryptionAlgorithm -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO EncryptionAlgorithm -readAesGcmCtrV1 v@(AesGcmCtrV1 aadPrefix aadFileUnique supplyAadPrefix) buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return v - Just (elemType, identifier) -> case identifier of - 1 -> do - aadPrefix <- readByteString buf pos - readAesGcmCtrV1 (v {aadPrefix = aadPrefix}) buf pos lastFieldId fieldStack - 2 -> do - aadFileUnique <- readByteString buf pos - readAesGcmCtrV1 (v {aadFileUnique = aadFileUnique}) buf pos lastFieldId fieldStack - 3 -> do - supplyAadPrefix <- readAndAdvance pos buf - readAesGcmCtrV1 (v {supplyAadPrefix = supplyAadPrefix == compactBooleanTrue}) buf pos lastFieldId fieldStack - _ -> return ENCRYPTION_ALGORITHM_UNKNOWN - -readColumnOrder :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnOrder -readColumnOrder buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return COLUMN_ORDER_UNKNOWN - Just (elemType, identifier) -> case identifier of - 1 -> do - _ <- replicateM_ 2 (readTypeOrder buf pos 0 []) - return TYPE_ORDER - _ -> return COLUMN_ORDER_UNKNOWN - -readTypeOrder :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnOrder -readTypeOrder buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return TYPE_ORDER - Just (elemType, identifier) -> if elemType == STOP - then return TYPE_ORDER - else readTypeOrder buf pos identifier fieldStack - -data SchemaElement = SchemaElement - { elementName :: T.Text, - elementType :: TType, - typeLength :: Int32, - numChildren :: Int32, - fieldId :: Int32, - repetitionType :: RepetitionType, - convertedType :: Int32, - scale :: Int32, - precision :: Int32, - logicalType :: LogicalType - } - deriving (Show, Eq) - -data RepetitionType = REQUIRED | OPTIONAL | REPEATED | UNKNOWN_REPETITION_TYPE deriving (Eq, Show) - -data LogicalType - = STRING_TYPE - | MAP_TYPE - | LIST_TYPE - | ENUM_TYPE - | DECIMAL_TYPE - | DATE_TYPE - | DecimalType {decimalTypePrecision :: Int32, decimalTypeScale :: Int32} - | TimeType {isAdjustedToUTC :: Bool, unit :: TimeUnit} - | -- This should probably have a different, more constrained TimeUnit type. - TimestampType {isAdjustedToUTC :: Bool, unit :: TimeUnit} - | IntType {bitWidth :: Int8, intIsSigned :: Bool} - | LOGICAL_TYPE_UNKNOWN - | JSON_TYPE - | BSON_TYPE - | UUID_TYPE - | FLOAT16_TYPE - | VariantType {specificationVersion :: Int8} - | GeometryType {crs :: T.Text} - | GeographyType {crs :: T.Text, algorithm :: EdgeInterpolationAlgorithm} - deriving (Eq, Show) - -data TimeUnit - = MILLISECONDS - | MICROSECONDS - | NANOSECONDS - | TIME_UNIT_UNKNOWN - deriving (Eq, Show) - -data EdgeInterpolationAlgorithm - = SPHERICAL - | VINCENTY - | THOMAS - | ANDOYER - | KARNEY - deriving (Eq, Show) - -repetitionTypeFromInt :: Int32 -> RepetitionType -repetitionTypeFromInt 0 = REQUIRED -repetitionTypeFromInt 1 = OPTIONAL -repetitionTypeFromInt 2 = REPEATED -repetitionTypeFromInt _ = UNKNOWN_REPETITION_TYPE - -defaultSchemaElement :: SchemaElement -defaultSchemaElement = SchemaElement "" STOP 0 0 (-1) UNKNOWN_REPETITION_TYPE 0 0 0 LOGICAL_TYPE_UNKNOWN - -toIntegralType :: Int32 -> TType -toIntegralType n - | n == 0 = BOOL - | n == 1 = I32 - | n == 2 = I64 - | n == 3 = I64 - | n == 4 = DOUBLE - | n == 5 = DOUBLE - | n == 6 = STRING - | n == 7 = STRING - | otherwise = error $ "Unknown integral type: " ++ show n - -readSchemaElement :: SchemaElement -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO SchemaElement -readSchemaElement schemaElement buf pos lastFieldId fieldStack = do - t <- readAndAdvance pos buf - if t .&. 0x0f == 0 - then return schemaElement - else do - let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16 - identifier <- - if modifier == 0 - then readIntFromBuffer @Int16 buf pos - else return (lastFieldId + modifier) - let elemType = toTType (t .&. 0x0f) - case identifier of - 1 -> do - schemaElemType <- toIntegralType <$> readInt32FromBuffer buf pos - readSchemaElement (schemaElement {elementType = schemaElemType}) buf pos identifier fieldStack - 2 -> do - typeLength <- readInt32FromBuffer buf pos - readSchemaElement (schemaElement {typeLength = typeLength}) buf pos identifier fieldStack - 3 -> do - fieldRepetitionType <- readInt32FromBuffer buf pos - readSchemaElement (schemaElement {repetitionType = repetitionTypeFromInt fieldRepetitionType}) buf pos identifier fieldStack - 4 -> do - nameSize <- readVarIntFromBuffer @Int buf pos - contents <- replicateM nameSize (readAndAdvance pos buf) - readSchemaElement (schemaElement {elementName = T.pack (map (chr . fromIntegral) contents)}) buf pos identifier fieldStack - 5 -> do - numChildren <- readInt32FromBuffer buf pos - readSchemaElement (schemaElement {numChildren = numChildren}) buf pos identifier fieldStack - 6 -> do - convertedType <- readInt32FromBuffer buf pos - readSchemaElement (schemaElement {convertedType = convertedType}) buf pos identifier fieldStack - 7 -> do - scale <- readInt32FromBuffer buf pos - readSchemaElement (schemaElement {scale = scale}) buf pos identifier fieldStack - 8 -> do - precision <- readInt32FromBuffer buf pos - readSchemaElement (schemaElement {precision = precision}) buf pos identifier fieldStack - 9 -> do - fieldId <- readInt32FromBuffer buf pos - readSchemaElement (schemaElement {fieldId = fieldId}) buf pos identifier fieldStack - 10 -> do - logicalType <- readLogicalType buf pos 0 [] - readSchemaElement (schemaElement {logicalType = logicalType}) buf pos identifier fieldStack - _ -> error $ show identifier -- return schemaElement - -readLogicalType :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType -readLogicalType buf pos lastFieldId fieldStack = do - t <- readAndAdvance pos buf - if t .&. 0x0f == 0 - then return LOGICAL_TYPE_UNKNOWN - else do - let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16 - identifier <- - if modifier == 0 - then readIntFromBuffer @Int16 buf pos - else return (lastFieldId + modifier) - let elemType = toTType (t .&. 0x0f) - case identifier of - 1 -> do - replicateM_ 2 (readField buf pos 0 []) - return STRING_TYPE - 2 -> do - replicateM_ 2 (readField buf pos 0 []) - return MAP_TYPE - 3 -> do - replicateM_ 2 (readField buf pos 0 []) - return LIST_TYPE - 4 -> do - replicateM_ 2 (readField buf pos 0 []) - return ENUM_TYPE - 5 -> do - _ <- readField buf pos 0 [] - readDecimalType (DecimalType {decimalTypeScale = 0, decimalTypePrecision = 0}) buf pos 0 [] - 6 -> do - replicateM_ 2 (readField buf pos 0 []) - return DATE_TYPE - 7 -> do - _ <- readField buf pos 0 [] - readTimeType (TimeType {isAdjustedToUTC = False, unit = MILLISECONDS}) buf pos 0 [] - 8 -> do - _ <- readField buf pos 0 [] - readTimeType (TimestampType {isAdjustedToUTC = False, unit = MILLISECONDS}) buf pos 0 [] - -- Apparently reserved for interval types - 9 -> return LOGICAL_TYPE_UNKNOWN - 10 -> do - _ <- readField buf pos 0 [] - readIntType (IntType {intIsSigned = False, bitWidth = 0}) buf pos 0 [] - 11 -> do - replicateM_ 2 (readField buf pos 0 []) - return LOGICAL_TYPE_UNKNOWN - 12 -> do - replicateM_ 2 (readField buf pos 0 []) - return JSON_TYPE - 13 -> do - replicateM_ 2 (readField buf pos 0 []) - return BSON_TYPE - 14 -> do - replicateM_ 2 (readField buf pos 0 []) - return UUID_TYPE - 15 -> do - replicateM_ 2 (readField buf pos 0 []) - return FLOAT16_TYPE - 16 -> do - _ <- readField buf pos 0 [] - return VariantType {specificationVersion = 1} - 17 -> do - _ <- readField buf pos 0 [] - return GeometryType {crs = ""} - 18 -> do - _ <- readField buf pos 0 [] - return GeographyType {crs = "", algorithm = SPHERICAL} - _ -> return LOGICAL_TYPE_UNKNOWN - -readIntType :: LogicalType -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType -readIntType v@(IntType bitWidth intIsSigned) buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return v - Just (elemType, identifier) -> case identifier of - 1 -> do - bitWidth <- readAndAdvance pos buf - readIntType (v {bitWidth = fromIntegral bitWidth}) buf pos lastFieldId fieldStack - 2 -> do - -- TODO: Check for empty - intIsSigned <- readAndAdvance pos buf - readIntType (v {intIsSigned = intIsSigned == compactBooleanTrue}) buf pos lastFieldId fieldStack - _ -> error $ "UNKNOWN field ID for IntType" ++ show identifier - -readDecimalType :: LogicalType -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType -readDecimalType v@(DecimalType p s) buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return v - Just (elemType, identifier) -> case identifier of - 1 -> do - s' <- readInt32FromBuffer buf pos - readDecimalType (v {decimalTypeScale = s'}) buf pos lastFieldId fieldStack - 2 -> do - p' <- readInt32FromBuffer buf pos - readDecimalType (v {decimalTypePrecision = p'}) buf pos lastFieldId fieldStack - _ -> error $ "UNKNOWN field ID for DecimalType" ++ show identifier - -readTimeType :: LogicalType -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType -readTimeType v@(TimeType _ _) buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return v - Just (elemType, identifier) -> case identifier of - 1 -> do - -- TODO: Check for empty - isAdjustedToUTC <- readAndAdvance pos buf - readTimeType (v {isAdjustedToUTC = isAdjustedToUTC == compactBooleanTrue}) buf pos lastFieldId fieldStack - 2 -> do - u <- readUnit buf pos 0 [] - readTimeType (v {unit = u}) buf pos lastFieldId fieldStack - _ -> error $ "UNKNOWN field ID for TimeType" ++ show identifier -readTimeType v@(TimestampType _ _) buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return v - Just (elemType, identifier) -> case identifier of - 1 -> do - -- TODO: Check for empty - isAdjustedToUTC <- readAndAdvance pos buf - readTimeType (v {isAdjustedToUTC = isAdjustedToUTC == compactBooleanTrue}) buf pos lastFieldId fieldStack - 2 -> do - u <- readUnit buf pos 0 [] - readTimeType (v {unit = u}) buf pos lastFieldId fieldStack - _ -> error $ "UNKNOWN field ID for TimestampType" ++ show identifier - -readUnit :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO TimeUnit -readUnit buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return TIME_UNIT_UNKNOWN - Just (elemType, identifier) -> case identifier of - 1 -> do - _ <- readField buf pos 0 [] - return MILLISECONDS - 2 -> do - _ <- readField buf pos 0 [] - return MICROSECONDS - 3 -> do - _ <- readField buf pos 0 [] - return NANOSECONDS - _ -> return TIME_UNIT_UNKNOWN - -readRowGroup :: RowGroup -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO RowGroup -readRowGroup r buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return r - Just (elemType, identifier) -> case identifier of - 1 -> do - sizeAndType <- readAndAdvance pos buf - let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int - -- type of the contents of the list. - let elemType = toTType sizeAndType - columnChunks <- replicateM sizeOnly (readColumnChunk emptyColumnChunk buf pos 0 []) - readRowGroup (r {rowGroupColumns = columnChunks}) buf pos identifier fieldStack - 2 -> do - totalBytes <- readIntFromBuffer @Int64 buf pos - readRowGroup (r {totalByteSize = totalBytes}) buf pos identifier fieldStack - 3 -> do - nRows <- readIntFromBuffer @Int64 buf pos - readRowGroup (r {rowGroupNumRows = nRows}) buf pos identifier fieldStack - 4 -> return r - 5 -> do - offset <- readIntFromBuffer @Int64 buf pos - readRowGroup (r {fileOffset = offset}) buf pos identifier fieldStack - 6 -> do - compressedSize <- readIntFromBuffer @Int64 buf pos - readRowGroup (r {totalCompressedSize = compressedSize}) buf pos identifier fieldStack - 7 -> do - ordinal <- readIntFromBuffer @Int16 buf pos - readRowGroup (r {ordinal = ordinal}) buf pos identifier fieldStack - _ -> error $ "Unknown row group field: " ++ show identifier - -readColumnChunk :: ColumnChunk -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnChunk -readColumnChunk c buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return c - Just (elemType, identifier) -> case identifier of - 1 -> do - stringSize <- readVarIntFromBuffer @Int buf pos - contents <- map (chr . fromIntegral) <$> replicateM stringSize (readAndAdvance pos buf) - readColumnChunk (c {columnChunkFilePath = contents}) buf pos identifier fieldStack - 2 -> do - columnChunkMetadataFileOffset <- readIntFromBuffer @Int64 buf pos - readColumnChunk (c {columnChunkMetadataFileOffset = columnChunkMetadataFileOffset}) buf pos identifier fieldStack - 3 -> do - columnMetadata <- readColumnMetadata emptyColumnMetadata buf pos 0 [] - readColumnChunk (c {columnMetaData = columnMetadata}) buf pos identifier fieldStack - 4 -> do - columnOffsetIndexOffset <- readIntFromBuffer @Int64 buf pos - readColumnChunk (c {columnChunkOffsetIndexOffset = columnOffsetIndexOffset}) buf pos identifier fieldStack - 5 -> do - columnOffsetIndexLength <- readInt32FromBuffer buf pos - readColumnChunk (c {columnChunkOffsetIndexLength = columnOffsetIndexLength}) buf pos identifier fieldStack - 6 -> do - columnChunkColumnIndexOffset <- readIntFromBuffer @Int64 buf pos - readColumnChunk (c {columnChunkColumnIndexOffset = columnChunkColumnIndexOffset}) buf pos identifier fieldStack - 7 -> do - columnChunkColumnIndexLength <- readInt32FromBuffer buf pos - readColumnChunk (c {columnChunkColumnIndexLength = columnChunkColumnIndexLength}) buf pos identifier fieldStack - _ -> return c - -readColumnMetadata :: ColumnMetaData -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnMetaData -readColumnMetadata cm buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return cm - Just (elemType, identifier) -> case identifier of - 1 -> do - cType <- parquetTypeFromInt <$> readInt32FromBuffer buf pos - readColumnMetadata (cm {columnType = cType}) buf pos identifier [] - 2 -> do - sizeAndType <- readAndAdvance pos buf - let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int - let elemType = toTType sizeAndType - encodings <- replicateM sizeOnly (readParquetEncoding buf pos 0 []) - readColumnMetadata (cm {columnEncodings = encodings}) buf pos identifier fieldStack - 3 -> do - sizeAndType <- readAndAdvance pos buf - let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int - let elemType = toTType sizeAndType - paths <- replicateM sizeOnly (readString buf pos) - readColumnMetadata (cm {columnPathInSchema = paths}) buf pos identifier fieldStack - 4 -> do - cType <- compressionCodecFromInt <$> readInt32FromBuffer buf pos - readColumnMetadata (cm {columnCodec = cType}) buf pos identifier [] - 5 -> do - numValues <- readIntFromBuffer @Int64 buf pos - readColumnMetadata (cm {columnNumValues = numValues}) buf pos identifier [] - 6 -> do - columnTotalUncompressedSize <- readIntFromBuffer @Int64 buf pos - readColumnMetadata (cm {columnTotalUncompressedSize = columnTotalUncompressedSize}) buf pos identifier [] - 7 -> do - columnTotalCompressedSize <- readIntFromBuffer @Int64 buf pos - readColumnMetadata (cm {columnTotalCompressedSize = columnTotalCompressedSize}) buf pos identifier [] - 8 -> do - sizeAndType <- readAndAdvance pos buf - let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int - let elemType = toTType sizeAndType - columnKeyValueMetadata <- replicateM sizeOnly (readKeyValue emptyKeyValue buf pos 0 []) - readColumnMetadata (cm {columnKeyValueMetadata = columnKeyValueMetadata}) buf pos identifier fieldStack - 9 -> do - columnDataPageOffset <- readIntFromBuffer @Int64 buf pos - readColumnMetadata (cm {columnDataPageOffset = columnDataPageOffset}) buf pos identifier [] - 10 -> do - columnIndexPageOffset <- readIntFromBuffer @Int64 buf pos - readColumnMetadata (cm {columnIndexPageOffset = columnIndexPageOffset}) buf pos identifier [] - 11 -> do - columnDictionaryPageOffset <- readIntFromBuffer @Int64 buf pos - readColumnMetadata (cm {columnDictionaryPageOffset = columnDictionaryPageOffset}) buf pos identifier [] - 12 -> do - stats <- readStatistics emptyColumnStatistics buf pos 0 [] - readColumnMetadata (cm {columnStatistics = stats}) buf pos identifier fieldStack - 13 -> do - sizeAndType <- readAndAdvance pos buf - let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int - let elemType = toTType sizeAndType - pageEncodingStats <- replicateM sizeOnly (readPageEncodingStats emptyPageEncodingStats buf pos 0 []) - readColumnMetadata (cm {columnEncodingStats = pageEncodingStats}) buf pos identifier fieldStack - 14 -> do - bloomFilterOffset <- readIntFromBuffer @Int64 buf pos - readColumnMetadata (cm {bloomFilterOffset = bloomFilterOffset}) buf pos identifier [] - 15 -> do - bloomFilterLength <- readInt32FromBuffer buf pos - readColumnMetadata (cm {bloomFilterLength = bloomFilterLength}) buf pos identifier [] - 16 -> do - stats <- readSizeStatistics emptySizeStatistics buf pos 0 [] - readColumnMetadata (cm {columnSizeStatistics = stats}) buf pos identifier fieldStack - 17 -> return $ error "UNIMPLEMENTED" - _ -> return cm - -readParquetEncoding :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ParquetEncoding -readParquetEncoding buf pos lastFieldId fieldStack = parquetEncodingFromInt <$> readInt32FromBuffer buf pos - -readPageEncodingStats :: PageEncodingStats -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO PageEncodingStats -readPageEncodingStats pes buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return pes - Just (elemType, identifier) -> case identifier of - 1 -> do - pType <- pageTypeFromInt <$> readInt32FromBuffer buf pos - readPageEncodingStats (pes {pageEncodingPageType = pType}) buf pos identifier [] - 2 -> do - pEnc <- parquetEncodingFromInt <$> readInt32FromBuffer buf pos - readPageEncodingStats (pes {pageEncoding = pEnc}) buf pos identifier [] - 3 -> do - encodedCount <- readInt32FromBuffer buf pos - readPageEncodingStats (pes {pagesWithEncoding = encodedCount}) buf pos identifier [] - _ -> pure pes - -readStatistics :: ColumnStatistics -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnStatistics -readStatistics cs buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return cs - Just (elemType, identifier) -> case identifier of - 1 -> do - maxInBytes <- readByteString buf pos - readStatistics (cs {columnMax = maxInBytes}) buf pos identifier fieldStack - 2 -> do - minInBytes <- readByteString buf pos - readStatistics (cs {columnMin = minInBytes}) buf pos identifier fieldStack - 3 -> do - nullCount <- readIntFromBuffer @Int64 buf pos - readStatistics (cs {columnNullCount = nullCount}) buf pos identifier fieldStack - 4 -> do - distinctCount <- readIntFromBuffer @Int64 buf pos - readStatistics (cs {columnDistictCount = distinctCount}) buf pos identifier fieldStack - 5 -> do - maxInBytes <- readByteString buf pos - readStatistics (cs {columnMaxValue = maxInBytes}) buf pos identifier fieldStack - 6 -> do - minInBytes <- readByteString buf pos - readStatistics (cs {columnMinValue = minInBytes}) buf pos identifier fieldStack - 7 -> do - isMaxValueExact <- readAndAdvance pos buf - readStatistics (cs {isColumnMaxValueExact = isMaxValueExact == compactBooleanTrue}) buf pos identifier fieldStack - 8 -> do - isMinValueExact <- readAndAdvance pos buf - readStatistics (cs {isColumnMinValueExact = isMinValueExact == compactBooleanTrue}) buf pos identifier fieldStack - _ -> pure cs - -readSizeStatistics :: SizeStatistics -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO SizeStatistics -readSizeStatistics ss buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return ss - Just (elemType, identifier) -> case identifier of - 1 -> do - unencodedByteArrayDataTypes <- readIntFromBuffer @Int64 buf pos - readSizeStatistics (ss {unencodedByteArrayDataTypes = unencodedByteArrayDataTypes}) buf pos identifier fieldStack - 2 -> do - sizeAndType <- readAndAdvance pos buf - let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int - let elemType = toTType sizeAndType - repetitionLevelHistogram <- replicateM sizeOnly (readIntFromBuffer @Int64 buf pos) - readSizeStatistics (ss {repetitionLevelHistogram = repetitionLevelHistogram}) buf pos identifier fieldStack - 3 -> do - sizeAndType <- readAndAdvance pos buf - let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int - let elemType = toTType sizeAndType - definitionLevelHistogram <- replicateM sizeOnly (readIntFromBuffer @Int64 buf pos) - readSizeStatistics (ss {definitionLevelHistogram = definitionLevelHistogram}) buf pos identifier fieldStack - _ -> pure ss - -readKeyValue :: KeyValue -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO KeyValue -readKeyValue kv buf pos lastFieldId fieldStack = do - fieldContents <- readField buf pos lastFieldId fieldStack - case fieldContents of - Nothing -> return kv - Just (elemType, identifier) -> case identifier of - 1 -> do - k <- readString buf pos - readKeyValue (kv {key = k}) buf pos identifier fieldStack - 2 -> do - v <- readString buf pos - readKeyValue (kv {key = v}) buf pos identifier fieldStack - _ -> return kv - -readString :: Ptr Word8 -> IORef Int -> IO String -readString buf pos = do - nameSize <- readVarIntFromBuffer @Int buf pos - map (chr . fromIntegral) <$> replicateM nameSize (readAndAdvance pos buf) - -readByteStringFromBytes :: [Word8] -> ([Word8], [Word8]) -readByteStringFromBytes xs = let - (size, rem) = readVarIntFromBytes @Int xs - in (take size rem, drop size rem) - -readByteString :: Ptr Word8 -> IORef Int -> IO [Word8] -readByteString buf pos = do - size <- readVarIntFromBuffer @Int buf pos - replicateM size (readAndAdvance pos buf) - -readByteString' :: Ptr Word8 -> Int64 -> IO [Word8] -readByteString' buf size = mapM (`readSingleByte` buf) [0..(size - 1)] - -readField :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO (Maybe (TType, Int16)) -readField buf pos lastFieldId fieldStack = do - t <- readAndAdvance pos buf - if t .&. 0x0f == 0 - then return Nothing - else do - let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16 - identifier <- - if modifier == 0 - then readIntFromBuffer @Int16 buf pos - else return (lastFieldId + modifier) - let elemType = toTType (t .&. 0x0f) - pure $ Just (elemType, identifier) - -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) - -readAndAdvance :: IORef Int -> Ptr b -> IO Word8 -readAndAdvance bufferPos buffer = do - pos <- readIORef bufferPos - b <- peekByteOff buffer pos :: IO Word8 - modifyIORef bufferPos (+ 1) - return b - -readSingleByte :: Int64 -> Ptr b -> IO Word8 -readSingleByte pos buffer = peekByteOff buffer (fromIntegral pos) - -readNoAdvance :: IORef Int -> Ptr b -> IO Word8 -readNoAdvance bufferPos buffer = do - pos <- readIORef bufferPos - peekByteOff buffer pos :: IO Word8 - -compactBooleanTrue :: Word8 -compactBooleanTrue = 0x01 - -compactBooleanFalse :: Word8 -compactBooleanFalse = 0x02 - -compactByte :: Word8 -compactByte = 0x03 - -compactI16 :: Word8 -compactI16 = 0x04 - -compactI32 :: Word8 -compactI32 = 0x05 - -compactI64 :: Word8 -compactI64 = 0x06 - -compactDouble :: Word8 -compactDouble = 0x07 - -compactBinary :: Word8 -compactBinary = 0x08 - -compactList :: Word8 -compactList = 0x09 - -compactSet :: Word8 -compactSet = 0x0A - -compactMap :: Word8 -compactMap = 0x0B - -compactStruct :: Word8 -compactStruct = 0x0C - -compactUuid :: Word8 -compactUuid = 0x0D - -data TType - = STOP - | BOOL - | BYTE - | I16 - | I32 - | I64 - | DOUBLE - | STRING - | LIST - | SET - | MAP - | STRUCT - | UUID - deriving (Show, Eq) - -toTType :: Word8 -> TType -toTType t = - fromMaybe STOP $ - M.lookup (t .&. 0x0f) $ - M.fromList - [ (compactBooleanTrue, BOOL), - (compactBooleanFalse, BOOL), - (compactByte, BYTE), - (compactI16, I16), - (compactI32, I32), - (compactI64, I64), - (compactDouble, DOUBLE), - (compactBinary, STRING), - (compactList, LIST), - (compactSet, SET), - (compactMap, MAP), - (compactStruct, STRUCT), - (compactUuid, UUID) - ] - -readIntFromBuffer :: (Integral a) => Ptr b -> IORef Int -> IO a -readIntFromBuffer buf bufferPos = do - n <- readVarIntFromBuffer buf bufferPos - let u = fromIntegral n :: Word32 - return $ fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)) - -readIntFromBytes :: (Integral a) => [Word8] -> (a, [Word8]) -readIntFromBytes bs = let - (n, rem) = readVarIntFromBytes bs - u = fromIntegral n :: Word32 - in (fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), rem) - -readInt32FromBuffer :: Ptr b -> IORef Int -> IO Int32 -readInt32FromBuffer buf bufferPos = do - n <- (fromIntegral <$> readVarIntFromBuffer @Int64 buf bufferPos) :: IO Int32 - let u = fromIntegral n :: Word32 - return $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)) - -readInt32FromBytes :: [Word8] -> (Int32, [Word8]) -readInt32FromBytes bs = let - (n', rem) = readVarIntFromBytes @Int64 bs - n = fromIntegral n' :: Int32 - u = fromIntegral n :: Word32 - in ((fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), rem) - -readVarIntFromBuffer :: (Integral a) => Ptr b -> IORef Int -> IO a -readVarIntFromBuffer buf bufferPos = do - start <- readIORef bufferPos - let loop i shift result = do - b <- readAndAdvance bufferPos buf - let res = result .|. ((fromIntegral (b .&. 0x7f) :: Integer) `shiftL` shift) - if (b .&. 0x80) /= 0x80 - then return res - else loop (i + 1) (shift + 7) res - fromIntegral <$> loop start 0 0 - -readVarIntFromBytes :: (Integral a) => [Word8] -> (a, [Word8]) -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 - -littleEndianWord32 :: [Word8] -> Word32 -littleEndianWord32 bytes - | length bytes == 4 = foldr (\v acc -> acc .|. v) 0 (zipWith (\b i -> (fromIntegral b) `shiftL` i) bytes [0,8..]) - | length bytes < 4 = littleEndianWord32 (take 4 $ bytes ++ (cycle[0])) - | otherwise = error $ "Expected exactly 4 bytes for Word32 but got " ++ (show bytes) - -littleEndianWord64 :: [Word8] -> Word64 -littleEndianWord64 bytes - | length bytes == 8 = foldr (\v acc -> acc .|. v) 0 (zipWith (\b i -> (fromIntegral b) `shiftL` i) bytes [0,8..]) - | otherwise = error "Expected exactly 8 bytes" - -littleEndianInt32 :: [Word8] -> Int32 -littleEndianInt32 bytes - | length bytes == 4 = foldr (\v acc -> acc .|. v) 0 (zipWith (\b i -> (fromIntegral b) `shiftL` i) bytes [0,8..]) - | otherwise = error "Expected exactly 4 bytes for Int32" - -readUVarInt :: [Word8] -> (Word64, [Word8]) -readUVarInt xs = loop xs 0 0 0 - where loop bs x _ 10 = (x, bs) - loop (b:bs) x s i - | b < 0x80 = (x .|. ((fromIntegral b) `shiftL` s), bs) - | otherwise = loop bs (x .|. (fromIntegral ((b .&. 0x7f) `shiftL` s))) (s + 7) (i + 1) - -bitStream :: [Word8] -> [[Word8]] -bitStream xs = map (reverse . toBits) xs - -toBits :: Word8 -> [Word8] -toBits b = go 1 b - where - go 8 n = [(n .&. 1)] - go i n = (n .&. 1) : go (i + 1) (n .>>. 1) - -bitStreamToInt :: Word8 -> [Word8] -> [Int32] -bitStreamToInt _ [] = [] -bitStreamToInt bitWidth bits = let - currBits = take (fromIntegral bitWidth) bits - remaining = drop (fromIntegral bitWidth) bits - in bitsToInt32 bitWidth currBits : bitStreamToInt bitWidth remaining - -bitsToInt32 :: Word8 -> [Word8] -> Int32 -bitsToInt32 bitWidth bits = fromIntegral $ foldr (.|.) 0 (zipWith (\s b -> b .<<. s) [0..] (reverse bits)) +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} + +module DataFrame.IO.Parquet ( + readParquet, +) where + +import Control.Monad +import qualified Data.ByteString as BSO +import Data.Char +import Data.Foldable +import Data.IORef +import Data.List +import qualified Data.Map as M +import Data.Maybe +import qualified Data.Text as T +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.DataFrame (DataFrame) +import qualified DataFrame.Internal.DataFrame as DI +import qualified DataFrame.Operations.Core as DI +import Foreign +import System.IO + +import DataFrame.IO.Parquet.Binary +import DataFrame.IO.Parquet.Dictionary +import DataFrame.IO.Parquet.Levels +import DataFrame.IO.Parquet.Page +import DataFrame.IO.Parquet.Thrift +import DataFrame.IO.Parquet.Types + +{- | Read a parquet file from path and load it into a dataframe. + +==== __Example__ +@ +ghci> D.readParquet "./data/mtcars.parquet" df + +@ +-} +readParquet :: String -> IO DataFrame +readParquet path = withBinaryFile path ReadMode $ \handle -> do + (size, magicString) <- readMetadataSizeFromFooter handle + when (magicString /= "PAR1") $ error "Invalid Parquet file" + + fileMetadata <- readMetadata handle size + + let columnPaths = getColumnPaths (drop 1 $ schema fileMetadata) + let columnNames = map fst columnPaths + + colMap <- newIORef (M.empty :: M.Map T.Text DI.Column) + + let schemaElements = schema fileMetadata + let getTypeLength :: [String] -> Maybe Int32 + getTypeLength path = findTypeLength schemaElements path 0 + where + findTypeLength [] _ _ = Nothing + findTypeLength (s : ss) targetPath depth + | map T.unpack (pathToElement s ss depth) == targetPath + && elementType s == STRING + && typeLength s > 0 = + Just (typeLength s) + | otherwise = findTypeLength ss targetPath (if numChildren s > 0 then depth + 1 else depth) + + pathToElement _ _ _ = [] + + forM_ (rowGroups fileMetadata) $ \rowGroup -> do + forM_ (zip (rowGroupColumns rowGroup) [0 ..]) $ \(colChunk, colIdx) -> do + let metadata = columnMetaData colChunk + let colPath = columnPathInSchema metadata + let colName = + if null colPath + then T.pack $ "col_" ++ show colIdx + else T.pack $ last colPath + + let colDataPageOffset = columnDataPageOffset metadata + let colDictionaryPageOffset = columnDictionaryPageOffset metadata + let colStart = + if colDictionaryPageOffset > 0 && colDataPageOffset > colDictionaryPageOffset + then colDictionaryPageOffset + else colDataPageOffset + let colLength = columnTotalCompressedSize metadata + + columnBytes <- readBytes handle colStart colLength + + pages <- readAllPages (columnCodec metadata) columnBytes + + let maybeTypeLength = + if columnType metadata == PFIXED_LEN_BYTE_ARRAY + then getTypeLength colPath + else Nothing + + let primaryEncoding = fromMaybe EPLAIN (fmap fst (uncons (columnEncodings metadata))) + + let schemaTail = drop 1 (schema fileMetadata) + let colPath = columnPathInSchema (columnMetaData colChunk) + let (maxDef, maxRep) = levelsForPath schemaTail colPath + column <- processColumnPages (maxDef, maxRep) pages (columnType metadata) primaryEncoding maybeTypeLength + + modifyIORef colMap (M.insert colName column) + + finalColMap <- readIORef colMap + let orderedColumns = + map + (\name -> (name, finalColMap M.! name)) + (filter (`M.member` finalColMap) columnNames) + + pure $ DI.fromNamedColumns orderedColumns + +readMetadataSizeFromFooter :: Handle -> IO (Integer, BSO.ByteString) +readMetadataSizeFromFooter handle = do + footerOffSet <- numBytesInFile handle + buf <- mallocBytes 8 :: IO (Ptr Word8) + hSeek handle AbsoluteSeek (fromIntegral $ footerOffSet - 8) + _ <- hGetBuf handle buf 8 + + sizeBytes <- mapM (\i -> fromIntegral <$> (peekElemOff buf i :: IO Word8) :: IO Int32) [0 .. 3] + let size = fromIntegral $ foldl' (.|.) 0 $ zipWith shift sizeBytes [0, 8, 16, 24] + + magicStringBytes <- mapM (\i -> peekElemOff buf i :: IO Word8) [4 .. 7] + let magicString = BSO.pack magicStringBytes + free buf + return (size, magicString) + +getColumnPaths :: [SchemaElement] -> [(T.Text, Int)] +getColumnPaths schema = extractLeafPaths schema 0 [] + where + extractLeafPaths :: [SchemaElement] -> Int -> [T.Text] -> [(T.Text, Int)] + extractLeafPaths [] _ _ = [] + extractLeafPaths (s : ss) idx path + | numChildren s == 0 = + let fullPath = T.intercalate "." (path ++ [elementName s]) + in (fullPath, idx) : extractLeafPaths ss (idx + 1) path + | otherwise = + let newPath = if T.null (elementName s) then path else path ++ [elementName s] + childrenCount = fromIntegral (numChildren s) + (children, remaining) = splitAt childrenCount ss + childResults = extractLeafPaths children idx newPath + in childResults ++ extractLeafPaths remaining (idx + length childResults) path + +processColumnPages :: (Int, Int) -> [Page] -> ParquetType -> ParquetEncoding -> Maybe Int32 -> IO DI.Column +processColumnPages (maxDef, maxRep) pages pType _ maybeTypeLength = do + let dictPages = filter isDictionaryPage pages + let dataPages = filter isDataPage pages + + let dictValsM = + case dictPages of + [] -> Nothing + (dictPage : _) -> + case pageTypeHeader (pageHeader dictPage) of + DictionaryPageHeader{..} -> + let countForBools = + if pType == PBOOLEAN + then error "is bool" Just dictionaryPageHeaderNumValues + else maybeTypeLength + in Just (readDictVals pType (pageBytes dictPage) countForBools) + _ -> Nothing + + cols <- forM dataPages $ \page -> do + case pageTypeHeader (pageHeader page) of + DataPageHeader{..} -> do + let n = fromIntegral dataPageHeaderNumValues + let bs0 = pageBytes page + let (defLvls, _repLvls, afterLvls) = readLevelsV1 n maxDef maxRep bs0 + let nPresent = length (filter (== maxDef) defLvls) + + case dataPageHeaderEncoding of + EPLAIN -> + case pType of + PBOOLEAN -> + let (vals, _) = readNBool nPresent afterLvls + in pure (toMaybeBool maxDef defLvls vals) + PINT32 -> + let (vals, _) = readNInt32 nPresent afterLvls + in pure (toMaybeInt32 maxDef defLvls vals) + PINT64 -> + let (vals, _) = readNInt64 nPresent afterLvls + in pure (toMaybeInt64 maxDef defLvls vals) + PINT96 -> + let (vals, _) = readNInt96Times nPresent afterLvls + in pure (toMaybeUTCTime maxDef defLvls vals) + PFLOAT -> + let (vals, _) = readNFloat nPresent afterLvls + in pure (toMaybeFloat maxDef defLvls vals) + PDOUBLE -> + let (vals, _) = readNDouble nPresent afterLvls + in pure (toMaybeDouble maxDef defLvls vals) + PBYTE_ARRAY -> + let (raws, _) = readNByteArrays nPresent afterLvls + texts = map (T.pack . map (chr . fromIntegral)) 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 + in pure (toMaybeText maxDef defLvls texts) + Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type length" + PARQUET_TYPE_UNKNOWN -> error "Cannot read unknown Parquet type" + ERLE_DICTIONARY -> decodeDictV1 dictValsM maxDef defLvls nPresent afterLvls + EPLAIN_DICTIONARY -> decodeDictV1 dictValsM maxDef defLvls nPresent afterLvls + other -> error ("Unsupported v1 encoding: " ++ show other) + DataPageHeaderV2{..} -> do + let n = fromIntegral dataPageHeaderV2NumValues + let bs0 = pageBytes page + let (defLvls, _repLvls, afterLvls) = + readLevelsV2 n maxDef maxRep definitionLevelByteLength repetitionLevelByteLength bs0 + let nPresent = + if dataPageHeaderV2NumNulls > 0 + then fromIntegral (dataPageHeaderV2NumValues - dataPageHeaderV2NumNulls) + else length (filter (== maxDef) defLvls) + + case dataPageHeaderV2Encoding of + EPLAIN -> + case pType of + PBOOLEAN -> + let (vals, _) = readNBool nPresent afterLvls + in pure (toMaybeBool maxDef defLvls vals) + PINT32 -> + let (vals, _) = readNInt32 nPresent afterLvls + in pure (toMaybeInt32 maxDef defLvls vals) + PINT64 -> + let (vals, _) = readNInt64 nPresent afterLvls + in pure (toMaybeInt64 maxDef defLvls vals) + PINT96 -> + let (vals, _) = readNInt96Times nPresent afterLvls + in pure (toMaybeUTCTime maxDef defLvls vals) + PFLOAT -> + let (vals, _) = readNFloat nPresent afterLvls + in pure (toMaybeFloat maxDef defLvls vals) + PDOUBLE -> + let (vals, _) = readNDouble nPresent afterLvls + in pure (toMaybeDouble maxDef defLvls vals) + PBYTE_ARRAY -> + let (raws, _) = readNByteArrays nPresent afterLvls + texts = map (T.pack . map (chr . fromIntegral)) 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 + in pure (toMaybeText maxDef defLvls texts) + Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type length" + PARQUET_TYPE_UNKNOWN -> error "Cannot read unknown Parquet type" + ERLE_DICTIONARY -> decodeDictV1 dictValsM maxDef defLvls nPresent afterLvls + EPLAIN_DICTIONARY -> decodeDictV1 dictValsM maxDef defLvls nPresent afterLvls + other -> error ("Unsupported v2 encoding: " ++ show other) + + case cols of + [] -> pure $ DI.fromList ([] :: [Maybe Int]) + (c : cs) -> pure $ foldl' (\l r -> fromMaybe (error "concat failed") (DI.concatColumns l r)) c cs diff --git a/src/DataFrame/IO/Parquet/Binary.hs b/src/DataFrame/IO/Parquet/Binary.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Binary.hs @@ -0,0 +1,130 @@ +{-# LANGUAGE TypeApplications #-} + +module DataFrame.IO.Parquet.Binary where + +import Control.Monad +import Data.Bits +import Data.Char +import Data.IORef +import Data.Word +import Foreign +import System.IO + +littleEndianWord32 :: [Word8] -> Word32 +littleEndianWord32 bytes + | 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) + +littleEndianWord64 :: [Word8] -> Word64 +littleEndianWord64 bytes = foldr (.|.) 0 (zipWith (\b i -> (fromIntegral b) `shiftL` i) (take 8 bytes) [0, 8 ..]) + +littleEndianInt32 :: [Word8] -> Int32 +littleEndianInt32 = fromIntegral . littleEndianWord32 + +word64ToLittleEndian :: Word64 -> [Word8] +word64ToLittleEndian w = map (\i -> fromIntegral (w `shiftR` i)) [0, 8, 16, 24, 32, 40, 48, 56] + +word32ToLittleEndian :: Word32 -> [Word8] +word32ToLittleEndian w = map (\i -> fromIntegral (w `shiftR` i)) [0, 8, 16, 24] + +readUVarInt :: [Word8] -> (Word64, [Word8]) +readUVarInt xs = loop xs 0 0 0 + where + loop bs x _ 10 = (x, bs) + loop (b : bs) x s i + | b < 0x80 = (x .|. ((fromIntegral b) `shiftL` s), bs) + | otherwise = loop bs (x .|. (fromIntegral ((b .&. 0x7f) `shiftL` s))) (s + 7) (i + 1) + +readVarIntFromBytes :: (Integral a) => [Word8] -> (a, [Word8]) +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 + +readIntFromBytes :: (Integral a) => [Word8] -> (a, [Word8]) +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 = + let (n', rem) = readVarIntFromBytes @Int64 bs + n = fromIntegral n' :: Int32 + u = fromIntegral n :: Word32 + in ((fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), rem) + +readAndAdvance :: IORef Int -> Ptr b -> IO Word8 +readAndAdvance bufferPos buffer = do + pos <- readIORef bufferPos + b <- peekByteOff buffer pos :: IO Word8 + modifyIORef bufferPos (+ 1) + return b + +readVarIntFromBuffer :: (Integral a) => Ptr b -> IORef Int -> IO a +readVarIntFromBuffer buf bufferPos = do + start <- readIORef bufferPos + let loop i shift result = do + b <- readAndAdvance bufferPos buf + let res = result .|. ((fromIntegral (b .&. 0x7f) :: Integer) `shiftL` shift) + if (b .&. 0x80) /= 0x80 + then return res + else loop (i + 1) (shift + 7) res + fromIntegral <$> loop start 0 0 + +readIntFromBuffer :: (Integral a) => Ptr b -> IORef Int -> IO a +readIntFromBuffer buf bufferPos = do + n <- readVarIntFromBuffer buf bufferPos + let u = fromIntegral n :: Word32 + return $ fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)) + +readInt32FromBuffer :: Ptr b -> IORef Int -> IO Int32 +readInt32FromBuffer buf bufferPos = do + n <- (fromIntegral <$> readVarIntFromBuffer @Int64 buf bufferPos) :: IO Int32 + let u = fromIntegral n :: Word32 + return $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)) + +readString :: Ptr Word8 -> IORef Int -> IO String +readString buf pos = do + nameSize <- readVarIntFromBuffer @Int buf pos + map (chr . fromIntegral) <$> replicateM nameSize (readAndAdvance pos buf) + +readBytes :: Handle -> Int64 -> Int64 -> IO [Word8] +readBytes handle colStart colLen = do + buf <- mallocBytes (fromIntegral colLen) :: IO (Ptr Word8) + hSeek handle AbsoluteSeek (fromIntegral colStart) + _ <- hGetBuf handle buf (fromIntegral colLen) + columnBytes <- mapM (peekByteOff buf) [0 .. fromIntegral colLen - 1] + free buf + pure columnBytes + +numBytesInFile :: Handle -> IO Integer +numBytesInFile handle = do + hSeek handle SeekFromEnd 0 + hTell handle + +readByteStringFromBytes :: [Word8] -> ([Word8], [Word8]) +readByteStringFromBytes xs = + let + (size, rem) = readVarIntFromBytes @Int xs + in + (take size rem, drop size rem) + +readByteString :: Ptr Word8 -> IORef Int -> IO [Word8] +readByteString buf pos = do + size <- readVarIntFromBuffer @Int buf pos + replicateM size (readAndAdvance pos buf) + +readByteString' :: Ptr Word8 -> Int64 -> IO [Word8] +readByteString' buf size = mapM (`readSingleByte` buf) [0 .. (size - 1)] + +readSingleByte :: Int64 -> Ptr b -> IO Word8 +readSingleByte pos buffer = peekByteOff buffer (fromIntegral pos) + +readNoAdvance :: IORef Int -> Ptr b -> IO Word8 +readNoAdvance bufferPos buffer = do + pos <- readIORef bufferPos + peekByteOff buffer pos :: IO Word8 diff --git a/src/DataFrame/IO/Parquet/ColumnStatistics.hs b/src/DataFrame/IO/Parquet/ColumnStatistics.hs --- a/src/DataFrame/IO/Parquet/ColumnStatistics.hs +++ b/src/DataFrame/IO/Parquet/ColumnStatistics.hs @@ -4,16 +4,16 @@ import Data.Word data ColumnStatistics = ColumnStatistics - { columnMin :: [Word8], - columnMax :: [Word8], - columnNullCount :: Int64, - columnDistictCount :: Int64, - columnMinValue :: [Word8], - columnMaxValue :: [Word8], - isColumnMaxValueExact :: Bool, - isColumnMinValueExact :: Bool - } - deriving (Show, Eq) + { columnMin :: [Word8] + , columnMax :: [Word8] + , columnNullCount :: Int64 + , columnDistictCount :: Int64 + , columnMinValue :: [Word8] + , columnMaxValue :: [Word8] + , isColumnMaxValueExact :: Bool + , isColumnMinValueExact :: Bool + } + deriving (Show, Eq) emptyColumnStatistics :: ColumnStatistics emptyColumnStatistics = ColumnStatistics [] [] 0 0 [] [] False False diff --git a/src/DataFrame/IO/Parquet/Compression.hs b/src/DataFrame/IO/Parquet/Compression.hs --- a/src/DataFrame/IO/Parquet/Compression.hs +++ b/src/DataFrame/IO/Parquet/Compression.hs @@ -3,16 +3,16 @@ import Data.Int data CompressionCodec - = UNCOMPRESSED - | SNAPPY - | GZIP - | LZO - | BROTLI - | LZ4 - | ZSTD - | LZ4_RAW - | COMPRESSION_CODEC_UNKNOWN - deriving (Show, Eq) + = UNCOMPRESSED + | SNAPPY + | GZIP + | LZO + | BROTLI + | LZ4 + | ZSTD + | LZ4_RAW + | COMPRESSION_CODEC_UNKNOWN + deriving (Show, Eq) compressionCodecFromInt :: Int32 -> CompressionCodec compressionCodecFromInt 0 = UNCOMPRESSED diff --git a/src/DataFrame/IO/Parquet/Dictionary.hs b/src/DataFrame/IO/Parquet/Dictionary.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Dictionary.hs @@ -0,0 +1,176 @@ +{-# LANGUAGE OverloadedStrings #-} + +module DataFrame.IO.Parquet.Dictionary where + +import Control.Monad +import Data.Char +import Data.Int +import Data.Maybe +import qualified Data.Text as T +import Data.Time +import Data.Word +import DataFrame.IO.Parquet.Binary +import DataFrame.IO.Parquet.Encoding +import DataFrame.IO.Parquet.Levels +import DataFrame.IO.Parquet.Time +import DataFrame.IO.Parquet.Types +import qualified DataFrame.Internal.Column as DI +import Foreign +import GHC.Float +import GHC.IO (unsafePerformIO) +import Text.Printf + +dictCardinality :: DictVals -> Int +dictCardinality (DBool ds) = length ds +dictCardinality (DInt32 ds) = length ds +dictCardinality (DInt64 ds) = length ds +dictCardinality (DInt96 ds) = length ds +dictCardinality (DFloat ds) = length ds +dictCardinality (DDouble ds) = length ds +dictCardinality (DText ds) = length ds + +readDictVals :: ParquetType -> [Word8] -> 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) +readDictVals PINT96 bs _ = DInt96 (readPageInt96Times bs) +readDictVals PFLOAT bs _ = DFloat (readPageFloat bs) +readDictVals PDOUBLE bs _ = DDouble (readPageWord64 bs) +readDictVals PBYTE_ARRAY bs _ = DText (readPageBytes bs) +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) + +readPageWord64 :: [Word8] -> [Double] +readPageWord64 [] = [] +readPageWord64 xs = castWord64ToDouble (littleEndianWord64 (take 8 xs)) : readPageWord64 (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) + +readPageBool :: [Word8] -> [Bool] +readPageBool [] = [] +readPageBool bs = + let totalBits = length bs * 8 + bits = concatMap (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7]) bs + in bits + +readPageInt64 :: [Word8] -> [Int64] +readPageInt64 [] = [] +readPageInt64 xs = fromIntegral (littleEndianWord64 (take 8 xs)) : readPageInt64 (drop 8 xs) + +readPageFloat :: [Word8] -> [Float] +readPageFloat [] = [] +readPageFloat xs = castWord32ToFloat (littleEndianWord32 (take 4 xs)) : readPageFloat (drop 4 xs) + +readNInt96Times :: Int -> [Word8] -> ([UTCTime], [Word8]) +readNInt96Times 0 bs = ([], bs) +readNInt96Times k bs = + let timestamp96 = take 12 bs + utcTime = int96ToUTCTime timestamp96 + bs' = drop 12 bs + (times, rest) = readNInt96Times (k - 1) bs' + in (utcTime : times, rest) + +readPageInt96Times :: [Word8] -> [UTCTime] +readPageInt96Times [] = [] +readPageInt96Times bs = + let (times, rest) = readNInt96Times (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 + +decodeDictV1 :: Maybe DictVals -> Int -> [Int] -> Int -> [Word8] -> IO DI.Column +decodeDictV1 dictValsM maxDef defLvls nPresent bytes = + case dictValsM of + Nothing -> error "Dictionary-encoded page but dictionary is missing" + Just dictVals -> + let (idxs, _rest) = decodeDictIndicesV1 nPresent (dictCardinality dictVals) bytes + in do + when (length idxs /= nPresent) $ + error $ + "dict index count mismatch: got " ++ show (length idxs) ++ ", expected " ++ show nPresent + case dictVals of + DBool ds -> do + let values = [ds !! i | i <- idxs] + pure (toMaybeBool maxDef defLvls values) + DInt32 ds -> do + let values = [ds !! i | i <- idxs] + pure (toMaybeInt32 maxDef defLvls values) + DInt64 ds -> do + let values = [ds !! i | i <- idxs] + pure (toMaybeInt64 maxDef defLvls values) + DInt96 ds -> do + let values = [ds !! i | i <- idxs] + pure (toMaybeUTCTime maxDef defLvls values) + DFloat ds -> do + let values = [ds !! i | i <- idxs] + pure (toMaybeFloat maxDef defLvls values) + DDouble ds -> do + let values = [ds !! i | i <- idxs] + pure (toMaybeDouble maxDef defLvls values) + DText ds -> do + let values = [ds !! i | i <- idxs] + pure (toMaybeText maxDef defLvls values) + +toMaybeInt32 :: Int -> [Int] -> [Int32] -> DI.Column +toMaybeInt32 maxDef def xs = + let filled = stitchNullable maxDef def xs + in if all isJust filled + then DI.fromList (map (fromMaybe 0) filled) + else DI.fromList filled + +toMaybeDouble :: Int -> [Int] -> [Double] -> DI.Column +toMaybeDouble maxDef def xs = + let filled = stitchNullable maxDef def xs + in if all isJust filled + then DI.fromList (map (fromMaybe 0) filled) + else DI.fromList filled + +toMaybeText :: Int -> [Int] -> [T.Text] -> DI.Column +toMaybeText maxDef def xs = + let filled = stitchNullable maxDef def xs + in if all isJust filled + then DI.fromList (map (fromMaybe "") filled) + else DI.fromList filled + +toMaybeBool :: Int -> [Int] -> [Bool] -> DI.Column +toMaybeBool maxDef def xs = + let filled = stitchNullable maxDef def xs + in if all isJust filled + then DI.fromList (map (fromMaybe False) filled) + else DI.fromList filled + +toMaybeInt64 :: Int -> [Int] -> [Int64] -> DI.Column +toMaybeInt64 maxDef def xs = + let filled = stitchNullable maxDef def xs + in if all isJust filled + then DI.fromList (map (fromMaybe 0) filled) + else DI.fromList filled + +toMaybeFloat :: Int -> [Int] -> [Float] -> DI.Column +toMaybeFloat maxDef def xs = + let filled = stitchNullable maxDef def xs + in if all isJust filled + then DI.fromList (map (fromMaybe 0.0) filled) + else DI.fromList filled + +toMaybeUTCTime :: Int -> [Int] -> [UTCTime] -> DI.Column +toMaybeUTCTime maxDef def times = + let filled = stitchNullable maxDef def times + defaultTime = UTCTime (fromGregorian 1970 1 1) (secondsToDiffTime 0) + in if all isJust filled + then DI.fromList (map (fromMaybe defaultTime) filled) + else DI.fromList filled diff --git a/src/DataFrame/IO/Parquet/Encoding.hs b/src/DataFrame/IO/Parquet/Encoding.hs --- a/src/DataFrame/IO/Parquet/Encoding.hs +++ b/src/DataFrame/IO/Parquet/Encoding.hs @@ -1,28 +1,98 @@ module DataFrame.IO.Parquet.Encoding where -import Data.Int +import Data.Bits +import Data.Foldable +import Data.List (mapAccumL) +import Data.Word +import DataFrame.IO.Parquet.Binary -data ParquetEncoding - = EPLAIN - | EPLAIN_DICTIONARY - | ERLE - | EBIT_PACKED - | EDELTA_BINARY_PACKED - | EDELTA_LENGTH_BYTE_ARRAY - | EDELTA_BYTE_ARRAY - | ERLE_DICTIONARY - | EBYTE_STREAM_SPLIT - | PARQUET_ENCODING_UNKNOWN - deriving (Show, Eq) +ceilLog2 :: Int -> Int +ceilLog2 x + | x <= 1 = 0 + | otherwise = 1 + ceilLog2 ((x + 1) `div` 2) -parquetEncodingFromInt :: Int32 -> ParquetEncoding -parquetEncodingFromInt 0 = EPLAIN -parquetEncodingFromInt 2 = EPLAIN_DICTIONARY -parquetEncodingFromInt 3 = ERLE -parquetEncodingFromInt 4 = EBIT_PACKED -parquetEncodingFromInt 5 = EDELTA_BINARY_PACKED -parquetEncodingFromInt 6 = EDELTA_LENGTH_BYTE_ARRAY -parquetEncodingFromInt 7 = EDELTA_BYTE_ARRAY -parquetEncodingFromInt 8 = ERLE_DICTIONARY -parquetEncodingFromInt 9 = EBYTE_STREAM_SPLIT -parquetEncodingFromInt _ = PARQUET_ENCODING_UNKNOWN +bitWidthForMaxLevel :: Int -> Int +bitWidthForMaxLevel maxLevel = ceilLog2 (maxLevel + 1) + +bytesForBW :: Int -> Int +bytesForBW bw = (bw + 7) `div` 8 + +unpackBitPacked :: Int -> Int -> [Word8] -> ([Word32], [Word8]) +unpackBitPacked bw count bs + | count <= 0 = ([], bs) + | null bs = ([], bs) + | otherwise = + let totalBits = bw * count + totalBytes = (totalBits + 7) `div` 8 + chunk = take totalBytes bs + rest = drop totalBytes bs + in if length chunk < totalBytes + then error $ "Not enough bytes: need " ++ show totalBytes ++ ", got " ++ show (length chunk) + else + let 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) + + extractValues _ [] = [] + extractValues n bitsLeft + | n <= 0 = [] + | length bitsLeft < bw = [] + | otherwise = + let (this, bitsLeft') = splitAt bw bitsLeft + in toN this : extractValues (n - 1) bitsLeft' + + vals = extractValues count bits + in (map fromIntegral vals, rest) + +decodeRLEBitPackedHybrid :: Int -> Int -> [Word8] -> ([Word32], [Word8]) +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 0 rest acc = (reverse acc, rest) + go n rest acc + | null rest = error $ "Premature end of stream: need " ++ show n ++ " more values" + | otherwise = + let (hdr64, afterHdr) = readUVarInt rest + isPacked = (hdr64 .&. 1) == 1 + in if isPacked + then + let groups = fromIntegral (hdr64 `shiftR` 1) :: Int + totalVals = groups * 8 + (valsAll, afterRun) = unpackBitPacked bw totalVals afterHdr + takeN = min n totalVals + actualTaken = take takeN valsAll + in go (n - takeN) afterRun (reverse actualTaken ++ acc) + else + let runLen = fromIntegral (hdr64 `shiftR` 1) :: Int + nbytes = bytesForBW bw + word32 = littleEndianWord32 (take 4 afterHdr) + afterV = 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 need dictCard bs = + case bs of + [] -> error "empty dictionary index stream" + (w0 : rest0) -> + let widthFromDict = ceilLog2 dictCard + looksLikeWidth = w0 <= 32 && (w0 /= 0 || dictCard <= 1) + tryWithWidthByte = + let bw = fromIntegral w0 + (u32s, rest1) = decodeRLEBitPackedHybrid bw need rest0 + in (map fromIntegral u32s, rest1) + tryWithoutWidthByte = + let bw = widthFromDict + (u32s, rest1) = decodeRLEBitPackedHybrid bw need bs + in (map fromIntegral u32s, rest1) + (idxs, rest') = + if looksLikeWidth + then + let (xs, r) = tryWithWidthByte + in if length xs == need then (xs, r) else tryWithoutWidthByte + else + tryWithoutWidthByte + in (idxs, rest') diff --git a/src/DataFrame/IO/Parquet/Levels.hs b/src/DataFrame/IO/Parquet/Levels.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Levels.hs @@ -0,0 +1,106 @@ +module DataFrame.IO.Parquet.Levels where + +import Data.Int +import Data.List +import qualified Data.Text as T +import Data.Word + +import DataFrame.IO.Parquet.Binary +import DataFrame.IO.Parquet.Encoding +import DataFrame.IO.Parquet.Thrift +import DataFrame.IO.Parquet.Types + +readLevelsV1 :: Int -> Int -> Int -> [Word8] -> ([Int], [Int], [Word8]) +readLevelsV1 n maxDef maxRep bs = + let bwDef = bitWidthForMaxLevel maxDef + bwRep = bitWidthForMaxLevel maxRep + + (repLvlsU32, afterRep) = + 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 + (repVals, _) = decodeRLEBitPackedHybrid bwRep n repData + in (repVals, afterRepData) + + (defLvlsU32, afterDef) = + 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 + (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]) +readLevelsV2 n maxDef maxRep defLen repLen bs = + let (repBytes, afterRepBytes) = splitAt (fromIntegral repLen) bs + (defBytes, afterDefBytes) = splitAt (fromIntegral defLen) afterRepBytes + bwDef = bitWidthForMaxLevel maxDef + bwRep = bitWidthForMaxLevel maxRep + (repLvlsU32, _) = + if bwRep == 0 + then (replicate n 0, repBytes) + else decodeRLEBitPackedHybrid bwRep n repBytes + (defLvlsU32, _) = + if bwDef == 0 + then (replicate n 0, defBytes) + else decodeRLEBitPackedHybrid bwDef n defBytes + in (map fromIntegral defLvlsU32, map fromIntegral repLvlsU32, afterDefBytes) + +stitchNullable :: Int -> [Int] -> [a] -> [Maybe a] +stitchNullable maxDef defLvls vals = go defLvls vals + where + go [] _ = [] + go (d : ds) vs + | d == maxDef = case vs of + (v : vs') -> Just v : go ds vs' + [] -> error "value stream exhausted" + | otherwise = Nothing : go ds vs + +data SNode = SNode + { sName :: String + , sRep :: RepetitionType + , sChildren :: [SNode] + } + deriving (Show, Eq) + +parseOne :: [SchemaElement] -> (SNode, [SchemaElement]) +parseOne [] = error "parseOne: empty schema list" +parseOne (se : rest) = + let childCount = fromIntegral (numChildren se) + (kids, rest') = parseMany childCount rest + in ( SNode + { sName = T.unpack (elementName se) + , sRep = repetitionType se + , sChildren = kids + } + , rest' + ) + +parseMany :: Int -> [SchemaElement] -> ([SNode], [SchemaElement]) +parseMany 0 xs = ([], xs) +parseMany n xs = + let (node, xs') = parseOne xs + (nodes, xs'') = parseMany (n - 1) xs' + in (node : nodes, xs'') + +parseAll :: [SchemaElement] -> [SNode] +parseAll [] = [] +parseAll xs = let (n, xs') = parseOne xs in n : parseAll xs' + +levelsForPath :: [SchemaElement] -> [String] -> (Int, Int) +levelsForPath schemaTail path = go 0 0 (parseAll schemaTail) path + where + go defC repC _ [] = (defC, repC) + go defC repC nodes (p : ps) = + case find (\n -> sName n == p) nodes of + Nothing -> (defC, repC) + Just n -> + let defC' = defC + (if sRep n == OPTIONAL then 1 else 0) + repC' = repC + (if sRep n == REPEATED then 1 else 0) + in go defC' repC' (sChildren n) ps diff --git a/src/DataFrame/IO/Parquet/Page.hs b/src/DataFrame/IO/Parquet/Page.hs --- a/src/DataFrame/IO/Parquet/Page.hs +++ b/src/DataFrame/IO/Parquet/Page.hs @@ -1,19 +1,342 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TypeApplications #-} + module DataFrame.IO.Parquet.Page where +import Codec.Compression.Zstd.Streaming +import Control.Monad +import Data.Bits +import qualified Data.ByteString as BSO +import Data.Char +import Data.Foldable import Data.Int +import Data.List +import Data.Maybe +import qualified Data.Text as T +import Data.Word +import DataFrame.IO.Parquet.Binary +import DataFrame.IO.Parquet.Dictionary +import DataFrame.IO.Parquet.Levels +import DataFrame.IO.Parquet.Thrift +import DataFrame.IO.Parquet.Types +import qualified DataFrame.Internal.Column as DI +import GHC.Float +import qualified Snappy as Snappy +import Text.Printf -data PageType - = DATA_PAGE - | INDEX_PAGE - | DICTIONARY_PAGE - | DATA_PAGE_V2 - | PAGE_TYPE_UNKNOWN - deriving (Show, Eq) +isDataPage :: Page -> Bool +isDataPage page = case pageTypeHeader (pageHeader page) of + DataPageHeader{..} -> True + DataPageHeaderV2{..} -> True + _ -> False -pageTypeFromInt :: Int32 -> PageType -pageTypeFromInt 0 = DATA_PAGE -pageTypeFromInt 1 = INDEX_PAGE -pageTypeFromInt 2 = DICTIONARY_PAGE -pageTypeFromInt 3 = DATA_PAGE_V2 -pageTypeFromInt _ = PAGE_TYPE_UNKNOWN +isDictionaryPage :: Page -> Bool +isDictionaryPage page = case pageTypeHeader (pageHeader page) of + 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 + + 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) + other -> error ("Unsupported compression type: " ++ show other) + pure $ (Just $ Page hdr (BSO.unpack fullData), drop (fromIntegral $ compressedPageSize hdr) rem) + +readPageHeader :: PageHeader -> [Word8] -> Int16 -> (PageHeader, [Word8]) +readPageHeader hdr [] _ = (hdr, []) +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 + +readPageTypeHeader :: PageTypeHeader -> [Word8] -> Int16 -> (PageTypeHeader, [Word8]) +readPageTypeHeader hdr [] _ = (hdr, []) +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 : rem') = rem + in + readPageTypeHeader (hdr{dictionaryPageIsSorted = isSorted == compactBooleanTrue}) rem' identifier + n -> error $ 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 +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 = fromMaybe True $ fmap ((== compactBooleanTrue) . (.&. 0x0f)) (safeHead xs) + in + readPageTypeHeader (hdr{dataPageHeaderV2IsCompressed = dataPageHeaderV2IsCompressed}) rem identifier + 8 -> + let + (stats, rem') = readStatisticsFromBytes emptyColumnStatistics rem 0 + in + readPageTypeHeader (hdr{dataPageHeaderV2Statistics = stats}) rem' identifier + n -> error $ show n + +safeHead :: [a] -> Maybe a +safeHead [] = Nothing +safeHead (x : _) = Just x + +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) + +readAllPages :: CompressionCodec -> [Word8] -> 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) + +readNInt32 :: Int -> [Word8] -> ([Int32], [Word8]) +readNInt32 0 bs = ([], bs) +readNInt32 k bs = + let x = littleEndianInt32 (take 4 bs) + bs' = drop 4 bs + (xs, rest) = readNInt32 (k - 1) bs' + in (x : xs, rest) + +readNDouble :: Int -> [Word8] -> ([Double], [Word8]) +readNDouble 0 bs = ([], bs) +readNDouble k bs = + let x = castWord64ToDouble (littleEndianWord64 (take 8 bs)) + bs' = drop 8 bs + (xs, rest) = readNDouble (k - 1) bs' + in (x : xs, rest) + +readNByteArrays :: Int -> [Word8] -> ([[Word8]], [Word8]) +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 + (xs, rest) = readNByteArrays (k - 1) bs' + in (body : xs, rest) + +readNBool :: Int -> [Word8] -> ([Bool], [Word8]) +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 + bools = take count bits + in (bools, rest) + +readNInt64 :: Int -> [Word8] -> ([Int64], [Word8]) +readNInt64 0 bs = ([], bs) +readNInt64 k bs = + let x = fromIntegral (littleEndianWord64 (take 8 bs)) + bs' = drop 8 bs + (xs, rest) = readNInt64 (k - 1) bs' + in (x : xs, rest) + +readNFloat :: Int -> [Word8] -> ([Float], [Word8]) +readNFloat 0 bs = ([], bs) +readNFloat k bs = + let x = castWord32ToFloat (littleEndianWord32 (take 4 bs)) + bs' = drop 4 bs + (xs, rest) = readNFloat (k - 1) bs' + in (x : xs, rest) + +splitFixed :: Int -> Int -> [Word8] -> ([[Word8]], [Word8]) +splitFixed 0 _ bs = ([], bs) +splitFixed k len bs = + let body = take len bs + bs' = drop len bs + (xs, rest) = splitFixed (k - 1) len bs' + in (body : xs, rest) + +readStatisticsFromBytes :: ColumnStatistics -> [Word8] -> Int16 -> (ColumnStatistics, [Word8]) +readStatisticsFromBytes cs xs lastFieldId = + let + fieldContents = readField' xs lastFieldId + in + case fieldContents of + Nothing -> (cs, drop 1 xs) + Just (rem, elemType, identifier) -> case identifier of + 1 -> + let + (maxInBytes, rem') = readByteStringFromBytes rem + in + readStatisticsFromBytes (cs{columnMax = maxInBytes}) rem' identifier + 2 -> + let + (minInBytes, rem') = readByteStringFromBytes rem + in + readStatisticsFromBytes (cs{columnMin = minInBytes}) rem' identifier + 3 -> + let + (nullCount, rem') = readIntFromBytes @Int64 rem + in + readStatisticsFromBytes (cs{columnNullCount = nullCount}) rem' identifier + 4 -> + let + (distinctCount, rem') = readIntFromBytes @Int64 rem + in + readStatisticsFromBytes (cs{columnDistictCount = distinctCount}) rem' identifier + 5 -> + let + (maxInBytes, rem') = readByteStringFromBytes rem + in + readStatisticsFromBytes (cs{columnMaxValue = maxInBytes}) rem' identifier + 6 -> + let + (minInBytes, rem') = readByteStringFromBytes rem + in + readStatisticsFromBytes (cs{columnMinValue = minInBytes}) rem' identifier + 7 -> + let + (isMaxValueExact : rem') = rem + in + readStatisticsFromBytes (cs{isColumnMaxValueExact = isMaxValueExact == compactBooleanTrue}) rem' identifier + 8 -> + let + (isMinValueExact : rem') = rem + in + readStatisticsFromBytes (cs{isColumnMinValueExact = isMinValueExact == compactBooleanTrue}) rem' identifier + n -> error $ show n diff --git a/src/DataFrame/IO/Parquet/Thrift.hs b/src/DataFrame/IO/Parquet/Thrift.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Thrift.hs @@ -0,0 +1,774 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.IO.Parquet.Thrift where + +import Control.Monad +import Data.Bits +import Data.Char +import Data.IORef +import Data.Int +import qualified Data.Map as M +import Data.Maybe +import qualified Data.Text as T +import Data.Word +import DataFrame.IO.Parquet.Binary +import DataFrame.IO.Parquet.Types +import Foreign +import System.IO + +data SchemaElement = SchemaElement + { elementName :: T.Text + , elementType :: TType + , typeLength :: Int32 + , numChildren :: Int32 + , fieldId :: Int32 + , repetitionType :: RepetitionType + , convertedType :: Int32 + , scale :: Int32 + , precision :: Int32 + , logicalType :: LogicalType + } + deriving (Show, Eq) + +data KeyValue = KeyValue + { key :: String + , value :: String + } + deriving (Show, Eq) + +data FileMetadata = FileMetaData + { version :: Int32 + , schema :: [SchemaElement] + , numRows :: Integer + , rowGroups :: [RowGroup] + , keyValueMetadata :: [KeyValue] + , createdBy :: Maybe String + , columnOrders :: [ColumnOrder] + , encryptionAlgorithm :: EncryptionAlgorithm + , footerSigningKeyMetadata :: [Word8] + } + deriving (Show, Eq) + +data TType + = STOP + | BOOL + | BYTE + | I16 + | I32 + | I64 + | DOUBLE + | STRING + | LIST + | SET + | MAP + | STRUCT + | UUID + deriving (Show, Eq) + +defaultMetadata :: FileMetadata +defaultMetadata = + FileMetaData + { version = 0 + , schema = [] + , numRows = 0 + , rowGroups = [] + , keyValueMetadata = [] + , createdBy = Nothing + , columnOrders = [] + , encryptionAlgorithm = ENCRYPTION_ALGORITHM_UNKNOWN + , footerSigningKeyMetadata = [] + } + +data ColumnMetaData = ColumnMetaData + { columnType :: ParquetType + , columnEncodings :: [ParquetEncoding] + , columnPathInSchema :: [String] + , columnCodec :: CompressionCodec + , columnNumValues :: Int64 + , columnTotalUncompressedSize :: Int64 + , columnTotalCompressedSize :: Int64 + , columnKeyValueMetadata :: [KeyValue] + , columnDataPageOffset :: Int64 + , columnIndexPageOffset :: Int64 + , columnDictionaryPageOffset :: Int64 + , columnStatistics :: ColumnStatistics + , columnEncodingStats :: [PageEncodingStats] + , bloomFilterOffset :: Int64 + , bloomFilterLength :: Int32 + , columnSizeStatistics :: SizeStatistics + , columnGeospatialStatistics :: GeospatialStatistics + } + deriving (Show, Eq) + +data ColumnChunk = ColumnChunk + { columnChunkFilePath :: String + , columnChunkMetadataFileOffset :: Int64 + , columnMetaData :: ColumnMetaData + , columnChunkOffsetIndexOffset :: Int64 + , columnChunkOffsetIndexLength :: Int32 + , columnChunkColumnIndexOffset :: Int64 + , columnChunkColumnIndexLength :: Int32 + , cryptoMetadata :: ColumnCryptoMetadata + , encryptedColumnMetadata :: [Word8] + } + deriving (Show, Eq) + +data RowGroup = RowGroup + { rowGroupColumns :: [ColumnChunk] + , totalByteSize :: Int64 + , rowGroupNumRows :: Int64 + , rowGroupSortingColumns :: [SortingColumn] + , fileOffset :: Int64 + , totalCompressedSize :: Int64 + , ordinal :: Int16 + } + deriving (Show, Eq) + +defaultSchemaElement :: SchemaElement +defaultSchemaElement = SchemaElement "" STOP 0 0 (-1) UNKNOWN_REPETITION_TYPE 0 0 0 LOGICAL_TYPE_UNKNOWN + +emptyColumnMetadata :: ColumnMetaData +emptyColumnMetadata = ColumnMetaData PARQUET_TYPE_UNKNOWN [] [] COMPRESSION_CODEC_UNKNOWN 0 0 0 [] 0 0 0 emptyColumnStatistics [] 0 0 emptySizeStatistics emptyGeospatialStatistics + +emptyColumnChunk :: ColumnChunk +emptyColumnChunk = ColumnChunk "" 0 emptyColumnMetadata 0 0 0 0 COLUMN_CRYPTO_METADATA_UNKNOWN [] + +emptyKeyValue :: KeyValue +emptyKeyValue = KeyValue{key = "", value = ""} + +emptyRowGroup :: RowGroup +emptyRowGroup = RowGroup [] 0 0 [] 0 0 0 + +compactBooleanTrue, compactI32, compactI64, compactDouble, compactBinary, compactList, compactStruct :: Word8 +compactBooleanTrue = 0x01 +compactI32 = 0x05 +compactI64 = 0x06 +compactDouble = 0x07 +compactBinary = 0x08 +compactList = 0x09 +compactStruct = 0x0C + +toTType :: Word8 -> TType +toTType t = + fromMaybe STOP $ + M.lookup (t .&. 0x0f) $ + M.fromList + [ (compactBooleanTrue, BOOL) + , (compactI32, I32) + , (compactI64, I64) + , (compactDouble, DOUBLE) + , (compactBinary, STRING) + , (compactList, LIST) + , (compactStruct, STRUCT) + ] + +readField :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO (Maybe (TType, Int16)) +readField buf pos lastFieldId fieldStack = do + t <- readAndAdvance pos buf + if t .&. 0x0f == 0 + then return Nothing + else do + let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16 + identifier <- + if modifier == 0 + then readIntFromBuffer @Int16 buf pos + else return (lastFieldId + modifier) + let elemType = toTType (t .&. 0x0f) + pure $ Just (elemType, identifier) + +skipToStructEnd :: Ptr Word8 -> IORef Int -> IO () +skipToStructEnd buf pos = do + t <- readAndAdvance pos buf + if t .&. 0x0f == 0 + then return () + else do + let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16 + identifier <- + if modifier == 0 + then readIntFromBuffer @Int16 buf pos + else return 0 + let elemType = toTType (t .&. 0x0f) + skipFieldData elemType buf pos + skipToStructEnd buf pos + +skipFieldData :: TType -> Ptr Word8 -> IORef Int -> IO () +skipFieldData fieldType buf pos = case fieldType of + BOOL -> return () + I32 -> readIntFromBuffer @Int32 buf pos >> pure () + I64 -> readIntFromBuffer @Int64 buf pos >> pure () + DOUBLE -> readIntFromBuffer @Int64 buf pos >> pure () + STRING -> readByteString buf pos >> pure () + LIST -> skipList buf pos + STRUCT -> skipToStructEnd buf pos + _ -> return () + +skipList :: Ptr Word8 -> IORef Int -> IO () +skipList buf pos = do + sizeAndType <- readAndAdvance pos buf + let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int + let elemType = toTType sizeAndType + replicateM_ sizeOnly (skipFieldData elemType buf pos) + +readMetadata :: Handle -> Integer -> IO FileMetadata +readMetadata handle size = do + metaDataBuf <- mallocBytes (fromIntegral size) :: IO (Ptr Word8) + footerOffSet <- numBytesInFile handle + + hSeek handle AbsoluteSeek (fromIntegral $! footerOffSet - footerSize - size) + + metadataBytesRead <- hGetBuf handle metaDataBuf (fromIntegral size) + let lastFieldId = 0 + let fieldStack = [] + bufferPos <- newIORef (0 :: Int) + metadata <- readFileMetaData defaultMetadata metaDataBuf bufferPos lastFieldId fieldStack + free metaDataBuf + return metadata + +readFileMetaData :: FileMetadata -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO FileMetadata +readFileMetaData metadata metaDataBuf bufferPos lastFieldId fieldStack = do + fieldContents <- readField metaDataBuf bufferPos lastFieldId fieldStack + case fieldContents of + Nothing -> return metadata + Just (elemType, identifier) -> case identifier of + 1 -> do + version <- readIntFromBuffer @Int32 metaDataBuf bufferPos + readFileMetaData (metadata{version = version}) metaDataBuf bufferPos identifier fieldStack + 2 -> do + sizeAndType <- readAndAdvance bufferPos metaDataBuf + listSize <- + if (sizeAndType `shiftR` 4) .&. 0x0f == 15 + then readVarIntFromBuffer @Int metaDataBuf bufferPos + else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) + + let elemType = toTType sizeAndType + schemaElements <- replicateM listSize (readSchemaElement defaultSchemaElement metaDataBuf bufferPos 0 []) + readFileMetaData (metadata{schema = schemaElements}) metaDataBuf bufferPos identifier fieldStack + 3 -> do + numRows <- readIntFromBuffer @Int64 metaDataBuf bufferPos + readFileMetaData (metadata{numRows = fromIntegral numRows}) metaDataBuf bufferPos identifier fieldStack + 4 -> do + sizeAndType <- readAndAdvance bufferPos metaDataBuf + listSize <- + if (sizeAndType `shiftR` 4) .&. 0x0f == 15 + then readVarIntFromBuffer @Int metaDataBuf bufferPos + else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) + + let elemType = toTType sizeAndType + rowGroups <- replicateM listSize (readRowGroup emptyRowGroup metaDataBuf bufferPos 0 []) + readFileMetaData (metadata{rowGroups = rowGroups}) metaDataBuf bufferPos identifier fieldStack + 5 -> do + sizeAndType <- readAndAdvance bufferPos metaDataBuf + listSize <- + if (sizeAndType `shiftR` 4) .&. 0x0f == 15 + then readVarIntFromBuffer @Int metaDataBuf bufferPos + else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) + + let elemType = toTType sizeAndType + keyValueMetadata <- replicateM listSize (readKeyValue emptyKeyValue metaDataBuf bufferPos 0 []) + readFileMetaData (metadata{keyValueMetadata = keyValueMetadata}) metaDataBuf bufferPos identifier fieldStack + 6 -> do + createdBy <- readString metaDataBuf bufferPos + readFileMetaData (metadata{createdBy = Just createdBy}) metaDataBuf bufferPos identifier fieldStack + 7 -> do + sizeAndType <- readAndAdvance bufferPos metaDataBuf + listSize <- + if (sizeAndType `shiftR` 4) .&. 0x0f == 15 + then readVarIntFromBuffer @Int metaDataBuf bufferPos + else return $ fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) + + let elemType = toTType sizeAndType + columnOrders <- replicateM listSize (readColumnOrder metaDataBuf bufferPos 0 []) + readFileMetaData (metadata{columnOrders = columnOrders}) metaDataBuf bufferPos identifier fieldStack + 8 -> do + encryptionAlgorithm <- readEncryptionAlgorithm metaDataBuf bufferPos 0 [] + readFileMetaData (metadata{encryptionAlgorithm = encryptionAlgorithm}) metaDataBuf bufferPos identifier fieldStack + 9 -> do + footerSigningKeyMetadata <- readByteString metaDataBuf bufferPos + readFileMetaData (metadata{footerSigningKeyMetadata = footerSigningKeyMetadata}) metaDataBuf bufferPos identifier fieldStack + n -> return $ error $ "UNIMPLEMENTED " ++ show n + +readSchemaElement :: SchemaElement -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO SchemaElement +readSchemaElement schemaElement buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return schemaElement + Just (STOP, _) -> return schemaElement + Just (elemType, identifier) -> case identifier of + 1 -> do + schemaElemType <- toIntegralType <$> readInt32FromBuffer buf pos + readSchemaElement (schemaElement{elementType = schemaElemType}) buf pos identifier fieldStack + 2 -> do + typeLength <- readInt32FromBuffer buf pos + readSchemaElement (schemaElement{typeLength = typeLength}) buf pos identifier fieldStack + 3 -> do + fieldRepetitionType <- readInt32FromBuffer buf pos + readSchemaElement (schemaElement{repetitionType = repetitionTypeFromInt fieldRepetitionType}) buf pos identifier fieldStack + 4 -> do + nameSize <- readVarIntFromBuffer @Int buf pos + if nameSize <= 0 + then readSchemaElement schemaElement buf pos identifier fieldStack + else do + contents <- replicateM nameSize (readAndAdvance pos buf) + readSchemaElement (schemaElement{elementName = T.pack (map (chr . fromIntegral) contents)}) buf pos identifier fieldStack + 5 -> do + numChildren <- readInt32FromBuffer buf pos + readSchemaElement (schemaElement{numChildren = numChildren}) buf pos identifier fieldStack + 6 -> do + convertedType <- readInt32FromBuffer buf pos + readSchemaElement (schemaElement{convertedType = convertedType}) buf pos identifier fieldStack + 7 -> do + scale <- readInt32FromBuffer buf pos + readSchemaElement (schemaElement{scale = scale}) buf pos identifier fieldStack + 8 -> do + precision <- readInt32FromBuffer buf pos + readSchemaElement (schemaElement{precision = precision}) buf pos identifier fieldStack + 9 -> do + fieldId <- readInt32FromBuffer buf pos + readSchemaElement (schemaElement{fieldId = fieldId}) buf pos identifier fieldStack + 10 -> do + logicalType <- readLogicalType buf pos 0 [] + readSchemaElement (schemaElement{logicalType = logicalType}) buf pos identifier fieldStack + _ -> do + skipFieldData elemType buf pos + readSchemaElement schemaElement buf pos identifier fieldStack + +readRowGroup :: RowGroup -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO RowGroup +readRowGroup r buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return r + Just (elemType, identifier) -> case identifier of + 1 -> do + sizeAndType <- readAndAdvance pos buf + let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int + -- type of the contents of the list. + let elemType = toTType sizeAndType + columnChunks <- replicateM sizeOnly (readColumnChunk emptyColumnChunk buf pos 0 []) + readRowGroup (r{rowGroupColumns = columnChunks}) buf pos identifier fieldStack + 2 -> do + totalBytes <- readIntFromBuffer @Int64 buf pos + readRowGroup (r{totalByteSize = totalBytes}) buf pos identifier fieldStack + 3 -> do + nRows <- readIntFromBuffer @Int64 buf pos + readRowGroup (r{rowGroupNumRows = nRows}) buf pos identifier fieldStack + 4 -> return r + 5 -> do + offset <- readIntFromBuffer @Int64 buf pos + readRowGroup (r{fileOffset = offset}) buf pos identifier fieldStack + 6 -> do + compressedSize <- readIntFromBuffer @Int64 buf pos + readRowGroup (r{totalCompressedSize = compressedSize}) buf pos identifier fieldStack + 7 -> do + ordinal <- readIntFromBuffer @Int16 buf pos + readRowGroup (r{ordinal = ordinal}) buf pos identifier fieldStack + _ -> error $ "Unknown row group field: " ++ show identifier + +readColumnChunk :: ColumnChunk -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnChunk +readColumnChunk c buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return c + Just (elemType, identifier) -> case identifier of + 1 -> do + stringSize <- readVarIntFromBuffer @Int buf pos + contents <- map (chr . fromIntegral) <$> replicateM stringSize (readAndAdvance pos buf) + readColumnChunk (c{columnChunkFilePath = contents}) buf pos identifier fieldStack + 2 -> do + columnChunkMetadataFileOffset <- readIntFromBuffer @Int64 buf pos + readColumnChunk (c{columnChunkMetadataFileOffset = columnChunkMetadataFileOffset}) buf pos identifier fieldStack + 3 -> do + columnMetadata <- readColumnMetadata emptyColumnMetadata buf pos 0 [] + readColumnChunk (c{columnMetaData = columnMetadata}) buf pos identifier fieldStack + 4 -> do + columnOffsetIndexOffset <- readIntFromBuffer @Int64 buf pos + readColumnChunk (c{columnChunkOffsetIndexOffset = columnOffsetIndexOffset}) buf pos identifier fieldStack + 5 -> do + columnOffsetIndexLength <- readInt32FromBuffer buf pos + readColumnChunk (c{columnChunkOffsetIndexLength = columnOffsetIndexLength}) buf pos identifier fieldStack + 6 -> do + columnChunkColumnIndexOffset <- readIntFromBuffer @Int64 buf pos + readColumnChunk (c{columnChunkColumnIndexOffset = columnChunkColumnIndexOffset}) buf pos identifier fieldStack + 7 -> do + columnChunkColumnIndexLength <- readInt32FromBuffer buf pos + readColumnChunk (c{columnChunkColumnIndexLength = columnChunkColumnIndexLength}) buf pos identifier fieldStack + _ -> return c + +readColumnMetadata :: ColumnMetaData -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnMetaData +readColumnMetadata cm buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return cm + Just (elemType, identifier) -> case identifier of + 1 -> do + cType <- parquetTypeFromInt <$> readInt32FromBuffer buf pos + readColumnMetadata (cm{columnType = cType}) buf pos identifier [] + 2 -> do + sizeAndType <- readAndAdvance pos buf + let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int + let elemType = toTType sizeAndType + encodings <- replicateM sizeOnly (readParquetEncoding buf pos 0 []) + readColumnMetadata (cm{columnEncodings = encodings}) buf pos identifier fieldStack + 3 -> do + sizeAndType <- readAndAdvance pos buf + let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int + let elemType = toTType sizeAndType + paths <- replicateM sizeOnly (readString buf pos) + readColumnMetadata (cm{columnPathInSchema = paths}) buf pos identifier fieldStack + 4 -> do + cType <- compressionCodecFromInt <$> readInt32FromBuffer buf pos + readColumnMetadata (cm{columnCodec = cType}) buf pos identifier [] + 5 -> do + numValues <- readIntFromBuffer @Int64 buf pos + readColumnMetadata (cm{columnNumValues = numValues}) buf pos identifier [] + 6 -> do + columnTotalUncompressedSize <- readIntFromBuffer @Int64 buf pos + readColumnMetadata (cm{columnTotalUncompressedSize = columnTotalUncompressedSize}) buf pos identifier [] + 7 -> do + columnTotalCompressedSize <- readIntFromBuffer @Int64 buf pos + readColumnMetadata (cm{columnTotalCompressedSize = columnTotalCompressedSize}) buf pos identifier [] + 8 -> do + sizeAndType <- readAndAdvance pos buf + let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int + let elemType = toTType sizeAndType + columnKeyValueMetadata <- replicateM sizeOnly (readKeyValue emptyKeyValue buf pos 0 []) + readColumnMetadata (cm{columnKeyValueMetadata = columnKeyValueMetadata}) buf pos identifier fieldStack + 9 -> do + columnDataPageOffset <- readIntFromBuffer @Int64 buf pos + readColumnMetadata (cm{columnDataPageOffset = columnDataPageOffset}) buf pos identifier [] + 10 -> do + columnIndexPageOffset <- readIntFromBuffer @Int64 buf pos + readColumnMetadata (cm{columnIndexPageOffset = columnIndexPageOffset}) buf pos identifier [] + 11 -> do + columnDictionaryPageOffset <- readIntFromBuffer @Int64 buf pos + readColumnMetadata (cm{columnDictionaryPageOffset = columnDictionaryPageOffset}) buf pos identifier [] + 12 -> do + stats <- readStatistics emptyColumnStatistics buf pos 0 [] + readColumnMetadata (cm{columnStatistics = stats}) buf pos identifier fieldStack + 13 -> do + sizeAndType <- readAndAdvance pos buf + let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int + let elemType = toTType sizeAndType + pageEncodingStats <- replicateM sizeOnly (readPageEncodingStats emptyPageEncodingStats buf pos 0 []) + readColumnMetadata (cm{columnEncodingStats = pageEncodingStats}) buf pos identifier fieldStack + 14 -> do + bloomFilterOffset <- readIntFromBuffer @Int64 buf pos + readColumnMetadata (cm{bloomFilterOffset = bloomFilterOffset}) buf pos identifier [] + 15 -> do + bloomFilterLength <- readInt32FromBuffer buf pos + readColumnMetadata (cm{bloomFilterLength = bloomFilterLength}) buf pos identifier [] + 16 -> do + stats <- readSizeStatistics emptySizeStatistics buf pos 0 [] + readColumnMetadata (cm{columnSizeStatistics = stats}) buf pos identifier fieldStack + 17 -> return $ error "UNIMPLEMENTED" + _ -> return cm + +readEncryptionAlgorithm :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO EncryptionAlgorithm +readEncryptionAlgorithm buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return ENCRYPTION_ALGORITHM_UNKNOWN + Just (elemType, identifier) -> case identifier of + 1 -> do + readAesGcmV1 (AesGcmV1{aadPrefix = [], aadFileUnique = [], supplyAadPrefix = False}) buf pos 0 [] + 2 -> do + readAesGcmCtrV1 (AesGcmCtrV1{aadPrefix = [], aadFileUnique = [], supplyAadPrefix = False}) buf pos 0 [] + n -> return ENCRYPTION_ALGORITHM_UNKNOWN + +readColumnOrder :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnOrder +readColumnOrder buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return COLUMN_ORDER_UNKNOWN + Just (elemType, identifier) -> case identifier of + 1 -> do + _ <- replicateM_ 2 (readTypeOrder buf pos 0 []) + return TYPE_ORDER + _ -> return COLUMN_ORDER_UNKNOWN + +readAesGcmCtrV1 :: EncryptionAlgorithm -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO EncryptionAlgorithm +readAesGcmCtrV1 v@(AesGcmCtrV1 aadPrefix aadFileUnique supplyAadPrefix) buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return v + Just (elemType, identifier) -> case identifier of + 1 -> do + aadPrefix <- readByteString buf pos + readAesGcmCtrV1 (v{aadPrefix = aadPrefix}) buf pos lastFieldId fieldStack + 2 -> do + aadFileUnique <- readByteString buf pos + readAesGcmCtrV1 (v{aadFileUnique = aadFileUnique}) buf pos lastFieldId fieldStack + 3 -> do + supplyAadPrefix <- readAndAdvance pos buf + readAesGcmCtrV1 (v{supplyAadPrefix = supplyAadPrefix == compactBooleanTrue}) buf pos lastFieldId fieldStack + _ -> return ENCRYPTION_ALGORITHM_UNKNOWN + +readAesGcmV1 :: EncryptionAlgorithm -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO EncryptionAlgorithm +readAesGcmV1 v@(AesGcmV1 aadPrefix aadFileUnique supplyAadPrefix) buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return v + Just (elemType, identifier) -> case identifier of + 1 -> do + aadPrefix <- readByteString buf pos + readAesGcmV1 (v{aadPrefix = aadPrefix}) buf pos lastFieldId fieldStack + 2 -> do + aadFileUnique <- readByteString buf pos + readAesGcmV1 (v{aadFileUnique = aadFileUnique}) buf pos lastFieldId fieldStack + 3 -> do + supplyAadPrefix <- readAndAdvance pos buf + readAesGcmV1 (v{supplyAadPrefix = supplyAadPrefix == compactBooleanTrue}) buf pos lastFieldId fieldStack + _ -> return ENCRYPTION_ALGORITHM_UNKNOWN + +readTypeOrder :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnOrder +readTypeOrder buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return TYPE_ORDER + Just (elemType, identifier) -> + if elemType == STOP + then return TYPE_ORDER + else readTypeOrder buf pos identifier fieldStack + +readKeyValue :: KeyValue -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO KeyValue +readKeyValue kv buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return kv + Just (elemType, identifier) -> case identifier of + 1 -> do + k <- readString buf pos + readKeyValue (kv{key = k}) buf pos identifier fieldStack + 2 -> do + v <- readString buf pos + readKeyValue (kv{value = v}) buf pos identifier fieldStack + _ -> return kv + +readPageEncodingStats :: PageEncodingStats -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO PageEncodingStats +readPageEncodingStats pes buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return pes + Just (elemType, identifier) -> case identifier of + 1 -> do + pType <- pageTypeFromInt <$> readInt32FromBuffer buf pos + readPageEncodingStats (pes{pageEncodingPageType = pType}) buf pos identifier [] + 2 -> do + pEnc <- parquetEncodingFromInt <$> readInt32FromBuffer buf pos + readPageEncodingStats (pes{pageEncoding = pEnc}) buf pos identifier [] + 3 -> do + encodedCount <- readInt32FromBuffer buf pos + readPageEncodingStats (pes{pagesWithEncoding = encodedCount}) buf pos identifier [] + _ -> pure pes + +readParquetEncoding :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ParquetEncoding +readParquetEncoding buf pos lastFieldId fieldStack = parquetEncodingFromInt <$> readInt32FromBuffer buf pos + +readStatistics :: ColumnStatistics -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnStatistics +readStatistics cs buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return cs + Just (elemType, identifier) -> case identifier of + 1 -> do + maxInBytes <- readByteString buf pos + readStatistics (cs{columnMax = maxInBytes}) buf pos identifier fieldStack + 2 -> do + minInBytes <- readByteString buf pos + readStatistics (cs{columnMin = minInBytes}) buf pos identifier fieldStack + 3 -> do + nullCount <- readIntFromBuffer @Int64 buf pos + readStatistics (cs{columnNullCount = nullCount}) buf pos identifier fieldStack + 4 -> do + distinctCount <- readIntFromBuffer @Int64 buf pos + readStatistics (cs{columnDistictCount = distinctCount}) buf pos identifier fieldStack + 5 -> do + maxInBytes <- readByteString buf pos + readStatistics (cs{columnMaxValue = maxInBytes}) buf pos identifier fieldStack + 6 -> do + minInBytes <- readByteString buf pos + readStatistics (cs{columnMinValue = minInBytes}) buf pos identifier fieldStack + 7 -> do + isMaxValueExact <- readAndAdvance pos buf + readStatistics (cs{isColumnMaxValueExact = isMaxValueExact == compactBooleanTrue}) buf pos identifier fieldStack + 8 -> do + isMinValueExact <- readAndAdvance pos buf + readStatistics (cs{isColumnMinValueExact = isMinValueExact == compactBooleanTrue}) buf pos identifier fieldStack + _ -> pure cs + +readSizeStatistics :: SizeStatistics -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO SizeStatistics +readSizeStatistics ss buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return ss + Just (elemType, identifier) -> case identifier of + 1 -> do + unencodedByteArrayDataTypes <- readIntFromBuffer @Int64 buf pos + readSizeStatistics (ss{unencodedByteArrayDataTypes = unencodedByteArrayDataTypes}) buf pos identifier fieldStack + 2 -> do + sizeAndType <- readAndAdvance pos buf + let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int + let elemType = toTType sizeAndType + repetitionLevelHistogram <- replicateM sizeOnly (readIntFromBuffer @Int64 buf pos) + readSizeStatistics (ss{repetitionLevelHistogram = repetitionLevelHistogram}) buf pos identifier fieldStack + 3 -> do + sizeAndType <- readAndAdvance pos buf + let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int + let elemType = toTType sizeAndType + definitionLevelHistogram <- replicateM sizeOnly (readIntFromBuffer @Int64 buf pos) + readSizeStatistics (ss{definitionLevelHistogram = definitionLevelHistogram}) buf pos identifier fieldStack + _ -> pure ss + +footerSize :: Integer +footerSize = 8 + +toIntegralType :: Int32 -> TType +toIntegralType n + | n == 1 = I32 + | n == 2 = I64 + | n == 4 = DOUBLE + | n == 6 = STRING + | otherwise = STRING + +readLogicalType :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType +readLogicalType buf pos lastFieldId fieldStack = do + t <- readAndAdvance pos buf + if t .&. 0x0f == 0 + then return LOGICAL_TYPE_UNKNOWN + else do + let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16 + identifier <- + if modifier == 0 + then readIntFromBuffer @Int16 buf pos + else return (lastFieldId + modifier) + let elemType = toTType (t .&. 0x0f) + case identifier of + 1 -> do + replicateM_ 2 (readField buf pos 0 []) + return STRING_TYPE + 2 -> do + replicateM_ 2 (readField buf pos 0 []) + return MAP_TYPE + 3 -> do + replicateM_ 2 (readField buf pos 0 []) + return LIST_TYPE + 4 -> do + replicateM_ 2 (readField buf pos 0 []) + return ENUM_TYPE + 5 -> do + readDecimalType (DecimalType{decimalTypeScale = 0, decimalTypePrecision = 0}) buf pos 0 [] + 6 -> do + replicateM_ 2 (readField buf pos 0 []) + return DATE_TYPE + 7 -> do + readTimeType (TimeType{isAdjustedToUTC = False, unit = MILLISECONDS}) buf pos 0 [] + 8 -> do + readTimeType (TimestampType{isAdjustedToUTC = False, unit = MILLISECONDS}) buf pos 0 [] + -- Apparently reserved for interval types + 9 -> return LOGICAL_TYPE_UNKNOWN + 10 -> do + intType <- readIntType (IntType{intIsSigned = False, bitWidth = 0}) buf pos 0 [] + _ <- readField buf pos 0 [] + pure intType + 11 -> do + replicateM_ 2 (readField buf pos 0 []) + return LOGICAL_TYPE_UNKNOWN + 12 -> do + replicateM_ 2 (readField buf pos 0 []) + return JSON_TYPE + 13 -> do + replicateM_ 2 (readField buf pos 0 []) + return BSON_TYPE + 14 -> do + replicateM_ 2 (readField buf pos 0 []) + return UUID_TYPE + 15 -> do + replicateM_ 2 (readField buf pos 0 []) + return FLOAT16_TYPE + 16 -> do + return VariantType{specificationVersion = 1} + 17 -> do + return GeometryType{crs = ""} + 18 -> do + return GeographyType{crs = "", algorithm = SPHERICAL} + _ -> return LOGICAL_TYPE_UNKNOWN + +readIntType :: LogicalType -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType +readIntType v@(IntType bitWidth intIsSigned) buf pos lastFieldId fieldStack = do + t <- readAndAdvance pos buf + if t .&. 0x0f == 0 + then return v + else do + let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16 + identifier <- + if modifier == 0 + then readIntFromBuffer @Int16 buf pos + else return (lastFieldId + modifier) + + case identifier of + 1 -> do + bitWidthValue <- readAndAdvance pos buf + readIntType (v{bitWidth = fromIntegral bitWidthValue}) buf pos identifier fieldStack + 2 -> do + let isSigned = (t .&. 0x0f) == compactBooleanTrue + readIntType (v{intIsSigned = isSigned}) buf pos identifier fieldStack + _ -> error $ "UNKNOWN field ID for IntType: " ++ show identifier + +readDecimalType :: LogicalType -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType +readDecimalType v@(DecimalType p s) buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return v + Just (elemType, identifier) -> case identifier of + 1 -> do + s' <- readInt32FromBuffer buf pos + readDecimalType (v{decimalTypeScale = s'}) buf pos lastFieldId fieldStack + 2 -> do + p' <- readInt32FromBuffer buf pos + readDecimalType (v{decimalTypePrecision = p'}) buf pos lastFieldId fieldStack + _ -> error $ "UNKNOWN field ID for DecimalType" ++ show identifier + +readTimeType :: LogicalType -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType +readTimeType v@(TimeType _ _) buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return v + Just (elemType, identifier) -> case identifier of + 1 -> do + -- TODO: Check for empty + isAdjustedToUTC <- readAndAdvance pos buf + readTimeType (v{isAdjustedToUTC = isAdjustedToUTC == compactBooleanTrue}) buf pos lastFieldId fieldStack + 2 -> do + u <- readUnit buf pos 0 [] + readTimeType (v{unit = u}) buf pos lastFieldId fieldStack + _ -> error $ "UNKNOWN field ID for TimeType" ++ show identifier +readTimeType v@(TimestampType _ _) buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return v + Just (elemType, identifier) -> case identifier of + 1 -> do + -- TODO: Check for empty + isAdjustedToUTC <- readAndAdvance pos buf + readTimeType (v{isAdjustedToUTC = isAdjustedToUTC == compactBooleanTrue}) buf pos lastFieldId fieldStack + 2 -> do + u <- readUnit buf pos 0 [] + readTimeType (v{unit = u}) buf pos lastFieldId fieldStack + _ -> error $ "UNKNOWN field ID for TimestampType" ++ show identifier + +readUnit :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO TimeUnit +readUnit buf pos lastFieldId fieldStack = do + fieldContents <- readField buf pos lastFieldId fieldStack + case fieldContents of + Nothing -> return TIME_UNIT_UNKNOWN + Just (elemType, identifier) -> case identifier of + 1 -> do + _ <- readField buf pos 0 [] + return MILLISECONDS + 2 -> do + _ <- readField buf pos 0 [] + return MICROSECONDS + 3 -> do + _ <- readField buf pos 0 [] + return NANOSECONDS + _ -> return TIME_UNIT_UNKNOWN diff --git a/src/DataFrame/IO/Parquet/Time.hs b/src/DataFrame/IO/Parquet/Time.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Time.hs @@ -0,0 +1,56 @@ +{-# LANGUAGE NumericUnderscores #-} + +module DataFrame.IO.Parquet.Time where + +import Data.Int +import Data.Time +import Data.Word + +import DataFrame.IO.Parquet.Binary + +int96ToUTCTime :: [Word8] -> UTCTime +int96ToUTCTime bytes + | length bytes /= 12 = error "INT96 must be exactly 12 bytes" + | otherwise = + let (nanosBytes, julianBytes) = splitAt 8 bytes + nanosSinceMidnight = littleEndianWord64 nanosBytes + julianDay = littleEndianWord32 julianBytes + in julianDayAndNanosToUTCTime (fromIntegral julianDay) nanosSinceMidnight + +julianDayAndNanosToUTCTime :: Integer -> Word64 -> UTCTime +julianDayAndNanosToUTCTime julianDay nanosSinceMidnight = + let day = julianDayToDay julianDay + secondsSinceMidnight = fromIntegral nanosSinceMidnight / 1_000_000_000 + diffTime = secondsToDiffTime (floor secondsSinceMidnight) + in UTCTime day diffTime + +julianDayToDay :: Integer -> Day +julianDayToDay julianDay = + let a = julianDay + 32044 + b = (4 * a + 3) `div` 146097 + c = a - (146097 * b) `div` 4 + d = (4 * c + 3) `div` 1461 + e = c - (1461 * d) `div` 4 + m = (5 * e + 2) `div` 153 + day = e - (153 * m + 2) `div` 5 + 1 + month = m + 3 - 12 * (m `div` 10) + year = 100 * b + d - 4800 + m `div` 10 + in fromGregorian year (fromIntegral month) (fromIntegral day) + +-- 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 day diffTime) = + let julianDay = dayToJulianDay day + nanosSinceMidnight = floor (realToFrac diffTime * 1_000_000_000) + nanosBytes = word64ToLittleEndian nanosSinceMidnight + julianBytes = word32ToLittleEndian (fromIntegral julianDay) + in nanosBytes ++ julianBytes + +dayToJulianDay :: Day -> Integer +dayToJulianDay day = + let (year, month, dayOfMonth) = toGregorian day + a = fromIntegral $ (14 - (fromIntegral month)) `div` 12 + y = fromIntegral $ year + 4800 - a + m = fromIntegral $ month + 12 * (fromIntegral a) - 3 + in fromIntegral dayOfMonth + (153 * m + 2) `div` 5 + 365 * y + y `div` 4 - y `div` 100 + y `div` 400 - 32045 diff --git a/src/DataFrame/IO/Parquet/Types.hs b/src/DataFrame/IO/Parquet/Types.hs --- a/src/DataFrame/IO/Parquet/Types.hs +++ b/src/DataFrame/IO/Parquet/Types.hs @@ -1,18 +1,23 @@ +{-# LANGUAGE OverloadedStrings #-} + module DataFrame.IO.Parquet.Types where import Data.Int +import qualified Data.Text as T +import Data.Time +import Data.Word data ParquetType - = PBOOLEAN - | PINT32 - | PINT64 - | PINT96 - | PFLOAT - | PDOUBLE - | PBYTE_ARRAY - | PFIXED_LEN_BYTE_ARRAY - | PARQUET_TYPE_UNKNOWN - deriving (Show, Eq) + = PBOOLEAN + | PINT32 + | PINT64 + | PINT96 + | PFLOAT + | PDOUBLE + | PBYTE_ARRAY + | PFIXED_LEN_BYTE_ARRAY + | PARQUET_TYPE_UNKNOWN + deriving (Show, Eq) parquetTypeFromInt :: Int32 -> ParquetType parquetTypeFromInt 0 = PBOOLEAN @@ -24,3 +29,266 @@ parquetTypeFromInt 6 = PBYTE_ARRAY parquetTypeFromInt 7 = PFIXED_LEN_BYTE_ARRAY parquetTypeFromInt _ = PARQUET_TYPE_UNKNOWN + +data PageType + = DATA_PAGE + | INDEX_PAGE + | DICTIONARY_PAGE + | DATA_PAGE_V2 + | PAGE_TYPE_UNKNOWN + deriving (Show, Eq) + +pageTypeFromInt :: Int32 -> PageType +pageTypeFromInt 0 = DATA_PAGE +pageTypeFromInt 1 = INDEX_PAGE +pageTypeFromInt 2 = DICTIONARY_PAGE +pageTypeFromInt 3 = DATA_PAGE_V2 +pageTypeFromInt _ = PAGE_TYPE_UNKNOWN + +data ParquetEncoding + = EPLAIN + | EPLAIN_DICTIONARY + | ERLE + | EBIT_PACKED + | EDELTA_BINARY_PACKED + | EDELTA_LENGTH_BYTE_ARRAY + | EDELTA_BYTE_ARRAY + | ERLE_DICTIONARY + | EBYTE_STREAM_SPLIT + | PARQUET_ENCODING_UNKNOWN + deriving (Show, Eq) + +parquetEncodingFromInt :: Int32 -> ParquetEncoding +parquetEncodingFromInt 0 = EPLAIN +parquetEncodingFromInt 2 = EPLAIN_DICTIONARY +parquetEncodingFromInt 3 = ERLE +parquetEncodingFromInt 4 = EBIT_PACKED +parquetEncodingFromInt 5 = EDELTA_BINARY_PACKED +parquetEncodingFromInt 6 = EDELTA_LENGTH_BYTE_ARRAY +parquetEncodingFromInt 7 = EDELTA_BYTE_ARRAY +parquetEncodingFromInt 8 = ERLE_DICTIONARY +parquetEncodingFromInt 9 = EBYTE_STREAM_SPLIT +parquetEncodingFromInt _ = PARQUET_ENCODING_UNKNOWN + +data CompressionCodec + = UNCOMPRESSED + | SNAPPY + | GZIP + | LZO + | BROTLI + | LZ4 + | ZSTD + | LZ4_RAW + | COMPRESSION_CODEC_UNKNOWN + deriving (Show, Eq) + +data PageEncodingStats = PageEncodingStats + { pageEncodingPageType :: PageType + , pageEncoding :: ParquetEncoding + , pagesWithEncoding :: Int32 + } + deriving (Show, Eq) + +emptyPageEncodingStats :: PageEncodingStats +emptyPageEncodingStats = PageEncodingStats PAGE_TYPE_UNKNOWN PARQUET_ENCODING_UNKNOWN 0 + +data SizeStatistics = SizeStatisics + { unencodedByteArrayDataTypes :: Int64 + , repetitionLevelHistogram :: [Int64] + , definitionLevelHistogram :: [Int64] + } + deriving (Show, Eq) + +emptySizeStatistics :: SizeStatistics +emptySizeStatistics = SizeStatisics 0 [] [] + +data BoundingBox = BoundingBox + { xmin :: Double + , xmax :: Double + , ymin :: Double + , ymax :: Double + , zmin :: Double + , zmax :: Double + , mmin :: Double + , mmax :: Double + } + deriving (Show, Eq) + +emptyBoundingBox :: BoundingBox +emptyBoundingBox = BoundingBox 0 0 0 0 0 0 0 0 + +data GeospatialStatistics = GeospatialStatistics + { bbox :: BoundingBox + , geospatialTypes :: [Int32] + } + deriving (Show, Eq) + +emptyGeospatialStatistics :: GeospatialStatistics +emptyGeospatialStatistics = GeospatialStatistics emptyBoundingBox [] + +data ColumnStatistics = ColumnStatistics + { columnMin :: [Word8] + , columnMax :: [Word8] + , columnNullCount :: Int64 + , columnDistictCount :: Int64 + , columnMinValue :: [Word8] + , columnMaxValue :: [Word8] + , isColumnMaxValueExact :: Bool + , isColumnMinValueExact :: Bool + } + deriving (Show, Eq) + +emptyColumnStatistics :: ColumnStatistics +emptyColumnStatistics = ColumnStatistics [] [] 0 0 [] [] False False + +data ColumnCryptoMetadata + = COLUMN_CRYPTO_METADATA_UNKNOWN + | ENCRYPTION_WITH_FOOTER_KEY + | EncryptionWithColumnKey + { columnCryptPathInSchema :: [String] + , columnKeyMetadata :: [Word8] + } + deriving (Show, Eq) + +data SortingColumn = SortingColumn + { columnIndex :: Int32 + , columnOrderDescending :: Bool + , nullFirst :: Bool + } + deriving (Show, Eq) + +emptySortingColumn :: SortingColumn +emptySortingColumn = SortingColumn 0 False False + +data ColumnOrder + = TYPE_ORDER + | COLUMN_ORDER_UNKNOWN + deriving (Show, Eq) + +data EncryptionAlgorithm + = ENCRYPTION_ALGORITHM_UNKNOWN + | AesGcmV1 + { aadPrefix :: [Word8] + , aadFileUnique :: [Word8] + , supplyAadPrefix :: Bool + } + | AesGcmCtrV1 + { aadPrefix :: [Word8] + , aadFileUnique :: [Word8] + , supplyAadPrefix :: Bool + } + deriving (Show, Eq) + +data DictVals + = DBool [Bool] + | DInt32 [Int32] + | DInt64 [Int64] + | DInt96 [UTCTime] + | DFloat [Float] + | DDouble [Double] + | DText [T.Text] + deriving (Show, Eq) + +data Page = Page + { pageHeader :: PageHeader + , pageBytes :: [Word8] + } + deriving (Show, Eq) + +data PageHeader = PageHeader + { pageHeaderPageType :: PageType + , uncompressedPageSize :: Int32 + , compressedPageSize :: Int32 + , pageHeaderCrcChecksum :: Int32 + , pageTypeHeader :: PageTypeHeader + } + deriving (Show, Eq) + +emptyPageHeader = PageHeader PAGE_TYPE_UNKNOWN 0 0 0 PAGE_TYPE_HEADER_UNKNOWN + +data PageTypeHeader + = DataPageHeader + { dataPageHeaderNumValues :: Int32 + , dataPageHeaderEncoding :: ParquetEncoding + , definitionLevelEncoding :: ParquetEncoding + , repetitionLevelEncoding :: ParquetEncoding + , dataPageHeaderStatistics :: ColumnStatistics + } + | DataPageHeaderV2 + { dataPageHeaderV2NumValues :: Int32 + , dataPageHeaderV2NumNulls :: Int32 + , dataPageHeaderV2NumRows :: Int32 + , dataPageHeaderV2Encoding :: ParquetEncoding + , definitionLevelByteLength :: Int32 + , repetitionLevelByteLength :: Int32 + , dataPageHeaderV2IsCompressed :: Bool + , dataPageHeaderV2Statistics :: ColumnStatistics + } + | DictionaryPageHeader + { dictionaryPageHeaderNumValues :: Int32 + , dictionaryPageHeaderEncoding :: ParquetEncoding + , dictionaryPageIsSorted :: Bool + } + | INDEX_PAGE_HEADER + | PAGE_TYPE_HEADER_UNKNOWN + deriving (Show, Eq) + +emptyDictionaryPageHeader = DictionaryPageHeader 0 PARQUET_ENCODING_UNKNOWN False +emptyDataPageHeader = DataPageHeader 0 PARQUET_ENCODING_UNKNOWN PARQUET_ENCODING_UNKNOWN PARQUET_ENCODING_UNKNOWN emptyColumnStatistics +emptyDataPageHeaderV2 = DataPageHeaderV2 0 0 0 PARQUET_ENCODING_UNKNOWN 0 0 {- default for v2 is compressed -} True emptyColumnStatistics + +data RepetitionType = REQUIRED | OPTIONAL | REPEATED | UNKNOWN_REPETITION_TYPE deriving (Eq, Show) + +data LogicalType + = STRING_TYPE + | MAP_TYPE + | LIST_TYPE + | ENUM_TYPE + | DECIMAL_TYPE + | DATE_TYPE + | DecimalType {decimalTypePrecision :: Int32, decimalTypeScale :: Int32} + | TimeType {isAdjustedToUTC :: Bool, unit :: TimeUnit} + | -- This should probably have a different, more constrained TimeUnit type. + TimestampType {isAdjustedToUTC :: Bool, unit :: TimeUnit} + | IntType {bitWidth :: Int8, intIsSigned :: Bool} + | LOGICAL_TYPE_UNKNOWN + | JSON_TYPE + | BSON_TYPE + | UUID_TYPE + | FLOAT16_TYPE + | VariantType {specificationVersion :: Int8} + | GeometryType {crs :: T.Text} + | GeographyType {crs :: T.Text, algorithm :: EdgeInterpolationAlgorithm} + deriving (Eq, Show) + +data TimeUnit + = MILLISECONDS + | MICROSECONDS + | NANOSECONDS + | TIME_UNIT_UNKNOWN + deriving (Eq, Show) + +data EdgeInterpolationAlgorithm + = SPHERICAL + | VINCENTY + | THOMAS + | ANDOYER + | KARNEY + deriving (Eq, Show) + +repetitionTypeFromInt :: Int32 -> RepetitionType +repetitionTypeFromInt 0 = REQUIRED +repetitionTypeFromInt 1 = OPTIONAL +repetitionTypeFromInt 2 = REPEATED +repetitionTypeFromInt _ = UNKNOWN_REPETITION_TYPE + +compressionCodecFromInt :: Int32 -> CompressionCodec +compressionCodecFromInt 0 = UNCOMPRESSED +compressionCodecFromInt 1 = SNAPPY +compressionCodecFromInt 2 = GZIP +compressionCodecFromInt 3 = LZO +compressionCodecFromInt 4 = BROTLI +compressionCodecFromInt 5 = LZ4 +compressionCodecFromInt 6 = ZSTD +compressionCodecFromInt 7 = LZ4_RAW +compressionCodecFromInt _ = COMPRESSION_CODEC_UNKNOWN diff --git a/src/DataFrame/Internal/Column.hs b/src/DataFrame/Internal/Column.hs --- a/src/DataFrame/Internal/Column.hs +++ b/src/DataFrame/Internal/Column.hs @@ -1,72 +1,76 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE PolyKinds #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE LambdaCase #-} + module DataFrame.Internal.Column where import qualified Data.ByteString.Char8 as C import qualified Data.List as L import qualified Data.Set as S import qualified Data.Text as T +import qualified Data.Vector as VB import qualified Data.Vector.Algorithms.Merge as VA import qualified Data.Vector.Generic as VG -import qualified Data.Vector as VB import qualified Data.Vector.Mutable as VBM import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM +import Control.Exception (throw) import Control.Monad.ST (runST) -import DataFrame.Internal.Types -import DataFrame.Internal.Parsing import Data.Int +import Data.Kind (Constraint, Type) import Data.Maybe import Data.Proxy import Data.Text.Encoding (decodeUtf8Lenient) -import Data.Type.Equality (type (:~:)(Refl), TestEquality (..)) +import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) import Data.Typeable (Typeable, cast) import Data.Word +import DataFrame.Errors +import DataFrame.Internal.Parsing +import DataFrame.Internal.Types import Type.Reflection import Unsafe.Coerce (unsafeCoerce) -import DataFrame.Errors -import Control.Exception (throw) -import Data.Kind (Type, Constraint) --- | Our representation of a column is a GADT that can store data based on the underlying data. --- --- This allows us to pattern match on data kinds and limit some operations to only some --- kinds of vectors. E.g. operations for missing data only happen in an OptionalColumn. +{- | Our representation of a column is a GADT that can store data based on the underlying data. + +This allows us to pattern match on data kinds and limit some operations to only some +kinds of vectors. E.g. operations for missing data only happen in an OptionalColumn. +-} data Column where - BoxedColumn :: Columnable a => VB.Vector a -> Column - UnboxedColumn :: (Columnable a, VU.Unbox a) => VU.Vector a -> Column - OptionalColumn :: Columnable a => VB.Vector (Maybe a) -> Column + BoxedColumn :: (Columnable a) => VB.Vector a -> Column + UnboxedColumn :: (Columnable a, VU.Unbox a) => VU.Vector a -> Column + OptionalColumn :: (Columnable a) => VB.Vector (Maybe a) -> Column data MutableColumn where - MBoxedColumn :: Columnable a => VBM.IOVector a -> MutableColumn - MUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> MutableColumn + MBoxedColumn :: (Columnable a) => VBM.IOVector a -> MutableColumn + MUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> MutableColumn --- | A TypedColumn is a wrapper around our type-erased column. --- It is used to type check expressions on columns. +{- | A TypedColumn is a wrapper around our type-erased column. +It is used to type check expressions on columns. +-} data TypedColumn a where - TColumn :: Columnable a => Column -> TypedColumn a + TColumn :: (Columnable a) => Column -> TypedColumn a -- | Gets the underlying value from a TypedColumn. unwrapTypedColumn :: TypedColumn a -> Column unwrapTypedColumn (TColumn value) = value --- | An internal function that checks if a column can --- be used in missing value operations. +{- | An internal function that checks if a column can +be used in missing value operations. +-} isOptional :: Column -> Bool isOptional (OptionalColumn column) = True isOptional _ = False @@ -74,62 +78,70 @@ -- | An internal/debugging function to get the column type of a column. columnVersionString :: Column -> String columnVersionString column = case column of - BoxedColumn _ -> "Boxed" - UnboxedColumn _ -> "Unboxed" - OptionalColumn _ -> "Optional" + BoxedColumn _ -> "Boxed" + UnboxedColumn _ -> "Unboxed" + OptionalColumn _ -> "Optional" --- | An internal/debugging function to get the type stored in the outermost vector --- of a column. +{- | An internal/debugging function to get the type stored in the outermost vector +of a column. +-} columnTypeString :: Column -> String columnTypeString column = case column of - BoxedColumn (column :: VB.Vector a) -> show (typeRep @a) - UnboxedColumn (column :: VU.Vector a) -> show (typeRep @a) - OptionalColumn (column :: VB.Vector a) -> show (typeRep @a) + BoxedColumn (column :: VB.Vector a) -> show (typeRep @a) + UnboxedColumn (column :: VU.Vector a) -> show (typeRep @a) + OptionalColumn (column :: VB.Vector a) -> show (typeRep @a) instance (Show a) => Show (TypedColumn a) where - show (TColumn col) = show col + show (TColumn col) = show col instance Show Column where - show :: Column -> String - show (BoxedColumn column) = show column - show (UnboxedColumn column) = show column - show (OptionalColumn column) = show column + show :: Column -> String + show (BoxedColumn column) = show column + show (UnboxedColumn column) = show column + show (OptionalColumn column) = show column instance Eq Column where - (==) :: Column -> Column -> Bool - (==) (BoxedColumn (a :: VB.Vector t1)) (BoxedColumn (b :: VB.Vector t2)) = - case testEquality (typeRep @t1) (typeRep @t2) of - Nothing -> False - Just Refl -> a == b - (==) (OptionalColumn (a :: VB.Vector t1)) (OptionalColumn (b :: VB.Vector t2)) = - case testEquality (typeRep @t1) (typeRep @t2) of - Nothing -> False - Just Refl -> a == b - (==) (UnboxedColumn (a :: VU.Vector t1)) (UnboxedColumn (b :: VU.Vector t2)) = - case testEquality (typeRep @t1) (typeRep @t2) of - Nothing -> False - Just Refl -> a == b - (==) _ _ = False + (==) :: Column -> Column -> Bool + (==) (BoxedColumn (a :: VB.Vector t1)) (BoxedColumn (b :: VB.Vector t2)) = + case testEquality (typeRep @t1) (typeRep @t2) of + Nothing -> False + Just Refl -> a == b + (==) (OptionalColumn (a :: VB.Vector t1)) (OptionalColumn (b :: VB.Vector t2)) = + case testEquality (typeRep @t1) (typeRep @t2) of + Nothing -> False + Just Refl -> a == b + (==) (UnboxedColumn (a :: VU.Vector t1)) (UnboxedColumn (b :: VU.Vector t2)) = + case testEquality (typeRep @t1) (typeRep @t2) of + Nothing -> False + Just Refl -> a == b + (==) _ _ = False --- | A class for converting a vector to a column of the appropriate type. --- Given each Rep we tell the `toColumnRep` function which Column type to pick. +{- | A class for converting a vector to a column of the appropriate type. +Given each Rep we tell the `toColumnRep` function which Column type to pick. +-} class ColumnifyRep (r :: Rep) a where - toColumnRep :: VB.Vector a -> Column + toColumnRep :: VB.Vector a -> Column -- | Constraint synonym for what we can put into columns. -type Columnable a = (Columnable' a, ColumnifyRep (KindOf a) a, UnboxIf a, SBoolI (Unboxable a), SBoolI (Numeric a) ) +type Columnable a = (Columnable' a, ColumnifyRep (KindOf a) a, UnboxIf a, SBoolI (Unboxable a), SBoolI (Numeric a)) -instance (Columnable a, VU.Unbox a) - => ColumnifyRep 'RUnboxed a where - toColumnRep = UnboxedColumn . VU.convert +instance + (Columnable a, VU.Unbox a) => + ColumnifyRep 'RUnboxed a + where + toColumnRep = UnboxedColumn . VU.convert -instance Columnable a - => ColumnifyRep 'RBoxed a where - toColumnRep = BoxedColumn +instance + (Columnable a) => + ColumnifyRep 'RBoxed a + where + toColumnRep = BoxedColumn -instance Columnable a - => ColumnifyRep 'ROptional (Maybe a) where - toColumnRep = OptionalColumn +instance + (Columnable a) => + ColumnifyRep 'ROptional (Maybe a) + where + toColumnRep = OptionalColumn {- | O(n) Convert a vector to a column. Automatically picks the best representation of a vector to store the underlying data in. @@ -142,8 +154,9 @@ @ -} fromVector :: - forall a. (Columnable a, ColumnifyRep (KindOf a) a) - => VB.Vector a -> Column + forall a. + (Columnable a, ColumnifyRep (KindOf a) a) => + VB.Vector a -> Column fromVector = toColumnRep @(KindOf a) {- | O(n) Convert an unboxed vector to a column. This avoids the extra conversion if you already have the data in an unboxed vector. @@ -169,35 +182,36 @@ @ -} fromList :: - forall a. (Columnable a, ColumnifyRep (KindOf a) a) - => [a] -> Column + forall a. + (Columnable a, ColumnifyRep (KindOf a) a) => + [a] -> Column fromList = toColumnRep @(KindOf a) . VB.fromList -- | An internal function to map a function over the values of a column. -mapColumn - :: forall b c. - ( Columnable b - , Columnable c - , UnboxIf c) - => (b -> c) - -> Column - -> Maybe Column +mapColumn :: + forall b c. + ( Columnable b + , Columnable c + , UnboxIf c + ) => + (b -> c) -> + Column -> + Maybe Column mapColumn f = \case - BoxedColumn (col :: VB.Vector a) - | Just Refl <- testEquality (typeRep @a) (typeRep @b) - -> Just (fromVector @c (VB.map f col)) - | otherwise -> Nothing - OptionalColumn (col :: VB.Vector a) - | Just Refl <- testEquality (typeRep @a) (typeRep @b) - -> Just (fromVector @c (VB.map f col)) - | otherwise -> Nothing - UnboxedColumn (col :: VU.Vector a) - | Just Refl <- testEquality (typeRep @a) (typeRep @b) - -> Just $ case sUnbox @c of - STrue -> UnboxedColumn (VU.map f col) + BoxedColumn (col :: VB.Vector a) + | Just Refl <- testEquality (typeRep @a) (typeRep @b) -> + Just (fromVector @c (VB.map f col)) + | otherwise -> Nothing + OptionalColumn (col :: VB.Vector a) + | Just Refl <- testEquality (typeRep @a) (typeRep @b) -> + Just (fromVector @c (VB.map f col)) + | otherwise -> Nothing + UnboxedColumn (col :: VU.Vector a) + | Just Refl <- testEquality (typeRep @a) (typeRep @b) -> + Just $ case sUnbox @c of + STrue -> UnboxedColumn (VU.map f col) SFalse -> fromVector @c (VB.map f (VB.convert col)) - | otherwise -> Nothing - + | otherwise -> Nothing -- | O(1) Gets the number of elements in the column. columnLength :: Column -> Int @@ -242,7 +256,7 @@ -- | O(n) Selects the elements at a given set of indices. Does not change the order. atIndicesStable :: VU.Vector Int -> Column -> Column atIndicesStable indexes (BoxedColumn column) = BoxedColumn $ indexes `getIndices` column -atIndicesStable indexes (UnboxedColumn column) = UnboxedColumn $ indexes `getIndicesUnboxed` column +atIndicesStable indexes (UnboxedColumn column) = UnboxedColumn $ VU.unsafeBackpermute column indexes atIndicesStable indexes (OptionalColumn column) = OptionalColumn $ indexes `getIndices` column {-# INLINE atIndicesStable #-} @@ -256,119 +270,123 @@ getIndicesUnboxed indices xs = VU.generate (VU.length indices) (\i -> xs VU.! (indices VU.! i)) {-# INLINE getIndicesUnboxed #-} -findIndices :: forall a. (Columnable a) - => (a -> Bool) - -> Column - -> Maybe (VU.Vector Int) +findIndices :: + forall a. + (Columnable a) => + (a -> Bool) -> + Column -> + Maybe (VU.Vector Int) findIndices pred (BoxedColumn (column :: VB.Vector b)) = do - Refl <- testEquality (typeRep @a) (typeRep @b) - pure $ VG.convert (VG.findIndices pred column) + Refl <- testEquality (typeRep @a) (typeRep @b) + pure $ VG.convert (VG.findIndices pred column) findIndices pred (UnboxedColumn (column :: VU.Vector b)) = do - Refl <- testEquality (typeRep @a) (typeRep @b) - pure $ VG.findIndices pred column + Refl <- testEquality (typeRep @a) (typeRep @b) + pure $ VG.findIndices pred column findIndices pred (OptionalColumn (column :: VB.Vector (Maybe b))) = do - Refl <- testEquality (typeRep @a) (typeRep @(Maybe b)) - pure $ VG.convert (VG.findIndices pred column) + Refl <- testEquality (typeRep @a) (typeRep @(Maybe b)) + pure $ VG.convert (VG.findIndices pred column) -- | An internal function that returns a vector of how indexes change after a column is sorted. sortedIndexes :: Bool -> Column -> VU.Vector Int -sortedIndexes asc (BoxedColumn column ) = runST $ do - withIndexes <- VG.thaw $ VG.indexed column - VA.sortBy (\(a, b) (a', b') -> (if asc then compare else flip compare) b b') withIndexes - sorted <- VG.unsafeFreeze withIndexes - return $ VU.generate (VG.length column) (\i -> fst (sorted VG.! i)) +sortedIndexes asc (BoxedColumn column) = runST $ do + withIndexes <- VG.thaw $ VG.indexed column + VA.sortBy (\(a, b) (a', b') -> (if asc then compare else flip compare) b b') withIndexes + sorted <- VG.unsafeFreeze withIndexes + return $ VU.generate (VG.length column) (\i -> fst (sorted VG.! i)) sortedIndexes asc (UnboxedColumn column) = runST $ do - withIndexes <- VG.thaw $ VG.indexed column - VA.sortBy (\(a, b) (a', b') -> (if asc then compare else flip compare) b b') withIndexes - sorted <- VG.unsafeFreeze withIndexes - return $ VU.generate (VG.length column) (\i -> fst (sorted VG.! i)) -sortedIndexes asc (OptionalColumn column ) = runST $ do - withIndexes <- VG.thaw $ VG.indexed column - VA.sortBy (\(a, b) (a', b') -> (if asc then compare else flip compare) b b') withIndexes - sorted <- VG.unsafeFreeze withIndexes - return $ VU.generate (VG.length column) (\i -> fst (sorted VG.! i)) + withIndexes <- VG.thaw $ VG.indexed column + VA.sortBy (\(a, b) (a', b') -> (if asc then compare else flip compare) b b') withIndexes + sorted <- VG.unsafeFreeze withIndexes + return $ VU.generate (VG.length column) (\i -> fst (sorted VG.! i)) +sortedIndexes asc (OptionalColumn column) = runST $ do + withIndexes <- VG.thaw $ VG.indexed column + VA.sortBy (\(a, b) (a', b') -> (if asc then compare else flip compare) b b') withIndexes + sorted <- VG.unsafeFreeze withIndexes + return $ VU.generate (VG.length column) (\i -> fst (sorted VG.! i)) {-# INLINE sortedIndexes #-} -- | Applies a function that returns an unboxed result to an unboxed vector, storing the result in a column. -imapColumn - :: forall b c. (Columnable b, Columnable c) - => (Int -> b -> c) -> Column -> Maybe Column +imapColumn :: + forall b c. + (Columnable b, Columnable c) => + (Int -> b -> c) -> Column -> Maybe Column imapColumn f = \case - BoxedColumn (col :: VB.Vector a) - | Just Refl <- testEquality (typeRep @a) (typeRep @b) - -> Just (fromVector @c (VB.imap f col)) - | otherwise -> Nothing - UnboxedColumn (col :: VU.Vector a) - | Just Refl <- testEquality (typeRep @a) (typeRep @b) - -> Just $ - case sUnbox @c of - STrue -> UnboxedColumn (VU.imap f col) - SFalse -> fromVector @c (VB.imap f (VB.convert col)) - | otherwise -> Nothing - OptionalColumn (col :: VB.Vector a) - | Just Refl <- testEquality (typeRep @a) (typeRep @b) - -> Just (fromVector @c (VB.imap f col)) - | otherwise -> Nothing + BoxedColumn (col :: VB.Vector a) + | Just Refl <- testEquality (typeRep @a) (typeRep @b) -> + Just (fromVector @c (VB.imap f col)) + | otherwise -> Nothing + UnboxedColumn (col :: VU.Vector a) + | Just Refl <- testEquality (typeRep @a) (typeRep @b) -> + Just $ + case sUnbox @c of + STrue -> UnboxedColumn (VU.imap f col) + SFalse -> fromVector @c (VB.imap f (VB.convert col)) + | otherwise -> Nothing + OptionalColumn (col :: VB.Vector a) + | Just Refl <- testEquality (typeRep @a) (typeRep @b) -> + Just (fromVector @c (VB.imap f col)) + | otherwise -> Nothing -- | Filter column with index. -ifilterColumn :: forall a . (Columnable a) => (Int -> a -> Bool) -> Column -> Maybe Column +ifilterColumn :: forall a. (Columnable a) => (Int -> a -> Bool) -> Column -> Maybe Column ifilterColumn f c@(BoxedColumn (column :: VB.Vector b)) = do - Refl <- testEquality (typeRep @a) (typeRep @b) - return $ BoxedColumn $ VG.ifilter f column + Refl <- testEquality (typeRep @a) (typeRep @b) + return $ BoxedColumn $ VG.ifilter f column ifilterColumn f c@(UnboxedColumn (column :: VU.Vector b)) = do - Refl <- testEquality (typeRep @a) (typeRep @b) - return $ UnboxedColumn $ VG.ifilter f column + Refl <- testEquality (typeRep @a) (typeRep @b) + return $ UnboxedColumn $ VG.ifilter f column ifilterColumn _ _ = Nothing -- | Fold (right) column with index. ifoldrColumn :: forall a b. (Columnable a, Columnable b) => (Int -> a -> b -> b) -> b -> Column -> Maybe b ifoldrColumn f acc c@(BoxedColumn (column :: VB.Vector d)) = do - Refl <- testEquality (typeRep @a) (typeRep @d) - return $ VG.ifoldr f acc column + Refl <- testEquality (typeRep @a) (typeRep @d) + return $ VG.ifoldr f acc column ifoldrColumn f acc c@(OptionalColumn (column :: VB.Vector d)) = do - Refl <- testEquality (typeRep @a) (typeRep @d) - return $ VG.ifoldr f acc column + Refl <- testEquality (typeRep @a) (typeRep @d) + return $ VG.ifoldr f acc column ifoldrColumn f acc c@(UnboxedColumn (column :: VU.Vector d)) = do - Refl <- testEquality (typeRep @a) (typeRep @d) - return $ VG.ifoldr f acc column + Refl <- testEquality (typeRep @a) (typeRep @d) + return $ VG.ifoldr f acc column -- | Fold (left) column with index. -ifoldlColumn :: forall a b . (Columnable a, Columnable b) => (b -> Int -> a -> b) -> b -> Column -> Maybe b +ifoldlColumn :: forall a b. (Columnable a, Columnable b) => (b -> Int -> a -> b) -> b -> Column -> Maybe b ifoldlColumn f acc c@(BoxedColumn (column :: VB.Vector d)) = do - Refl <- testEquality (typeRep @a) (typeRep @d) - return $ VG.ifoldl' f acc column + Refl <- testEquality (typeRep @a) (typeRep @d) + return $ VG.ifoldl' f acc column ifoldlColumn f acc c@(OptionalColumn (column :: VB.Vector d)) = do - Refl <- testEquality (typeRep @a) (typeRep @d) - return $ VG.ifoldl' f acc column + Refl <- testEquality (typeRep @a) (typeRep @d) + return $ VG.ifoldl' f acc column ifoldlColumn f acc c@(UnboxedColumn (column :: VU.Vector d)) = do - Refl <- testEquality (typeRep @a) (typeRep @d) - return $ VG.ifoldl' f acc column + Refl <- testEquality (typeRep @a) (typeRep @d) + return $ VG.ifoldl' f acc column -headColumn :: forall a . Columnable a => Column -> Maybe a +headColumn :: forall a. (Columnable a) => Column -> Maybe a headColumn (BoxedColumn (col :: VB.Vector b)) = do - Refl <- testEquality (typeRep @a) (typeRep @b) - pure (VG.head col) + Refl <- testEquality (typeRep @a) (typeRep @b) + pure (VG.head col) headColumn (UnboxedColumn (col :: VU.Vector b)) = do - Refl <- testEquality (typeRep @a) (typeRep @b) - pure (VG.head col) + Refl <- testEquality (typeRep @a) (typeRep @b) + pure (VG.head col) headColumn (OptionalColumn (col :: VB.Vector b)) = do - Refl <- testEquality (typeRep @a) (typeRep @b) - pure (VG.head col) + Refl <- testEquality (typeRep @a) (typeRep @b) + pure (VG.head col) -- | Generic reduce function for all Column types. -reduceColumn :: forall a b. Columnable a => (a -> b) -> Column -> Maybe b +reduceColumn :: forall a b. (Columnable a) => (a -> b) -> Column -> Maybe b {-# SPECIALIZE reduceColumn :: - (VU.Vector (Double, Double) -> Double) -> Column -> Maybe Double, - (VU.Vector Double -> Double) -> Column -> Maybe Double #-} + (VU.Vector (Double, Double) -> Double) -> Column -> Maybe Double + , (VU.Vector Double -> Double) -> Column -> Maybe Double + #-} reduceColumn f (BoxedColumn (column :: c)) = do - Refl <- testEquality (typeRep @c) (typeRep @a) - pure $ f column + Refl <- testEquality (typeRep @c) (typeRep @a) + pure $ f column reduceColumn f (UnboxedColumn (column :: c)) = do - Refl <- testEquality (typeRep @c) (typeRep @a) - pure $ f column + Refl <- testEquality (typeRep @c) (typeRep @a) + pure $ f column reduceColumn f (OptionalColumn (column :: c)) = do - Refl <- testEquality (typeRep @c) (typeRep @a) - pure $ f column + Refl <- testEquality (typeRep @c) (typeRep @a) + pure $ f column {-# INLINE reduceColumn #-} -- | An internal, column version of zip. @@ -376,95 +394,98 @@ zipColumns (BoxedColumn column) (BoxedColumn other) = BoxedColumn (VG.zip column other) zipColumns (BoxedColumn column) (UnboxedColumn other) = BoxedColumn (VB.generate (min (VG.length column) (VG.length other)) (\i -> (column VG.! i, other VG.! i))) zipColumns (BoxedColumn column) (OptionalColumn optcolumn) = BoxedColumn (VG.zip (VB.convert column) optcolumn) - zipColumns (UnboxedColumn column) (BoxedColumn other) = BoxedColumn (VB.generate (min (VG.length column) (VG.length other)) (\i -> (column VG.! i, other VG.! i))) zipColumns (UnboxedColumn column) (UnboxedColumn other) = UnboxedColumn (VG.zip column other) zipColumns (UnboxedColumn column) (OptionalColumn optcolumn) = BoxedColumn (VG.zip (VB.convert column) optcolumn) - zipColumns (OptionalColumn optcolumn) (BoxedColumn column) = BoxedColumn (VG.zip optcolumn (VB.convert column)) zipColumns (OptionalColumn optcolumn) (UnboxedColumn column) = BoxedColumn (VG.zip optcolumn (VB.convert column)) zipColumns (OptionalColumn optcolumn) (OptionalColumn optother) = BoxedColumn (VG.zip optcolumn optother) {-# INLINE zipColumns #-} -- | An internal, column version of zipWith. -zipWithColumns :: forall a b c . (Columnable a, Columnable b, Columnable c) => (a -> b -> c) -> Column -> Column -> Maybe Column +zipWithColumns :: forall a b c. (Columnable a, Columnable b, Columnable c) => (a -> b -> c) -> Column -> Column -> Maybe Column zipWithColumns f (UnboxedColumn (column :: VU.Vector d)) (UnboxedColumn (other :: VU.Vector e)) = case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> case testEquality (typeRep @b) (typeRep @e) of - Just Refl -> pure $ case sUnbox @c of - STrue -> fromUnboxedVector (VU.zipWith f column other) - SFalse -> fromVector $ VB.zipWith f (VG.convert column) (VG.convert other) + Just Refl -> case testEquality (typeRep @b) (typeRep @e) of + Just Refl -> pure $ case sUnbox @c of + STrue -> fromUnboxedVector (VU.zipWith f column other) + SFalse -> fromVector $ VB.zipWith f (VG.convert column) (VG.convert other) + Nothing -> Nothing Nothing -> Nothing - Nothing -> Nothing -zipWithColumns f left right = let - left' = toVector @a left - right' = toVector @b right - in pure $ fromVector $ VB.zipWith f left' right' +zipWithColumns f left right = + let + left' = toVector @a left + right' = toVector @b right + in + pure $ fromVector $ VB.zipWith f left' right' {-# INLINE zipWithColumns #-} -- Functions for mutable columns (intended for IO). writeColumn :: Int -> T.Text -> MutableColumn -> IO (Either T.Text Bool) -writeColumn i value (MBoxedColumn (col :: VBM.IOVector a)) = let - in case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> (if isNullish value +writeColumn i value (MBoxedColumn (col :: VBM.IOVector a)) = + let + in case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> + ( if isNullish value then VBM.unsafeWrite col i "" >> return (Left $! value) - else VBM.unsafeWrite col i value >> return (Right True)) - Nothing -> return (Left value) + else VBM.unsafeWrite col i value >> return (Right True) + ) + Nothing -> return (Left value) writeColumn i value (MUnboxedColumn (col :: VUM.IOVector a)) = - case testEquality (typeRep @a) (typeRep @Int) of - Just Refl -> case readInt value of - Just v -> VUM.unsafeWrite col i v >> return (Right True) - Nothing -> VUM.unsafeWrite col i 0 >> return (Left value) - Nothing -> case testEquality (typeRep @a) (typeRep @Double) of - Nothing -> return (Left $! value) - Just Refl -> case readDouble value of + case testEquality (typeRep @a) (typeRep @Int) of + Just Refl -> case readInt value of Just v -> VUM.unsafeWrite col i v >> return (Right True) - Nothing -> VUM.unsafeWrite col i 0 >> return (Left $! value) + Nothing -> VUM.unsafeWrite col i 0 >> return (Left value) + Nothing -> case testEquality (typeRep @a) (typeRep @Double) of + Nothing -> return (Left $! value) + Just Refl -> case readDouble value of + Just v -> VUM.unsafeWrite col i v >> return (Right True) + Nothing -> VUM.unsafeWrite col i 0 >> return (Left $! value) {-# INLINE writeColumn #-} freezeColumn' :: [(Int, T.Text)] -> MutableColumn -> IO Column freezeColumn' nulls (MBoxedColumn col) - | null nulls = BoxedColumn <$> VB.unsafeFreeze col - | all (isNullish . snd) nulls = OptionalColumn . VB.imap (\i v -> if i `elem` map fst nulls then Nothing else Just v) <$> VB.unsafeFreeze col - | otherwise = BoxedColumn . VB.imap (\i v -> if i `elem` map fst nulls then Left (fromMaybe (error "") (lookup i nulls)) else Right v) <$> VB.unsafeFreeze col + | null nulls = BoxedColumn <$> VB.unsafeFreeze col + | all (isNullish . snd) nulls = OptionalColumn . VB.imap (\i v -> if i `elem` map fst nulls then Nothing else Just v) <$> VB.unsafeFreeze col + | otherwise = BoxedColumn . VB.imap (\i v -> if i `elem` map fst nulls then Left (fromMaybe (error "") (lookup i nulls)) else Right v) <$> VB.unsafeFreeze col freezeColumn' nulls (MUnboxedColumn col) - | null nulls = UnboxedColumn <$> VU.unsafeFreeze col - | all (isNullish . snd) nulls = VU.unsafeFreeze col >>= \c -> return $ OptionalColumn $ VB.generate (VU.length c) (\i -> if i `elem` map fst nulls then Nothing else Just (c VU.! i)) - | otherwise = VU.unsafeFreeze col >>= \c -> return $ BoxedColumn $ VB.generate (VU.length c) (\i -> if i `elem` map fst nulls then Left (fromMaybe (error "") (lookup i nulls)) else Right (c VU.! i)) + | null nulls = UnboxedColumn <$> VU.unsafeFreeze col + | all (isNullish . snd) nulls = VU.unsafeFreeze col >>= \c -> return $ OptionalColumn $ VB.generate (VU.length c) (\i -> if i `elem` map fst nulls then Nothing else Just (c VU.! i)) + | otherwise = VU.unsafeFreeze col >>= \c -> return $ BoxedColumn $ VB.generate (VU.length c) (\i -> if i `elem` map fst nulls then Left (fromMaybe (error "") (lookup i nulls)) else Right (c VU.! i)) {-# INLINE freezeColumn' #-} -- | Fills the end of a column, up to n, with Nothing. Does nothing if column has length greater than n. expandColumn :: Int -> Column -> Column expandColumn n (OptionalColumn col) = OptionalColumn $ col <> VB.replicate (n - VG.length col) Nothing expandColumn n column@(BoxedColumn col) - | n > VG.length col = OptionalColumn $ VB.map Just col <> VB.replicate (n - VG.length col) Nothing - | otherwise = column + | n > VG.length col = OptionalColumn $ VB.map Just col <> VB.replicate (n - VG.length col) Nothing + | otherwise = column expandColumn n column@(UnboxedColumn col) - | n > VG.length col = OptionalColumn $ VB.map Just (VU.convert col) <> VB.replicate (n - VG.length col) Nothing - | otherwise = column + | n > VG.length col = OptionalColumn $ VB.map Just (VU.convert col) <> VB.replicate (n - VG.length col) Nothing + | otherwise = column -- | Fills the beginning of a column, up to n, with Nothing. Does nothing if column has length greater than n. leftExpandColumn :: Int -> Column -> Column leftExpandColumn n column@(OptionalColumn col) - | n > VG.length col = OptionalColumn $ VG.replicate (n - VG.length col) Nothing <> col - | otherwise = column + | n > VG.length col = OptionalColumn $ VG.replicate (n - VG.length col) Nothing <> col + | otherwise = column leftExpandColumn n column@(BoxedColumn col) - | n > VG.length col = OptionalColumn $ VG.replicate (n - VG.length col) Nothing <> VG.map Just col - | otherwise = column + | n > VG.length col = OptionalColumn $ VG.replicate (n - VG.length col) Nothing <> VG.map Just col + | otherwise = column leftExpandColumn n column@(UnboxedColumn col) - | n > VG.length col = OptionalColumn $ VG.replicate (n - VG.length col) Nothing <> VG.map Just (VU.convert col) - | otherwise = column + | n > VG.length col = OptionalColumn $ VG.replicate (n - VG.length col) Nothing <> VG.map Just (VU.convert col) + | otherwise = column -- | Concatenates two columns. concatColumns :: Column -> Column -> Maybe Column concatColumns (OptionalColumn left) (OptionalColumn right) = case testEquality (typeOf left) (typeOf right) of - Nothing -> Nothing - Just Refl -> Just (OptionalColumn $ left <> right) + Nothing -> Nothing + Just Refl -> Just (OptionalColumn $ left <> right) concatColumns (BoxedColumn left) (BoxedColumn right) = case testEquality (typeOf left) (typeOf right) of - Nothing -> Nothing - Just Refl -> Just (BoxedColumn $ left <> right) + Nothing -> Nothing + Just Refl -> Just (BoxedColumn $ left <> right) concatColumns (UnboxedColumn left) (UnboxedColumn right) = case testEquality (typeOf left) (typeOf right) of - Nothing -> Nothing - Just Refl -> Just (UnboxedColumn $ left <> right) + Nothing -> Nothing + Just Refl -> Just (UnboxedColumn $ left <> right) concatColumns _ _ = Nothing {- | O(n) Converts a column to a boxed vector. Throws an exception if the wrong type is specified. @@ -477,11 +498,12 @@ [1,2,3,4] > toVector @Double column exception: ... +@ -} -toVector :: forall a . Columnable a => Column -> VB.Vector a +toVector :: forall a. (Columnable a) => Column -> VB.Vector a toVector xs = case toVectorSafe xs of - Left err -> throw err - Right val -> val + Left err -> throw err + Right val -> val {- | O(n) Converts a column to a list. Throws an exception if the wrong type is specified. @@ -493,32 +515,51 @@ [1,2,3,4] > toList @Double column exception: ... +@ -} -toList :: forall a . Columnable a => Column -> [a] +toList :: forall a. (Columnable a) => Column -> [a] toList xs = case toVectorSafe @a xs of - Left err -> throw err - Right val -> VB.toList val + Left err -> throw err + Right val -> VB.toList val -- | A safe version of toVector that returns an Either type. -toVectorSafe :: forall a v . (VG.Vector v a, Columnable a) => Column -> Either DataFrameException (v a) +toVectorSafe :: forall a v. (VG.Vector v a, Columnable a) => Column -> Either DataFrameException (v a) toVectorSafe column@(OptionalColumn (col :: VB.Vector b)) = - case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> Right $ VG.convert col - Nothing -> Left $ TypeMismatchException (MkTypeErrorContext { userType = Right (typeRep @a) - , expectedType = Right (typeRep @b) - , callingFunctionName = Just "toVectorSafe" - , errorColumnName = Nothing}) + case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> Right $ VG.convert col + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @a) + , expectedType = Right (typeRep @b) + , callingFunctionName = Just "toVectorSafe" + , errorColumnName = Nothing + } + ) toVectorSafe (BoxedColumn (col :: VB.Vector b)) = - case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> Right $ VG.convert col - Nothing -> Left $ TypeMismatchException (MkTypeErrorContext { userType = Right (typeRep @a) - , expectedType = Right (typeRep @b) - , callingFunctionName = Just "toVectorSafe" - , errorColumnName = Nothing}) + case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> Right $ VG.convert col + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @a) + , expectedType = Right (typeRep @b) + , callingFunctionName = Just "toVectorSafe" + , errorColumnName = Nothing + } + ) toVectorSafe (UnboxedColumn (col :: VU.Vector b)) = - case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> Right $ VG.convert col - Nothing -> Left $ TypeMismatchException (MkTypeErrorContext { userType = Right (typeRep @a) - , expectedType = Right (typeRep @b) - , callingFunctionName = Just "toVectorSafe" - , errorColumnName = Nothing}) + case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> Right $ VG.convert col + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @a) + , expectedType = Right (typeRep @b) + , callingFunctionName = Just "toVectorSafe" + , errorColumnName = Nothing + } + ) diff --git a/src/DataFrame/Internal/DataFrame.hs b/src/DataFrame/Internal/DataFrame.hs --- a/src/DataFrame/Internal/DataFrame.hs +++ b/src/DataFrame/Internal/DataFrame.hs @@ -1,98 +1,110 @@ {-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE FlexibleContexts #-} + module DataFrame.Internal.DataFrame where import qualified Data.Map as M import qualified Data.Text as T import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Unboxed as VU import Control.Monad (join) -import DataFrame.Display.Terminal.PrettyPrint -import DataFrame.Internal.Column import Data.Function (on) import Data.List (sortBy, transpose, (\\)) import Data.Maybe (isJust) -import Data.Type.Equality (type (:~:)(Refl), TestEquality (testEquality)) +import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) +import DataFrame.Display.Terminal.PrettyPrint +import DataFrame.Internal.Column import Text.Printf import Type.Reflection (typeRep) data DataFrame = DataFrame - { -- | Our main data structure stores a dataframe as - -- a vector of columns. This improv - columns :: V.Vector Column, - -- | Keeps the column names in the order they were inserted in. - columnIndices :: M.Map T.Text Int, - dataframeDimensions :: (Int, Int) - } + { columns :: V.Vector Column + {- ^ Our main data structure stores a dataframe as + a vector of columns. This improv + -} + , columnIndices :: M.Map T.Text Int + -- ^ Keeps the column names in the order they were inserted in. + , dataframeDimensions :: (Int, Int) + } -data GroupedDataFrame = Grouped { - fullDataframe :: DataFrame, - groupedColumns :: [T.Text], - valueIndices :: VU.Vector Int, - offsets :: VU.Vector Int -} +{- | A record that contains information about how and what +rows are grouped in the dataframe. This can only be used with +`aggregate`. +-} +data GroupedDataFrame = Grouped + { fullDataframe :: DataFrame + , groupedColumns :: [T.Text] + , valueIndices :: VU.Vector Int + , offsets :: VU.Vector Int + } instance Show GroupedDataFrame where - show (Grouped df cols indices os) = printf "{ keyColumns: %s groupedColumns: %s }" - (show cols) - (show ((M.keys (columnIndices df)) \\ cols)) + show (Grouped df cols indices os) = + printf + "{ keyColumns: %s groupedColumns: %s }" + (show cols) + (show ((M.keys (columnIndices df)) \\ cols)) instance Eq GroupedDataFrame where - (==) (Grouped df cols indices os) (Grouped df' cols' indices' os') = (df == df') && (cols == cols') - + (==) (Grouped df cols indices os) (Grouped df' cols' indices' os') = (df == df') && (cols == cols') instance Eq DataFrame where - (==) :: DataFrame -> DataFrame -> Bool - a == b = map fst (M.toList $ columnIndices a) == map fst (M.toList $ columnIndices b) && - foldr (\(name, index) acc -> acc && (columns a V.!? index == (columns b V.!? (columnIndices b M.! name)))) True (M.toList $ columnIndices a) + (==) :: DataFrame -> DataFrame -> Bool + a == b = + map fst (M.toList $ columnIndices a) == map fst (M.toList $ columnIndices b) + && foldr (\(name, index) acc -> acc && (columns a V.!? index == (columns b V.!? (columnIndices b M.! name)))) True (M.toList $ columnIndices a) instance Show DataFrame where - show :: DataFrame -> String - show d = T.unpack (asText d False) + show :: DataFrame -> String + show d = T.unpack (asText d False) asText :: DataFrame -> Bool -> T.Text asText d properMarkdown = - let header = "index" : map fst (sortBy (compare `on` snd) $ M.toList (columnIndices d)) - types = V.toList $ V.filter (/= "") $ V.map getType (columns d) - getType :: Column -> T.Text - getType (BoxedColumn (column :: V.Vector a)) = T.pack $ show (typeRep @a) - getType (UnboxedColumn (column :: VU.Vector a)) = T.pack $ show (typeRep @a) - getType (OptionalColumn (column :: V.Vector a)) = T.pack $ show (typeRep @a) - -- Separate out cases dynamically so we don't end up making round trip string - -- copies. - get :: Maybe Column -> V.Vector T.Text - get (Just (BoxedColumn (column :: V.Vector a))) = case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> column - Nothing -> case testEquality (typeRep @a) (typeRep @String) of + let header = "index" : map fst (sortBy (compare `on` snd) $ M.toList (columnIndices d)) + types = V.toList $ V.filter (/= "") $ V.map getType (columns d) + getType :: Column -> T.Text + getType (BoxedColumn (column :: V.Vector a)) = T.pack $ show (typeRep @a) + getType (UnboxedColumn (column :: VU.Vector a)) = T.pack $ show (typeRep @a) + getType (OptionalColumn (column :: V.Vector a)) = T.pack $ show (typeRep @a) + -- Separate out cases dynamically so we don't end up making round trip string + -- copies. + get :: Maybe Column -> V.Vector T.Text + get (Just (BoxedColumn (column :: V.Vector a))) = case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> column + Nothing -> case testEquality (typeRep @a) (typeRep @String) of Just Refl -> V.map T.pack column Nothing -> V.map (T.pack . show) column - get (Just (UnboxedColumn column)) = V.map (T.pack . show) (V.convert column) - get (Just (OptionalColumn column)) = V.map (T.pack . show) column - getTextColumnFromFrame df (i, name) = if i == 0 - then V.fromList (map (T.pack . show) [0..(fst (dataframeDimensions df) - 1)]) - else get $ (V.!?) (columns d) ((M.!) (columnIndices d) name) - rows = - transpose $ - zipWith (curry (V.toList . getTextColumnFromFrame d)) [0..] header - in showTable properMarkdown header ("Int":types) rows + get (Just (UnboxedColumn column)) = V.map (T.pack . show) (V.convert column) + get (Just (OptionalColumn column)) = V.map (T.pack . show) column + getTextColumnFromFrame df (i, name) = + if i == 0 + then V.fromList (map (T.pack . show) [0 .. (fst (dataframeDimensions df) - 1)]) + else get $ (V.!?) (columns d) ((M.!) (columnIndices d) name) + rows = + transpose $ + zipWith (curry (V.toList . getTextColumnFromFrame d)) [0 ..] header + in showTable properMarkdown header ("Int" : types) rows -- | O(1) Creates an empty dataframe empty :: DataFrame -empty = DataFrame {columns = V.empty, - columnIndices = M.empty, - dataframeDimensions = (0, 0) } +empty = + DataFrame + { columns = V.empty + , columnIndices = M.empty + , dataframeDimensions = (0, 0) + } getColumn :: T.Text -> DataFrame -> Maybe Column getColumn name df = do - i <- columnIndices df M.!? name - columns df V.!? i + i <- columnIndices df M.!? name + columns df V.!? i unsafeGetColumn :: T.Text -> DataFrame -> Column unsafeGetColumn name df = columns df V.! (columnIndices df M.! name) @@ -100,19 +112,30 @@ null :: DataFrame -> Bool null df = V.null (columns df) +{- | Returns a dataframe as a two dimentions vector of floats. + +All entries in the dataframe must be doubles. +This is useful for handing data over into ML systems. +-} toMatrix :: DataFrame -> V.Vector (VU.Vector Float) -toMatrix df = let - m = V.map (toVector @Double) (columns df) - in V.generate (fst (dataframeDimensions df)) (\i -> foldl (\acc j -> acc `VU.snoc` (realToFrac ((m V.! j) V.! i))) VU.empty [0..(V.length m - 1)]) +toMatrix df = + let + m = V.map (toVector @Double) (columns df) + in + V.generate (fst (dataframeDimensions df)) (\i -> foldl (\acc j -> acc `VU.snoc` (realToFrac ((m V.! j) V.! i))) VU.empty [0 .. (V.length m - 1)]) -columnAsVector :: forall a . Columnable a => T.Text -> DataFrame -> V.Vector a +{- | Get a specific column as a vector. + +You must specify the type via type applications. +-} +columnAsVector :: forall a. (Columnable a) => T.Text -> DataFrame -> V.Vector a columnAsVector name df = case unsafeGetColumn name df of - (BoxedColumn (col :: V.Vector b)) -> case testEquality (typeRep @a) (typeRep @b) of - Nothing -> error "Type error" - Just Refl -> col - (OptionalColumn (col :: V.Vector b)) -> case testEquality (typeRep @a) (typeRep @b) of - Nothing -> error "Type error" - Just Refl -> col - (UnboxedColumn (col :: VU.Vector b)) -> case testEquality (typeRep @a) (typeRep @b) of - Nothing -> error "Type error" - Just Refl -> VG.convert col + (BoxedColumn (col :: V.Vector b)) -> case testEquality (typeRep @a) (typeRep @b) of + Nothing -> error "Type error" + Just Refl -> col + (OptionalColumn (col :: V.Vector b)) -> case testEquality (typeRep @a) (typeRep @b) of + Nothing -> error "Type error" + Just Refl -> col + (UnboxedColumn (col :: VU.Vector b)) -> case testEquality (typeRep @a) (typeRep @b) of + Nothing -> error "Type error" + Just Refl -> VG.convert col diff --git a/src/DataFrame/Internal/Expression.hs b/src/DataFrame/Internal/Expression.hs --- a/src/DataFrame/Internal/Expression.hs +++ b/src/DataFrame/Internal/Expression.hs @@ -1,221 +1,302 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE BangPatterns #-} + module DataFrame.Internal.Expression where -import qualified Data.Map as M -import Data.Type.Equality (type (:~:)(Refl), TestEquality (testEquality)) +import Control.Exception (throw) import Data.Data (Typeable) -import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame -import DataFrame.Internal.Types +import qualified Data.Map as M +import Data.Maybe (fromMaybe) import qualified Data.Text as T +import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Unboxed as VU +import DataFrame.Errors (DataFrameException (ColumnNotFoundException)) +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame +import DataFrame.Internal.Types import Type.Reflection (typeRep) -import DataFrame.Errors (DataFrameException(ColumnNotFoundException)) -import Control.Exception (throw) -import Data.Maybe (fromMaybe) data Expr a where - Col :: Columnable a => T.Text -> Expr a - Lit :: Columnable a => a -> Expr a - Apply :: (Columnable a, - Columnable b) - => T.Text -- Operation name - -> (b -> a) - -> Expr b - -> Expr a - BinOp :: (Columnable c, - Columnable b, - Columnable a) - => T.Text -- operation name - -> (c -> b -> a) - -> Expr c - -> Expr b - -> Expr a - GeneralAggregate :: (Columnable a) - => T.Text -- Column name - -> T.Text -- Operation name - -> (forall v b. (VG.Vector v b, Columnable b) => v b -> a) - -> Expr a - ReductionAggregate :: (Columnable a) - => T.Text -- Column name - -> T.Text -- Operation name - -> (forall a . Columnable a => a -> a -> a) - -> Expr a - NumericAggregate :: (Columnable a, - Columnable b, - VU.Unbox a, - VU.Unbox b, - Num a, - Num b) - => T.Text -- Column name - -> T.Text -- Operation name - -> (VU.Vector b -> a) - -> Expr a + Col :: (Columnable a) => T.Text -> Expr a + Lit :: (Columnable a) => a -> Expr a + Apply :: + ( Columnable a + , Columnable b + ) => + T.Text -> -- Operation name + (b -> a) -> + Expr b -> + Expr a + BinOp :: + ( Columnable c + , Columnable b + , Columnable a + ) => + T.Text -> -- operation name + (c -> b -> a) -> + Expr c -> + Expr b -> + Expr a + GeneralAggregate :: + (Columnable a) => + T.Text -> -- Column name + T.Text -> -- Operation name + (forall v b. (VG.Vector v b, Columnable b) => v b -> a) -> + Expr a + ReductionAggregate :: + (Columnable a) => + T.Text -> -- Column name + T.Text -> -- Operation name + (forall a. (Columnable a) => a -> a -> a) -> + Expr a + NumericAggregate :: + ( Columnable a + , Columnable b + , VU.Unbox a + , VU.Unbox b + , Num a + , Num b + ) => + T.Text -> -- Column name + T.Text -> -- Operation name + (VU.Vector b -> a) -> + Expr a data UExpr where - Wrap :: Columnable a => Expr a -> UExpr + Wrap :: (Columnable a) => Expr a -> UExpr -interpret :: forall a . (Columnable a) => DataFrame -> Expr a -> TypedColumn a +interpret :: forall a. (Columnable a) => DataFrame -> Expr a -> TypedColumn a interpret df (Lit value) = TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value interpret df (Col name) = case getColumn name df of Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df) Just col -> TColumn col -interpret df (Apply _ (f :: c -> d) value) = let +interpret df (Apply _ (f :: c -> d) value) = + let (TColumn value') = interpret @c df value - -- TODO: Handle this gracefully. - in TColumn $ fromMaybe (error "mapColumn returned nothing") (mapColumn f value') -interpret df (BinOp _ (f :: c -> d -> e) left right) = let + in + -- TODO: Handle this gracefully. + TColumn $ fromMaybe (error "mapColumn returned nothing") (mapColumn f value') +interpret df (BinOp _ (f :: c -> d -> e) (Lit left) right) = + let + (TColumn right') = interpret @d df right + in + TColumn $ fromMaybe (error "mapColumn returned nothing") (mapColumn (f left) right') +interpret df (BinOp _ (f :: c -> d -> e) left (Lit right)) = + let (TColumn left') = interpret @c df left + in + TColumn $ fromMaybe (error "mapColumn returned nothing") (mapColumn (flip f right) left') +interpret df (BinOp _ (f :: c -> d -> e) left right) = + let + (TColumn left') = interpret @c df left (TColumn right') = interpret @d df right - in TColumn $ fromMaybe (error "mapColumn returned nothing") (zipWithColumns f left' right') -interpret df (ReductionAggregate name op (f :: forall a . Columnable a => a -> a -> a)) = let + in + TColumn $ fromMaybe (error "mapColumn returned nothing") (zipWithColumns f left' right') +interpret df (ReductionAggregate name op (f :: forall a. (Columnable a) => a -> a -> a)) = + let (TColumn column) = interpret @a df (Col name) - in case headColumn @a column of - Nothing -> error "Invalid operation" - Just h -> case ifoldlColumn (\acc _ v -> f acc v) h column of - Nothing -> error "Invalid operation" - Just value -> TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value + in + case headColumn @a column of + Nothing -> error "Invalid operation" + Just h -> case ifoldlColumn (\acc _ v -> f acc v) h column of + Nothing -> error "Invalid operation" + Just value -> TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value interpret _ expr = error ("Invalid operation for dataframe: " ++ show expr) -interpretAggregation :: forall a . (Columnable a) => GroupedDataFrame -> Expr a -> TypedColumn a +interpretAggregation :: forall a. (Columnable a) => GroupedDataFrame -> Expr a -> TypedColumn a interpretAggregation gdf (Lit value) = TColumn $ fromVector $ V.replicate (VG.length (offsets gdf) - 1) value interpretAggregation gdf@(Grouped df names indices os) (Col name) = case getColumn name df of Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df) Just col -> TColumn $ atIndicesStable (VG.map (indices `VG.unsafeIndex`) (VG.init os)) col -interpretAggregation gdf (Apply _ (f :: c -> d) expr) = let +interpretAggregation gdf (Apply _ (f :: c -> d) expr) = + let (TColumn value) = interpretAggregation @c gdf expr - in case mapColumn f value of - Nothing -> error "Type error in interpretation" - Just col -> TColumn col -interpretAggregation gdf (BinOp _ (f :: c -> d -> e) left right) = let + in + case mapColumn f value of + Nothing -> error "Type error in interpretation" + Just col -> TColumn col +interpretAggregation gdf (BinOp _ (f :: c -> d -> e) left right) = + let (TColumn left') = interpretAggregation @c gdf left (TColumn right') = interpretAggregation @d gdf right - in case zipWithColumns f left' right' of - Nothing -> error "Type error in binary operation" - Just col -> TColumn col + in + case zipWithColumns f left' right' of + Nothing -> error "Type error in binary operation" + Just col -> TColumn col interpretAggregation gdf@(Grouped df names indices os) (GeneralAggregate name op (f :: forall v b. (VG.Vector v b, Columnable b) => v b -> c)) = case getColumn name df of Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df) - Just (BoxedColumn col) -> TColumn $ fromVector $ - V.generate (VG.length os - 1) - (\i -> f (V.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) - (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) - ) - ) + Just (BoxedColumn col) -> + TColumn $ + fromVector $ + V.generate + (VG.length os - 1) + ( \i -> + f + ( V.generate + (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) + (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) + ) + ) Just (UnboxedColumn col) -> case sUnbox @c of - SFalse -> TColumn $ fromVector $ - V.generate (VG.length os - 1) - (\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) - (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) - ) - ) - STrue -> TColumn $ fromUnboxedVector $ - VU.generate (VG.length os - 1) - (\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) - (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) - ) - ) - Just (OptionalColumn col) -> TColumn $ fromVector $ - V.generate (VG.length os - 1) - (\i -> f (V.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) - (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) - ) - ) -interpretAggregation gdf@(Grouped df names indices os) (ReductionAggregate name op (f :: forall a . Columnable a => a -> a -> a)) = case getColumn name df of + SFalse -> + TColumn $ + fromVector $ + V.generate + (VG.length os - 1) + ( \i -> + f + ( VU.generate + (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) + (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) + ) + ) + STrue -> + TColumn $ + fromUnboxedVector $ + VU.generate + (VG.length os - 1) + ( \i -> + f + ( VU.generate + (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) + (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) + ) + ) + Just (OptionalColumn col) -> + TColumn $ + fromVector $ + V.generate + (VG.length os - 1) + ( \i -> + f + ( V.generate + (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) + (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) + ) + ) +interpretAggregation gdf@(Grouped df names indices os) (ReductionAggregate name op (f :: forall a. (Columnable a) => a -> a -> a)) = case getColumn name df of Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df) - Just (BoxedColumn col) -> TColumn $ fromVector $ - VG.generate (VG.length os - 1) $ \g -> - let !start = os `VG.unsafeIndex` g - !end = os `VG.unsafeIndex` (g+1) - in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end - where - {-# INLINE go #-} - go !acc j e - | j == e = acc - | otherwise = - let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j) - in go (f acc x) (j + 1) e + Just (BoxedColumn col) -> TColumn $ + fromVector $ + VG.generate (VG.length os - 1) $ \g -> + let !start = os `VG.unsafeIndex` g + !end = os `VG.unsafeIndex` (g + 1) + in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end + where + {-# INLINE go #-} + go !acc j e + | j == e = acc + | otherwise = + let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j) + in go (f acc x) (j + 1) e Just (UnboxedColumn col) -> case sUnbox @a of - SFalse -> TColumn $ fromVector $ - VG.generate (VG.length os - 1) $ \g -> - let !start = os `VG.unsafeIndex` g - !end = os `VG.unsafeIndex` (g+1) - in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end - where - {-# INLINE go #-} - go !acc j e - | j == e = acc - | otherwise = - let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j) - in go (f acc x) (j + 1) e - STrue -> TColumn $ fromVector $ - VG.generate (VG.length os - 1) $ \g -> - let !start = os `VG.unsafeIndex` g - !end = os `VG.unsafeIndex` (g+1) - in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end - where - {-# INLINE go #-} - go !acc j e - | j == e = acc - | otherwise = - let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j) - in go (f acc x) (j + 1) e - Just (OptionalColumn col) -> TColumn $ fromVector $ - VG.generate (VG.length os - 1) $ \g -> - let !start = os `VG.unsafeIndex` g - !end = os `VG.unsafeIndex` (g+1) - in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end - where - {-# INLINE go #-} - go !acc j e - | j == e = acc - | otherwise = - let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j) - in go (f acc x) (j + 1) e + SFalse -> TColumn $ + fromVector $ + VG.generate (VG.length os - 1) $ \g -> + let !start = os `VG.unsafeIndex` g + !end = os `VG.unsafeIndex` (g + 1) + in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end + where + {-# INLINE go #-} + go !acc j e + | j == e = acc + | otherwise = + let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j) + in go (f acc x) (j + 1) e + STrue -> TColumn $ + fromVector $ + VG.generate (VG.length os - 1) $ \g -> + let !start = os `VG.unsafeIndex` g + !end = os `VG.unsafeIndex` (g + 1) + in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end + where + {-# INLINE go #-} + go !acc j e + | j == e = acc + | otherwise = + let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j) + in go (f acc x) (j + 1) e + Just (OptionalColumn col) -> TColumn $ + fromVector $ + VG.generate (VG.length os - 1) $ \g -> + let !start = os `VG.unsafeIndex` g + !end = os `VG.unsafeIndex` (g + 1) + in go (col `VG.unsafeIndex` (indices `VG.unsafeIndex` start)) (start + 1) end + where + {-# INLINE go #-} + go !acc j e + | j == e = acc + | otherwise = + let !x = col `VG.unsafeIndex` (indices `VG.unsafeIndex` j) + in go (f acc x) (j + 1) e interpretAggregation gdf@(Grouped df names indices os) (NumericAggregate name op (f :: VU.Vector b -> c)) = case getColumn name df of Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df) Just (UnboxedColumn (col :: VU.Vector d)) -> case testEquality (typeRep @b) (typeRep @d) of - Nothing -> case testEquality (typeRep @d) (typeRep @Int) of + Nothing -> case testEquality (typeRep @d) (typeRep @Int) of Just Refl -> case sUnbox @c of - SFalse -> TColumn $ fromVector $ - V.generate (VG.length os - 1) - (\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) - (\j -> fromIntegral (col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i))))) - ) - ) - STrue -> TColumn $ fromUnboxedVector $ - VU.generate (VG.length os - 1) - (\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) - (\j -> fromIntegral (col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i))))) - ) - ) + SFalse -> + TColumn $ + fromVector $ + V.generate + (VG.length os - 1) + ( \i -> + f + ( VU.generate + (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) + (\j -> fromIntegral (col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i))))) + ) + ) + STrue -> + TColumn $ + fromUnboxedVector $ + VU.generate + (VG.length os - 1) + ( \i -> + f + ( VU.generate + (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) + (\j -> fromIntegral (col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i))))) + ) + ) Just Refl -> case sNumeric @d of SFalse -> error $ "Cannot apply numeric aggregation to non-numeric column: " ++ (T.unpack name) - STrue -> case sUnbox @c of - SFalse -> TColumn $ fromVector $ - V.generate (VG.length os - 1) - (\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) - (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) - ) + STrue -> case sUnbox @c of + SFalse -> + TColumn $ + fromVector $ + V.generate + (VG.length os - 1) + ( \i -> + f + ( VU.generate + (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) + (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) ) - STrue -> TColumn $ fromUnboxedVector $ - VU.generate (VG.length os - 1) - (\i -> f (VU.generate (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) - (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) - ) + ) + STrue -> + TColumn $ + fromUnboxedVector $ + VU.generate + (VG.length os - 1) + ( \i -> + f + ( VU.generate + (os `VG.unsafeIndex` (i + 1) - (os `VG.unsafeIndex` i)) + (\j -> col `VG.unsafeIndex` (indices `VG.unsafeIndex` (j + (os `VG.unsafeIndex` i)))) ) + ) _ -> error $ "Cannot apply numeric aggregation to non-numeric column: " ++ (T.unpack name) instance (Num a, Columnable a) => Num (Expr a) where @@ -226,15 +307,15 @@ (*) = BinOp "mult" (*) fromInteger :: Integer -> Expr a - fromInteger = Lit . fromInteger + fromInteger = Lit . fromInteger negate :: Expr a -> Expr a negate = Apply "negate" negate - abs :: Num a => Expr a -> Expr a + abs :: (Num a) => Expr a -> Expr a abs = Apply "abs" abs - signum :: Num a => Expr a -> Expr a + signum :: (Num a) => Expr a -> Expr a signum = Apply "signum" signum instance (Fractional a, Columnable a) => Fractional (Expr a) where @@ -258,7 +339,7 @@ asin :: (Floating a, Columnable a) => Expr a -> Expr a asin = Apply "asin" asin acos :: (Floating a, Columnable a) => Expr a -> Expr a - acos = Apply "acos" acos + acos = Apply "acos" acos atan :: (Floating a, Columnable a) => Expr a -> Expr a atan = Apply "atan" atan sinh :: (Floating a, Columnable a) => Expr a -> Expr a @@ -272,10 +353,9 @@ atanh :: (Floating a, Columnable a) => Expr a -> Expr a atanh = Apply "atanh" atanh - instance (Show a) => Show (Expr a) where - show :: forall a . Show a => Expr a -> String + show :: forall a. (Show a) => Expr a -> String show (Col name) = "col@" ++ show (typeRep @a) ++ "(" ++ T.unpack name ++ ")" show (Lit value) = show value show (Apply name f value) = T.unpack name ++ "(" ++ show value ++ ")" - show (BinOp name f a b) = T.unpack name ++ "(" ++ show a ++ ", " ++ show b ++ ")" + show (BinOp name f a b) = T.unpack name ++ "(" ++ show a ++ ", " ++ show b ++ ")" diff --git a/src/DataFrame/Internal/Parsing.hs b/src/DataFrame/Internal/Parsing.hs --- a/src/DataFrame/Internal/Parsing.hs +++ b/src/DataFrame/Internal/Parsing.hs @@ -1,12 +1,14 @@ {-# LANGUAGE OverloadedStrings #-} + module DataFrame.Internal.Parsing where import qualified Data.ByteString.Char8 as C import qualified Data.Set as S import qualified Data.Text as T -import Data.Text.Read +import Data.ByteString.Lex.Fractional import Data.Maybe (fromMaybe) +import Data.Text.Read import GHC.Stack (HasCallStack) import Text.Read (readMaybe) @@ -15,57 +17,64 @@ readValue :: (HasCallStack, Read a) => T.Text -> a readValue s = case readMaybe (T.unpack s) of - Nothing -> error $ "Could not read value: " ++ T.unpack s - Just value -> value + Nothing -> error $ "Could not read value: " ++ T.unpack s + Just value -> value readInteger :: (HasCallStack) => T.Text -> Maybe Integer readInteger s = case signed decimal (T.strip s) of - Left _ -> Nothing - Right (value, "") -> Just value - Right (value, _) -> Nothing + Left _ -> Nothing + Right (value, "") -> Just value + Right (value, _) -> Nothing readInt :: (HasCallStack) => T.Text -> Maybe Int readInt s = case signed decimal (T.strip s) of - Left _ -> Nothing - Right (value, "") -> Just value - Right (value, _) -> Nothing + Left _ -> Nothing + Right (value, "") -> Just value + Right (value, _) -> Nothing {-# INLINE readInt #-} readByteStringInt :: (HasCallStack) => C.ByteString -> Maybe Int readByteStringInt s = case C.readInt (C.strip s) of - Nothing -> Nothing - Just (value, "") -> Just value - Just (value, _) -> Nothing + Nothing -> Nothing + Just (value, "") -> Just value + Just (value, _) -> Nothing {-# INLINE readByteStringInt #-} +readByteStringDouble :: (HasCallStack) => C.ByteString -> Maybe Double +readByteStringDouble s = case readSigned readDecimal (C.strip s) of + Nothing -> Nothing + Just (value, "") -> Just value + Just (value, _) -> Nothing +{-# INLINE readByteStringDouble #-} + readDouble :: (HasCallStack) => T.Text -> Maybe Double readDouble s = - case signed double s of - Left _ -> Nothing - Right (value, "") -> Just value - Right (value, _) -> Nothing + case signed double s of + Left _ -> Nothing + Right (value, "") -> Just value + Right (value, _) -> Nothing {-# INLINE readDouble #-} readIntegerEither :: (HasCallStack) => T.Text -> Either T.Text Integer readIntegerEither s = case signed decimal (T.strip s) of - Left _ -> Left s - Right (value, "") -> Right value - Right (value, _) -> Left s + Left _ -> Left s + Right (value, "") -> Right value + Right (value, _) -> Left s {-# INLINE readIntegerEither #-} readIntEither :: (HasCallStack) => T.Text -> Either T.Text Int readIntEither s = case signed decimal (T.strip s) of - Left _ -> Left s - Right (value, "") -> Right value - Right (value, _) -> Left s + Left _ -> Left s + Right (value, "") -> Right value + Right (value, _) -> Left s {-# INLINE readIntEither #-} readDoubleEither :: (HasCallStack) => T.Text -> Either T.Text Double readDoubleEither s = - case signed double s of - Left _ -> Left s - Right (value, "") -> Right value - Right (value, _) -> Left s + case signed double s of + Left _ -> Left s + Right (value, "") -> Right value + Right (value, _) -> Left s {-# INLINE readDoubleEither #-} safeReadValue :: (Read a) => T.Text -> Maybe a diff --git a/src/DataFrame/Internal/Row.hs b/src/DataFrame/Internal/Row.hs --- a/src/DataFrame/Internal/Row.hs +++ b/src/DataFrame/Internal/Row.hs @@ -1,10 +1,11 @@ -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ExistentialQuantification #-} {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} -{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + module DataFrame.Internal.Row where import qualified Data.List as L @@ -12,25 +13,24 @@ import qualified Data.Set as S import qualified Data.Text as T import qualified Data.Vector as V +import qualified Data.Vector.Algorithms.Merge as VA import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Algorithms.Merge as VA import Control.Exception (throw) import Control.Monad.ST (runST) -import DataFrame.Errors (DataFrameException(..)) -import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame -import DataFrame.Internal.Types import Data.Function (on) import Data.Maybe (fromMaybe) +import Data.Type.Equality (TestEquality (..)) import Data.Typeable (Typeable, type (:~:) (..)) -import Data.Word ( Word8, Word16, Word32, Word64 ) -import Type.Reflection (TypeRep, typeOf, typeRep) -import Data.Type.Equality (TestEquality(..)) -import Text.ParserCombinators.ReadPrec(ReadPrec) +import Data.Word (Word16, Word32, Word64, Word8) +import DataFrame.Errors (DataFrameException (..)) +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame +import DataFrame.Internal.Types +import Text.ParserCombinators.ReadPrec (ReadPrec) import Text.Read (Lexeme (Ident), lexP, parens, readListPrec, readListPrecDefault, readPrec) - +import Type.Reflection (TypeRep, typeOf, typeRep) data Any where Value :: (Columnable' a) => a -> Any @@ -51,71 +51,76 @@ show :: Any -> String show (Value a) = T.unpack (showValue a) -showValue :: forall a . (Columnable' a) => a -> T.Text +showValue :: forall a. (Columnable' a) => a -> T.Text showValue v = case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> v - Nothing -> case testEquality (typeRep @a) (typeRep @String) of - Just Refl -> T.pack v - Nothing -> (T.pack . show) v + Just Refl -> v + Nothing -> case testEquality (typeRep @a) (typeRep @String) of + Just Refl -> T.pack v + Nothing -> (T.pack . show) v instance Read Any where - readListPrec :: ReadPrec [Any] - readListPrec = readListPrecDefault - - readPrec :: ReadPrec Any - readPrec = parens $ do - Ident "Value" <- lexP - readPrec + readListPrec :: ReadPrec [Any] + readListPrec = readListPrecDefault + readPrec :: ReadPrec Any + readPrec = parens $ do + Ident "Value" <- lexP + readPrec -toAny :: forall a . (Columnable' a) => a -> Any -toAny = Value +toAny :: forall a. (Columnable' a) => a -> Any +toAny = Value type Row = V.Vector Any toRowList :: [T.Text] -> DataFrame -> [Row] -toRowList names df = let - nameSet = S.fromList names - in map (mkRowRep df nameSet) [0..(fst (dataframeDimensions df) - 1)] +toRowList names df = + let + nameSet = S.fromList names + in + map (mkRowRep df nameSet) [0 .. (fst (dataframeDimensions df) - 1)] toRowVector :: [T.Text] -> DataFrame -> V.Vector Row -toRowVector names df = let - nameSet = S.fromList names - in V.generate (fst (dataframeDimensions df)) (mkRowRep df nameSet) +toRowVector names df = + let + nameSet = S.fromList names + in + V.generate (fst (dataframeDimensions df)) (mkRowRep df nameSet) mkRowFromArgs :: [T.Text] -> DataFrame -> Int -> Row mkRowFromArgs names df i = V.map get (V.fromList names) where get name = case getColumn name df of - Nothing -> throw $ ColumnNotFoundException name "[INTERNAL] mkRowFromArgs" (map fst $ M.toList $ columnIndices df) - Just (BoxedColumn column) -> toAny (column V.! i) - Just (UnboxedColumn column) -> toAny (column VU.! i) - Just (OptionalColumn column) -> toAny (column V.! i) + Nothing -> throw $ ColumnNotFoundException name "[INTERNAL] mkRowFromArgs" (map fst $ M.toList $ columnIndices df) + Just (BoxedColumn column) -> toAny (column V.! i) + Just (UnboxedColumn column) -> toAny (column VU.! i) + Just (OptionalColumn column) -> toAny (column V.! i) mkRowRep :: DataFrame -> S.Set T.Text -> Int -> Row mkRowRep df names i = V.generate (S.size names) (\index -> get (names' V.! index)) where inOrderIndexes = map fst $ L.sortBy (compare `on` snd) $ M.toList (columnIndices df) names' = V.fromList [n | n <- inOrderIndexes, S.member n names] - throwError name = error $ "Column " + throwError name = + error $ + "Column " ++ T.unpack name ++ " has less items than " ++ "the other columns at index " ++ show i get name = case getColumn name df of - Just (BoxedColumn c) -> case c V.!? i of - Just e -> toAny e - Nothing -> throwError name - Just (OptionalColumn c) -> case c V.!? i of - Just e -> toAny e - Nothing -> throwError name - Just (UnboxedColumn c) -> case c VU.!? i of - Just e -> toAny e - Nothing -> throwError name + Just (BoxedColumn c) -> case c V.!? i of + Just e -> toAny e + Nothing -> throwError name + Just (OptionalColumn c) -> case c V.!? i of + Just e -> toAny e + Nothing -> throwError name + Just (UnboxedColumn c) -> case c VU.!? i of + Just e -> toAny e + Nothing -> throwError name sortedIndexes' :: Bool -> V.Vector Row -> VU.Vector Int sortedIndexes' asc rows = runST $ do - withIndexes <- VG.thaw (V.indexed rows) - VA.sortBy ((if asc then compare else flip compare) `on` snd) withIndexes - sorted <- VG.unsafeFreeze withIndexes - return $ VU.generate (VG.length rows) (\i -> fst (sorted VG.! i)) + withIndexes <- VG.thaw (V.indexed rows) + VA.sortBy ((if asc then compare else flip compare) `on` snd) withIndexes + sorted <- VG.unsafeFreeze withIndexes + return $ VU.generate (VG.length rows) (\i -> fst (sorted VG.! i)) diff --git a/src/DataFrame/Internal/Schema.hs b/src/DataFrame/Internal/Schema.hs --- a/src/DataFrame/Internal/Schema.hs +++ b/src/DataFrame/Internal/Schema.hs @@ -1,23 +1,24 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} + module DataFrame.Internal.Schema where import qualified Data.Map as M import qualified Data.Proxy as P import qualified Data.Text as T -import DataFrame.Internal.Column import Data.Maybe -import Data.Type.Equality (type (:~:)(Refl), TestEquality (..)) +import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) +import DataFrame.Internal.Column import Type.Reflection (typeRep) data SchemaType where - SType :: Columnable a => P.Proxy a -> SchemaType + SType :: (Columnable a) => P.Proxy a -> SchemaType instance Show SchemaType where show (SType (_ :: P.Proxy a)) = show (typeRep @a) @@ -25,9 +26,10 @@ instance Eq SchemaType where (==) (SType (_ :: P.Proxy a)) (SType (_ :: P.Proxy b)) = isJust (testEquality (typeRep @a) (typeRep @b)) -schemaType :: forall a . Columnable a => SchemaType +schemaType :: forall a. (Columnable a) => SchemaType schemaType = SType (P.Proxy @a) -data Schema = Schema { - elements :: M.Map T.Text SchemaType -} deriving (Show, Eq) +data Schema = Schema + { elements :: M.Map T.Text SchemaType + } + deriving (Show, Eq) diff --git a/src/DataFrame/Internal/Types.hs b/src/DataFrame/Internal/Types.hs --- a/src/DataFrame/Internal/Types.hs +++ b/src/DataFrame/Internal/Types.hs @@ -1,101 +1,103 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} +{-# LANGUAGE PolyKinds #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE UndecidableInstances #-} + module DataFrame.Internal.Types where -import Data.Int ( Int8, Int16, Int32, Int64 ) -import Data.Kind (Type, Constraint) +import Data.Int (Int16, Int32, Int64, Int8) +import Data.Kind (Constraint, Type) import Data.Maybe (fromMaybe) +import Data.Type.Equality (TestEquality (..)) import Data.Typeable (Typeable, type (:~:) (..)) -import Data.Word ( Word8, Word16, Word32, Word64 ) -import Type.Reflection (TypeRep, typeOf, typeRep) -import Data.Type.Equality (TestEquality(..)) import qualified Data.Vector as VB import qualified Data.Vector.Unboxed as VU +import Data.Word (Word16, Word32, Word64, Word8) +import Type.Reflection (TypeRep, typeOf, typeRep) type Columnable' a = (Typeable a, Show a, Ord a, Eq a, Read a) --- | A type with column representations used to select the --- "right" representation when specializing the `toColumn` function. +{- | A type with column representations used to select the +"right" representation when specializing the `toColumn` function. +-} data Rep - = RBoxed - | RUnboxed - | ROptional + = RBoxed + | RUnboxed + | ROptional -- | Type-level if statement. type family If (cond :: Bool) (yes :: k) (no :: k) :: k where - If 'True yes _ = yes - If 'False _ no = no + If 'True yes _ = yes + If 'False _ no = no -- | All unboxable types (according to the `vector` package). type family Unboxable (a :: Type) :: Bool where - Unboxable Int = 'True - Unboxable Int8 = 'True - Unboxable Int16 = 'True - Unboxable Int32 = 'True - Unboxable Int64 = 'True - Unboxable Word = 'True - Unboxable Word8 = 'True - Unboxable Word16 = 'True - Unboxable Word32 = 'True - Unboxable Word64 = 'True - Unboxable Char = 'True - Unboxable Bool = 'True - Unboxable Double = 'True - Unboxable Float = 'True - Unboxable _ = 'False + Unboxable Int = 'True + Unboxable Int8 = 'True + Unboxable Int16 = 'True + Unboxable Int32 = 'True + Unboxable Int64 = 'True + Unboxable Word = 'True + Unboxable Word8 = 'True + Unboxable Word16 = 'True + Unboxable Word32 = 'True + Unboxable Word64 = 'True + Unboxable Char = 'True + Unboxable Bool = 'True + Unboxable Double = 'True + Unboxable Float = 'True + Unboxable _ = 'False -- | All unboxable types (according to the `vector` package). type family Numeric (a :: Type) :: Bool where - Numeric Int = 'True - Numeric Int8 = 'True - Numeric Int16 = 'True - Numeric Int32 = 'True - Numeric Int64 = 'True - Numeric Word = 'True - Numeric Word8 = 'True - Numeric Word16 = 'True - Numeric Word32 = 'True - Numeric Word64 = 'True - Numeric Double = 'True - Numeric Float = 'True - Numeric _ = 'False + Numeric Int = 'True + Numeric Int8 = 'True + Numeric Int16 = 'True + Numeric Int32 = 'True + Numeric Int64 = 'True + Numeric Word = 'True + Numeric Word8 = 'True + Numeric Word16 = 'True + Numeric Word32 = 'True + Numeric Word64 = 'True + Numeric Double = 'True + Numeric Float = 'True + Numeric _ = 'False -- | Compute the column representation tag for any ‘a’. type family KindOf a :: Rep where - KindOf (Maybe a) = 'ROptional - KindOf a = If (Unboxable a) 'RUnboxed 'RBoxed + KindOf (Maybe a) = 'ROptional + KindOf a = If (Unboxable a) 'RUnboxed 'RBoxed -- | Type-level boolean for constraint/type comparison. data SBool (b :: Bool) where - STrue :: SBool 'True - SFalse :: SBool 'False + STrue :: SBool 'True + SFalse :: SBool 'False -- | The runtime witness for our type-level branching. class SBoolI (b :: Bool) where - sbool :: SBool b + sbool :: SBool b -instance SBoolI 'True where sbool = STrue +instance SBoolI 'True where sbool = STrue instance SBoolI 'False where sbool = SFalse -- | Type-level function to determine whether or not a type is unboxa -sUnbox :: forall a. SBoolI (Unboxable a) => SBool (Unboxable a) +sUnbox :: forall a. (SBoolI (Unboxable a)) => SBool (Unboxable a) sUnbox = sbool @(Unboxable a) -sNumeric :: forall a. SBoolI (Numeric a) => SBool (Numeric a) +sNumeric :: forall a. (SBoolI (Numeric a)) => SBool (Numeric a) sNumeric = sbool @(Numeric a) type family When (flag :: Bool) (c :: Constraint) :: Constraint where - When 'True c = c - When 'False c = () -- empty constraint + When 'True c = c + When 'False c = () -- empty constraint type UnboxIf a = When (Unboxable a) (VU.Unbox a) diff --git a/src/DataFrame/Lazy/IO/CSV.hs b/src/DataFrame/Lazy/IO/CSV.hs --- a/src/DataFrame/Lazy/IO/CSV.hs +++ b/src/DataFrame/Lazy/IO/CSV.hs @@ -1,10 +1,11 @@ {-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE GADTs #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE RankNTypes #-} + module DataFrame.Lazy.IO.CSV where import qualified Data.ByteString.Char8 as C @@ -12,62 +13,65 @@ import qualified Data.Map as M import qualified Data.Set as S import qualified Data.Text as T +import qualified Data.Text.IO as TIO import qualified Data.Text.Lazy as TL import qualified Data.Text.Lazy.IO as TLIO -import qualified Data.Text.IO as TIO import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Mutable as VM +import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM -import Control.Applicative ((<$>), (<|>), (<*>), (<*), (*>), many) -import Control.Monad (forM_, zipWithM_, unless, when, void, replicateM_) +import Control.Applicative (many, (*>), (<$>), (<*), (<*>), (<|>)) +import Control.Monad (forM_, replicateM_, unless, void, when, zipWithM_) import Data.Attoparsec.Text import Data.Char -import DataFrame.Internal.Column (Column(..), MutableColumn(..), freezeColumn', writeColumn, columnLength) -import DataFrame.Internal.DataFrame (DataFrame(..)) -import DataFrame.Internal.Parsing -import DataFrame.Operations.Typing import Data.Foldable (fold) import Data.Function (on) import Data.IORef import Data.Maybe import Data.Text.Encoding (decodeUtf8Lenient) -import Data.Type.Equality - ( TestEquality (testEquality), - type (:~:) (Refl) - ) +import Data.Type.Equality ( + TestEquality (testEquality), + type (:~:) (Refl), + ) +import DataFrame.Internal.Column (Column (..), MutableColumn (..), columnLength, freezeColumn', writeColumn) +import DataFrame.Internal.DataFrame (DataFrame (..)) +import DataFrame.Internal.Parsing +import DataFrame.Operations.Typing import GHC.IO.Handle (Handle) -import Prelude hiding (concat, takeWhile) import System.IO import Type.Reflection +import Prelude hiding (concat, takeWhile) -- | Record for CSV read options. -data ReadOptions = ReadOptions { - hasHeader :: Bool, - inferTypes :: Bool, - safeRead :: Bool, - rowRange :: !(Maybe (Int, Int)), -- (start, length) - seekPos :: !(Maybe Integer), - totalRows :: !(Maybe Int), - leftOver :: !T.Text, - rowsRead :: !Int -} +data ReadOptions = ReadOptions + { hasHeader :: Bool + , inferTypes :: Bool + , safeRead :: Bool + , rowRange :: !(Maybe (Int, Int)) -- (start, length) + , seekPos :: !(Maybe Integer) + , totalRows :: !(Maybe Int) + , leftOver :: !T.Text + , rowsRead :: !Int + } --- | By default we assume the file has a header, we infer the types on read --- and we convert any rows with nullish objects into Maybe (safeRead). +{- | By default we assume the file has a header, we infer the types on read +and we convert any rows with nullish objects into Maybe (safeRead). +-} defaultOptions :: ReadOptions -defaultOptions = ReadOptions { hasHeader = True, inferTypes = True, safeRead = True, rowRange = Nothing, seekPos = Nothing, totalRows = Nothing, leftOver = "", rowsRead = 0 } +defaultOptions = ReadOptions{hasHeader = True, inferTypes = True, safeRead = True, rowRange = Nothing, seekPos = Nothing, totalRows = Nothing, leftOver = "", rowsRead = 0} --- | Reads a CSV file from the given path. --- Note this file stores intermediate temporary files --- while converting the CSV from a row to a columnar format. +{- | Reads a CSV file from the given path. +Note this file stores intermediate temporary files +while converting the CSV from a row to a columnar format. +-} readCsv :: String -> IO DataFrame readCsv path = fst <$> readSeparated ',' defaultOptions path --- | Reads a tab separated file from the given path. --- Note this file stores intermediate temporary files --- while converting the CSV from a row to a columnar format. +{- | Reads a tab separated file from the given path. +Note this file stores intermediate temporary files +while converting the CSV from a row to a columnar format. +-} readTsv :: String -> IO DataFrame readTsv path = fst <$> readSeparated '\t' defaultOptions path @@ -78,13 +82,14 @@ Nothing -> countRows c path >>= \total -> if hasHeader opts then return (total - 1) else return total Just n -> if hasHeader opts then return (n - 1) else return n let (begin, len) = case rowRange opts of - Nothing -> (0, totalRows) + Nothing -> (0, totalRows) Just (start, len) -> (start, min len (totalRows - rowsRead opts)) withFile path ReadMode $ \handle -> do firstRow <- map T.strip . parseSep c <$> TIO.hGetLine handle - let columnNames = if hasHeader opts - then map (T.filter (/= '\"')) firstRow - else map (T.singleton . intToDigit) [0..(length firstRow - 1)] + let columnNames = + if hasHeader opts + then map (T.filter (/= '\"')) firstRow + else map (T.singleton . intToDigit) [0 .. (length firstRow - 1)] -- If there was no header rewind the file cursor. unless (hasHeader opts) $ hSeek handle AbsoluteSeek 0 @@ -115,64 +120,80 @@ cols <- V.mapM (freezeColumn mutableCols nulls' opts) (V.generate numColumns id) pos <- hTell handle - return (DataFrame { - columns = cols, - columnIndices = M.fromList (zip columnNames [0..]), - dataframeDimensions = (maybe 0 columnLength (cols V.!? 0), V.length cols) - }, (pos, unconsumed, r + 1)) + return + ( DataFrame + { columns = cols + , columnIndices = M.fromList (zip columnNames [0 ..]) + , dataframeDimensions = (maybe 0 columnLength (cols V.!? 0), V.length cols) + } + , (pos, unconsumed, r + 1) + ) {-# INLINE readSeparated #-} getInitialDataVectors :: Int -> VM.IOVector MutableColumn -> [T.Text] -> IO () getInitialDataVectors n mCol xs = do - forM_ (zip [0..] xs) $ \(i, x) -> do + forM_ (zip [0 ..] xs) $ \(i, x) -> do col <- case inferValueType x of - "Int" -> MUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c) - "Double" -> MUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c) - _ -> MBoxedColumn <$> ((VM.unsafeNew n :: IO (VM.IOVector T.Text)) >>= \c -> VM.unsafeWrite c 0 x >> return c) + "Int" -> MUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c) + "Double" -> MUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c) + _ -> MBoxedColumn <$> ((VM.unsafeNew n :: IO (VM.IOVector T.Text)) >>= \c -> VM.unsafeWrite c 0 x >> return c) VM.unsafeWrite mCol i col {-# INLINE getInitialDataVectors #-} inferValueType :: T.Text -> T.Text -inferValueType s = let +inferValueType s = + let example = s - in case readInt example of - Just _ -> "Int" - Nothing -> case readDouble example of - Just _ -> "Double" - Nothing -> "Other" + in + case readInt example of + Just _ -> "Int" + Nothing -> case readDouble example of + Just _ -> "Double" + Nothing -> "Other" {-# INLINE inferValueType #-} readSingleLine :: Char -> T.Text -> Handle -> IO ([T.Text], T.Text) -readSingleLine c unused handle = parseWith (TIO.hGetChunk handle) (parseRow c) unused >>= \case - Fail unconsumed ctx er -> do - erpos <- hTell handle - fail $ "Failed to parse CSV file around " <> show erpos <> " byte; due: " - <> show er <> "; context: " <> show ctx - Partial c -> do - fail "Partial handler is called" - Done (unconsumed :: T.Text) (row :: [T.Text]) -> do - return (row, unconsumed) +readSingleLine c unused handle = + parseWith (TIO.hGetChunk handle) (parseRow c) unused >>= \case + Fail unconsumed ctx er -> do + erpos <- hTell handle + fail $ + "Failed to parse CSV file around " + <> show erpos + <> " byte; due: " + <> show er + <> "; context: " + <> show ctx + Partial c -> do + fail "Partial handler is called" + Done (unconsumed :: T.Text) (row :: [T.Text]) -> do + return (row, unconsumed) -- | Reads rows from the handle and stores values in mutable vectors. fillColumns :: Int -> Char -> VM.IOVector MutableColumn -> VM.IOVector [(Int, T.Text)] -> T.Text -> Handle -> IO (T.Text, Int) fillColumns n c mutableCols nullIndices unused handle = do input <- newIORef unused rowsRead <- newIORef (0 :: Int) - forM_ [1..(n - 1)] $ \i -> do + forM_ [1 .. (n - 1)] $ \i -> do isEOF <- hIsEOF handle input' <- readIORef input unless (isEOF && input' == mempty) $ do - parseWith (TIO.hGetChunk handle) (parseRow c) input' >>= \case + parseWith (TIO.hGetChunk handle) (parseRow c) input' >>= \case Fail unconsumed ctx er -> do - erpos <- hTell handle - fail $ "Failed to parse CSV file around " <> show erpos <> " byte; due: " - <> show er <> "; context: " <> show ctx + erpos <- hTell handle + fail $ + "Failed to parse CSV file around " + <> show erpos + <> " byte; due: " + <> show er + <> "; context: " + <> show ctx Partial c -> do - fail "Partial handler is called" + fail "Partial handler is called" Done (unconsumed :: T.Text) (row :: [T.Text]) -> do - writeIORef input unconsumed - modifyIORef rowsRead (+1) - zipWithM_ (writeValue mutableCols nullIndices i) [0..] row + writeIORef input unconsumed + modifyIORef rowsRead (+ 1) + zipWithM_ (writeValue mutableCols nullIndices i) [0 ..] row l <- readIORef input r <- readIORef rowsRead pure (l, r) @@ -200,17 +221,17 @@ record :: Char -> Parser [T.Text] record c = - field c `sepBy1` char c - "record" + field c `sepBy1` char c + "record" {-# INLINE record #-} parseRow :: Char -> Parser [T.Text] -parseRow c = (record c <* lineEnd) "record-new-line" +parseRow c = (record c <* lineEnd) "record-new-line" field :: Char -> Parser T.Text field c = - quotedField <|> unquotedField c - "field" + quotedField <|> unquotedField c + "field" {-# INLINE field #-} unquotedTerminators :: Char -> S.Set Char @@ -218,59 +239,71 @@ unquotedField :: Char -> Parser T.Text unquotedField sep = - takeWhile (not . (`S.member` terminators)) "unquoted field" - where terminators = unquotedTerminators sep + takeWhile (not . (`S.member` terminators)) "unquoted field" + where + terminators = unquotedTerminators sep {-# INLINE unquotedField #-} quotedField :: Parser T.Text quotedField = char '"' *> contents <* char '"' "quoted field" - where - contents = fold <$> many (unquote <|> unescape) - where - unquote = takeWhile1 (notInClass "\"\\") - unescape = char '\\' *> do - T.singleton <$> do - char '\\' <|> char '"' + where + contents = fold <$> many (unquote <|> unescape) + where + unquote = takeWhile1 (notInClass "\"\\") + unescape = + char '\\' *> do + T.singleton <$> do + char '\\' <|> char '"' {-# INLINE quotedField #-} lineEnd :: Parser () lineEnd = - (endOfLine <|> endOfInput) - "end of line" + (endOfLine <|> endOfInput) + "end of line" {-# INLINE lineEnd #-} -- | First pass to count rows for exact allocation countRows :: Char -> FilePath -> IO Int countRows c path = withFile path ReadMode $! go 0 "" - where - go n input h = do - isEOF <- hIsEOF h - if isEOF && input == mempty + where + go n input h = do + isEOF <- hIsEOF h + if isEOF && input == mempty then pure n else - parseWith (TIO.hGetChunk h) (parseRow c) input >>= \case - Fail unconsumed ctx er -> do - erpos <- hTell h - fail $ "Failed to parse CSV file around " <> show erpos <> " byte; due: " - <> show er <> "; context: " <> show ctx <> " " <> show unconsumed - Partial c -> do - fail $ "Partial handler is called; n = " <> show n - Done (unconsumed :: T.Text) _ -> - go (n + 1) unconsumed h + parseWith (TIO.hGetChunk h) (parseRow c) input >>= \case + Fail unconsumed ctx er -> do + erpos <- hTell h + fail $ + "Failed to parse CSV file around " + <> show erpos + <> " byte; due: " + <> show er + <> "; context: " + <> show ctx + <> " " + <> show unconsumed + Partial c -> do + fail $ "Partial handler is called; n = " <> show n + Done (unconsumed :: T.Text) _ -> + go (n + 1) unconsumed h {-# INLINE countRows #-} writeCsv :: String -> DataFrame -> IO () writeCsv = writeSeparated ',' -writeSeparated :: Char -- ^ Separator - -> String -- ^ Path to write to - -> DataFrame - -> IO () -writeSeparated c filepath df = withFile filepath WriteMode $ \handle ->do +writeSeparated :: + -- | Separator + Char -> + -- | Path to write to + String -> + DataFrame -> + IO () +writeSeparated c filepath df = withFile filepath WriteMode $ \handle -> do let (rows, columns) = dataframeDimensions df let headers = map fst (L.sortBy (compare `on` snd) (M.toList (columnIndices df))) TIO.hPutStrLn handle (T.intercalate ", " headers) - forM_ [0..(rows - 1)] $ \i -> do + forM_ [0 .. (rows - 1)] $ \i -> do let row = getRowAsText df i TIO.hPutStrLn handle (T.intercalate ", " row) @@ -280,43 +313,46 @@ indexMap = M.fromList (map (\(a, b) -> (b, a)) $ M.toList (columnIndices df)) go k (BoxedColumn (c :: V.Vector a)) acc = case c V.!? i of Just e -> textRep : acc - where textRep = case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> e - Nothing -> case typeRep @a of - App t1 t2 -> case eqTypeRep t1 (typeRep @Maybe) of - Just HRefl -> case testEquality t2 (typeRep @T.Text) of - Just Refl -> fromMaybe "null" e - Nothing -> (fromOptional . (T.pack . show)) e - where fromOptional s - | T.isPrefixOf "Just " s = T.drop (T.length "Just ") s - | otherwise = "null" - Nothing -> (T.pack . show) e - _ -> (T.pack . show) e + where + textRep = case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> e + Nothing -> case typeRep @a of + App t1 t2 -> case eqTypeRep t1 (typeRep @Maybe) of + Just HRefl -> case testEquality t2 (typeRep @T.Text) of + Just Refl -> fromMaybe "null" e + Nothing -> (fromOptional . (T.pack . show)) e + where + fromOptional s + | T.isPrefixOf "Just " s = T.drop (T.length "Just ") s + | otherwise = "null" + Nothing -> (T.pack . show) e + _ -> (T.pack . show) e Nothing -> error $ "Column " - ++ T.unpack (indexMap M.! k) - ++ " has less items than " - ++ "the other columns at index " - ++ show i + ++ T.unpack (indexMap M.! k) + ++ " has less items than " + ++ "the other columns at index " + ++ show i go k (UnboxedColumn c) acc = case c VU.!? i of Just e -> T.pack (show e) : acc Nothing -> error $ "Column " - ++ T.unpack (indexMap M.! k) - ++ " has less items than " - ++ "the other columns at index " - ++ show i + ++ T.unpack (indexMap M.! k) + ++ " has less items than " + ++ "the other columns at index " + ++ show i go k (OptionalColumn (c :: V.Vector (Maybe a))) acc = case c V.!? i of Just e -> textRep : acc - where textRep = case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> fromMaybe "Nothing" e - Nothing -> (T.pack . show) e + where + textRep = case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> fromMaybe "Nothing" e + Nothing -> (T.pack . show) e Nothing -> error $ "Column " - ++ T.unpack (indexMap M.! k) - ++ " has less items than " - ++ "the other columns at index " - ++ show i + ++ T.unpack (indexMap M.! k) + ++ " has less items than " + ++ "the other columns at index " + ++ show i diff --git a/src/DataFrame/Lazy/Internal/DataFrame.hs b/src/DataFrame/Lazy/Internal/DataFrame.hs --- a/src/DataFrame/Lazy/Internal/DataFrame.hs +++ b/src/DataFrame/Lazy/Internal/DataFrame.hs @@ -1,90 +1,103 @@ -{-# LANGUAGE GADTs #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ExistentialQuantification #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} -{-# LANGUAGE ExistentialQuantification #-} -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE NumericUnderscores #-} +{-# LANGUAGE OverloadedStrings #-} + module DataFrame.Lazy.Internal.DataFrame where -import Control.Monad (forM, foldM) -import Data.IORef -import Data.Kind +import Control.Monad (foldM, forM) +import Data.IORef +import Data.Kind import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T import qualified Data.Vector as V -import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Column as C +import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E +import qualified DataFrame.Lazy.IO.CSV as D import qualified DataFrame.Operations.Core as D -import DataFrame.Operations.Merge +import DataFrame.Operations.Merge import qualified DataFrame.Operations.Subset as D import qualified DataFrame.Operations.Transformations as D -import qualified DataFrame.Lazy.IO.CSV as D -import System.FilePath +import System.FilePath data LazyOperation where - Derive :: C.Columnable a => T.Text -> E.Expr a -> LazyOperation - Select :: [T.Text] -> LazyOperation - Filter :: E.Expr Bool -> LazyOperation + Derive :: (C.Columnable a) => T.Text -> E.Expr a -> LazyOperation + Select :: [T.Text] -> LazyOperation + Filter :: E.Expr Bool -> LazyOperation instance Show LazyOperation where - show :: LazyOperation -> String - show (Derive name expr) = T.unpack name ++ " := " ++ show expr - show (Select columns) = "select(" ++ show columns ++ ")" - show (Filter expr) = "filter(" ++ show expr ++ ")" + show :: LazyOperation -> String + show (Derive name expr) = T.unpack name ++ " := " ++ show expr + show (Select columns) = "select(" ++ show columns ++ ")" + show (Filter expr) = "filter(" ++ show expr ++ ")" -data InputType = ICSV deriving Show +data InputType = ICSV deriving (Show) data LazyDataFrame = LazyDataFrame - { inputPath :: FilePath - , inputType :: InputType - , operations :: [LazyOperation] - , batchSize :: Int - } deriving Show + { inputPath :: FilePath + , inputType :: InputType + , operations :: [LazyOperation] + , batchSize :: Int + } + deriving (Show) eval :: LazyOperation -> D.DataFrame -> D.DataFrame eval (Derive name expr) = D.derive name expr eval (Select columns) = D.select columns eval (Filter expr) = D.filterWhere expr -runDataFrame :: forall a . (C.Columnable a) => LazyDataFrame -> IO D.DataFrame +runDataFrame :: forall a. (C.Columnable a) => LazyDataFrame -> IO D.DataFrame runDataFrame df = do - let path = inputPath df - totalRows <- D.countRows ',' path - let batches = batchRanges totalRows (batchSize df) - (df', _) <- foldM (\(accDf, (pos, unused, r)) (start, end) -> do - mapM_ putStr ["Scanning: ", show start, " to ", show end, " rows out of ", show totalRows, "\n"] + let path = inputPath df + totalRows <- D.countRows ',' path + let batches = batchRanges totalRows (batchSize df) + (df', _) <- + foldM + ( \(accDf, (pos, unused, r)) (start, end) -> do + mapM_ putStr ["Scanning: ", show start, " to ", show end, " rows out of ", show totalRows, "\n"] - (sdf, (pos', unconsumed, rowsRead)) <- D.readSeparated ',' ( - D.defaultOptions { D.rowRange = Just (start, batchSize df) - , D.totalRows = Just totalRows - , D.seekPos = pos - , D.rowsRead = r - , D.leftOver = unused}) path - let rdf = L.foldl' (flip eval) sdf (operations df) - return (accDf <> rdf, (Just pos', unconsumed, rowsRead + r)) ) (D.empty, (Nothing, "", 0)) batches - return df' + (sdf, (pos', unconsumed, rowsRead)) <- + D.readSeparated + ',' + ( D.defaultOptions + { D.rowRange = Just (start, batchSize df) + , D.totalRows = Just totalRows + , D.seekPos = pos + , D.rowsRead = r + , D.leftOver = unused + } + ) + path + let rdf = L.foldl' (flip eval) sdf (operations df) + return (accDf <> rdf, (Just pos', unconsumed, rowsRead + r)) + ) + (D.empty, (Nothing, "", 0)) + batches + return df' batchRanges :: Int -> Int -> [(Int, Int)] -batchRanges n inc = go n [0,inc..n] +batchRanges n inc = go n [0, inc .. n] where - go _ [] = [] - go n [x] = [(x, n)] - go n (f:s:rest) =(f, s) : go n (s:rest) + go _ [] = [] + go n [x] = [(x, n)] + go n (f : s : rest) = (f, s) : go n (s : rest) scanCsv :: T.Text -> LazyDataFrame scanCsv path = LazyDataFrame (T.unpack path) ICSV [] 512_000 addOperation :: LazyOperation -> LazyDataFrame -> LazyDataFrame -addOperation op df = df { operations = operations df ++ [op] } +addOperation op df = df{operations = operations df ++ [op]} -derive :: C.Columnable a => T.Text -> E.Expr a -> LazyDataFrame -> LazyDataFrame +derive :: (C.Columnable a) => T.Text -> E.Expr a -> LazyDataFrame -> LazyDataFrame derive name expr = addOperation (Derive name expr) -select :: C.Columnable a => [T.Text] -> LazyDataFrame -> LazyDataFrame +select :: (C.Columnable a) => [T.Text] -> LazyDataFrame -> LazyDataFrame select columns = addOperation (Select columns) -filter :: C.Columnable a => E.Expr Bool -> LazyDataFrame -> LazyDataFrame +filter :: (C.Columnable a) => E.Expr Bool -> LazyDataFrame -> LazyDataFrame filter cond = addOperation (Filter cond) diff --git a/src/DataFrame/Operations/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs --- a/src/DataFrame/Operations/Aggregation.hs +++ b/src/DataFrame/Operations/Aggregation.hs @@ -1,10 +1,11 @@ {-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE FlexibleContexts #-} + module DataFrame.Operations.Aggregation where import qualified Data.Set as S @@ -14,114 +15,129 @@ import qualified Data.Map as M import qualified Data.Map.Strict as MS import qualified Data.Text as T -import qualified Data.Vector.Generic as VG import qualified Data.Vector as V +import qualified Data.Vector.Algorithms.Merge as VA +import qualified Data.Vector.Generic as VG import qualified Data.Vector.Mutable as VM import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Algorithms.Merge as VA -import qualified Statistics.Quantile as SS -import qualified Statistics.Sample as SS import Control.Exception (throw) import Control.Monad (foldM_) import Control.Monad.ST (runST) -import DataFrame.Internal.Column (Column(..), fromVector, - getIndicesUnboxed, getIndices, - atIndicesStable, - Columnable, unwrapTypedColumn, - columnVersionString) -import DataFrame.Internal.DataFrame (GroupedDataFrame(..), DataFrame(..), empty, getColumn, unsafeGetColumn) +import Data.Function ((&)) +import Data.Hashable +import Data.List ((\\)) +import Data.Maybe +import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) +import DataFrame.Errors +import DataFrame.Internal.Column ( + Column (..), + Columnable, + atIndicesStable, + columnVersionString, + fromVector, + getIndices, + getIndicesUnboxed, + unwrapTypedColumn, + ) +import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..), empty, getColumn, unsafeGetColumn) import DataFrame.Internal.Expression import DataFrame.Internal.Parsing import DataFrame.Internal.Types -import DataFrame.Errors import DataFrame.Operations.Core import DataFrame.Operations.Merge import DataFrame.Operations.Subset -import Data.Function ((&)) -import Data.Hashable -import Data.List ((\\)) -import Data.Maybe -import Data.Type.Equality (type (:~:)(Refl), TestEquality(..)) -import Type.Reflection (typeRep, typeOf) +import Type.Reflection (typeOf, typeRep) --- | O(k * n) groups the dataframe by the given rows aggregating the remaining rows --- into vector that should be reduced later. +{- | O(k * n) groups the dataframe by the given rows aggregating the remaining rows +into vector that should be reduced later. +-} groupBy :: - [T.Text] -> - DataFrame -> - GroupedDataFrame + [T.Text] -> + DataFrame -> + GroupedDataFrame groupBy names df - | any (`notElem` columnNames df) names = throw $ ColumnNotFoundException (T.pack $ show $ names L.\\ columnNames df) "groupBy" (columnNames df) - | otherwise = Grouped df names (VG.map fst valueIndices) (VU.fromList (reverse (changingPoints valueIndices))) + | any (`notElem` columnNames df) names = throw $ ColumnNotFoundException (T.pack $ show $ names L.\\ columnNames df) "groupBy" (columnNames df) + | otherwise = Grouped df names (VG.map fst valueIndices) (VU.fromList (reverse (changingPoints valueIndices))) where indicesToGroup = M.elems $ M.filterWithKey (\k _ -> k `elem` names) (columnIndices df) rowRepresentations = VU.generate (fst (dimensions df)) (mkRowRep indicesToGroup df) valueIndices = runST $ do - withIndexes <- VG.thaw $ VG.indexed rowRepresentations - VA.sortBy (\(a, b) (a', b') -> compare b b') withIndexes - VG.unsafeFreeze withIndexes + withIndexes <- VG.thaw $ VG.indexed rowRepresentations + VA.sortBy (\(a, b) (a', b') -> compare b b') withIndexes + VG.unsafeFreeze withIndexes changingPoints :: (Eq a, VU.Unbox a) => VU.Vector (Int, a) -> [Int] changingPoints vs = VG.length vs : (fst (VU.ifoldl findChangePoints initialState vs)) where initialState = ([0], snd (VG.head vs)) findChangePoints (offsets, currentVal) index (_, newVal) - | currentVal == newVal = (offsets, currentVal) - | otherwise = (index : offsets, newVal) + | currentVal == newVal = (offsets, currentVal) + | otherwise = (index : offsets, newVal) mkRowRep :: [Int] -> DataFrame -> Int -> Int mkRowRep groupColumnIndices df i = case h of - (x:[]) -> x - xs -> hash h + (x : []) -> x + xs -> hash h where h = (map mkHash groupColumnIndices) getHashedElem :: Column -> Int -> Int getHashedElem (BoxedColumn (c :: V.Vector a)) j = hash' @a (c V.! j) getHashedElem (UnboxedColumn (c :: VU.Vector a)) j = hash' @a (c VU.! j) getHashedElem (OptionalColumn (c :: V.Vector a)) j = hash' @a (c V.! j) - mkHash j = getHashedElem ((V.!) (columns df) j) i + mkHash j = getHashedElem ((V.!) (columns df) j) i --- | This hash function returns the hash when given a non numeric type but --- the value when given a numeric. -hash' :: Columnable a => a -> Int +{- | This hash function returns the hash when given a non numeric type but +the value when given a numeric. +-} +hash' :: (Columnable a) => a -> Int hash' value = case testEquality (typeOf value) (typeRep @Double) of - Just Refl -> round $ value * 1000 - Nothing -> case testEquality (typeOf value) (typeRep @Int) of - Just Refl -> value - Nothing -> case testEquality (typeOf value) (typeRep @T.Text) of - Just Refl -> hash value - Nothing -> hash (show value) + Just Refl -> round $ value * 1000 + Nothing -> case testEquality (typeOf value) (typeRep @Int) of + Just Refl -> value + Nothing -> case testEquality (typeOf value) (typeRep @T.Text) of + Just Refl -> hash value + Nothing -> hash (show value) mkGroupedColumns :: VU.Vector Int -> DataFrame -> DataFrame -> T.Text -> DataFrame mkGroupedColumns indices df acc name = - case (V.!) (columns df) (columnIndices df M.! name) of - BoxedColumn column -> - let vs = indices `getIndices` column - in insertVector name vs acc - OptionalColumn column -> - let vs = indices `getIndices` column - in insertVector name vs acc - UnboxedColumn column -> - let vs = indices `getIndicesUnboxed` column - in insertUnboxedVector name vs acc + case (V.!) (columns df) (columnIndices df M.! name) of + BoxedColumn column -> + let vs = indices `getIndices` column + in insertVector name vs acc + OptionalColumn column -> + let vs = indices `getIndices` column + in insertVector name vs acc + UnboxedColumn column -> + let vs = indices `getIndicesUnboxed` column + in insertUnboxedVector name vs acc +{- | Aggregate a grouped dataframe using the expressions give. +All ungrouped columns will be dropped. +-} aggregate :: [(T.Text, UExpr)] -> GroupedDataFrame -> DataFrame --- aggregate aggs df = undefined -aggregate aggs gdf@(Grouped df groupingColumns valueIndices offsets) = let - df' = selectIndices (VG.map (valueIndices VG.!) (VG.init offsets)) (select groupingColumns df) - groupedColumns = columnNames df L.\\ groupingColumns - f (name, Wrap (expr :: Expr a)) d = let - value = interpretAggregation @a gdf expr - in insertColumn name (unwrapTypedColumn value) d - in fold f aggs df' +aggregate aggs gdf@(Grouped df groupingColumns valueIndices offsets) = + let + df' = selectIndices (VG.map (valueIndices VG.!) (VG.init offsets)) (select groupingColumns df) + groupedColumns = columnNames df L.\\ groupingColumns + f (name, Wrap (expr :: Expr a)) d = + let + value = interpretAggregation @a gdf expr + in + insertColumn name (unwrapTypedColumn value) d + in + fold f aggs df' selectIndices :: VU.Vector Int -> DataFrame -> DataFrame -selectIndices xs df = df { columns = VG.map (atIndicesStable xs) (columns df) - , dataframeDimensions = (VG.length xs, VG.length (columns df)) - } +selectIndices xs df = + df + { columns = VG.map (atIndicesStable xs) (columns df) + , dataframeDimensions = (VG.length xs, VG.length (columns df)) + } +-- | Filter out all non-unique values in a dataframe. distinct :: DataFrame -> DataFrame distinct df = selectIndices (VG.map (indices VG.!) (VG.init os)) df - where (Grouped _ _ indices os) = groupBy (columnNames df) df + where + (Grouped _ _ indices os) = groupBy (columnNames df) df diff --git a/src/DataFrame/Operations/Core.hs b/src/DataFrame/Operations/Core.hs --- a/src/DataFrame/Operations/Core.hs +++ b/src/DataFrame/Operations/Core.hs @@ -5,6 +5,7 @@ {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} + module DataFrame.Operations.Core where import qualified Data.List as L @@ -12,225 +13,529 @@ import qualified Data.Map.Strict as MS import qualified Data.Set as S import qualified Data.Text as T -import qualified Data.Vector.Generic as VG import qualified Data.Vector as V +import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU -import Control.Exception ( throw ) -import DataFrame.Errors -import DataFrame.Internal.Column ( Column(..), fromVector, fromList, columnLength, columnTypeString, expandColumn, Columnable) -import DataFrame.Internal.DataFrame (DataFrame(..), getColumn, null, empty) -import DataFrame.Internal.Parsing (isNullish) +import Control.Exception (throw) import Data.Either import Data.Function (on, (&)) import Data.Maybe -import Data.Type.Equality (type (:~:)(Refl), TestEquality(..)) +import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) +import DataFrame.Errors +import DataFrame.Internal.Column (Column (..), Columnable, columnLength, columnTypeString, expandColumn, fromList, fromVector) +import DataFrame.Internal.DataFrame (DataFrame (..), empty, getColumn, null) +import DataFrame.Internal.Parsing (isNullish) import Type.Reflection import Prelude hiding (null) --- | O(1) Get DataFrame dimensions i.e. (rows, columns) +{- | O(1) Get DataFrame dimensions i.e. (rows, columns) + +==== __Example__ +@ +ghci> D.dimensions df + +(100, 3) +@ +-} dimensions :: DataFrame -> (Int, Int) dimensions = dataframeDimensions {-# INLINE dimensions #-} --- | O(k) Get column names of the DataFrame in order of insertion. +{- | O(k) Get column names of the DataFrame in order of insertion. + +==== __Example__ +@ +ghci> D.columnNames df + +["col_a", "col_b", "col_c"] +@ +-} columnNames :: DataFrame -> [T.Text] -columnNames = map fst . L.sortBy (compare `on` snd). M.toList . columnIndices +columnNames = map fst . L.sortBy (compare `on` snd) . M.toList . columnIndices {-# INLINE columnNames #-} --- | /O(n)/ Adds a vector to the dataframe. +{- | Adds a vector to the dataframe. If the vector has less elements than the dataframe and the dataframe is not empty +the vector is converted to type `Maybe a` filled with `Nothing` to match the size of the dataframe. Similarly, +if the vector has more elements than what's currently in the dataframe, the other columns in the dataframe are +change to `Maybe ` and filled with `Nothing`. + +==== __Example__ +@ +ghci> import qualified Data.Vector as V + +ghci> D.insertVector "numbers" (V.fromList [1..10]) D.empty + +--------------- +index | numbers +------|-------- + Int | Int +------|-------- +0 | 1 +1 | 2 +2 | 3 +3 | 4 +4 | 5 +5 | 6 +6 | 7 +7 | 8 +8 | 9 +9 | 10 + +@ +-} insertVector :: - forall a. - Columnable a => - -- | Column Name - T.Text -> - -- | Vector to add to column - V.Vector a -> - -- | DataFrame to add column to - DataFrame -> - DataFrame + forall a. + (Columnable a) => + -- | Column Name + T.Text -> + -- | Vector to add to column + V.Vector a -> + -- | DataFrame to add column to + DataFrame -> + DataFrame insertVector name xs = insertColumn name (fromVector xs) {-# INLINE insertVector #-} -cloneColumn :: T.Text -> T.Text -> DataFrame -> DataFrame -cloneColumn original new df = fromMaybe (throw $ ColumnNotFoundException original "cloneColumn" (map fst $ M.toList $ columnIndices df)) $ do - column <- getColumn original df - return $ insertColumn new column df +{- | /O(k)/ Add a column to the dataframe providing a default. +This constructs a new vector and also may convert it +to an unboxed vector if necessary. Since columns are usually +large the runtime is dominated by the length of the list, k. +-} +insertVectorWithDefault :: + forall a. + (Columnable a) => + -- | Default Value + a -> + -- | Column name + T.Text -> + -- | Data to add to column + V.Vector a -> + -- | DataFrame to add to column + DataFrame -> + DataFrame +insertVectorWithDefault defaultValue name xs d = + let (rows, _) = dataframeDimensions d + values = xs V.++ V.replicate (rows - V.length xs) defaultValue + in insertColumn name (fromVector values) d --- | /O(n)/ Adds an unboxed vector to the dataframe. +{- | /O(n)/ Adds an unboxed vector to the dataframe. + +Same as insertVector but takes an unboxed vector. If you insert a vector of numbers through insertVector it will either way be converted +into an unboxed vector so this function saves that extra work/conversion. +-} insertUnboxedVector :: - forall a. - (Columnable a, VU.Unbox a) => - -- | Column Name - T.Text -> - -- | Unboxed vector to add to column - VU.Vector a -> - -- | DataFrame to add to column - DataFrame -> - DataFrame + forall a. + (Columnable a, VU.Unbox a) => + -- | Column Name + T.Text -> + -- | Unboxed vector to add to column + VU.Vector a -> + -- | DataFrame to add to column + DataFrame -> + DataFrame insertUnboxedVector name xs = insertColumn name (UnboxedColumn xs) --- -- | /O(n)/ Add a column to the dataframe. Not meant for external use. +{- | /O(n)/ Add a column to the dataframe. + +==== __Example__ +@ +ghci> D.insertColumn "numbers" (D.fromList [1..10]) D.empty + +--------------- +index | numbers +------|-------- + Int | Int +------|-------- +0 | 1 +1 | 2 +2 | 3 +3 | 4 +4 | 5 +5 | 6 +6 | 7 +7 | 8 +8 | 9 +9 | 10 + +@ +-} insertColumn :: - -- | Column Name - T.Text -> - -- | Column to add - Column -> - -- | DataFrame to add to column - DataFrame -> - DataFrame -insertColumn name column d = let - (r, c) = dataframeDimensions d - n = max (columnLength column) r - in case M.lookup name (columnIndices d) of - Just i -> DataFrame (V.map (expandColumn n) (columns d V.// [(i, column)])) (columnIndices d) (n, c) - Nothing -> DataFrame (V.map (expandColumn n) (columns d `V.snoc` column)) (M.insert name c (columnIndices d)) (n, c + 1) + -- | Column Name + T.Text -> + -- | Column to add + Column -> + -- | DataFrame to add to column + DataFrame -> + DataFrame +insertColumn name column d = + let + (r, c) = dataframeDimensions d + n = max (columnLength column) r + in + case M.lookup name (columnIndices d) of + Just i -> DataFrame (V.map (expandColumn n) (columns d V.// [(i, column)])) (columnIndices d) (n, c) + Nothing -> DataFrame (V.map (expandColumn n) (columns d `V.snoc` column)) (M.insert name c (columnIndices d)) (n, c + 1) --- | /O(k)/ Add a column to the dataframe providing a default. --- This constructs a new vector and also may convert it --- to an unboxed vector if necessary. Since columns are usually --- large the runtime is dominated by the length of the list, k. -insertVectorWithDefault :: - forall a. - (Columnable a) => - -- | Default Value - a -> - -- | Column name - T.Text -> - -- | Data to add to column - V.Vector a -> - -- | DataFrame to add to column - DataFrame -> - DataFrame -insertVectorWithDefault defaultValue name xs d = - let (rows, _) = dataframeDimensions d - values = xs V.++ V.replicate (rows - V.length xs) defaultValue - in insertColumn name (fromVector values) d +{- | /O(n)/ Clones a column and places it under a new name in the dataframe. +==== __Example__ +@ +ghci> import qualified Data.Vector as V + +ghci> df = insertVector "numbers" (V.fromList [1..10]) D.empty + +ghci> D.cloneColumn "numbers" "others" df + +------------------------ +index | numbers | others +------|---------|------- + Int | Int | Int +------|---------|------- +0 | 1 | 1 +1 | 2 | 2 +2 | 3 | 3 +3 | 4 | 4 +4 | 5 | 5 +5 | 6 | 6 +6 | 7 | 7 +7 | 8 | 8 +8 | 9 | 9 +9 | 10 | 10 + +@ +-} +cloneColumn :: T.Text -> T.Text -> DataFrame -> DataFrame +cloneColumn original new df = fromMaybe (throw $ ColumnNotFoundException original "cloneColumn" (map fst $ M.toList $ columnIndices df)) $ do + column <- getColumn original df + return $ insertColumn new column df + +{- | /O(n)/ Renames a single column. + +==== __Example__ +@ +ghci> import qualified Data.Vector as V + +ghci> df = insertVector "numbers" (V.fromList [1..10]) D.empty + +ghci> D.rename "numbers" "others" df + +-------------- +index | others +------|------- + Int | Int +------|------- +0 | 1 +1 | 2 +2 | 3 +3 | 4 +4 | 5 +5 | 6 +6 | 7 +7 | 8 +8 | 9 +9 | 10 + +@ +-} rename :: T.Text -> T.Text -> DataFrame -> DataFrame rename orig new df = either throw id (renameSafe orig new df) +{- | /O(n)/ Renames many columns. + +==== __Example__ +@ +ghci> import qualified Data.Vector as V + +ghci> df = D.insertVector "others" (V.fromList [11..20]) (D.insertVector "numbers" (V.fromList [1..10]) D.empty) + +ghci> df + +------------------------ +index | numbers | others +------|---------|------- + Int | Int | Int +------|---------|------- +0 | 1 | 11 +1 | 2 | 12 +2 | 3 | 13 +3 | 4 | 14 +4 | 5 | 15 +5 | 6 | 16 +6 | 7 | 17 +7 | 8 | 18 +8 | 9 | 19 +9 | 10 | 20 + +ghci> D.renameMany [("numbers", "first_10"), ("others", "next_10")] df + +-------------------------- +index | first_10 | next_10 +------|----------|-------- + Int | Int | Int +------|----------|-------- +0 | 1 | 11 +1 | 2 | 12 +2 | 3 | 13 +3 | 4 | 14 +4 | 5 | 15 +5 | 6 | 16 +6 | 7 | 17 +7 | 8 | 18 +8 | 9 | 19 +9 | 10 | 20 + +@ +-} renameMany :: [(T.Text, T.Text)] -> DataFrame -> DataFrame renameMany replacements df = fold (uncurry rename) replacements df renameSafe :: T.Text -> T.Text -> DataFrame -> Either DataFrameException DataFrame renameSafe orig new df = fromMaybe (Left $ ColumnNotFoundException orig "rename" (map fst $ M.toList $ columnIndices df)) $ do - columnIndex <- M.lookup orig (columnIndices df) - let origRemoved = M.delete orig (columnIndices df) - let newAdded = M.insert new columnIndex origRemoved - return (Right df { columnIndices = newAdded }) + columnIndex <- M.lookup orig (columnIndices df) + let origRemoved = M.delete orig (columnIndices df) + let newAdded = M.insert new columnIndex origRemoved + return (Right df{columnIndices = newAdded}) --- | O(1) Get the number of elements in a given column. -columnSize :: T.Text -> DataFrame -> Maybe Int -columnSize name df = columnLength <$> getColumn name df +data ColumnInfo = ColumnInfo + { nameOfColumn :: !T.Text + , nonNullValues :: !Int + , nullValues :: !Int + , partiallyParsedValues :: !Int + , uniqueValues :: !Int + , typeOfColumn :: !T.Text + } -data ColumnInfo = ColumnInfo { - nameOfColumn :: !T.Text, - nonNullValues :: !Int, - nullValues :: !Int, - partiallyParsedValues :: !Int, - uniqueValues :: !Int, - typeOfColumn :: !T.Text - } +{- | O(n * k ^ 2) Returns the number of non-null columns in the dataframe and the type associated with each column. --- | O(n) Returns the number of non-null columns in the dataframe and the type associated --- with each column. +==== __Example__ +@ +ghci> import qualified Data.Vector as V + +ghci> df = D.insertVector "others" (V.fromList [11..20]) (D.insertVector "numbers" (V.fromList [1..10]) D.empty) + +ghci> D.describeColumns df + +----------------------------------------------------------------------------------------------------- +index | Column Name | # Non-null Values | # Null Values | # Partially parsed | # Unique Values | Type +------|-------------|-------------------|---------------|--------------------|-----------------|----- + Int | Text | Int | Int | Int | Int | Text +------|-------------|-------------------|---------------|--------------------|-----------------|----- +0 | others | 10 | 0 | 0 | 10 | Int +1 | numbers | 10 | 0 | 0 | 10 | Int + +@ +-} describeColumns :: DataFrame -> DataFrame -describeColumns df = empty & insertColumn "Column Name" (fromList (map nameOfColumn infos)) - & insertColumn "# Non-null Values" (fromList (map nonNullValues infos)) - & insertColumn "# Null Values" (fromList (map nullValues infos)) - & insertColumn "# Partially parsed" (fromList (map partiallyParsedValues infos)) - & insertColumn "# Unique Values" (fromList (map uniqueValues infos)) - & insertColumn "Type" (fromList (map typeOfColumn infos)) +describeColumns df = + empty + & insertColumn "Column Name" (fromList (map nameOfColumn infos)) + & insertColumn "# Non-null Values" (fromList (map nonNullValues infos)) + & insertColumn "# Null Values" (fromList (map nullValues infos)) + & insertColumn "# Partially parsed" (fromList (map partiallyParsedValues infos)) + & insertColumn "# Unique Values" (fromList (map uniqueValues infos)) + & insertColumn "Type" (fromList (map typeOfColumn infos)) where infos = L.sortBy (compare `on` nonNullValues) (V.ifoldl' go [] (columns df)) :: [ColumnInfo] indexMap = M.fromList (map (\(a, b) -> (b, a)) $ M.toList (columnIndices df)) columnName i = M.lookup i indexMap - go acc i col@(OptionalColumn (c :: V.Vector a)) = let - cname = columnName i - countNulls = nulls col - countPartial = partiallyParsed col - columnType = T.pack $ show $ typeRep @a - unique = S.size $ VG.foldr S.insert S.empty c - in if isNothing cname then acc else ColumnInfo (fromMaybe "" cname) (columnLength col - countNulls) countNulls countPartial unique columnType : acc - go acc i col@(BoxedColumn (c :: V.Vector a)) = let - cname = columnName i - countPartial = partiallyParsed col - columnType = T.pack $ show $ typeRep @a - unique = S.size $ VG.foldr S.insert S.empty c - in if isNothing cname then acc else ColumnInfo (fromMaybe "" cname) (columnLength col) 0 countPartial unique columnType : acc - go acc i col@(UnboxedColumn c) = let - cname = columnName i - columnType = T.pack $ columnTypeString col - unique = S.size $ VG.foldr S.insert S.empty c - -- Unboxed columns cannot have nulls since Maybe - -- is not an instance of Unbox a - in if isNothing cname then acc else ColumnInfo (fromMaybe "" cname) (columnLength col) 0 0 unique columnType : acc + go acc i col@(OptionalColumn (c :: V.Vector a)) = + let + cname = columnName i + countNulls = nulls col + countPartial = partiallyParsed col + columnType = T.pack $ show $ typeRep @a + unique = S.size $ VG.foldr S.insert S.empty c + in + if isNothing cname then acc else ColumnInfo (fromMaybe "" cname) (columnLength col - countNulls) countNulls countPartial unique columnType : acc + go acc i col@(BoxedColumn (c :: V.Vector a)) = + let + cname = columnName i + countPartial = partiallyParsed col + columnType = T.pack $ show $ typeRep @a + unique = S.size $ VG.foldr S.insert S.empty c + in + if isNothing cname then acc else ColumnInfo (fromMaybe "" cname) (columnLength col) 0 countPartial unique columnType : acc + go acc i col@(UnboxedColumn c) = + let + cname = columnName i + columnType = T.pack $ columnTypeString col + unique = S.size $ VG.foldr S.insert S.empty c + in + -- Unboxed columns cannot have nulls since Maybe + -- is not an instance of Unbox a + if isNothing cname then acc else ColumnInfo (fromMaybe "" cname) (columnLength col) 0 0 unique columnType : acc nulls :: Column -> Int nulls (OptionalColumn xs) = VG.length $ VG.filter isNothing xs nulls (BoxedColumn (xs :: V.Vector a)) = case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> VG.length $ VG.filter isNullish xs - Nothing -> case testEquality (typeRep @a) (typeRep @String) of - Just Refl -> VG.length $ VG.filter (isNullish . T.pack) xs - Nothing -> case typeRep @a of - App t1 t2 -> case eqTypeRep t1 (typeRep @Maybe) of - Just HRefl -> VG.length $ VG.filter isNothing xs - Nothing -> 0 - _ -> 0 + Just Refl -> VG.length $ VG.filter isNullish xs + Nothing -> case testEquality (typeRep @a) (typeRep @String) of + Just Refl -> VG.length $ VG.filter (isNullish . T.pack) xs + Nothing -> case typeRep @a of + App t1 t2 -> case eqTypeRep t1 (typeRep @Maybe) of + Just HRefl -> VG.length $ VG.filter isNothing xs + Nothing -> 0 + _ -> 0 nulls _ = 0 partiallyParsed :: Column -> Int partiallyParsed (BoxedColumn (xs :: V.Vector a)) = - case typeRep @a of - App (App tycon t1) t2 -> case eqTypeRep tycon (typeRep @Either) of - Just HRefl -> VG.length $ VG.filter isLeft xs - Nothing -> 0 - _ -> 0 + case typeRep @a of + App (App tycon t1) t2 -> case eqTypeRep tycon (typeRep @Either) of + Just HRefl -> VG.length $ VG.filter isLeft xs + Nothing -> 0 + _ -> 0 partiallyParsed _ = 0 +{- | Creates a dataframe from a list of tuples with name and column. + +==== __Example__ +@ +ghci> df = D.fromNamedColumns [("numbers", D.fromList [1..10]), ("others", D.fromList [11..20])] + +ghci> df + +------------------------ +index | numbers | others +------|---------|------- + Int | Int | Int +------|---------|------- +0 | 1 | 11 +1 | 2 | 12 +2 | 3 | 13 +3 | 4 | 14 +4 | 5 | 15 +5 | 6 | 16 +6 | 7 | 17 +7 | 8 | 18 +8 | 9 | 19 +9 | 10 | 20 + +@ +-} fromNamedColumns :: [(T.Text, Column)] -> DataFrame fromNamedColumns = L.foldl' (\df (name, column) -> insertColumn name column df) empty +{- | Create a dataframe from a list of columns. The column names are "0", "1"... etc. +Useful for quick exploration but you should probably alwyas rename the columns after +or drop the ones you don't want. + +==== __Example__ +@ +ghci> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]] + +ghci> df + +----------------- +index | 0 | 1 +------|-----|---- + Int | Int | Int +------|-----|---- +0 | 1 | 11 +1 | 2 | 12 +2 | 3 | 13 +3 | 4 | 14 +4 | 5 | 15 +5 | 6 | 16 +6 | 7 | 17 +7 | 8 | 18 +8 | 9 | 19 +9 | 10 | 20 + +@ +-} fromUnnamedColumns :: [Column] -> DataFrame -fromUnnamedColumns = fromNamedColumns . zip (map (T.pack . show) [0..]) +fromUnnamedColumns = fromNamedColumns . zip (map (T.pack . show) [0 ..]) --- | O (k * n) Counts the occurences of each value in a given column. +{- | O (k * n) Counts the occurences of each value in a given column. + +==== __Example__ +@ +ghci> df = D.fromUnnamedColumns [D.fromList [1..10], D.fromList [11..20]] + +ghci> D.valueCounts @Int "0" df + +[(1,1),(2,1),(3,1),(4,1),(5,1),(6,1),(7,1),(8,1),(9,1),(10,1)] + +@ +-} valueCounts :: forall a. (Columnable a) => T.Text -> DataFrame -> [(a, Int)] valueCounts columnName df = case getColumn columnName df of - Nothing -> throw $ ColumnNotFoundException columnName "valueCounts" (map fst $ M.toList $ columnIndices df) - Just (BoxedColumn (column' :: V.Vector c)) -> + Nothing -> throw $ ColumnNotFoundException columnName "valueCounts" (map fst $ M.toList $ columnIndices df) + Just (BoxedColumn (column' :: V.Vector c)) -> let - column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column' - in case (typeRep @a) `testEquality` (typeRep @c) of - Nothing -> throw $ TypeMismatchException (MkTypeErrorContext - { userType = Right $ typeRep @a - , expectedType = Right $ typeRep @c - , errorColumnName = Just (T.unpack columnName) - , callingFunctionName = Just "valueCounts" - }) - Just Refl -> M.toAscList column - Just (OptionalColumn (column' :: V.Vector c)) -> + column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column' + in + case (typeRep @a) `testEquality` (typeRep @c) of + Nothing -> + throw $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right $ typeRep @a + , expectedType = Right $ typeRep @c + , errorColumnName = Just (T.unpack columnName) + , callingFunctionName = Just "valueCounts" + } + ) + Just Refl -> M.toAscList column + Just (OptionalColumn (column' :: V.Vector c)) -> let - column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column' - in case (typeRep @a) `testEquality` (typeRep @c) of - Nothing -> throw $ TypeMismatchException (MkTypeErrorContext - { userType = Right $ typeRep @a - , expectedType = Right $ typeRep @c - , errorColumnName = Just (T.unpack columnName) - , callingFunctionName = Just "valueCounts" - }) - Just Refl -> M.toAscList column - Just (UnboxedColumn (column' :: VU.Vector c)) -> let - column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty (V.convert column') - in case (typeRep @a) `testEquality` (typeRep @c) of - Nothing -> throw $ TypeMismatchException (MkTypeErrorContext - { userType = Right $ typeRep @a - , expectedType = Right $ typeRep @c - , errorColumnName = Just (T.unpack columnName) - , callingFunctionName = Just "valueCounts" - }) - Just Refl -> M.toAscList column + column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column' + in + case (typeRep @a) `testEquality` (typeRep @c) of + Nothing -> + throw $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right $ typeRep @a + , expectedType = Right $ typeRep @c + , errorColumnName = Just (T.unpack columnName) + , callingFunctionName = Just "valueCounts" + } + ) + Just Refl -> M.toAscList column + Just (UnboxedColumn (column' :: VU.Vector c)) -> + let + column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty (V.convert column') + in + case (typeRep @a) `testEquality` (typeRep @c) of + Nothing -> + throw $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right $ typeRep @a + , expectedType = Right $ typeRep @c + , errorColumnName = Just (T.unpack columnName) + , callingFunctionName = Just "valueCounts" + } + ) + Just Refl -> M.toAscList column +{- | A left fold for dataframes that takes the dataframe as the last object. +this makes it easier to chain operations. + +==== __Example__ +@ +ghci> D.fold (const id) [1..5] df + +----------------- +index | 0 | 1 +------|-----|---- + Int | Int | Int +------|-----|---- +0 | 1 | 11 +1 | 2 | 12 +2 | 3 | 13 +3 | 4 | 14 +4 | 5 | 15 +5 | 6 | 16 +6 | 7 | 17 +7 | 8 | 18 +8 | 9 | 19 +9 | 10 | 20 + +@ +-} fold :: (a -> DataFrame -> DataFrame) -> [a] -> DataFrame -> DataFrame fold f xs acc = L.foldl' (flip f) acc xs diff --git a/src/DataFrame/Operations/Join.hs b/src/DataFrame/Operations/Join.hs --- a/src/DataFrame/Operations/Join.hs +++ b/src/DataFrame/Operations/Join.hs @@ -1,42 +1,52 @@ {-# LANGUAGE OverloadedStrings #-} + module DataFrame.Operations.Join where import qualified Data.Map.Strict as M +import qualified Data.Text as T import qualified Data.Vector as VB import qualified Data.Vector.Unboxed as VU -import qualified Data.Text as T -import DataFrame.Internal.Column as D -import DataFrame.Internal.DataFrame as D -import DataFrame.Operations.Aggregation as D -import DataFrame.Operations.Core as D +import DataFrame.Internal.Column as D +import DataFrame.Internal.DataFrame as D +import DataFrame.Operations.Aggregation as D +import DataFrame.Operations.Core as D -data JoinType = INNER - | LEFT - | RIGHT - | FULL_OUTER +-- | Equivalent to SQL join types. +data JoinType + = INNER + | LEFT + | RIGHT + | FULL_OUTER -join :: JoinType - -> [T.Text] - -> DataFrame -- Right hand side - -> DataFrame -- Left hand side - -> DataFrame +{- | Join two dataframes using SQL join semantics. + +Only inner join is implemented for now. +-} +join :: + JoinType -> + [T.Text] -> + DataFrame -> -- Right hand side + DataFrame -> -- Left hand side + DataFrame join INNER xs right = innerJoin xs right join LEFT xs right = error "UNIMPLEMENTED" join RIGHT xs right = error "UNIMPLEMENTED" join FULL_OUTER xs right = error "UNIMPLEMENTED" --- Create row representation for each of the two dataframes --- get the product of left and right counts for each key +{- | Inner join of two dataframes. Note: for chaining, the left dataframe is actually +on the right side. +-} innerJoin :: [T.Text] -> DataFrame -> DataFrame -> DataFrame -innerJoin cs right left = let +innerJoin cs right left = + let leftIndicesToGroup = M.elems $ M.filterWithKey (\k _ -> k `elem` cs) (D.columnIndices left) leftRowRepresentations = VU.generate (fst (D.dimensions left)) (D.mkRowRep leftIndicesToGroup left) -- key -> [index0, index1] - leftKeyCountsAndIndices = VU.foldr (\(i, v) acc -> M.insertWith (++) v [i] acc) M.empty (VU.indexed leftRowRepresentations) + leftKeyCountsAndIndices = VU.foldr (\(i, v) acc -> M.insertWith (++) v [i] acc) M.empty (VU.indexed leftRowRepresentations) -- key -> [index0, index1] rightIndicesToGroup = M.elems $ M.filterWithKey (\k _ -> k `elem` cs) (D.columnIndices right) rightRowRepresentations = VU.generate (fst (D.dimensions right)) (D.mkRowRep rightIndicesToGroup right) - rightKeyCountsAndIndices = VU.foldr (\(i, v) acc -> M.insertWith (++) v [i] acc) M.empty (VU.indexed rightRowRepresentations) + rightKeyCountsAndIndices = VU.foldr (\(i, v) acc -> M.insertWith (++) v [i] acc) M.empty (VU.indexed rightRowRepresentations) -- key -> [(left_indexes0, right_indexes1)] mergedKeyCountsAndIndices = M.foldrWithKey (\k v m -> if k `M.member` rightKeyCountsAndIndices then M.insert k (VU.fromList v, VU.fromList (rightKeyCountsAndIndices M.! k)) m else m) M.empty leftKeyCountsAndIndices -- [(ints, ints)] @@ -46,13 +56,14 @@ expandedLeftIndicies = mconcat (map fst expandedIndices) expandedRightIndicies = mconcat (map snd expandedIndices) -- df - expandedLeft = left { columns = VB.map (D.atIndicesStable expandedLeftIndicies) (D.columns left), dataframeDimensions = (VU.length expandedLeftIndicies, snd (D.dataframeDimensions left))} - -- df - expandedRight = right { columns = VB.map (D.atIndicesStable expandedRightIndicies) (D.columns right), dataframeDimensions = (VU.length expandedRightIndicies, snd (D.dataframeDimensions right))} + expandedLeft = left{columns = VB.map (D.atIndicesStable expandedLeftIndicies) (D.columns left), dataframeDimensions = (VU.length expandedLeftIndicies, snd (D.dataframeDimensions left))} + -- df + expandedRight = right{columns = VB.map (D.atIndicesStable expandedRightIndicies) (D.columns right), dataframeDimensions = (VU.length expandedRightIndicies, snd (D.dataframeDimensions right))} -- [string] leftColumns = D.columnNames left - rightColumns = D.columnNames right + rightColumns = D.columnNames right initDf = expandedLeft insertIfPresent _ Nothing df = df insertIfPresent name (Just c) df = D.insertColumn name c df - in D.fold (\name df -> if name `elem` cs then df else (if name `elem` leftColumns then insertIfPresent ("Right_" <> name) (D.getColumn name expandedRight) df else insertIfPresent name (D.getColumn name expandedRight) df)) rightColumns initDf + in + D.fold (\name df -> if name `elem` cs then df else (if name `elem` leftColumns then insertIfPresent ("Right_" <> name) (D.getColumn name expandedRight) df else insertIfPresent name (D.getColumn name expandedRight) df)) rightColumns initDf diff --git a/src/DataFrame/Operations/Merge.hs b/src/DataFrame/Operations/Merge.hs --- a/src/DataFrame/Operations/Merge.hs +++ b/src/DataFrame/Operations/Merge.hs @@ -1,4 +1,5 @@ {-# LANGUAGE InstanceSigs #-} + module DataFrame.Operations.Merge where import qualified Data.List as L @@ -11,8 +12,10 @@ import Data.Maybe instance Semigroup D.DataFrame where + -- \| Vertically merge two dataframes using shared columns. (<>) :: D.DataFrame -> D.DataFrame -> D.DataFrame - (<>) a b = let + (<>) a b = + let columnsInBOnly = filter (\c -> c `notElem` D.columnNames b) (D.columnNames b) columnsInA = D.columnNames a addColumns a' b' df name @@ -23,25 +26,29 @@ | fst (D.dimensions b') == 0 = fromMaybe df $ do col <- D.getColumn name a' pure $ D.insertColumn name col df - | otherwise = let + | otherwise = + let numColumnsA = (fst $ D.dimensions a') numColumnsB = (fst $ D.dimensions b') numColumns = max numColumnsA numColumnsB optA = D.getColumn name a' optB = D.getColumn name b' - in case optB of - Nothing -> case optA of - Nothing -> D.insertColumn name (D.fromList ([] :: [T.Text])) df - Just a'' -> D.insertColumn name (D.expandColumn numColumnsB a'') df - Just b'' -> case optA of - Nothing -> D.insertColumn name (D.leftExpandColumn numColumnsA b'') df - Just a'' -> fromMaybe df $ do - concatedColumns <- D.concatColumns a'' b'' - pure $ D.insertColumn name concatedColumns df - in L.foldl' (addColumns a b) D.empty (D.columnNames a `L.union` D.columnNames b) + in + case optB of + Nothing -> case optA of + Nothing -> D.insertColumn name (D.fromList ([] :: [T.Text])) df + Just a'' -> D.insertColumn name (D.expandColumn numColumnsB a'') df + Just b'' -> case optA of + Nothing -> D.insertColumn name (D.leftExpandColumn numColumnsA b'') df + Just a'' -> fromMaybe df $ do + concatedColumns <- D.concatColumns a'' b'' + pure $ D.insertColumn name concatedColumns df + in + L.foldl' (addColumns a b) D.empty (D.columnNames a `L.union` D.columnNames b) instance Monoid D.DataFrame where - mempty = D.empty + mempty = D.empty +-- | Add two dataframes side by side/horizontally. (|||) :: D.DataFrame -> D.DataFrame -> D.DataFrame (|||) a b = D.fold (\name acc -> D.insertColumn name (D.unsafeGetColumn name b) acc) (D.columnNames b) a diff --git a/src/DataFrame/Operations/Sorting.hs b/src/DataFrame/Operations/Sorting.hs --- a/src/DataFrame/Operations/Sorting.hs +++ b/src/DataFrame/Operations/Sorting.hs @@ -1,4 +1,5 @@ {-# LANGUAGE OverloadedStrings #-} + module DataFrame.Operations.Sorting where import qualified Data.List as L @@ -6,28 +7,31 @@ import qualified Data.Vector as V import Control.Exception (throw) -import DataFrame.Errors (DataFrameException(..)) +import DataFrame.Errors (DataFrameException (..)) import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame (DataFrame(..), getColumn) +import DataFrame.Internal.DataFrame (DataFrame (..), getColumn) import DataFrame.Internal.Row import DataFrame.Operations.Core -- | Sort order taken as a parameter by the sortby function. data SortOrder = Ascending | Descending deriving (Eq) --- | O(k log n) Sorts the dataframe by a given row. --- --- > sortBy "Age" df +{- | O(k log n) Sorts the dataframe by a given row. + +> sortBy "Age" df +-} sortBy :: - SortOrder -> - [T.Text] -> - DataFrame -> - DataFrame + SortOrder -> + [T.Text] -> + DataFrame -> + DataFrame sortBy order names df - | any (`notElem` columnNames df) names = throw $ ColumnNotFoundException (T.pack $ show $ names L.\\ columnNames df) "sortBy" (columnNames df) - | otherwise = let - -- TODO: Remove the SortOrder defintion from operations so we can share it between here and internal and - -- we don't have to do this Bool mapping. - indexes = sortedIndexes' (order == Ascending) (toRowVector names df) - pick idxs col = atIndicesStable idxs col - in df {columns = V.map (pick indexes) (columns df)} + | any (`notElem` columnNames df) names = throw $ ColumnNotFoundException (T.pack $ show $ names L.\\ columnNames df) "sortBy" (columnNames df) + | otherwise = + let + -- TODO: Remove the SortOrder defintion from operations so we can share it between here and internal and + -- we don't have to do this Bool mapping. + indexes = sortedIndexes' (order == Ascending) (toRowVector names df) + pick idxs col = atIndicesStable idxs col + in + df{columns = V.map (pick indexes) (columns df)} diff --git a/src/DataFrame/Operations/Statistics.hs b/src/DataFrame/Operations/Statistics.hs --- a/src/DataFrame/Operations/Statistics.hs +++ b/src/DataFrame/Operations/Statistics.hs @@ -1,171 +1,257 @@ -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + module DataFrame.Operations.Statistics where import Data.Bifunctor (second) import qualified Data.List as L import qualified Data.Map as M import qualified Data.Text as T -import qualified Data.Vector.Generic as VG import qualified Data.Vector as V +import qualified Data.Vector.Algorithms.Intro as VA +import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM import qualified Statistics.Quantile as SS import qualified Statistics.Sample as SS import Prelude as P import Control.Exception (throw) -import DataFrame.Errors (DataFrameException(..)) +import Control.Monad.ST (runST) +import qualified Data.Bifunctor as Data +import Data.Foldable (asum) +import Data.Function ((&)) +import Data.Maybe (fromMaybe, isJust) +import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) +import DataFrame.Errors (DataFrameException (..)) import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame (DataFrame(..), getColumn, empty, unsafeGetColumn) +import DataFrame.Internal.DataFrame (DataFrame (..), empty, getColumn, unsafeGetColumn) +import DataFrame.Internal.Row (showValue, toAny) import DataFrame.Operations.Core import DataFrame.Operations.Subset (filterJust) -import Data.Foldable (asum) -import Data.Maybe (isJust, fromMaybe) -import Data.Function ((&)) -import Data.Type.Equality (type (:~:)(Refl), TestEquality (testEquality)) -import Type.Reflection (typeRep) -import qualified Data.Bifunctor as Data -import DataFrame.Internal.Row (toAny, showValue) import GHC.Float (int2Double) import Text.Printf (printf) +import Type.Reflection (typeRep) +{- | Show a frequency table for a categorical feaure. + +__Examples:__ + +@ +ghci> df <- D.readCsv "./data/housing.csv" + +ghci> D.frequencies "ocean_proximity" df + +---------------------------------------------------------------------------- +index | Statistic | <1H OCEAN | INLAND | ISLAND | NEAR BAY | NEAR OCEAN +------|----------------|-----------|--------|--------|----------|----------- + Int | Text | Any | Any | Any | Any | Any +------|----------------|-----------|--------|--------|----------|----------- +0 | Count | 9136 | 6551 | 5 | 2290 | 2658 +1 | Percentage (%) | 44.26% | 31.74% | 0.02% | 11.09% | 12.88% +@ +-} frequencies :: T.Text -> DataFrame -> DataFrame -frequencies name df = let - counts :: forall a . Columnable a => [(a, Int)] - counts = valueCounts name 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]) d) initDf counts - in case getColumn name df of - Nothing -> throw $ ColumnNotFoundException name "frequencies" (map fst $ M.toList $ 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 +frequencies name df = + let + counts :: forall a. (Columnable a) => [(a, Int)] + counts = valueCounts name 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]) d) initDf counts + in + case getColumn name df of + Nothing -> throw $ ColumnNotFoundException name "frequencies" (map fst $ M.toList $ 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 +-- | Calculates the mean of a given column as a standalone value. mean :: T.Text -> DataFrame -> Maybe Double mean = applyStatistic mean' +-- | Calculates the median of a given column as a standalone value. median :: T.Text -> DataFrame -> Maybe Double -median = applyStatistic (SS.median SS.medianUnbiased) +median = applyStatistic median' +-- | Calculates the standard deviation of a given column as a standalone value. standardDeviation :: T.Text -> DataFrame -> Maybe Double -standardDeviation = applyStatistic SS.fastStdDev +standardDeviation = applyStatistic (sqrt . variance') +-- | Calculates the skewness of a given column as a standalone value. skewness :: T.Text -> DataFrame -> Maybe Double skewness = applyStatistic SS.skewness +-- | Calculates the variance of a given column as a standalone value. variance :: T.Text -> DataFrame -> Maybe Double variance = applyStatistic variance' +-- | Calculates the inter-quartile range of a given column as a standalone value. interQuartileRange :: T.Text -> DataFrame -> Maybe Double interQuartileRange = applyStatistic (SS.midspread SS.medianUnbiased 4) +-- | Calculates the Pearson's correlation coefficient between two given columns as a standalone value. correlation :: T.Text -> T.Text -> DataFrame -> Maybe Double correlation first second df = do - f <- _getColumnAsDouble first df - s <- _getColumnAsDouble second df - return $ SS.correlation (VG.zip f s) + f <- _getColumnAsDouble first df + s <- _getColumnAsDouble second df + correlation' f s _getColumnAsDouble :: T.Text -> DataFrame -> Maybe (VU.Vector Double) _getColumnAsDouble name df = case getColumn name df of - Just (UnboxedColumn (f :: VU.Vector a)) -> case testEquality (typeRep @a) (typeRep @Double) of - Just Refl -> Just f - Nothing -> case testEquality (typeRep @a) (typeRep @Int) of - Just Refl -> Just $ VU.map fromIntegral f - Nothing -> Nothing - _ -> Nothing + Just (UnboxedColumn (f :: VU.Vector a)) -> case testEquality (typeRep @a) (typeRep @Double) of + Just Refl -> Just f + Nothing -> case testEquality (typeRep @a) (typeRep @Int) of + Just Refl -> Just $ VU.map fromIntegral f + Nothing -> case testEquality (typeRep @a) (typeRep @Float) of + Just Refl -> Just $ VU.map realToFrac f + Nothing -> Nothing + _ -> Nothing +{-# INLINE _getColumnAsDouble #-} +-- | Calculates the sum of a given column as a standalone value. sum :: forall a. (Columnable a, Num a, VU.Unbox a) => T.Text -> DataFrame -> Maybe a sum name df = case getColumn name df of - Nothing -> throw $ ColumnNotFoundException name "sum" (map fst $ M.toList $ columnIndices df) - Just ((UnboxedColumn (column :: VU.Vector a'))) -> case testEquality (typeRep @a') (typeRep @a) of - Just Refl -> Just $ VG.sum column - Nothing -> Nothing + Nothing -> throw $ ColumnNotFoundException name "sum" (map fst $ M.toList $ columnIndices df) + Just ((UnboxedColumn (column :: VU.Vector a'))) -> case testEquality (typeRep @a') (typeRep @a) of + Just Refl -> Just $ VG.sum column + Nothing -> Nothing applyStatistic :: (VU.Vector Double -> Double) -> T.Text -> DataFrame -> Maybe Double applyStatistic f name df = case getColumn name (filterJust name df) of - Nothing -> throw $ ColumnNotFoundException name "applyStatistic" (map fst $ M.toList $ columnIndices df) - Just column@(UnboxedColumn (col :: VU.Vector a)) -> case testEquality (typeRep @a) (typeRep @Double) of - Just Refl -> reduceColumn f column + Nothing -> throw $ ColumnNotFoundException name "applyStatistic" (map fst $ M.toList $ columnIndices df) + Just column@(UnboxedColumn (col :: VU.Vector a)) -> case testEquality (typeRep @a) (typeRep @Double) of + Just Refl -> + let + res = (f col) + in + if isNaN res then Nothing else pure res Nothing -> do - matching <- asum [mapColumn (fromIntegral :: Int -> Double) column, - mapColumn (fromIntegral :: Integer -> Double) column, - mapColumn (realToFrac :: Float -> Double) column, - Just column ] - reduceColumn f matching - _ -> Nothing + col' <- _getColumnAsDouble name df + let res = (f col') + if isNaN res then Nothing else pure res + _ -> Nothing +{-# INLINE applyStatistic #-} applyStatistics :: (VU.Vector Double -> VU.Vector Double) -> T.Text -> DataFrame -> Maybe (VU.Vector Double) applyStatistics f name df = case getColumn name (filterJust name df) of - Just ((UnboxedColumn (column :: VU.Vector a'))) -> case testEquality (typeRep @a') (typeRep @Int) of - Just Refl -> Just $! f (VU.map fromIntegral column) - Nothing -> case testEquality (typeRep @a') (typeRep @Double) of - Just Refl -> Just $! f column - Nothing -> case testEquality (typeRep @a') (typeRep @Float) of - Just Refl -> Just $! f (VG.map realToFrac column) - Nothing -> Nothing - _ -> Nothing + Just ((UnboxedColumn (column :: VU.Vector a'))) -> case testEquality (typeRep @a') (typeRep @Int) of + Just Refl -> Just $! f (VU.map fromIntegral column) + Nothing -> case testEquality (typeRep @a') (typeRep @Double) of + Just Refl -> Just $! f column + Nothing -> case testEquality (typeRep @a') (typeRep @Float) of + Just Refl -> Just $! f (VG.map realToFrac column) + Nothing -> Nothing + _ -> Nothing +-- | Descriprive statistics of the numeric columns. summarize :: DataFrame -> DataFrame -summarize df = fold columnStats (columnNames df) (fromNamedColumns [("Statistic", fromList ["Count" :: T.Text, "Mean", "Minimum", "25%" ,"Median", "75%", "Max", "StdDev", "IQR", "Skewness"])]) - where columnStats name d = if all isJust (stats name) then insertUnboxedVector name (VU.fromList (map (roundTo 2 . fromMaybe 0) $ stats name)) d else d - stats name = let +summarize df = fold columnStats (columnNames df) (fromNamedColumns [("Statistic", fromList ["Count" :: T.Text, "Mean", "Minimum", "25%", "Median", "75%", "Max", "StdDev", "IQR", "Skewness"])]) + where + columnStats name d = if all isJust (stats name) then insertUnboxedVector name (VU.fromList (map (roundTo 2 . fromMaybe 0) $ stats name)) d else d + stats name = + let count = fromIntegral . numElements <$> getColumn name df - quantiles = applyStatistics (SS.quantilesVec SS.medianUnbiased (VU.fromList [0,1,2,3,4]) 4) name df + quantiles = applyStatistics (SS.quantilesVec SS.medianUnbiased (VU.fromList [0, 1, 2, 3, 4]) 4) name df min' = flip (VG.!) 0 <$> quantiles quartile1 = flip (VG.!) 1 <$> quantiles median' = flip (VG.!) 2 <$> quantiles quartile3 = flip (VG.!) 3 <$> quantiles max' = flip (VG.!) 4 <$> quantiles iqr = (-) <$> quartile3 <*> quartile1 - in [count, - mean name df, - min', - quartile1, - median', - quartile3, - max', - standardDeviation name df, - iqr, - skewness name df] + in + [ count + , mean name df + , min' + , quartile1 + , median' + , quartile3 + , max' + , standardDeviation name df + , iqr + , skewness name df + ] -- | Round a @Double@ to Specified Precision roundTo :: Int -> Double -> Double -roundTo n x = fromInteger (round $ x * (10^n)) / (10.0^^n) +roundTo n x = fromInteger (round $ x * (10 ^ n)) / (10.0 ^^ n) toPct2dp :: Double -> String toPct2dp x - | x < 0.00005 = "<0.01%" - | otherwise = printf "%.2f%%" (x * 100) + | x < 0.00005 = "<0.01%" + | otherwise = printf "%.2f%%" (x * 100) mean' :: VU.Vector Double -> Double -mean' samp = let - (!total, !n) = VG.foldl' (\(!total, !n) v -> (total + v, n + 1)) (0 :: Double, 0 :: Int) samp - in total / fromIntegral n +mean' samp = VU.sum samp / fromIntegral (VU.length samp) +{-# INLINE mean #-} +median' :: VU.Vector Double -> Double +median' samp + | VU.null samp = throw $ EmptyDataSetException "median" + | otherwise = runST $ do + mutableSamp <- VU.thaw samp + VA.sort mutableSamp + let len = VU.length samp + middleIndex = len `div` 2 + middleElement <- VUM.read mutableSamp middleIndex + if odd len + then pure middleElement + else do + prev <- VUM.read mutableSamp (middleIndex - 1) + pure ((middleElement + prev) / 2) +{-# INLINE median' #-} + -- accumulator: count, mean, m2 -data VarAcc = VarAcc !Int !Double !Double deriving Show +data VarAcc = VarAcc !Int !Double !Double deriving (Show) step :: VarAcc -> Double -> VarAcc step (VarAcc !n !mean !m2) !x = - let !n' = n + 1 - !delta = x - mean - !mean' = mean + delta / fromIntegral n' - !m2' = m2 + delta * (x - mean') - in VarAcc n' mean' m2' + let !n' = n + 1 + !delta = x - mean + !mean' = mean + delta / fromIntegral n' + !m2' = m2 + delta * (x - mean') + in VarAcc n' mean' m2' +{-# INLINE step #-} computeVariance :: VarAcc -> Double -computeVariance (VarAcc n _ m2) - | n < 2 = 0 -- or error "variance of <2 samples" - | otherwise = m2 / fromIntegral (n - 1) +computeVariance (VarAcc !n _ !m2) + | n < 2 = 0 -- or error "variance of <2 samples" + | otherwise = m2 / fromIntegral (n - 1) +{-# INLINE computeVariance #-} variance' :: VU.Vector Double -> Double -variance' = computeVariance . VG.foldl' step (VarAcc 0 0 0) +variance' = computeVariance . VU.foldl' step (VarAcc 0 0 0) +{-# INLINE variance' #-} + +correlation' :: VU.Vector Double -> VU.Vector Double -> Maybe Double +correlation' xs ys + | VU.length xs /= VU.length ys = Nothing + | nI < 2 = Nothing + | otherwise = + let !nf = fromIntegral nI + (!sumX, !sumY, !sumSquaredX, !sumSquaredY, !sumXY) = go 0 0 0 0 0 0 + !num = nf * sumXY - sumX * sumY + !den = sqrt ((nf * sumSquaredX - sumX * sumX) * (nf * sumSquaredY - sumY * sumY)) + in pure (num / den) + where + !nI = VU.length xs + go !i !sumX !sumY !sumSquaredX !sumSquaredY !sumXY + | i < nI = + let !x = VU.unsafeIndex xs i + !y = VU.unsafeIndex ys i + !sumX' = sumX + x + !sumY' = sumY + y + !sumSquaredX' = sumSquaredX + x * x + !sumSquaredY' = sumSquaredY + y * y + !sumXY' = sumXY + x * y + in go (i + 1) sumX' sumY' sumSquaredX' sumSquaredY' sumXY' + | otherwise = (sumX, sumY, sumSquaredX, sumSquaredY, sumXY) +{-# INLINE correlation' #-} diff --git a/src/DataFrame/Operations/Subset.hs b/src/DataFrame/Operations/Subset.hs --- a/src/DataFrame/Operations/Subset.hs +++ b/src/DataFrame/Operations/Subset.hs @@ -1,11 +1,12 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ExplicitNamespaces #-} + module DataFrame.Operations.Subset where import qualified Data.List as L @@ -13,54 +14,57 @@ import qualified Data.Set as S import qualified Data.Text as T import qualified Data.Vector as V +import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM -import qualified Data.Vector.Generic as VG import qualified Prelude import Control.Exception (throw) -import DataFrame.Errors (DataFrameException(..), TypeErrorContext(..)) +import Control.Monad.ST +import Data.Function ((&)) +import Data.Maybe (fromJust, fromMaybe, isJust) +import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) +import DataFrame.Errors (DataFrameException (..), TypeErrorContext (..)) import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame (DataFrame(..), getColumn, empty) +import DataFrame.Internal.DataFrame (DataFrame (..), empty, getColumn) import DataFrame.Internal.Expression -import DataFrame.Internal.Row (mkRowFromArgs, Any, toAny) -import Data.Type.Equality (type (:~:)(Refl), TestEquality (..)) +import DataFrame.Internal.Row (Any, mkRowFromArgs, toAny) import DataFrame.Operations.Core import DataFrame.Operations.Transformations (apply) -import Data.Function ((&)) -import Data.Maybe (isJust, fromJust, fromMaybe) -import Prelude hiding (filter, take) import Type.Reflection -import Control.Monad.ST +import Prelude hiding (filter, take) -- | O(k * n) Take the first n rows of a DataFrame. take :: Int -> DataFrame -> DataFrame -take n d = d {columns = V.map (takeColumn n') (columns d), dataframeDimensions = (n', c)} +take n d = d{columns = V.map (takeColumn n') (columns d), dataframeDimensions = (n', c)} where (r, c) = dataframeDimensions d n' = clip n 0 r +-- | O(k * n) Take the last n rows of a DataFrame. takeLast :: Int -> DataFrame -> DataFrame -takeLast n d = d {columns = V.map (takeLastColumn n') (columns d), dataframeDimensions = (n', c)} +takeLast n d = d{columns = V.map (takeLastColumn n') (columns d), dataframeDimensions = (n', c)} where (r, c) = dataframeDimensions d n' = clip n 0 r +-- | O(k * n) Drop the first n rows of a DataFrame. drop :: Int -> DataFrame -> DataFrame -drop n d = d {columns = V.map (sliceColumn n' (max (r - n') 0)) (columns d), dataframeDimensions = (max (r - n') 0, c)} +drop n d = d{columns = V.map (sliceColumn n' (max (r - n') 0)) (columns d), dataframeDimensions = (max (r - n') 0, c)} where (r, c) = dataframeDimensions d n' = clip n 0 r +-- | O(k * n) Drop the last n rows of a DataFrame. dropLast :: Int -> DataFrame -> DataFrame -dropLast n d = d {columns = V.map (sliceColumn 0 n') (columns d), dataframeDimensions = (n', c)} +dropLast n d = d{columns = V.map (sliceColumn 0 n') (columns d), dataframeDimensions = (n', c)} where (r, c) = dataframeDimensions d n' = clip (r - n) 0 r -- | O(k * n) Take a range of rows of a DataFrame. range :: (Int, Int) -> DataFrame -> DataFrame -range (start, end) d = d {columns = V.map (sliceColumn (clip start 0 r) n') (columns d), dataframeDimensions = (n', c)} +range (start, end) d = d{columns = V.map (sliceColumn (clip start 0 r) n') (columns d), dataframeDimensions = (n', c)} where (r, c) = dataframeDimensions d n' = clip (end - start) 0 r @@ -68,106 +72,126 @@ clip :: Int -> Int -> Int -> Int clip n left right = min right $ max n left --- | O(n * k) Filter rows by a given condition. --- --- filter "x" even df +{- | O(n * k) Filter rows by a given condition. + +filter "x" even df +-} filter :: - forall a. - (Columnable a) => - -- | Column to filter by - T.Text -> - -- | Filter condition - (a -> Bool) -> - -- | Dataframe to filter - DataFrame -> - DataFrame + forall a. + (Columnable a) => + -- | Column to filter by + T.Text -> + -- | Filter condition + (a -> Bool) -> + -- | Dataframe to filter + DataFrame -> + DataFrame filter filterColumnName condition df = case getColumn filterColumnName df of - Nothing -> throw $ ColumnNotFoundException filterColumnName "filter" (map fst $ M.toList $ columnIndices df) - Just (BoxedColumn (column :: V.Vector b)) -> filterByVector filterColumnName column condition df - Just (OptionalColumn (column :: V.Vector b)) -> filterByVector filterColumnName column condition df - Just (UnboxedColumn (column :: VU.Vector b)) -> filterByVector filterColumnName column condition df + Nothing -> throw $ ColumnNotFoundException filterColumnName "filter" (map fst $ M.toList $ columnIndices df) + Just (BoxedColumn (column :: V.Vector b)) -> filterByVector filterColumnName column condition df + Just (OptionalColumn (column :: V.Vector b)) -> filterByVector filterColumnName column condition df + Just (UnboxedColumn (column :: VU.Vector b)) -> filterByVector filterColumnName column condition df -filterByVector :: forall a b v . (VG.Vector v b, Columnable a, Columnable b) => T.Text -> v b -> (a -> Bool) -> DataFrame -> DataFrame +filterByVector :: forall a b v. (VG.Vector v b, Columnable a, Columnable b) => T.Text -> v b -> (a -> Bool) -> DataFrame -> DataFrame filterByVector filterColumnName column condition df = case testEquality (typeRep @a) (typeRep @b) of - Nothing -> throw $ TypeMismatchException (MkTypeErrorContext - { userType = Right $ typeRep @a - , expectedType = Right $ typeRep @b - , errorColumnName = Just (T.unpack filterColumnName) - , callingFunctionName = Just "filter" - }) - Just Refl -> let - ixs = indexes condition column - in df {columns = V.map (atIndicesStable ixs) (columns df), dataframeDimensions = (VG.length ixs, snd (dataframeDimensions df))} + Nothing -> + throw $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right $ typeRep @a + , expectedType = Right $ typeRep @b + , errorColumnName = Just (T.unpack filterColumnName) + , callingFunctionName = Just "filter" + } + ) + Just Refl -> + let + ixs = indexes condition column + in + df{columns = V.map (atIndicesStable ixs) (columns df), dataframeDimensions = (VG.length ixs, snd (dataframeDimensions df))} indexes :: (VG.Vector v a) => (a -> Bool) -> v a -> VU.Vector Int indexes condition cols = runST $ do - ixs <- VUM.new 8192 - (!icount, _, _, !ixs') <- VG.foldM (\(!icount, !vcount, !cap, mv) v -> do - if not (condition v) then - pure (icount, vcount + 1, cap, mv) - else do - let shouldGrow = icount == cap - mv' <- if shouldGrow then VUM.grow mv cap else pure mv - VUM.write mv' icount vcount - pure (icount + 1, vcount + 1, cap + (cap * fromEnum shouldGrow), mv')) (0, 0, 8192, ixs) cols - VU.freeze (VUM.slice 0 icount ixs') + ixs <- VUM.new 8192 + (!icount, _, _, !ixs') <- + VG.foldM + ( \(!icount, !vcount, !cap, mv) v -> do + if not (condition v) + then + pure (icount, vcount + 1, cap, mv) + else do + let shouldGrow = icount == cap + mv' <- if shouldGrow then VUM.grow mv cap else pure mv + VUM.write mv' icount vcount + pure (icount + 1, vcount + 1, cap + (cap * fromEnum shouldGrow), mv') + ) + (0, 0, 8192, ixs) + cols + VU.freeze (VUM.slice 0 icount ixs') --- | O(k) a version of filter where the predicate comes first. --- --- > filterBy even "x" df +{- | O(k) a version of filter where the predicate comes first. + +> filterBy even "x" df +-} filterBy :: (Columnable a) => (a -> Bool) -> T.Text -> DataFrame -> DataFrame filterBy = flip filter --- | O(k) filters the dataframe with a row predicate. The arguments in the function --- must appear in the same order as they do in the list. --- --- > filterWhere (["x", "y"], func (\x y -> x + y > 5)) df +{- | O(k) filters the dataframe with a row predicate. The arguments in the function + must appear in the same order as they do in the list. + +> filterWhere (["x", "y"], func (\x y -> x + y > 5)) df +-} filterWhere :: Expr Bool -> DataFrame -> DataFrame -filterWhere expr df = let - (TColumn col) = interpret @Bool df expr - (Just indexes) = findIndices (==True) col - c' = snd $ dataframeDimensions df - pick idxs col = atIndicesStable idxs col - in df {columns = V.map (pick indexes) (columns df), dataframeDimensions = (VU.length indexes, c')} +filterWhere expr df = + let + (TColumn col) = interpret @Bool df expr + (Just indexes) = findIndices (== True) col + c' = snd $ dataframeDimensions df + pick idxs col = atIndicesStable idxs col + in + df{columns = V.map (pick indexes) (columns df), dataframeDimensions = (VU.length indexes, c')} +{- | O(k) removes all rows with `Nothing` in a given column from the dataframe. --- | O(k) removes all rows with `Nothing` in a given column from the dataframe. --- --- > filterJust df +> filterJust df +-} filterJust :: T.Text -> DataFrame -> DataFrame filterJust name df = case getColumn name df of - Nothing -> throw $ ColumnNotFoundException name "filterJust" (map fst $ M.toList $ columnIndices df) - Just column@(OptionalColumn (col :: V.Vector (Maybe a))) -> filter @(Maybe a) name isJust df & apply @(Maybe a) fromJust name - Just column -> df + Nothing -> throw $ ColumnNotFoundException name "filterJust" (map fst $ M.toList $ columnIndices df) + Just column@(OptionalColumn (col :: V.Vector (Maybe a))) -> filter @(Maybe a) name isJust df & apply @(Maybe a) fromJust name + Just column -> df --- | O(n * k) removes all rows with `Nothing` from the dataframe. --- --- > filterJust df +{- | O(n * k) removes all rows with `Nothing` from the dataframe. + +> filterJust df +-} filterAllJust :: DataFrame -> DataFrame filterAllJust df = foldr filterJust df (columnNames df) --- | O(k) cuts the dataframe in a cube of size (a, b) where --- a is the length and b is the width. --- --- > cube (10, 5) df +{- | O(k) cuts the dataframe in a cube of size (a, b) where + a is the length and b is the width. + +> cube (10, 5) df +-} cube :: (Int, Int) -> DataFrame -> DataFrame cube (length, width) = take length . selectIntRange (0, width - 1) --- | O(n) Selects a number of columns in a given dataframe. --- --- > select ["name", "age"] df +{- | O(n) Selects a number of columns in a given dataframe. + +> select ["name", "age"] df +-} select :: - [T.Text] -> - DataFrame -> - DataFrame + [T.Text] -> + DataFrame -> + DataFrame select cs df - | L.null cs = empty - | any (`notElem` columnNames df) cs = throw $ ColumnNotFoundException (T.pack $ show $ cs L.\\ columnNames df) "select" (columnNames df) - | otherwise = L.foldl' addKeyValue empty cs + | L.null cs = empty + | any (`notElem` columnNames df) cs = throw $ ColumnNotFoundException (T.pack $ show $ cs L.\\ columnNames df) "select" (columnNames df) + | otherwise = L.foldl' addKeyValue empty cs where addKeyValue d k = fromMaybe df $ do - col <- getColumn k df - pure $ insertColumn k col d + col <- getColumn k df + pure $ insertColumn k col d -- | O(n) select columns by index range of column names. selectIntRange :: (Int, Int) -> DataFrame -> DataFrame @@ -181,13 +205,14 @@ selectBy :: (T.Text -> Bool) -> DataFrame -> DataFrame selectBy f df = select (L.filter f (columnNames df)) df --- | O(n) inverse of select --- --- > exclude ["Name"] df +{- | O(n) inverse of select + +> exclude ["Name"] df +-} exclude :: - [T.Text] -> - DataFrame -> - DataFrame + [T.Text] -> + DataFrame -> + DataFrame exclude cs df = - let keysToKeep = columnNames df L.\\ cs - in select keysToKeep df + let keysToKeep = columnNames df L.\\ cs + in select keysToKeep df diff --git a/src/DataFrame/Operations/Transformations.hs b/src/DataFrame/Operations/Transformations.hs --- a/src/DataFrame/Operations/Transformations.hs +++ b/src/DataFrame/Operations/Transformations.hs @@ -1,166 +1,186 @@ +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE FlexibleContexts #-} + module DataFrame.Operations.Transformations where import qualified Data.List as L -import qualified Data.Text as T import qualified Data.Map as M -import qualified Data.Vector.Generic as VG +import qualified Data.Text as T import qualified Data.Vector as V +import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU import Control.Exception (throw) -import DataFrame.Errors (DataFrameException(..), TypeErrorContext(..)) -import DataFrame.Internal.Column (Column(..), columnTypeString, imapColumn, ifoldrColumn, TypedColumn (TColumn), Columnable, mapColumn, unwrapTypedColumn) -import DataFrame.Internal.DataFrame (DataFrame(..), getColumn) +import Data.Maybe +import DataFrame.Errors (DataFrameException (..), TypeErrorContext (..)) +import DataFrame.Internal.Column (Column (..), Columnable, TypedColumn (TColumn), columnTypeString, ifoldrColumn, imapColumn, mapColumn, unwrapTypedColumn) +import DataFrame.Internal.DataFrame (DataFrame (..), getColumn) import DataFrame.Internal.Expression -import DataFrame.Internal.Row (mkRowFromArgs, Any, toAny) +import DataFrame.Internal.Row (Any, mkRowFromArgs, toAny) import DataFrame.Operations.Core -import Data.Maybe -import Type.Reflection (typeRep, typeOf, TypeRep) +import Type.Reflection (TypeRep, typeOf, typeRep) -- | O(k) Apply a function to a given column in a dataframe. apply :: - forall b c. - (Columnable b, Columnable c) => - -- | function to apply - (b -> c) -> - -- | Column name - T.Text -> - -- | DataFrame to apply operation to - DataFrame -> - DataFrame + forall b c. + (Columnable b, Columnable c) => + -- | function to apply + (b -> c) -> + -- | Column name + T.Text -> + -- | DataFrame to apply operation to + DataFrame -> + DataFrame apply f columnName d = case safeApply f columnName d of - Left exception -> throw exception - Right df -> df + Left exception -> throw exception + Right df -> df -- | O(k) Safe version of the apply function. Returns (instead of throwing) the error. safeApply :: - forall b c. - (Columnable b, Columnable c) => - -- | function to apply - (b -> c) -> - -- | Column name - T.Text -> - -- | DataFrame to apply operation to - DataFrame -> - Either DataFrameException DataFrame + forall b c. + (Columnable b, Columnable c) => + -- | function to apply + (b -> c) -> + -- | Column name + T.Text -> + -- | DataFrame to apply operation to + DataFrame -> + Either DataFrameException DataFrame safeApply f columnName d = case getColumn columnName d of - Nothing -> Left $ ColumnNotFoundException columnName "apply" (map fst $ M.toList $ columnIndices d) - Just column -> case mapColumn f column of - Nothing -> Left $ TypeMismatchException (MkTypeErrorContext - { userType = Right $ typeRep @b - , expectedType = Left (columnTypeString column) :: Either String (TypeRep ()) - , errorColumnName = Just (T.unpack columnName) - , callingFunctionName = Just "apply" - }) - Just column' -> Right $ insertColumn columnName column' d - --- | O(k) Apply a function to a combination of columns in a dataframe and --- add the result into `alias` column. -derive :: forall a . Columnable a => T.Text -> Expr a -> DataFrame -> DataFrame -derive name expr df = let - value = interpret @a df expr - in insertColumn name (unwrapTypedColumn value) df + Nothing -> Left $ ColumnNotFoundException columnName "apply" (map fst $ M.toList $ columnIndices d) + Just column -> case mapColumn f column of + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right $ typeRep @b + , expectedType = Left (columnTypeString column) :: Either String (TypeRep ()) + , errorColumnName = Just (T.unpack columnName) + , callingFunctionName = Just "apply" + } + ) + Just column' -> Right $ insertColumn columnName column' d +{- | O(k) Apply a function to a combination of columns in a dataframe and +add the result into `alias` column. +-} +derive :: forall a. (Columnable a) => T.Text -> Expr a -> DataFrame -> DataFrame +derive name expr df = + let + value = interpret @a df expr + in + insertColumn name (unwrapTypedColumn value) df -- | O(k * n) Apply a function to given column names in a dataframe. applyMany :: - (Columnable b, Columnable c) => - (b -> c) -> - [T.Text] -> - DataFrame -> - DataFrame + (Columnable b, Columnable c) => + (b -> c) -> + [T.Text] -> + DataFrame -> + DataFrame applyMany f names df = L.foldl' (flip (apply f)) df names -- | O(k) Convenience function that applies to an int column. applyInt :: - (Columnable b) => - -- | Column name - -- | function to apply - (Int -> b) -> - T.Text -> - -- | DataFrame to apply operation to - DataFrame -> - DataFrame + (Columnable b) => + {- | Column name + | function to apply + -} + (Int -> b) -> + T.Text -> + -- | DataFrame to apply operation to + DataFrame -> + DataFrame applyInt = apply -- | O(k) Convenience function that applies to an double column. applyDouble :: - (Columnable b) => - -- | Column name - -- | function to apply - (Double -> b) -> - T.Text -> - -- | DataFrame to apply operation to - DataFrame -> - DataFrame + (Columnable b) => + {- | Column name + | function to apply + -} + (Double -> b) -> + T.Text -> + -- | DataFrame to apply operation to + DataFrame -> + DataFrame applyDouble = apply --- | O(k * n) Apply a function to a column only if there is another column --- value that matches the given criterion. --- --- > applyWhere "Age" (<20) "Generation" (const "Gen-Z") +{- | O(k * n) Apply a function to a column only if there is another column +value that matches the given criterion. + +> applyWhere "Age" (<20) "Generation" (const "Gen-Z") +-} applyWhere :: - forall a b . - (Columnable a, Columnable b) => - (a -> Bool) -> -- Filter condition - T.Text -> -- Criterion Column - (b -> b) -> -- function to apply - T.Text -> -- Column name - DataFrame -> -- DataFrame to apply operation to - DataFrame + forall a b. + (Columnable a, Columnable b) => + (a -> Bool) -> -- Filter condition + T.Text -> -- Criterion Column + (b -> b) -> -- function to apply + T.Text -> -- Column name + DataFrame -> -- DataFrame to apply operation to + DataFrame applyWhere condition filterColumnName f columnName df = case getColumn filterColumnName df of - Nothing -> throw $ ColumnNotFoundException filterColumnName "applyWhere" (map fst $ M.toList $ columnIndices df) - Just column -> case ifoldrColumn (\i val acc -> if condition val then V.cons i acc else acc) V.empty column of - Nothing -> throw $ TypeMismatchException (MkTypeErrorContext - { userType = Right $ typeRep @a - , expectedType = Left (columnTypeString column) :: Either String (TypeRep ()) - , errorColumnName = Just (T.unpack columnName) - , callingFunctionName = Just "applyWhere" - }) - Just indexes -> if V.null indexes - then df - else L.foldl' (\d i -> applyAtIndex i f columnName d) df indexes + Nothing -> throw $ ColumnNotFoundException filterColumnName "applyWhere" (map fst $ M.toList $ columnIndices df) + Just column -> case ifoldrColumn (\i val acc -> if condition val then V.cons i acc else acc) V.empty column of + Nothing -> + throw $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right $ typeRep @a + , expectedType = Left (columnTypeString column) :: Either String (TypeRep ()) + , errorColumnName = Just (T.unpack columnName) + , callingFunctionName = Just "applyWhere" + } + ) + Just indexes -> + if V.null indexes + then df + else L.foldl' (\d i -> applyAtIndex i f columnName d) df indexes -- | O(k) Apply a function to the column at a given index. applyAtIndex :: - forall a. - (Columnable a) => - -- | Index - Int -> - -- | function to apply - (a -> a) -> - -- | Column name - T.Text -> - -- | DataFrame to apply operation to - DataFrame -> - DataFrame + forall a. + (Columnable a) => + -- | Index + Int -> + -- | function to apply + (a -> a) -> + -- | Column name + T.Text -> + -- | DataFrame to apply operation to + DataFrame -> + DataFrame applyAtIndex i f columnName df = case getColumn columnName df of - Nothing -> throw $ ColumnNotFoundException columnName "applyAtIndex" (map fst $ M.toList $ columnIndices df) - Just column -> case imapColumn (\index value -> if index == i then f value else value) column of - Nothing -> throw $ TypeMismatchException (MkTypeErrorContext - { userType = Right $ typeRep @a - , expectedType = Left (columnTypeString column) :: Either String (TypeRep ()) - , errorColumnName = Just (T.unpack columnName) - , callingFunctionName = Just "applyAtIndex" - }) - Just column' -> insertColumn columnName column' df + Nothing -> throw $ ColumnNotFoundException columnName "applyAtIndex" (map fst $ M.toList $ columnIndices df) + Just column -> case imapColumn (\index value -> if index == i then f value else value) column of + Nothing -> + throw $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right $ typeRep @a + , expectedType = Left (columnTypeString column) :: Either String (TypeRep ()) + , errorColumnName = Just (T.unpack columnName) + , callingFunctionName = Just "applyAtIndex" + } + ) + Just column' -> insertColumn columnName column' df +-- | Replace all instances of `Nothing` in a column with the given value. impute :: - forall b . - (Columnable b) => - T.Text -> - b -> - DataFrame -> - DataFrame + forall b. + (Columnable b) => + T.Text -> + b -> + DataFrame -> + DataFrame impute columnName value df = case getColumn columnName df of - Nothing -> throw $ ColumnNotFoundException columnName "impute" (map fst $ M.toList $ columnIndices df) - Just (OptionalColumn _) -> case safeApply (fromMaybe value) columnName df of - Left (TypeMismatchException context) -> throw $ TypeMismatchException (context { callingFunctionName = Just "impute" }) - Left exception -> throw exception - Right res -> res - _ -> error "Cannot impute to a non-Empty column" + Nothing -> throw $ ColumnNotFoundException columnName "impute" (map fst $ M.toList $ columnIndices df) + Just (OptionalColumn _) -> case safeApply (fromMaybe value) columnName df of + Left (TypeMismatchException context) -> throw $ TypeMismatchException (context{callingFunctionName = Just "impute"}) + Left exception -> throw exception + Right res -> res + _ -> error "Cannot impute to a non-Empty column" diff --git a/src/DataFrame/Operations/Typing.hs b/src/DataFrame/Operations/Typing.hs --- a/src/DataFrame/Operations/Typing.hs +++ b/src/DataFrame/Operations/Typing.hs @@ -3,6 +3,7 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} + module DataFrame.Operations.Typing where import qualified Data.Set as S @@ -10,64 +11,74 @@ import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU -import DataFrame.Internal.Column (Column(..)) -import DataFrame.Internal.DataFrame (DataFrame(..)) -import DataFrame.Internal.Parsing import Data.Either import Data.Maybe import Data.Time -import Data.Type.Equality (type (:~:)(Refl), TestEquality(..)) +import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) +import DataFrame.Internal.Column (Column (..)) +import DataFrame.Internal.DataFrame (DataFrame (..)) +import DataFrame.Internal.Parsing import Type.Reflection (typeRep) parseDefaults :: Bool -> DataFrame -> DataFrame -parseDefaults safeRead df = df {columns = V.map (parseDefault safeRead) (columns df)} +parseDefaults safeRead df = df{columns = V.map (parseDefault safeRead) (columns df)} parseDefault :: Bool -> Column -> Column -parseDefault safeRead (BoxedColumn (c :: V.Vector a)) = let - parseTimeOpt s = parseTimeM {- Accept leading/trailing whitespace -} True defaultTimeLocale "%Y-%m-%d" (T.unpack s) :: Maybe Day - unsafeParseTime s = parseTimeOrError {- Accept leading/trailing whitespace -} True defaultTimeLocale "%Y-%m-%d" (T.unpack s) :: Day - in case (typeRep @a) `testEquality` (typeRep @T.Text) of - Nothing -> case (typeRep @a) `testEquality` (typeRep @String) of - Just Refl -> let - emptyToNothing v = if isNullish (T.pack v) then Nothing else Just v - safeVector = V.map emptyToNothing c - hasNulls = V.foldl' (\acc v -> if isNothing v then acc || True else acc) False safeVector - in if safeRead && hasNulls then BoxedColumn safeVector else BoxedColumn c - Nothing -> BoxedColumn c - Just Refl -> - let example = T.strip (V.head c) - emptyToNothing v = if isNullish v then Nothing else Just v - in case readInt example of - Just _ -> - let safeVector = V.map ((=<<) readInt . emptyToNothing) c - hasNulls = V.elem Nothing safeVector - in if safeRead && hasNulls then BoxedColumn safeVector else UnboxedColumn (VU.generate (V.length c) (fromMaybe 0 . (safeVector V.!))) - Nothing -> case readDouble example of - Just _ -> - let safeVector = V.map ((=<<) readDouble . emptyToNothing) c - hasNulls = V.elem Nothing safeVector - in if safeRead && hasNulls then BoxedColumn safeVector else UnboxedColumn (VU.generate (V.length c) (fromMaybe 0 . (safeVector V.!))) - Nothing -> case parseTimeOpt example of - Just d -> let - -- failed parse should be Either, nullish should be Maybe - emptyToNothing' v = if isNullish v then Left v else Right v - parseTimeEither v = case parseTimeOpt v of - Just v' -> Right v' - Nothing -> Left v - safeVector = V.map ((=<<) parseTimeEither . emptyToNothing') c - toMaybe (Left _) = Nothing - toMaybe (Right value) = Just value - lefts = V.filter isLeft safeVector - onlyNulls = (not (V.null lefts) && V.all (isNullish . fromLeft "non-null") lefts) - in if safeRead - then if onlyNulls - then BoxedColumn (V.map toMaybe safeVector) - else if V.any isLeft safeVector - then BoxedColumn safeVector - else BoxedColumn (V.map unsafeParseTime c) - else BoxedColumn (V.map unsafeParseTime c) - Nothing -> let +parseDefault safeRead (BoxedColumn (c :: V.Vector a)) = + let + parseTimeOpt s = parseTimeM {- Accept leading/trailing whitespace -} True defaultTimeLocale "%Y-%m-%d" (T.unpack s) :: Maybe Day + unsafeParseTime s = parseTimeOrError {- Accept leading/trailing whitespace -} True defaultTimeLocale "%Y-%m-%d" (T.unpack s) :: Day + in + case (typeRep @a) `testEquality` (typeRep @T.Text) of + Nothing -> case (typeRep @a) `testEquality` (typeRep @String) of + Just Refl -> + let + emptyToNothing v = if isNullish (T.pack v) then Nothing else Just v safeVector = V.map emptyToNothing c - hasNulls = V.any isNullish c - in if safeRead && hasNulls then BoxedColumn safeVector else BoxedColumn c + hasNulls = V.foldl' (\acc v -> if isNothing v then acc || True else acc) False safeVector + in + if safeRead && hasNulls then BoxedColumn safeVector else BoxedColumn c + Nothing -> BoxedColumn c + Just Refl -> + let example = T.strip (V.head c) + emptyToNothing v = if isNullish v then Nothing else Just v + in case readInt example of + Just _ -> + let safeVector = V.map ((=<<) readInt . emptyToNothing) c + hasNulls = V.elem Nothing safeVector + in if safeRead && hasNulls then BoxedColumn safeVector else UnboxedColumn (VU.generate (V.length c) (fromMaybe 0 . (safeVector V.!))) + Nothing -> case readDouble example of + Just _ -> + let safeVector = V.map ((=<<) readDouble . emptyToNothing) c + hasNulls = V.elem Nothing safeVector + in if safeRead && hasNulls then BoxedColumn safeVector else UnboxedColumn (VU.generate (V.length c) (fromMaybe 0 . (safeVector V.!))) + Nothing -> case parseTimeOpt example of + Just d -> + let + -- failed parse should be Either, nullish should be Maybe + emptyToNothing' v = if isNullish v then Left v else Right v + parseTimeEither v = case parseTimeOpt v of + Just v' -> Right v' + Nothing -> Left v + safeVector = V.map ((=<<) parseTimeEither . emptyToNothing') c + toMaybe (Left _) = Nothing + toMaybe (Right value) = Just value + lefts = V.filter isLeft safeVector + onlyNulls = (not (V.null lefts) && V.all (isNullish . fromLeft "non-null") lefts) + in + if safeRead + then + if onlyNulls + then BoxedColumn (V.map toMaybe safeVector) + else + if V.any isLeft safeVector + then BoxedColumn safeVector + else BoxedColumn (V.map unsafeParseTime c) + else BoxedColumn (V.map unsafeParseTime c) + Nothing -> + let + safeVector = V.map emptyToNothing c + hasNulls = V.any isNullish c + in + if safeRead && hasNulls then BoxedColumn safeVector else BoxedColumn c parseDefault safeRead column = column diff --git a/tests/Assertions.hs b/tests/Assertions.hs --- a/tests/Assertions.hs +++ b/tests/Assertions.hs @@ -1,4 +1,5 @@ {-# LANGUAGE ScopedTypeVariables #-} + module Assertions where import qualified Data.List as L @@ -7,20 +8,28 @@ import Test.HUnit -- Adapted from: https://github.com/BartMassey/chunk/blob/1ee4bd6545e0db6b8b5f4935d97e7606708eacc9/hunit.hs#L29 -assertExpectException :: String -> String -> - IO a -> Assertion +assertExpectException :: + String -> + String -> + IO a -> + Assertion assertExpectException preface expected action = do - r <- catch - (action >> (return . Just) "no exception thrown") - (\(e::SomeException) -> - return (checkForExpectedException e)) - case r of - Nothing -> return () - Just msg -> assertFailure $ preface ++ ": " ++ msg + r <- + catch + (action >> (return . Just) "no exception thrown") + ( \(e :: SomeException) -> + return (checkForExpectedException e) + ) + case r of + Nothing -> return () + Just msg -> assertFailure $ preface ++ ": " ++ msg where checkForExpectedException :: SomeException -> Maybe String checkForExpectedException e | expected `L.isInfixOf` show e = Nothing | otherwise = - Just $ "wrong exception detail, expected " ++ - expected ++ ", got: " ++ show e + Just $ + "wrong exception detail, expected " + ++ expected + ++ ", got: " + ++ show e diff --git a/tests/Functions.hs b/tests/Functions.hs --- a/tests/Functions.hs +++ b/tests/Functions.hs @@ -1,42 +1,62 @@ {-# LANGUAGE OverloadedStrings #-} + module Functions where import Control.Exception -import Test.HUnit -import DataFrame.Functions (sanitize) import qualified Data.Text as T +import DataFrame.Functions (sanitize) +import Test.HUnit -- Test cases for the sanitize function sanitizeIdentifiers :: Test -sanitizeIdentifiers = TestList - [ TestCase $ assertEqual "Reserved word 'Data' should become '_data_'" - "_data_" (sanitize "Data") - - , TestCase $ assertEqual "Spaces should become underscores" - "my_data" (sanitize "My Data") - - , TestCase $ assertEqual "Punctuation and parentheses should be handled" - "distance_km_h" (sanitize "Distance (km/h)") - - , TestCase $ assertEqual "Leading digit should be wrapped" - "_0_age_" (sanitize "0 Age") - - , TestCase $ assertEqual "Valid camelCase should be unchanged" - "camelCaseStr" (sanitize "camelCaseStr") - - , TestCase $ assertEqual "Valid camelCase with invalid characters mixed in" - "camelcase_str" (sanitize "camelCase$Str") - - , TestCase $ assertEqual "Valid snake_case should remain unchanged" - "snake_case_str" (sanitize "snake_case_str") - - , TestCase $ assertEqual "Leading digit with snake_case should be wrapped" - "_12_snake_case_" (sanitize "12_snake_case") - - , TestCase $ assertEqual "All symbols should become underscores" - "_____" (sanitize "***") - ] +sanitizeIdentifiers = + TestList + [ TestCase $ + assertEqual + "Reserved word 'Data' should become '_data_'" + "_data_" + (sanitize "Data") + , TestCase $ + assertEqual + "Spaces should become underscores" + "my_data" + (sanitize "My Data") + , TestCase $ + assertEqual + "Punctuation and parentheses should be handled" + "distance_km_h" + (sanitize "Distance (km/h)") + , TestCase $ + assertEqual + "Leading digit should be wrapped" + "_0_age_" + (sanitize "0 Age") + , TestCase $ + assertEqual + "Valid camelCase should be unchanged" + "camelCaseStr" + (sanitize "camelCaseStr") + , TestCase $ + assertEqual + "Valid camelCase with invalid characters mixed in" + "camelcase_str" + (sanitize "camelCase$Str") + , TestCase $ + assertEqual + "Valid snake_case should remain unchanged" + "snake_case_str" + (sanitize "snake_case_str") + , TestCase $ + assertEqual + "Leading digit with snake_case should be wrapped" + "_12_snake_case_" + (sanitize "12_snake_case") + , TestCase $ + assertEqual + "All symbols should become underscores" + "_____" + (sanitize "***") + ] tests :: [Test] tests = pure $ TestLabel "sanitizeIdentifiers" sanitizeIdentifiers - diff --git a/tests/Main.hs b/tests/Main.hs --- a/tests/Main.hs +++ b/tests/Main.hs @@ -1,13 +1,18 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} + module Main where -import qualified DataFrame as D -import qualified DataFrame as DI import qualified Data.List as L import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU +import qualified DataFrame as D +import qualified DataFrame.Internal.Column as D +import qualified DataFrame.Internal.Column as DI +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Internal.DataFrame as DI +import qualified DataFrame.Operations.Typing as D import qualified System.Exit as Exit import Control.Exception @@ -23,13 +28,16 @@ import qualified Operations.GroupBy import qualified Operations.InsertColumn import qualified Operations.Sort +import qualified Operations.Statistics import qualified Operations.Take import qualified Parquet testData :: D.DataFrame -testData = D.fromNamedColumns [ ("test1", DI.fromList ([1..26] :: [Int])) - , ("test2", DI.fromList ['a'..'z']) - ] +testData = + D.fromNamedColumns + [ ("test1", DI.fromList ([1 .. 26] :: [Int])) + , ("test2", DI.fromList ['a' .. 'z']) + ] -- Dimensions correctDimensions :: Test @@ -39,48 +47,58 @@ emptyDataframeDimensions = TestCase (assertEqual "should be (0, 0)" (0, 0) (D.dimensions D.empty)) dimensionsTest :: [Test] -dimensionsTest = [ TestLabel "dimensions_correctDimensions" correctDimensions - , TestLabel "dimensions_emptyDataframeDimensions" emptyDataframeDimensions - ] +dimensionsTest = + [ TestLabel "dimensions_correctDimensions" correctDimensions + , TestLabel "dimensions_emptyDataframeDimensions" emptyDataframeDimensions + ] -- parsing. parseDate :: Test -parseDate = let - expected = DI.BoxedColumn (V.fromList [fromGregorian 2020 02 14, fromGregorian 2021 02 14, fromGregorian 2022 02 14]) - actual = D.parseDefault True (DI.fromVector (V.fromList ["2020-02-14" :: T.Text, "2021-02-14", "2022-02-14"])) - in TestCase (assertEqual "Correctly parses gregorian date" expected actual) +parseDate = + let + expected = DI.BoxedColumn (V.fromList [fromGregorian 2020 02 14, fromGregorian 2021 02 14, fromGregorian 2022 02 14]) + actual = D.parseDefault True (DI.fromVector (V.fromList ["2020-02-14" :: T.Text, "2021-02-14", "2022-02-14"])) + in + TestCase (assertEqual "Correctly parses gregorian date" expected actual) incompleteDataParseEither :: Test -incompleteDataParseEither = let - expected = DI.BoxedColumn (V.fromList [Right $ fromGregorian 2020 02 14, Left ("2021-02-" :: T.Text), Right $ fromGregorian 2022 02 14]) - actual = D.parseDefault True (DI.fromVector (V.fromList ["2020-02-14" :: T.Text, "2021-02-", "2022-02-14"])) - in TestCase (assertEqual "Parses Either for gregorian date" expected actual) +incompleteDataParseEither = + let + expected = DI.BoxedColumn (V.fromList [Right $ fromGregorian 2020 02 14, Left ("2021-02-" :: T.Text), Right $ fromGregorian 2022 02 14]) + actual = D.parseDefault True (DI.fromVector (V.fromList ["2020-02-14" :: T.Text, "2021-02-", "2022-02-14"])) + in + TestCase (assertEqual "Parses Either for gregorian date" expected actual) incompleteDataParseMaybe :: Test -incompleteDataParseMaybe = let - expected = DI.BoxedColumn (V.fromList [Just $ fromGregorian 2020 02 14, Nothing, Just $ fromGregorian 2022 02 14]) - actual = D.parseDefault True (DI.fromVector (V.fromList ["2020-02-14" :: T.Text, "", "2022-02-14"])) - in TestCase (assertEqual "Parses Maybe for gregorian date with null/empty" expected actual) +incompleteDataParseMaybe = + let + expected = DI.BoxedColumn (V.fromList [Just $ fromGregorian 2020 02 14, Nothing, Just $ fromGregorian 2022 02 14]) + actual = D.parseDefault True (DI.fromVector (V.fromList ["2020-02-14" :: T.Text, "", "2022-02-14"])) + in + TestCase (assertEqual "Parses Maybe for gregorian date with null/empty" expected actual) parseTests :: [Test] -parseTests = [ - TestLabel "parseDate" parseDate, - TestLabel "incompleteDataParseMaybe" incompleteDataParseMaybe, - TestLabel "incompleteDataParseEither" incompleteDataParseEither - ] +parseTests = + [ TestLabel "parseDate" parseDate + , TestLabel "incompleteDataParseMaybe" incompleteDataParseMaybe + , TestLabel "incompleteDataParseEither" incompleteDataParseEither + ] tests :: Test -tests = TestList $ dimensionsTest - ++ Operations.Apply.tests - ++ Operations.Derive.tests - ++ Operations.Filter.tests - ++ Operations.GroupBy.tests - ++ Operations.InsertColumn.tests - ++ Operations.Sort.tests - ++ Operations.Take.tests - ++ Functions.tests - ++ Parquet.tests - ++ parseTests +tests = + TestList $ + dimensionsTest + ++ Operations.Apply.tests + ++ Operations.Derive.tests + ++ Operations.Filter.tests + ++ Operations.GroupBy.tests + ++ Operations.InsertColumn.tests + ++ Operations.Sort.tests + ++ Operations.Statistics.tests + ++ Operations.Take.tests + ++ Functions.tests + ++ Parquet.tests + ++ parseTests main :: IO () main = do diff --git a/tests/Operations/Apply.hs b/tests/Operations/Apply.hs --- a/tests/Operations/Apply.hs +++ b/tests/Operations/Apply.hs @@ -1,122 +1,203 @@ -{-# LANGUAGE TypeApplications #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} + module Operations.Apply where -import qualified DataFrame as D -import qualified DataFrame as DI -import qualified DataFrame as DE import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU +import qualified DataFrame as D +import qualified DataFrame as DE +import qualified DataFrame.Internal.Column as D +import qualified DataFrame.Internal.Column as DI +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Internal.DataFrame as DI import Assertions import Test.HUnit import Type.Reflection (typeRep) values :: [(T.Text, DI.Column)] -values = [ ("test1", DI.fromList ([1..26] :: [Int])) - , ("test2", DI.fromList (map show ['a'..'z'])) - , ("test3", DI.fromList ([1..26] :: [Int])) - , ("test4", DI.fromList ['a'..'z']) - , ("test5", DI.fromList ([1..26] :: [Int])) - , ("test6", DI.fromList ['a'..'z']) - , ("test7", DI.fromList ([1..26] :: [Int])) - , ("test8", DI.fromList ['a'..'z']) - ] +values = + [ ("test1", DI.fromList ([1 .. 26] :: [Int])) + , ("test2", DI.fromList (map show ['a' .. 'z'])) + , ("test3", DI.fromList ([1 .. 26] :: [Int])) + , ("test4", DI.fromList ['a' .. 'z']) + , ("test5", DI.fromList ([1 .. 26] :: [Int])) + , ("test6", DI.fromList ['a' .. 'z']) + , ("test7", DI.fromList ([1 .. 26] :: [Int])) + , ("test8", DI.fromList ['a' .. 'z']) + ] testData :: D.DataFrame testData = D.fromNamedColumns values applyBoxedToUnboxed :: Test -applyBoxedToUnboxed = TestCase (assertEqual "Boxed apply unboxed when result is unboxed" - (Just $ DI.UnboxedColumn (VU.fromList (replicate 26 (1 :: Int)))) - (DI.getColumn "test2" $ D.apply @String (const (1::Int)) "test2" testData)) +applyBoxedToUnboxed = + TestCase + ( assertEqual + "Boxed apply unboxed when result is unboxed" + (Just $ DI.UnboxedColumn (VU.fromList (replicate 26 (1 :: Int)))) + (DI.getColumn "test2" $ D.apply @String (const (1 :: Int)) "test2" testData) + ) applyBoxedToBoxed :: Test -applyBoxedToBoxed = TestCase (assertEqual "Boxed apply remains in boxed vector" - (Just $ DI.BoxedColumn (V.fromList (replicate 26 (1 :: Integer)))) - (DI.getColumn "test2" $ D.apply @String (const (1::Integer)) "test2" testData)) +applyBoxedToBoxed = + TestCase + ( assertEqual + "Boxed apply remains in boxed vector" + (Just $ DI.BoxedColumn (V.fromList (replicate 26 (1 :: Integer)))) + (DI.getColumn "test2" $ D.apply @String (const (1 :: Integer)) "test2" testData) + ) applyWrongType :: Test -applyWrongType = TestCase (assertExpectException "[Error Case]" - (DE.typeMismatchError (show $ typeRep @Char) (show $ typeRep @[Char])) - (print $ DI.getColumn "test2" $ D.apply @Char (const (1::Int)) "test2" testData)) +applyWrongType = + TestCase + ( assertExpectException + "[Error Case]" + (DE.typeMismatchError (show $ typeRep @Char) (show $ typeRep @[Char])) + (print $ DI.getColumn "test2" $ D.apply @Char (const (1 :: Int)) "test2" testData) + ) applyUnknownColumn :: Test -applyUnknownColumn = TestCase (assertExpectException "[Error Case]" - (DE.columnNotFound "test9" "apply" (D.columnNames testData)) - (print $ D.apply @[Char] (const (1::Int)) "test9" testData)) +applyUnknownColumn = + TestCase + ( assertExpectException + "[Error Case]" + (DE.columnNotFound "test9" "apply" (D.columnNames testData)) + (print $ D.apply @[Char] (const (1 :: Int)) "test9" testData) + ) applyManyOnlyGivenFields :: Test -applyManyOnlyGivenFields = TestCase (assertEqual "Applies function to many fields" - (D.fromNamedColumns (map (, D.fromList $ replicate 26 (1 :: Integer)) ["test4", "test6"] ++ - -- All other fields should have their original values. - filter (\(name, col) -> name /= "test4" && name /= "test6") values)) - (D.applyMany @Char (const (1::Integer)) - ["test4", "test6"] testData)) +applyManyOnlyGivenFields = + TestCase + ( assertEqual + "Applies function to many fields" + ( D.fromNamedColumns + ( map (,D.fromList $ replicate 26 (1 :: Integer)) ["test4", "test6"] + ++ + -- All other fields should have their original values. + filter (\(name, col) -> name /= "test4" && name /= "test6") values + ) + ) + ( D.applyMany @Char + (const (1 :: Integer)) + ["test4", "test6"] + testData + ) + ) applyManyBoxedToBoxed :: Test -applyManyBoxedToBoxed = TestCase (assertEqual "Applies function to many fields" - (D.fromNamedColumns (map (, D.fromList $ replicate 26 (1 :: Integer)) ["test4", "test6", "test8"])) - (D.select ["test4", "test6", "test8"] $ D.applyMany @Char (const (1::Integer)) - ["test4", "test6", "test8"] testData)) +applyManyBoxedToBoxed = + TestCase + ( assertEqual + "Applies function to many fields" + (D.fromNamedColumns (map (,D.fromList $ replicate 26 (1 :: Integer)) ["test4", "test6", "test8"])) + ( D.select ["test4", "test6", "test8"] $ + D.applyMany @Char + (const (1 :: Integer)) + ["test4", "test6", "test8"] + testData + ) + ) applyManyBoxedToUnboxed :: Test -applyManyBoxedToUnboxed = TestCase (assertEqual "Unboxes fields when necessary" - (D.fromNamedColumns (map (, D.fromList $ replicate 26 (1 :: Int)) ["test4", "test6", "test8"])) - (D.select ["test4", "test6", "test8"] $ D.applyMany @Char (const (1::Int)) - ["test4", "test6", "test8"] testData)) +applyManyBoxedToUnboxed = + TestCase + ( assertEqual + "Unboxes fields when necessary" + (D.fromNamedColumns (map (,D.fromList $ replicate 26 (1 :: Int)) ["test4", "test6", "test8"])) + ( D.select ["test4", "test6", "test8"] $ + D.applyMany @Char + (const (1 :: Int)) + ["test4", "test6", "test8"] + testData + ) + ) applyManyColumnNotFound :: Test -applyManyColumnNotFound = TestCase (assertExpectException "[Error Case]" - (DE.columnNotFound "test0" "apply" (D.columnNames testData)) - (print $ D.applyMany @Char (const (1::Integer)) - ["test0", "test6", "test8"] testData)) +applyManyColumnNotFound = + TestCase + ( assertExpectException + "[Error Case]" + (DE.columnNotFound "test0" "apply" (D.columnNames testData)) + ( print $ + D.applyMany @Char + (const (1 :: Integer)) + ["test0", "test6", "test8"] + testData + ) + ) applyManyWrongType :: Test -applyManyWrongType = TestCase (assertExpectException "[Error Case]" - (DE.typeMismatchError (show $ typeRep @Char) (show $ typeRep @[Char])) - (print $ DI.getColumn "test2" $ D.applyMany @Char (const (1::Int)) ["test2"] testData)) +applyManyWrongType = + TestCase + ( assertExpectException + "[Error Case]" + (DE.typeMismatchError (show $ typeRep @Char) (show $ typeRep @[Char])) + (print $ DI.getColumn "test2" $ D.applyMany @Char (const (1 :: Int)) ["test2"] testData) + ) applyWhereWrongConditionType :: Test -applyWhereWrongConditionType = TestCase (assertExpectException "[Error Case]" - (DE.typeMismatchError (show $ typeRep @Integer) (show $ typeRep @Int)) - (print $ D.applyWhere (even @Integer) "test1" ((+1) :: Int -> Int) "test5" testData)) +applyWhereWrongConditionType = + TestCase + ( assertExpectException + "[Error Case]" + (DE.typeMismatchError (show $ typeRep @Integer) (show $ typeRep @Int)) + (print $ D.applyWhere (even @Integer) "test1" ((+ 1) :: Int -> Int) "test5" testData) + ) applyWhereWrongTargetType :: Test -applyWhereWrongTargetType = TestCase (assertExpectException "[Error Case]" - (DE.typeMismatchError (show $ typeRep @Float) (show $ typeRep @Int)) - (print $ D.applyWhere (even @Int) "test1" ((+1) :: Float -> Float) "test5" testData)) +applyWhereWrongTargetType = + TestCase + ( assertExpectException + "[Error Case]" + (DE.typeMismatchError (show $ typeRep @Float) (show $ typeRep @Int)) + (print $ D.applyWhere (even @Int) "test1" ((+ 1) :: Float -> Float) "test5" testData) + ) applyWhereConditionColumnNotFound :: Test -applyWhereConditionColumnNotFound = TestCase (assertExpectException "[Error Case]" - (DE.columnNotFound "test0" "applyWhere" (D.columnNames testData)) - (print $ D.applyWhere (even @Int) "test0" ((+1) :: Int -> Int) "test5" testData)) +applyWhereConditionColumnNotFound = + TestCase + ( assertExpectException + "[Error Case]" + (DE.columnNotFound "test0" "applyWhere" (D.columnNames testData)) + (print $ D.applyWhere (even @Int) "test0" ((+ 1) :: Int -> Int) "test5" testData) + ) applyWhereTargetColumnNotFound :: Test -applyWhereTargetColumnNotFound = TestCase (assertExpectException "[Error Case]" - (DE.columnNotFound "test0" "applyAtIndex" (D.columnNames testData)) - (print $ D.applyWhere (even @Int) "test1" ((+1) :: Int -> Int) "test0" testData)) +applyWhereTargetColumnNotFound = + TestCase + ( assertExpectException + "[Error Case]" + (DE.columnNotFound "test0" "applyAtIndex" (D.columnNames testData)) + (print $ D.applyWhere (even @Int) "test1" ((+ 1) :: Int -> Int) "test0" testData) + ) applyWhereWAI :: Test -applyWhereWAI = TestCase (assertEqual "applyWhere works as intended" - (Just $ DI.UnboxedColumn (VU.fromList (zipWith ($) (cycle [id, (+1)]) [(1 :: Int)..26]))) - (D.getColumn "test5" $ D.applyWhere (even @Int) "test1" ((+1) :: Int -> Int) "test5" testData)) +applyWhereWAI = + TestCase + ( assertEqual + "applyWhere works as intended" + (Just $ DI.UnboxedColumn (VU.fromList (zipWith ($) (cycle [id, (+ 1)]) [(1 :: Int) .. 26]))) + (D.getColumn "test5" $ D.applyWhere (even @Int) "test1" ((+ 1) :: Int -> Int) "test5" testData) + ) tests :: [Test] -tests = [ TestLabel "applyBoxedToUnboxed" applyBoxedToUnboxed - , TestLabel "applyWrongType" applyWrongType - , TestLabel "applyUnknownColumn" applyUnknownColumn - , TestLabel "applyBoxedToBoxed" applyBoxedToBoxed - , TestLabel "applyManyBoxedToBoxed" applyManyBoxedToBoxed - , TestLabel "applyManyOnlyGivenFields" applyManyOnlyGivenFields - , TestLabel "applyManyBoxedToUnboxed" applyManyBoxedToUnboxed - , TestLabel "applyManyColumnNotFound" applyManyColumnNotFound - , TestLabel "applyManyWrongType" applyManyWrongType - , TestLabel "applyWhereWrongConditionType" applyWhereWrongConditionType - , TestLabel "applyWhereWrongTargetType" applyWhereWrongTargetType - , TestLabel "applyWhereConditionColumnNotFound" applyWhereConditionColumnNotFound - , TestLabel "applyWhereTargetColumnNotFound" applyWhereTargetColumnNotFound - , TestLabel "applyWhereWAI" applyWhereWAI - ] +tests = + [ TestLabel "applyBoxedToUnboxed" applyBoxedToUnboxed + , TestLabel "applyWrongType" applyWrongType + , TestLabel "applyUnknownColumn" applyUnknownColumn + , TestLabel "applyBoxedToBoxed" applyBoxedToBoxed + , TestLabel "applyManyBoxedToBoxed" applyManyBoxedToBoxed + , TestLabel "applyManyOnlyGivenFields" applyManyOnlyGivenFields + , TestLabel "applyManyBoxedToUnboxed" applyManyBoxedToUnboxed + , TestLabel "applyManyColumnNotFound" applyManyColumnNotFound + , TestLabel "applyManyWrongType" applyManyWrongType + , TestLabel "applyWhereWrongConditionType" applyWhereWrongConditionType + , TestLabel "applyWhereWrongTargetType" applyWhereWrongTargetType + , TestLabel "applyWhereConditionColumnNotFound" applyWhereConditionColumnNotFound + , TestLabel "applyWhereTargetColumnNotFound" applyWhereTargetColumnNotFound + , TestLabel "applyWhereWAI" applyWhereWAI + ] diff --git a/tests/Operations/Derive.hs b/tests/Operations/Derive.hs --- a/tests/Operations/Derive.hs +++ b/tests/Operations/Derive.hs @@ -1,36 +1,48 @@ -{-# LANGUAGE TypeApplications #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + module Operations.Derive where -import qualified DataFrame as D -import qualified DataFrame.Functions as F -import qualified DataFrame as DI -import qualified DataFrame as DE import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU +import qualified DataFrame as D +import qualified DataFrame.Functions as F +import qualified DataFrame.Internal.Column as D +import qualified DataFrame.Internal.Column as DI +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Internal.DataFrame as DI import Assertions import Test.HUnit import Type.Reflection (typeRep) values :: [(T.Text, DI.Column)] -values = [ ("test1", DI.fromList ([1..26] :: [Int])) - , ("test2", DI.fromList (map show ['a'..'z'])) - , ("test3", DI.fromList ['a'..'z']) - ] +values = + [ ("test1", DI.fromList ([1 .. 26] :: [Int])) + , ("test2", DI.fromList (map show ['a' .. 'z'])) + , ("test3", DI.fromList ['a' .. 'z']) + ] testData :: D.DataFrame testData = D.fromNamedColumns values deriveWAI :: Test -deriveWAI = TestCase (assertEqual "derive works with column expression" - (Just $ DI.BoxedColumn (V.fromList (zipWith (\n c -> show n ++ [c]) [1..26] ['a'..'z']))) - (DI.getColumn "test4" $ D.derive "test4" ( - F.lift2 (++) (F.lift show (F.col @Int "test1")) (F.lift (: ([] :: [Char])) (F.col @Char "test3")) - ) testData)) +deriveWAI = + TestCase + ( assertEqual + "derive works with column expression" + (Just $ DI.BoxedColumn (V.fromList (zipWith (\n c -> show n ++ [c]) [1 .. 26] ['a' .. 'z']))) + ( DI.getColumn "test4" $ + D.derive + "test4" + (F.lift2 (++) (F.lift show (F.col @Int "test1")) (F.lift (: ([] :: [Char])) (F.col @Char "test3"))) + testData + ) + ) tests :: [Test] -tests = [ TestLabel "deriveWAI" deriveWAI - ] +tests = + [ TestLabel "deriveWAI" deriveWAI + ] diff --git a/tests/Operations/Filter.hs b/tests/Operations/Filter.hs --- a/tests/Operations/Filter.hs +++ b/tests/Operations/Filter.hs @@ -1,73 +1,104 @@ -{-# LANGUAGE TypeApplications #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + module Operations.Filter where -import qualified DataFrame as D -import qualified DataFrame as DI -import qualified DataFrame as DE import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU +import qualified DataFrame as D +import qualified DataFrame as DE +import qualified DataFrame as DI import Assertions import Test.HUnit import Type.Reflection (typeRep) values :: [(T.Text, DI.Column)] -values = [ ("test1", DI.fromList ([1..26] :: [Int])) - , ("test2", DI.fromList (map show ['a'..'z'])) - , ("test3", DI.fromList ([1..26] :: [Int])) - , ("test4", DI.fromList ['a'..'z']) - , ("test5", DI.fromList ([1..26] :: [Int])) - , ("test6", DI.fromList ['a'..'z']) - , ("test7", DI.fromList ([1..26] :: [Int])) - , ("test8", DI.fromList ['a'..'z']) - ] +values = + [ ("test1", DI.fromList ([1 .. 26] :: [Int])) + , ("test2", DI.fromList (map show ['a' .. 'z'])) + , ("test3", DI.fromList ([1 .. 26] :: [Int])) + , ("test4", DI.fromList ['a' .. 'z']) + , ("test5", DI.fromList ([1 .. 26] :: [Int])) + , ("test6", DI.fromList ['a' .. 'z']) + , ("test7", DI.fromList ([1 .. 26] :: [Int])) + , ("test8", DI.fromList ['a' .. 'z']) + ] testData :: D.DataFrame testData = D.fromNamedColumns values filterColumnDoesNotExist :: Test -filterColumnDoesNotExist = TestCase (assertExpectException "[Error Case]" - (DE.columnNotFound "test0" "filter" (D.columnNames testData)) - (print $ D.filter @Int "test0" even testData)) +filterColumnDoesNotExist = + TestCase + ( assertExpectException + "[Error Case]" + (DE.columnNotFound "test0" "filter" (D.columnNames testData)) + (print $ D.filter @Int "test0" even testData) + ) filterColumnWrongType :: Test -filterColumnWrongType = TestCase (assertExpectException "[Error Case]" - (DE.typeMismatchError (show $ typeRep @Integer) (show $ typeRep @Int)) - (print $ D.filter @Integer "test1" even testData)) +filterColumnWrongType = + TestCase + ( assertExpectException + "[Error Case]" + (DE.typeMismatchError (show $ typeRep @Integer) (show $ typeRep @Int)) + (print $ D.filter @Integer "test1" even testData) + ) filterByColumnDoesNotExist :: Test -filterByColumnDoesNotExist = TestCase (assertExpectException "[Error Case]" - (DE.columnNotFound "test0" "filter" (D.columnNames testData)) - (print $ D.filterBy @Int even "test0" testData)) +filterByColumnDoesNotExist = + TestCase + ( assertExpectException + "[Error Case]" + (DE.columnNotFound "test0" "filter" (D.columnNames testData)) + (print $ D.filterBy @Int even "test0" testData) + ) filterByColumnWrongType :: Test -filterByColumnWrongType = TestCase (assertExpectException "[Error Case]" - (DE.typeMismatchError (show $ typeRep @Integer) (show $ typeRep @Int)) - (print $ D.filterBy @Integer even "test1" testData)) +filterByColumnWrongType = + TestCase + ( assertExpectException + "[Error Case]" + (DE.typeMismatchError (show $ typeRep @Integer) (show $ typeRep @Int)) + (print $ D.filterBy @Integer even "test1" testData) + ) filterColumnInexistentValues :: Test -filterColumnInexistentValues = TestCase (assertEqual "Non existent filter value returns no rows" - (0, 8) - (D.dimensions $ D.filter @Int "test1" (<0) testData)) +filterColumnInexistentValues = + TestCase + ( assertEqual + "Non existent filter value returns no rows" + (0, 8) + (D.dimensions $ D.filter @Int "test1" (< 0) testData) + ) filterColumnAllValues :: Test -filterColumnAllValues = TestCase (assertEqual "Filters all columns" - (26, 8) - (D.dimensions $ D.filter @Int "test1" (const True) testData)) +filterColumnAllValues = + TestCase + ( assertEqual + "Filters all columns" + (26, 8) + (D.dimensions $ D.filter @Int "test1" (const True) testData) + ) filterJustWAI :: Test -filterJustWAI = TestCase (assertEqual "Filters out Nothing and unwraps Maybe" - (D.fromNamedColumns [("test", D.fromList $ replicate 5 (1 :: Int))]) - (D.filterJust "test" (D.fromNamedColumns [("test", D.fromList$ take 10 $ cycle [Just (1 :: Int), Nothing])]))) +filterJustWAI = + TestCase + ( assertEqual + "Filters out Nothing and unwraps Maybe" + (D.fromNamedColumns [("test", D.fromList $ replicate 5 (1 :: Int))]) + (D.filterJust "test" (D.fromNamedColumns [("test", D.fromList $ take 10 $ cycle [Just (1 :: Int), Nothing])])) + ) tests :: [Test] -tests = [ TestLabel "filterColumnDoesNotExist" filterColumnDoesNotExist - , TestLabel "filterColumnWrongType" filterColumnWrongType - , TestLabel "filterByColumnDoesNotExist" filterByColumnDoesNotExist - , TestLabel "filterByColumnWrongType" filterByColumnWrongType - , TestLabel "filterColumnInexistentValues" filterColumnInexistentValues - , TestLabel "filterColumnAllValues" filterColumnAllValues - , TestLabel "filterJustWAI" filterJustWAI - ] +tests = + [ TestLabel "filterColumnDoesNotExist" filterColumnDoesNotExist + , TestLabel "filterColumnWrongType" filterColumnWrongType + , TestLabel "filterByColumnDoesNotExist" filterByColumnDoesNotExist + , TestLabel "filterByColumnWrongType" filterByColumnWrongType + , TestLabel "filterColumnInexistentValues" filterColumnInexistentValues + , TestLabel "filterColumnAllValues" filterColumnAllValues + , TestLabel "filterJustWAI" filterJustWAI + ] diff --git a/tests/Operations/GroupBy.hs b/tests/Operations/GroupBy.hs --- a/tests/Operations/GroupBy.hs +++ b/tests/Operations/GroupBy.hs @@ -1,46 +1,62 @@ {-# LANGUAGE OverloadedStrings #-} + module Operations.GroupBy where -import qualified DataFrame as D -import qualified DataFrame as DI -import qualified DataFrame as DE import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU +import qualified DataFrame as D +import qualified DataFrame.Internal.Column as D +import qualified DataFrame.Internal.Column as DI +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Internal.DataFrame as DI import Assertions import Test.HUnit values :: [(T.Text, DI.Column)] -values = [ ("test1", DI.fromList (concatMap (replicate 10) [1 :: Int, 2, 3, 4])) - , ("test2", DI.fromList (take 40 $ cycle [1 :: Int,2])) - , ("test3", DI.fromList [(1 :: Int)..40]) - , ("test4", DI.fromList (reverse [(1 :: Int)..40])) - ] +values = + [ ("test1", DI.fromList (concatMap (replicate 10) [1 :: Int, 2, 3, 4])) + , ("test2", DI.fromList (take 40 $ cycle [1 :: Int, 2])) + , ("test3", DI.fromList [(1 :: Int) .. 40]) + , ("test4", DI.fromList (reverse [(1 :: Int) .. 40])) + ] testData :: D.DataFrame testData = D.fromNamedColumns values groupBySingleRowWAI :: Test -groupBySingleRowWAI = TestCase (assertEqual "Groups by single column" - -- We don't yet compare offsets and indices - (D.Grouped testData ["test1"] VU.empty VU.empty) - (D.groupBy ["test1"] testData)) +groupBySingleRowWAI = + TestCase + ( assertEqual + "Groups by single column" + -- We don't yet compare offsets and indices + (D.Grouped testData ["test1"] VU.empty VU.empty) + (D.groupBy ["test1"] testData) + ) groupByMultipleRowsWAI :: Test -groupByMultipleRowsWAI = TestCase (assertEqual "Groups by single column" - -- We don't yet compare offsets and indices - (D.Grouped testData ["test1", "test2"] VU.empty VU.empty) - (D.groupBy ["test1", "test2"] testData)) +groupByMultipleRowsWAI = + TestCase + ( assertEqual + "Groups by single column" + -- We don't yet compare offsets and indices + (D.Grouped testData ["test1", "test2"] VU.empty VU.empty) + (D.groupBy ["test1", "test2"] testData) + ) groupByColumnDoesNotExist :: Test -groupByColumnDoesNotExist = TestCase (assertExpectException "[Error Case]" - (DE.columnNotFound "[\"test0\"]" "groupBy" (D.columnNames testData)) - (print $ D.groupBy ["test0"] testData)) +groupByColumnDoesNotExist = + TestCase + ( assertExpectException + "[Error Case]" + (D.columnNotFound "[\"test0\"]" "groupBy" (D.columnNames testData)) + (print $ D.groupBy ["test0"] testData) + ) tests :: [Test] -tests = [ TestLabel "groupBySingleRowWAI" groupBySingleRowWAI - , TestLabel "groupByMultipleRowsWAI" groupByMultipleRowsWAI - , TestLabel "groupByColumnDoesNotExist" groupByColumnDoesNotExist - ] - +tests = + [ TestLabel "groupBySingleRowWAI" groupBySingleRowWAI + , TestLabel "groupByMultipleRowsWAI" groupByMultipleRowsWAI + , TestLabel "groupByColumnDoesNotExist" groupByColumnDoesNotExist + ] diff --git a/tests/Operations/InsertColumn.hs b/tests/Operations/InsertColumn.hs --- a/tests/Operations/InsertColumn.hs +++ b/tests/Operations/InsertColumn.hs @@ -1,118 +1,171 @@ -{-# LANGUAGE TypeApplications #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + module Operations.InsertColumn where -import qualified DataFrame as D -import qualified DataFrame as DI import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU +import qualified DataFrame as D +import qualified DataFrame.Internal.Column as D +import qualified DataFrame.Internal.Column as DI +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Internal.DataFrame as DI import Assertions import Test.HUnit testData :: D.DataFrame -testData = D.fromNamedColumns [ ("test1", DI.fromList([1..26] :: [Int])) - , ("test2", DI.fromList ['a'..'z']) - , ("test3", DI.fromList ([1..26] :: [Int])) - , ("test4", DI.fromList ['a'..'z']) - , ("test5", DI.fromList ([1..26] :: [Int])) - , ("test6", DI.fromList ['a'..'z']) - , ("test7", DI.fromList ([1..26] :: [Int])) - , ("test8", DI.fromList ['a'..'z']) - ] +testData = + D.fromNamedColumns + [ ("test1", DI.fromList ([1 .. 26] :: [Int])) + , ("test2", DI.fromList ['a' .. 'z']) + , ("test3", DI.fromList ([1 .. 26] :: [Int])) + , ("test4", DI.fromList ['a' .. 'z']) + , ("test5", DI.fromList ([1 .. 26] :: [Int])) + , ("test6", DI.fromList ['a' .. 'z']) + , ("test7", DI.fromList ([1 .. 26] :: [Int])) + , ("test8", DI.fromList ['a' .. 'z']) + ] -- Adding a boxed vector to an empty dataframe creates a new column boxed containing the vector elements. addBoxedColumn :: Test -addBoxedColumn = TestCase (assertEqual "Two columns should be equal" - (Just $ DI.BoxedColumn (V.fromList ["Thuba" :: T.Text, "Zodwa", "Themba"])) - (DI.getColumn "new" $ D.insertVector "new" (V.fromList ["Thuba" :: T.Text, "Zodwa", "Themba"]) D.empty)) +addBoxedColumn = + TestCase + ( assertEqual + "Two columns should be equal" + (Just $ DI.BoxedColumn (V.fromList ["Thuba" :: T.Text, "Zodwa", "Themba"])) + (DI.getColumn "new" $ D.insertVector "new" (V.fromList ["Thuba" :: T.Text, "Zodwa", "Themba"]) D.empty) + ) addBoxedColumn' :: Test -addBoxedColumn' = TestCase (assertEqual "Two columns should be equal" - (Just $ DI.fromList["Thuba" :: T.Text, "Zodwa", "Themba"]) - (DI.getColumn "new" $ D.insertColumn "new" (DI.fromList["Thuba" :: T.Text, "Zodwa", "Themba"]) D.empty)) +addBoxedColumn' = + TestCase + ( assertEqual + "Two columns should be equal" + (Just $ DI.fromList ["Thuba" :: T.Text, "Zodwa", "Themba"]) + (DI.getColumn "new" $ D.insertColumn "new" (DI.fromList ["Thuba" :: T.Text, "Zodwa", "Themba"]) D.empty) + ) -- Adding an boxed vector with an unboxable type (Int/Double) to an empty dataframe creates a new column boxed containing the vector elements. addUnboxedColumn :: Test -addUnboxedColumn = TestCase (assertEqual "Value should be boxed" - (Just $ DI.UnboxedColumn (VU.fromList [1 :: Int, 2, 3])) - (DI.getColumn "new" $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3]) D.empty)) +addUnboxedColumn = + TestCase + ( assertEqual + "Value should be boxed" + (Just $ DI.UnboxedColumn (VU.fromList [1 :: Int, 2, 3])) + (DI.getColumn "new" $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3]) D.empty) + ) addUnboxedColumn' :: Test -addUnboxedColumn' = TestCase (assertEqual "Value should be boxed" - (Just $ DI.fromList[1 :: Int, 2, 3]) - (DI.getColumn "new" $ D.insertColumn "new" (DI.fromList[1 :: Int, 2, 3]) D.empty)) +addUnboxedColumn' = + TestCase + ( assertEqual + "Value should be boxed" + (Just $ DI.fromList [1 :: Int, 2, 3]) + (DI.getColumn "new" $ D.insertColumn "new" (DI.fromList [1 :: Int, 2, 3]) D.empty) + ) -- Adding a column with less values than the current DF dimensions adds column with optionals. addSmallerColumnBoxed :: Test -addSmallerColumnBoxed = TestCase ( - assertEqual "Missing values should be replaced with Nothing" - (Just $ DI.OptionalColumn (V.fromList [Just "a" :: Maybe T.Text, Just "b", Just "c", Nothing, Nothing])) - (DI.getColumn "newer" $ D.insertVector "newer" (V.fromList ["a" :: T.Text, "b", "c"]) $ D.insertVector "new" (V.fromList ["a" :: T.Text, "b", "c", "d", "e"]) D.empty) - ) +addSmallerColumnBoxed = + TestCase + ( assertEqual + "Missing values should be replaced with Nothing" + (Just $ DI.OptionalColumn (V.fromList [Just "a" :: Maybe T.Text, Just "b", Just "c", Nothing, Nothing])) + (DI.getColumn "newer" $ D.insertVector "newer" (V.fromList ["a" :: T.Text, "b", "c"]) $ D.insertVector "new" (V.fromList ["a" :: T.Text, "b", "c", "d", "e"]) D.empty) + ) addSmallerColumnUnboxed :: Test -addSmallerColumnUnboxed = TestCase ( - assertEqual "Missing values should be replaced with Nothing" - (Just $ DI.OptionalColumn (V.fromList [Just 1 :: Maybe Int, Just 2, Just 3, Nothing, Nothing])) - (DI.getColumn "newer" $ D.insertVector "newer" (V.fromList [1 :: Int, 2, 3]) $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3, 4, 5]) D.empty) - ) +addSmallerColumnUnboxed = + TestCase + ( assertEqual + "Missing values should be replaced with Nothing" + (Just $ DI.OptionalColumn (V.fromList [Just 1 :: Maybe Int, Just 2, Just 3, Nothing, Nothing])) + (DI.getColumn "newer" $ D.insertVector "newer" (V.fromList [1 :: Int, 2, 3]) $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3, 4, 5]) D.empty) + ) insertColumnWithDefaultFillsWithDefault :: Test -insertColumnWithDefaultFillsWithDefault = TestCase ( - assertEqual "Missing values should be replaced with Nothing" - (Just $ DI.UnboxedColumn (VU.fromList [1 :: Int, 2, 3, 0, 0])) - (DI.getColumn "newer" $ D.insertVectorWithDefault 0 "newer" (V.fromList [1 :: Int, 2, 3]) $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3, 4, 5]) D.empty) - ) +insertColumnWithDefaultFillsWithDefault = + TestCase + ( assertEqual + "Missing values should be replaced with Nothing" + (Just $ DI.UnboxedColumn (VU.fromList [1 :: Int, 2, 3, 0, 0])) + (DI.getColumn "newer" $ D.insertVectorWithDefault 0 "newer" (V.fromList [1 :: Int, 2, 3]) $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3, 4, 5]) D.empty) + ) insertColumnWithDefaultFillsLargerNoop :: Test -insertColumnWithDefaultFillsLargerNoop = TestCase ( - assertEqual "Lists should be the same size" - (Just $ DI.UnboxedColumn (VU.fromList [(6 :: Int)..10])) - (DI.getColumn "newer" $ D.insertVectorWithDefault 0 "newer" (V.fromList [(6 :: Int)..10]) $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3, 4, 5]) D.empty) - ) +insertColumnWithDefaultFillsLargerNoop = + TestCase + ( assertEqual + "Lists should be the same size" + (Just $ DI.UnboxedColumn (VU.fromList [(6 :: Int) .. 10])) + (DI.getColumn "newer" $ D.insertVectorWithDefault 0 "newer" (V.fromList [(6 :: Int) .. 10]) $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3, 4, 5]) D.empty) + ) addLargerColumnBoxed :: Test addLargerColumnBoxed = - TestCase (assertEqual "Smaller lists should grow and contain optionals" - (D.fromNamedColumns [("new", D.fromList [Just "a" :: Maybe T.Text, Just "b", Just "c", Nothing, Nothing]), - ("newer", D.fromList ["a" :: T.Text, "b", "c", "d", "e"])]) - (D.insertVector "newer" (V.fromList ["a" :: T.Text, "b", "c", "d", "e"]) - $ D.insertVector "new" (V.fromList ["a" :: T.Text, "b", "c"]) D.empty)) + TestCase + ( assertEqual + "Smaller lists should grow and contain optionals" + ( D.fromNamedColumns + [ ("new", D.fromList [Just "a" :: Maybe T.Text, Just "b", Just "c", Nothing, Nothing]) + , ("newer", D.fromList ["a" :: T.Text, "b", "c", "d", "e"]) + ] + ) + ( D.insertVector "newer" (V.fromList ["a" :: T.Text, "b", "c", "d", "e"]) $ + D.insertVector "new" (V.fromList ["a" :: T.Text, "b", "c"]) D.empty + ) + ) addLargerColumnUnboxed :: Test addLargerColumnUnboxed = - TestCase (assertEqual "Smaller lists should grow and contain optionals" - (D.fromNamedColumns [("old", D.fromList [Just 1 :: Maybe Int, Just 2, Nothing, Nothing, Nothing]), - ("new", D.fromList [Just 1 :: Maybe Int, Just 2, Just 3, Nothing, Nothing]), - ("newer", D.fromList [1 :: Int, 2, 3, 4, 5])]) - (D.insertVector "newer" (V.fromList [1 :: Int, 2, 3, 4, 5]) - $ D.insertVector "new" (V.fromList [1 :: Int, 2, 3]) $ - D.insertVector "old" (V.fromList [1 :: Int, 2]) D.empty)) + TestCase + ( assertEqual + "Smaller lists should grow and contain optionals" + ( D.fromNamedColumns + [ ("old", D.fromList [Just 1 :: Maybe Int, Just 2, Nothing, Nothing, Nothing]) + , ("new", D.fromList [Just 1 :: Maybe Int, Just 2, Just 3, Nothing, Nothing]) + , ("newer", D.fromList [1 :: Int, 2, 3, 4, 5]) + ] + ) + ( D.insertVector "newer" (V.fromList [1 :: Int, 2, 3, 4, 5]) $ + D.insertVector "new" (V.fromList [1 :: Int, 2, 3]) $ + D.insertVector "old" (V.fromList [1 :: Int, 2]) D.empty + ) + ) dimensionsChangeAfterAdd :: Test -dimensionsChangeAfterAdd = TestCase (assertEqual "should be (26, 3)" - (26, 9) - (D.dimensions $ D.insertVector @Int "new" (V.fromList [1..26]) testData)) +dimensionsChangeAfterAdd = + TestCase + ( assertEqual + "should be (26, 3)" + (26, 9) + (D.dimensions $ D.insertVector @Int "new" (V.fromList [1 .. 26]) testData) + ) dimensionsNotChangedAfterDuplicate :: Test -dimensionsNotChangedAfterDuplicate = TestCase (assertEqual "should be (26, 3)" - (26, 9) - (D.dimensions $ D.insertVector @Int "new" (V.fromList [1..26]) - $ D.insertVector @Int "new" (V.fromList [1..26]) testData)) - +dimensionsNotChangedAfterDuplicate = + TestCase + ( assertEqual + "should be (26, 3)" + (26, 9) + ( D.dimensions $ + D.insertVector @Int "new" (V.fromList [1 .. 26]) $ + D.insertVector @Int "new" (V.fromList [1 .. 26]) testData + ) + ) tests :: [Test] -tests = [ - TestLabel "dimensionsChangeAfterAdd" dimensionsChangeAfterAdd - , TestLabel "dimensionsNotChangedAfterDuplicate" dimensionsNotChangedAfterDuplicate - , TestLabel "addBoxedColunmToEmpty" addBoxedColumn - , TestLabel "addBoxedColumnAutoUnboxes" addBoxedColumn - , TestLabel "addSmallerColumnBoxed" addSmallerColumnBoxed - , TestLabel "addSmallerColumnUnboxed" addSmallerColumnUnboxed - , TestLabel "addLargerColumnBoxed" addLargerColumnBoxed - , TestLabel "addLargerColumnUnboxed" addLargerColumnUnboxed - , TestLabel "insertColumnWithDefaultFillsWithDefault" insertColumnWithDefaultFillsWithDefault - , TestLabel "insertColumnWithDefaultFillsLargerNoop" insertColumnWithDefaultFillsLargerNoop - ] +tests = + [ TestLabel "dimensionsChangeAfterAdd" dimensionsChangeAfterAdd + , TestLabel "dimensionsNotChangedAfterDuplicate" dimensionsNotChangedAfterDuplicate + , TestLabel "addBoxedColunmToEmpty" addBoxedColumn + , TestLabel "addBoxedColumnAutoUnboxes" addBoxedColumn + , TestLabel "addSmallerColumnBoxed" addSmallerColumnBoxed + , TestLabel "addSmallerColumnUnboxed" addSmallerColumnUnboxed + , TestLabel "addLargerColumnBoxed" addLargerColumnBoxed + , TestLabel "addLargerColumnUnboxed" addLargerColumnUnboxed + , TestLabel "insertColumnWithDefaultFillsWithDefault" insertColumnWithDefaultFillsWithDefault + , TestLabel "insertColumnWithDefaultFillsLargerNoop" insertColumnWithDefaultFillsLargerNoop + ] diff --git a/tests/Operations/Sort.hs b/tests/Operations/Sort.hs --- a/tests/Operations/Sort.hs +++ b/tests/Operations/Sort.hs @@ -1,50 +1,72 @@ {-# LANGUAGE OverloadedStrings #-} -module Operations.Sort where -import qualified DataFrame as D -import qualified DataFrame as DI -import qualified DataFrame as DE -import qualified Data.Text as T -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU +module Operations.Sort where import Assertions import Control.Monad import Data.Char +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import qualified DataFrame as D +import qualified DataFrame.Internal.Column as D +import qualified DataFrame.Internal.Column as DI +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Internal.DataFrame as DI import System.Random import System.Random.Shuffle (shuffle') import Test.HUnit values :: [(T.Text, DI.Column)] -values = let - ns = shuffle' [(1::Int)..26] 26 $ mkStdGen 252 - in [ ("test1", DI.fromList ns) - , ("test2", DI.fromList (map (chr . (+96)) ns)) - ] +values = + let + ns = shuffle' [(1 :: Int) .. 26] 26 $ mkStdGen 252 + in + [ ("test1", DI.fromList ns) + , ("test2", DI.fromList (map (chr . (+ 96)) ns)) + ] testData :: D.DataFrame testData = D.fromNamedColumns values sortByAscendingWAI :: Test -sortByAscendingWAI = TestCase (assertEqual "Sorting rows by ascending works as intended" - (D.fromNamedColumns [("test1", DI.fromList [(1::Int)..26]), - ("test2", DI.fromList ['a'..'z'])]) - (D.sortBy D.Ascending ["test1"] testData)) +sortByAscendingWAI = + TestCase + ( assertEqual + "Sorting rows by ascending works as intended" + ( D.fromNamedColumns + [ ("test1", DI.fromList [(1 :: Int) .. 26]) + , ("test2", DI.fromList ['a' .. 'z']) + ] + ) + (D.sortBy D.Ascending ["test1"] testData) + ) sortByDescendingWAI :: Test -sortByDescendingWAI = TestCase (assertEqual "Sorting rows by descending works as intended" - (D.fromNamedColumns [("test1", DI.fromList $ reverse [(1::Int)..26]), - ("test2", DI.fromList $ reverse ['a'..'z'])]) - (D.sortBy D.Descending ["test1"] testData)) +sortByDescendingWAI = + TestCase + ( assertEqual + "Sorting rows by descending works as intended" + ( D.fromNamedColumns + [ ("test1", DI.fromList $ reverse [(1 :: Int) .. 26]) + , ("test2", DI.fromList $ reverse ['a' .. 'z']) + ] + ) + (D.sortBy D.Descending ["test1"] testData) + ) sortByColumnDoesNotExist :: Test -sortByColumnDoesNotExist = TestCase (assertExpectException "[Error Case]" - (DE.columnNotFound "[\"test0\"]" "sortBy" (D.columnNames testData)) - (print $ D.sortBy D.Ascending ["test0"] testData)) +sortByColumnDoesNotExist = + TestCase + ( assertExpectException + "[Error Case]" + (D.columnNotFound "[\"test0\"]" "sortBy" (D.columnNames testData)) + (print $ D.sortBy D.Ascending ["test0"] testData) + ) tests :: [Test] -tests = [ TestLabel "sortByAscendingWAI" sortByAscendingWAI - , TestLabel "sortByDescendingWAI" sortByDescendingWAI - , TestLabel "sortByColumnDoesNotExist" sortByColumnDoesNotExist - ] - +tests = + [ TestLabel "sortByAscendingWAI" sortByAscendingWAI + , TestLabel "sortByDescendingWAI" sortByDescendingWAI + , TestLabel "sortByColumnDoesNotExist" sortByColumnDoesNotExist + ] diff --git a/tests/Operations/Statistics.hs b/tests/Operations/Statistics.hs new file mode 100644 --- /dev/null +++ b/tests/Operations/Statistics.hs @@ -0,0 +1,47 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Operations.Statistics where + +import qualified Data.Vector.Unboxed as VU +import qualified DataFrame as D +import qualified DataFrame.Internal.Column as D +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Operations.Statistics as D + +import Assertions +import Test.HUnit + +medianOfOddLengthDataSet :: Test +medianOfOddLengthDataSet = + TestCase + ( assertEqual + "Median of an odd length data set" + (D.median' (VU.fromList [179.94, 231.94, 839.06, 534.23, 248.94])) + 248.94 + ) + +medianOfEvenLengthDataSet :: Test +medianOfEvenLengthDataSet = + TestCase + ( assertEqual + "Median of an even length data set" + (D.median' (VU.fromList [179.94, 231.94, 839.06, 534.23, 248.94, 276.37])) + 262.655 + ) + +medianOfEmptyDataSet :: Test +medianOfEmptyDataSet = + TestCase + ( assertExpectException + "[Error Case]" + (D.emptyDataSetError "median") + (print $ D.median' (VU.fromList [])) + ) + +tests :: [Test] +tests = + [ TestLabel "medianOfOddLengthDataSet" medianOfOddLengthDataSet + , TestLabel "medianOfEvenLengthDataSet" medianOfEvenLengthDataSet + , TestLabel "medianOfEmptyDataSet" medianOfEmptyDataSet + ] diff --git a/tests/Operations/Take.hs b/tests/Operations/Take.hs --- a/tests/Operations/Take.hs +++ b/tests/Operations/Take.hs @@ -1,22 +1,27 @@ {-# LANGUAGE OverloadedStrings #-} + module Operations.Take where import qualified DataFrame as D -import qualified DataFrame as DI +import qualified DataFrame.Internal.Column as D +import qualified DataFrame.Internal.Column as DI +import qualified DataFrame.Internal.DataFrame as D +import qualified DataFrame.Internal.DataFrame as DI import Test.HUnit testData :: D.DataFrame -testData = D.fromNamedColumns [ ("test1", DI.fromList ([1..26] :: [Int])) - , ("test2", DI.fromList ['a'..'z']) - ] - +testData = + D.fromNamedColumns + [ ("test1", DI.fromList ([1 .. 26] :: [Int])) + , ("test2", DI.fromList ['a' .. 'z']) + ] takeWAI :: Test -takeWAI = TestCase (assertEqual "Gets first 10 numbers" (Just $ D.fromList [(1 :: Int)..10]) (D.getColumn "test1" $ D.take 10 testData)) +takeWAI = TestCase (assertEqual "Gets first 10 numbers" (Just $ D.fromList [(1 :: Int) .. 10]) (D.getColumn "test1" $ D.take 10 testData)) takeLastWAI :: Test -takeLastWAI = TestCase (assertEqual "Gets first 10 numbers" (Just $ D.fromList [(17 :: Int)..26]) (D.getColumn "test1" $ D.takeLast 10 testData)) +takeLastWAI = TestCase (assertEqual "Gets first 10 numbers" (Just $ D.fromList [(17 :: Int) .. 26]) (D.getColumn "test1" $ D.takeLast 10 testData)) lengthEqualsTakeParam :: Test lengthEqualsTakeParam = TestCase (assertEqual "should be (5, 2)" (5, 2) (D.dimensions $ D.take 5 testData)) @@ -43,14 +48,15 @@ negativeIsZeroTakeLast = TestCase (assertEqual "should be (0, 2)" (0, 2) (D.dimensions $ D.takeLast (-1) testData)) tests :: [Test] -tests = [ TestLabel "takeWAI" takeWAI - , TestLabel "takeLastWAI" takeLastWAI - , TestLabel "lengthEqualsTakeParam" lengthEqualsTakeParam - , TestLabel "lengthGreaterThanTakeParam" lengthGreaterThanTakeParam - , TestLabel "emptyIsZero" emptyIsZero - , TestLabel "negativeIsZero" negativeIsZero - , TestLabel "lengthEqualsTakeLastParam" lengthEqualsTakeLastParam - , TestLabel "lengthGreaterThanTakeLastParam" lengthGreaterThanTakeLastParam - , TestLabel "emptyIsZeroTakeLast" emptyIsZeroTakeLast - , TestLabel "negativeIsZeroTakeLast" negativeIsZeroTakeLast - ] +tests = + [ TestLabel "takeWAI" takeWAI + , TestLabel "takeLastWAI" takeLastWAI + , TestLabel "lengthEqualsTakeParam" lengthEqualsTakeParam + , TestLabel "lengthGreaterThanTakeParam" lengthGreaterThanTakeParam + , TestLabel "emptyIsZero" emptyIsZero + , TestLabel "negativeIsZero" negativeIsZero + , TestLabel "lengthEqualsTakeLastParam" lengthEqualsTakeLastParam + , TestLabel "lengthGreaterThanTakeLastParam" lengthGreaterThanTakeLastParam + , TestLabel "emptyIsZeroTakeLast" emptyIsZeroTakeLast + , TestLabel "negativeIsZeroTakeLast" negativeIsZeroTakeLast + ] diff --git a/tests/Parquet.hs b/tests/Parquet.hs --- a/tests/Parquet.hs +++ b/tests/Parquet.hs @@ -1,17 +1,73 @@ {-# LANGUAGE OverloadedStrings #-} + module Parquet where import qualified DataFrame as D -import Assertions -import GHC.IO (unsafePerformIO) -import Test.HUnit +import Assertions +import Data.Int +import Data.Text (Text) +import Data.Time +import Data.Time.Calendar +import GHC.IO (unsafePerformIO) +import Test.HUnit --- TODO: This currently fails +allTypes :: D.DataFrame +allTypes = + D.fromNamedColumns + [ ("id", D.fromList [(4 :: Int32), 5, 6, 7, 2, 3, 0, 1]) + , ("bool_col", D.fromList [True, False, True, False, True, False, True, False]) + , ("tinyint_col", D.fromList [(0 :: Int32), 1, 0, 1, 0, 1, 0, 1]) + , ("smallint_col", D.fromList [(0 :: Int32), 1, 0, 1, 0, 1, 0, 1]) + , ("int_col", D.fromList [(0 :: Int32), 1, 0, 1, 0, 1, 0, 1]) + , ("bigint_col", D.fromList [(0 :: Int64), 10, 0, 10, 0, 10, 0, 10]) + , ("float_col", D.fromList [(0 :: Float), 1.1, 0, 1.1, 0, 1.1, 0, 1.1]) + , ("double_col", D.fromList [(0 :: Double), 10.1, 0, 10.1, 0, 10.1, 0, 10.1]) + , ("date_string_col", D.fromList [("03/01/09" :: Text), "03/01/09", "04/01/09", "04/01/09", "02/01/09", "02/01/09", "01/01/09", "01/01/09"]) + , ("string_col", D.fromList (take 8 (cycle [("0" :: Text), "1"]))) + , + ( "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) + ] + ) + ] + allTypesPlain :: Test -allTypesPlain = TestCase (assertEqual "allTypesPlain" - (unsafePerformIO (D.readCsv "./data/mtcars.csv")) - (unsafePerformIO (D.readParquet "./data/mtcars.parquet"))) +allTypesPlain = + TestCase + ( assertEqual + "allTypesPlain" + allTypes + (unsafePerformIO (D.readParquet "./tests/data/alltypes_plain.parquet")) + ) +allTypesPlainSnappy :: Test +allTypesPlainSnappy = + TestCase + ( assertEqual + "allTypesPlainSnappy" + (D.filter "id" (`elem` [(6 :: Int32), 7]) allTypes) + (unsafePerformIO (D.readParquet "./tests/data/alltypes_plain.snappy.parquet")) + ) + +allTypesDictionary :: Test +allTypesDictionary = + TestCase + ( assertEqual + "allTypesPlainSnappy" + (D.filter "id" (`elem` [(0 :: Int32), 1]) allTypes) + (unsafePerformIO (D.readParquet "./tests/data/alltypes_dictionary.parquet")) + ) + +-- Uncomment to run parquet tests. +-- Currently commented because they don't run with github CI tests :: [Test] -tests = [] +tests = [] -- [allTypesPlain, allTypesPlainSnappy, allTypesDictionary]