diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,16 @@ # Revision history for dataframe +## 0.3.0.2 +* Re-enable Parquet. +* Change columnInfo to describeColumns +* We can now convert columns to lists. +* Fast reductions and groupings. GroupBys are now a dataframe construct not a column construct (thanks to @stites). +* Filter is now faster because we do mutation on the index vector. +* Frequencies table nnow correctly display percentages (thanks @kayvank) +* Show table implementations have been unified (thanks @metapho-re) +* We now compute statistics on null columns +* Drastic improvement in plotting since we now use granite. + ## 0.3.0.1 * Temporarily remove Parquet support. I think it'll be worth creating a spin off of snappy that doesn't rely on C bindings. Also I'll probably spin Parquet off into a separate library. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@
- + hackage Latest Release @@ -94,9 +94,9 @@ ## Installing ### Jupyter notebook -* We have a [hosted version of the Jupyter notebook](https://ihaskell-dataframe-crf7g5fvcpahdegz.westus2-01.azurewebsites.net/lab/) on azure sites. This is hosted on Azure's free tier so it slow and is only available at best effort. +* We have a [hosted version of the Jupyter notebook](https://ulwazi-exh9dbh2exbzgbc9.westus-01.azurewebsites.net/lab) on azure sites. This is hosted on Azure's free tier so it can only support 3 or 4 kernels at a time. * To get started quickly, use the Dockerfile in the [ihaskell-dataframe](https://github.com/mchav/ihaskell-dataframe) to build and run an image with dataframe integration. -* For a preview check out the [California Housing](https://ihaskell-dataframe-crf7g5fvcpahdegz.westus2-01.azurewebsites.net/lab/tree/California%20Housing.ipynb) notebook. +* For a preview check out the [California Housing](https://github.com/mchav/dataframe/blob/main/docs/California%20Housing.ipynb) notebook. ### CLI * Run the installation script `curl '=https' --tlsv1.2 -sSf https://raw.githubusercontent.com/mchav/dataframe/refs/heads/main/scripts/install.sh | sh` diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -2,22 +2,42 @@ -- 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 #-} 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 System.Random.Stateful + main :: IO () main = do - df <- D.readTsv "./data/chipotle.tsv" - let quantity = F.col "quantity" :: D.Expr Int -- A typed reference to a column. - print (df - |> D.select ["item_name", "quantity"] - |> D.groupBy ["item_name"] - |> D.aggregate [ (F.sum quantity) `F.as` "sum_quantity" - , (F.mean quantity) `F.as` "mean_quantity" - , (F.maximum quantity) `F.as` "maximum_quantity" - ] - |> D.sortBy D.Descending ["sum_quantity"] - |> D.take 10) + let n = 100_000_000 + g <- newIOGenM =<< newStdGen + let range = (-20.0 :: Double, 20.0 :: Double) + startGeneration <- getCurrentTime + ns <- VU.replicateM n (uniformRM range g) + xs <- VU.replicateM n (uniformRM range g) + ys <- VU.replicateM n (uniformRM range g) + let df = D.fromUnnamedColumns (map D.fromUnboxedVector [ns, xs, ys]) + endGeneration <- getCurrentTime + let generationTime = diffUTCTime endGeneration startGeneration + putStrLn $ "Data generation Time: " ++ (show generationTime) + startCalculation <- getCurrentTime + print $ D.mean "0" df + print $ D.variance "1" df + print $ D.correlation "1" "2" df + endCalculation <- getCurrentTime + let calculationTime = diffUTCTime endCalculation startCalculation + putStrLn $ "Calculation Time: " ++ (show calculationTime) + startFilter <- getCurrentTime + print $ D.filter "0" (>= (19.9 :: Double)) df D.|> D.take 10 + endFilter <- getCurrentTime + let filterTime = diffUTCTime endFilter startFilter + putStrLn $ "Filter Time: " ++ (show filterTime) + let totalTime = diffUTCTime endFilter startGeneration + putStrLn $ "Total Time: " ++ (show totalTime) 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.1 +version: 0.3.0.2 synopsis: A fast, safe, and intuitive DataFrame library. @@ -12,7 +12,7 @@ author: Michael Chavinda maintainer: mschavinda@gmail.com -copyright: (c) 2024-2024 Michael Chavinda +copyright: (c) 2024-2025 Michael Chavinda category: Data tested-with: GHC ==9.8.3 || ==9.6.6 || == 9.4.8 || ==9.10.1 || ==9.12.1 || ==9.12.2 extra-doc-files: CHANGELOG.md README.md @@ -34,6 +34,7 @@ DataFrame.Display.Terminal.Colours, DataFrame.Internal.DataFrame, DataFrame.Internal.Row, + DataFrame.Internal.Schema, DataFrame.Errors, DataFrame.Operations.Core, DataFrame.Operations.Join, @@ -46,6 +47,12 @@ DataFrame.Operations.Aggregation, DataFrame.Display.Terminal.Plot, DataFrame.IO.CSV, + DataFrame.IO.Parquet, + DataFrame.IO.Parquet.ColumnStatistics, + DataFrame.IO.Parquet.Compression, + DataFrame.IO.Parquet.Encoding, + DataFrame.IO.Parquet.Page, + DataFrame.IO.Parquet.Types, DataFrame.Lazy.IO.CSV, DataFrame.Lazy.Internal.DataFrame build-depends: base >= 4.17.2.0 && < 4.22, @@ -55,7 +62,9 @@ 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, hashable >= 1.2 && <= 1.5.0.0, + 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, @@ -66,135 +75,6 @@ hs-source-dirs: src default-language: Haskell2010 -executable chipotle - main-is: Chipotle.hs - other-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.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, - containers >= 0.6.7 && < 0.8, - directory >= 1.3.0.0 && <= 1.3.9.0, - hashable >= 1.2 && <= 1.5.0.0, - 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: examples, - src - default-language: Haskell2010 - -executable california_housing - main-is: CaliforniaHousing.hs - other-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.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, - containers >= 0.6.7 && < 0.8, - directory >= 1.3.0.0 && <= 1.3.9.0, - hashable >= 1.2 && <= 1.5.0.0, - 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: examples, - src - default-language: Haskell2010 - -executable one_billion_row_challenge - main-is: OneBillionRowChallenge.hs - other-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.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, - containers >= 0.6.7 && < 0.8, - directory >= 1.3.0.0 && <= 1.3.9.0, - hashable >= 1.2 && <= 1.5.0.0, - 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: examples, - src - default-language: Haskell2010 - executable dataframe main-is: Main.hs other-modules: DataFrame, @@ -216,6 +96,12 @@ DataFrame.Operations.Aggregation, DataFrame.Display.Terminal.Plot, DataFrame.IO.CSV, + DataFrame.IO.Parquet, + DataFrame.IO.Parquet.ColumnStatistics, + DataFrame.IO.Parquet.Compression, + DataFrame.IO.Parquet.Encoding, + DataFrame.IO.Parquet.Page, + DataFrame.IO.Parquet.Types, DataFrame.Operations.Join, DataFrame.Operations.Merge, DataFrame.Lazy.IO.CSV, @@ -226,8 +112,10 @@ bytestring >= 0.11 && <= 0.12.2.0, containers >= 0.6.7 && < 0.8, directory >= 1.3.0.0 && <= 1.3.9.0, + granite ^>= 0.1.0.2, hashable >= 1.2 && <= 1.5.0.0, 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, @@ -258,13 +146,15 @@ type: exitcode-stdio-1.0 main-is: Main.hs other-modules: Assertions, + Functions, Operations.Apply, Operations.Derive, Operations.Filter, Operations.GroupBy, Operations.InsertColumn, Operations.Sort, - Operations.Take + Operations.Take, + Parquet build-depends: base >= 4.17.2.0 && < 4.22, HUnit ^>= 1.6, random >= 1, diff --git a/examples/CaliforniaHousing.hs b/examples/CaliforniaHousing.hs deleted file mode 100644 --- a/examples/CaliforniaHousing.hs +++ /dev/null @@ -1,14 +0,0 @@ -{-# LANGUAGE OverloadedStrings #-} -module Main where - -import qualified DataFrame as D - -main :: IO () -main = do - parsed <- D.readCsv "./data/housing.csv" - - print $ D.describeColumns parsed - - print $ D.take 5 parsed - - D.plotHistograms D.PlotAll D.VerticalHistogram parsed diff --git a/examples/Chipotle.hs b/examples/Chipotle.hs deleted file mode 100644 --- a/examples/Chipotle.hs +++ /dev/null @@ -1,75 +0,0 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE TypeApplications #-} -module Main where - -import qualified DataFrame as D -import qualified DataFrame.Functions as F -import qualified Data.Text as T - -import DataFrame ((|>)) - -main :: IO () -main = do - raw <- D.readTsv "./data/chipotle.tsv" - print $ D.dimensions raw - - -- -- Sampling the dataframe - print $ D.take 5 raw - - -- Transform the data from a raw string into - -- respective types (throws error on failure) - let df = - raw - -- Change a specfic order ID - |> D.applyWhere (== (1 ::Int)) "order_id" (+ (2 :: Int)) "quantity" - -- Index based change. - |> D.applyAtIndex 0 ((\n -> n - 2) :: Int -> Int) "quantity" - -- Custom parsing: drop dollar sign and parse price as double - |> D.apply (D.readValue @Double . T.drop 1) "item_price" - - -- sample the dataframe. - print $ D.take 10 df - - -- Create a total_price column that is quantity * item_price - let withTotalPrice = D.derive "total_price" (F.lift fromIntegral (F.col @Int "quantity") * F.col @Double"item_price") df - - -- sample a filtered subset of the dataframe - putStrLn "Sample dataframe" - print $ - withTotalPrice - |> D.select ["quantity", "item_name", "item_price", "total_price"] - |> D.filter "total_price" ((100.0 :: Double) <) - |> D.take 10 - - -- Check how many chicken burritos were ordered. - -- There are two ways to checking how many chicken burritos - -- were ordered. - let searchTerm = "Chicken Burrito" :: T.Text - - print $ - df - |> D.select ["item_name", "quantity"] - -- It's more efficient to filter before grouping. - |> D.filter "item_name" (searchTerm ==) - |> D.groupBy ["item_name"] - |> D.aggregate [ (F.sum (F.col @Int "quantity")) `F.as` "sum" - , (F.maximum (F.col @Int "quantity")) `F.as` "max" - , (F.mean (F.col @Int "quantity")) `F.as` "mean"] - |> D.sortBy D.Descending ["sum"] - - -- Similarly, we can aggregate quantities by all rows. - print $ - df - |> D.select ["item_name", "quantity"] - |> D.groupBy ["item_name"] - |> D.aggregate [ (F.sum (F.col @Int "quantity")) `F.as` "sum" - , (F.maximum (F.col @Int "quantity")) `F.as` "maximum" - , (F.mean (F.col @Int "quantity")) `F.as` "mean"] - |> D.take 10 - - let firstOrder = - withTotalPrice - |> D.filterBy (maybe False (T.isInfixOf "Guacamole")) "choice_description" - |> D.filterBy (("Chicken Bowl" :: T.Text) ==) "item_name" - - print $ D.take 10 firstOrder diff --git a/examples/OneBillionRowChallenge.hs b/examples/OneBillionRowChallenge.hs deleted file mode 100644 --- a/examples/OneBillionRowChallenge.hs +++ /dev/null @@ -1,20 +0,0 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE TypeApplications #-} -module Main where - -import qualified DataFrame as D -import qualified DataFrame.Functions as F - -import DataFrame ((|>)) - -main :: IO () -main = do - parsed <- D.readSeparated ';' D.defaultOptions "./data/measurements.txt" - let measurement = (F.col @Double "Measurement") - print $ - parsed - |> D.groupBy ["City"] - |> D.aggregate [ (F.minimum measurement) `F.as` "minimum" - , (F.mean measurement) `F.as` "mean" - , (F.maximum measurement) `F.as` "maximum"] - |> D.sortBy D.Ascending ["City"] diff --git a/src/DataFrame.hs b/src/DataFrame.hs --- a/src/DataFrame.hs +++ b/src/DataFrame.hs @@ -24,6 +24,7 @@ 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 Data.Function import Data.List 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,338 +1,563 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE GADTs #-} -{-# LANGUAGE NumericUnderscores #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE FlexibleContexts #-} + 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 qualified Type.Reflection as Ref +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 Control.Monad ( forM_, forM ) -import Data.Bifunctor ( first ) -import Data.Char ( ord, chr ) -import DataFrame.Display.Terminal.Colours import DataFrame.Internal.Column (Column(..), Columnable) import DataFrame.Internal.DataFrame (DataFrame(..)) import DataFrame.Operations.Core -import Data.Maybe (fromMaybe) -import Data.Typeable (Typeable) -import Data.Type.Equality - ( type (:~:)(Refl), TestEquality(testEquality) ) -import GHC.Stack (HasCallStack) -import Text.Printf ( printf ) -import Type.Reflection (typeRep) +import Granite -data HistogramOrientation = VerticalHistogram | HorizontalHistogram +data PlotConfig = PlotConfig + { plotType :: PlotType + , plotTitle :: String + , plotSettings :: Plot + } -data PlotColumns = PlotAll | PlotSubset [T.Text] +data PlotType + = Histogram' + | Scatter' + | Line' + | Bar' + | BoxPlot' + | Pie' + | StackedBar' + | Heatmap' + deriving (Eq, Show) -plotHistograms :: HasCallStack => PlotColumns -> HistogramOrientation -> DataFrame -> IO () -plotHistograms plotSet orientation df = do - let cs = case plotSet of - PlotAll -> columnNames df - PlotSubset xs -> columnNames df `L.intersect` xs - forM_ cs $ \cname -> do - plotForColumn cname ((V.!) (columns df) (columnIndices df M.! cname)) orientation df +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 -plotHistogramsBy :: HasCallStack => T.Text -> PlotColumns -> HistogramOrientation -> DataFrame -> IO () -plotHistogramsBy col plotSet orientation df = do - let cs = case plotSet of - PlotAll -> columnNames df - PlotSubset xs -> columnNames df `L.intersect` xs - forM_ cs $ \cname -> do - let plotColumn = (V.!) (columns df) (columnIndices df M.! cname) - let byColumn = (V.!) (columns df) (columnIndices df M.! col) - plotForColumnBy col cname byColumn plotColumn orientation 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) --- Plot code adapted from: https://alexwlchan.net/2018/ascii-bar-charts/ -plotForColumnBy :: HasCallStack => T.Text -> T.Text -> Column -> Column -> HistogramOrientation -> DataFrame -> IO () -plotForColumnBy byCol cname (BoxedColumn (byColumn :: V.Vector a)) (BoxedColumn (plotColumn :: V.Vector b)) orientation df = do - let zipped = VG.zipWith (\left right -> (show left, show right)) plotColumn byColumn - let counts = countOccurrences zipped - if null counts || length counts > 20 - then pure () - else case orientation of - VerticalHistogram -> error "Vertical histograms aren't yet supported" - HorizontalHistogram -> plotGivenCounts' cname counts -plotForColumnBy byCol cname (UnboxedColumn byColumn) (BoxedColumn plotColumn) orientation df = do - let zipped = VG.zipWith (\left right -> (show left, show right)) plotColumn (V.convert byColumn) - let counts = countOccurrences zipped - if null counts || length counts > 20 - then pure () - else case orientation of - VerticalHistogram -> error "Vertical histograms aren't yet supported" - HorizontalHistogram -> plotGivenCounts' cname counts -plotForColumnBy byCol cname (BoxedColumn byColumn) (UnboxedColumn plotColumn) orientation df = do - let zipped = VG.zipWith (\left right -> (show left, show right)) (V.convert plotColumn) (V.convert byColumn) - let counts = countOccurrences zipped - if null counts || length counts > 20 - then pure () - else case orientation of - -- VerticalHistogram -> plotVerticalGivenCounts cname counts - HorizontalHistogram -> plotGivenCounts' cname counts -plotForColumnBy byCol cname (UnboxedColumn byColumn) (UnboxedColumn plotColumn) orientation df = do - let zipped = VG.zipWith (\left right -> (show left, show right)) (V.convert plotColumn) (V.convert byColumn) - let counts = countOccurrences zipped - if null counts || length counts > 20 - then pure () - else case orientation of - VerticalHistogram -> error "Vertical histograms aren't yet supported" - HorizontalHistogram -> plotGivenCounts' cname counts --- TODO: Add Optional columns -plotForColumnBy _ _ _ _ _ _ = return () +plotScatter :: HasCallStack => T.Text -> T.Text -> DataFrame -> IO () +plotScatter xCol yCol df = plotScatterWith xCol yCol (defaultPlotConfig Scatter') df --- Plot code adapted from: https://alexwlchan.net/2018/ascii-bar-charts/ -plotForColumn :: HasCallStack => T.Text -> Column -> HistogramOrientation -> DataFrame -> IO () -plotForColumn cname (BoxedColumn (column :: V.Vector a)) orientation df = do - let repa :: Ref.TypeRep a = Ref.typeRep @a - repText :: Ref.TypeRep T.Text = Ref.typeRep @T.Text - repString :: Ref.TypeRep String = Ref.typeRep @String - let counts = case repa `testEquality` repText of - Just Refl -> map (first T.unpack) $ valueCounts @T.Text cname df - Nothing -> case repa `testEquality` repString of - Just Refl -> valueCounts @String cname df - -- Support other scalar types. - Nothing -> [] -- numericHistogram column - if null counts || length counts > 20 - then putStrLn $ numericHistogram cname (V.convert column) - else case orientation of - VerticalHistogram -> plotVerticalGivenCounts cname counts - HorizontalHistogram -> plotGivenCounts cname counts -plotForColumn cname (UnboxedColumn (column :: VU.Vector a)) orientation df = do - let repa :: Ref.TypeRep a = Ref.typeRep @a - repText :: Ref.TypeRep T.Text = Ref.typeRep @T.Text - repString :: Ref.TypeRep String = Ref.typeRep @String - let counts = case repa `testEquality` repText of - Just Refl -> map (first show) $ valueCounts @T.Text cname df - Nothing -> case repa `testEquality` repString of - Just Refl -> valueCounts @String cname df - -- Support other scalar types. - Nothing -> [] - if null counts || length counts > 20 - then putStrLn $ numericHistogram cname (V.convert column) - else case orientation of - VerticalHistogram -> plotVerticalGivenCounts cname counts - HorizontalHistogram -> plotGivenCounts cname counts -plotForColumn _ _ _ _ = return () +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) -plotGivenCounts :: HasCallStack => T.Text -> [(String, Int)] -> IO () -plotGivenCounts cname counts = do - putStrLn $ "\nHistogram for " ++ show cname ++ "\n" - let n = 8 :: Int - let maxValue = maximum $ map snd counts - let increment = max 1 (maxValue `div` 50) - let longestLabelLength = maximum $ map (length . fst) counts - let longestBar = fromIntegral $ (maxValue * fromIntegral n `div` increment) `div` fromIntegral n + 1 - let border = "|" ++ replicate (longestLabelLength + length (show maxValue) + longestBar + 6) '-' ++ "|" - body <- forM counts $ \(label, count) -> do - let barChunks = fromIntegral $ (count * fromIntegral n `div` increment) `div` fromIntegral n - let remainder = fromIntegral $ (count * fromIntegral n `div` increment) `rem` fromIntegral n - let fractional = ([chr (ord '█' + n - remainder - 1) | remainder > 0]) - let bar = replicate barChunks '█' ++ fractional - let disp = if null bar then "| " else bar - let hist= "|" ++ brightGreen (leftJustify label longestLabelLength) ++ " | " ++ - leftJustify (show count) (length (show maxValue)) ++ " |" ++ - " " ++ brightBlue bar - return $ hist ++ "\n" ++ border - mapM_ putStrLn (border : body) - putChar '\n' +plotLines :: HasCallStack => [T.Text] -> DataFrame -> IO () +plotLines colNames df = plotLinesWith colNames (defaultPlotConfig Line') df -plotVerticalGivenCounts :: HasCallStack => T.Text -> [(String, Int)] -> IO () -plotVerticalGivenCounts cname counts' = do - putStrLn $ "\nHistogram for " ++ show cname ++ "\n" - let n = 8 :: Int - let clip s = if length s > n then take n s ++ ".." else s - let counts = map (first clip) counts' - let maxValue = maximum $ map snd counts - let increment = max 1 (maxValue `div` 10) - let longestLabelLength = 2 + maximum (map (length . fst) counts) - let longestBar = fromIntegral $ (maxValue * fromIntegral n `div` increment) `div` fromIntegral n + 1 - let border = "‾" ++ replicate (longestBar + 1) '|' ++ "+" - let maximumLineLength = length border - body <- forM counts $ \(label, count) -> do - let barChunks = fromIntegral $ (count * fromIntegral n `div` increment) `div` fromIntegral n - let remainder = fromIntegral $ (count * fromIntegral n `div` increment) `rem` fromIntegral n - let fractional = ([chr (ord '█' - (n - remainder - 1)) | remainder > 0]) - let bar = replicate barChunks '█' ++ fractional - let disp = if null bar then "| " else bar - let hist = "‾" ++ bar - return $ replicate longestLabelLength (leftJustify hist maximumLineLength) ++ [border] - let fullGraph = map brightBlue $ rotate $ border : concat body - let partition = smallestPartition increment intPlotRanges - let increments = reverse [0, maxValue `div` 2 , maxValue + partition] - let incString = reverse $ map (`leftJustify` (length (show maxValue) + 1)) $ show 0 : replicate (length fullGraph `div` 2 - 2) " " - ++ [show (maxValue `div` 2)] - ++ replicate (length fullGraph `div` 2 - 2) " " - ++ [show (maxValue + partition)] - ++ [""] - mapM_ putStrLn (zipWith (++) incString fullGraph) - putStrLn $ " " ++ replicate (length (show maxValue) + 1) ' ' ++ unwords (map (brightGreen . flip leftJustify longestLabelLength . fst) counts) - putChar '\n' +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) -leftJustify :: String -> Int -> String -leftJustify s n = s ++ replicate (max 0 (n - length s)) ' ' +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) -plotGivenCounts' :: HasCallStack => T.Text -> [((String, String), Int)] -> IO () -plotGivenCounts' cname counts = do - putStrLn $ "\nHistogram for " ++ show cname ++ "\n" - let n = 8 :: Int - let maxValue = maximum $ map snd counts - let increment = max 1 (maxValue `div` 50) - let longestLabelLength = maximum $ map (length. (\(a, b) -> a ++ " " ++ b) . fst) counts - let longestBar = fromIntegral $ (maxValue * fromIntegral n `div` increment) `div` fromIntegral n + 1 - let border = "|" ++ replicate (longestLabelLength + length (show maxValue) + longestBar + 6) '-' ++ "|" - body <- forM counts $ \((plotCol, byCol), count) -> do - let barChunks = fromIntegral $ (count * fromIntegral n `div` increment) `div` fromIntegral n - let remainder = fromIntegral $ (count * fromIntegral n `div` increment) `rem` fromIntegral n - let fractional = ([chr (ord '█' + n - remainder - 1) | remainder > 0]) - let bar = replicate barChunks '█' ++ fractional - let disp = if null bar then "| " else bar - let label = plotCol ++ " " ++ byCol - let hist= "|" ++ brightGreen (leftJustify label longestLabelLength) ++ " | " ++ - leftJustify (show count) (length (show maxValue)) ++ " |" ++ - " " ++ brightBlue bar - return $ hist ++ "\n" ++ border - mapM_ putStrLn (border : body) - putChar '\n' +plotStackedBars :: HasCallStack => T.Text -> [T.Text] -> DataFrame -> IO () +plotStackedBars categoryCol valueColumns df = + plotStackedBarsWith categoryCol valueColumns (defaultPlotConfig StackedBar') df -numericHistogram :: forall a . (HasCallStack, Columnable a) - => T.Text - -> V.Vector a - -> String -numericHistogram name xs = let - config = defaultConfig { - title = Just (T.unpack name), - width = 30, - height = 10 - } - in createHistogram config (V.toList xs') - where - xs' = case testEquality (typeRep @a) (typeRep @Double) of - Just Refl -> xs - Nothing -> case testEquality (typeRep @a) (typeRep @Int) of - Just Refl -> V.map fromIntegral xs - Nothing -> case testEquality (typeRep @a) (typeRep @Integer) of - Just Refl -> V.map fromIntegral xs - Nothing -> V.empty +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) -smallestPartition :: (Ord a) => a -> [a] -> a --- TODO: Find a more graceful way to handle this. -smallestPartition p [] = error "Data range too large to plot" -smallestPartition p (x:y:rest) - | p < y = x - | otherwise = smallestPartition p (y:rest) -smallestPartition p (x:rest) - | p < x = x - | otherwise = error "" +plotHeatmap :: HasCallStack => DataFrame -> IO () +plotHeatmap df = plotHeatmapWith (defaultPlotConfig Heatmap') df -largestPartition :: (Ord a) => a -> [a] -> a --- TODO: Find a more graceful way to handle this. -largestPartition p [] = error "Data range too large to plot" -largestPartition p (x:rest) - | p < x = x - | otherwise = largestPartition p rest +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) -intPlotRanges :: [Int] -intPlotRanges = [1, 5, - 10, 50, - 100, 500, - 1_000, 5_000, - 10_000, 50_000, - 100_000, 500_000, - 1_000_000, 5_000_000] +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 -rotate :: [String] -> [String] -rotate [] = [] -rotate xs - | head xs == "" = [] - | otherwise = map last xs : rotate (map init xs) +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) -countOccurrences :: Ord a => V.Vector a -> [(a, Int)] -countOccurrences xs = M.toList $ VG.foldr count initMap xs - where initMap = M.fromList (map (, 0) (V.toList xs)) - count k = M.insertWith (+) k 1 +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 -data HistogramConfig = HistogramConfig { - width :: Int, -- Width of the histogram in characters - height :: Int, -- Height of the histogram in rows - barChar :: Char, -- Character to use for bars - title :: Maybe String -- Optional title for the histogram -} +plotBars :: HasCallStack => T.Text -> DataFrame -> IO () +plotBars colName df = plotBarsWith colName Nothing (defaultPlotConfig Bar') df -defaultConfig :: HistogramConfig -defaultConfig = HistogramConfig { - width = 40, - height = 15, - barChar = '█', - title = Nothing -} +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 --- Calculate the histogram bins and counts -calculateBins :: [Double] -> Int -> [(Double, Int)] -calculateBins values numBins = - let minVal = minimum values - maxVal = maximum values - binWidth = (maxVal - minVal) / fromIntegral numBins - toBin x = floor ((x - minVal) / binWidth) - bins = map toBin values - counts = map length . L.group . L.sort $ bins - binValues = [minVal + (fromIntegral i * binWidth) | i <- [0..numBins-1]] - in zip binValues (counts ++ repeat 0) +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) --- Format a number with appropriate scaling (k, M, B, etc.) -formatNumber :: Double -> String -formatNumber n - | n >= 1e9 = printf "%.1fB" (n / 1e9) - | n >= 1e6 = printf "%.1fM" (n / 1e6) - | n >= 1e3 = printf "%.1fk" (n / 1e3) - | otherwise = printf "%.1f" n +plotBarsTopN :: HasCallStack => Int -> T.Text -> DataFrame -> IO () +plotBarsTopN n colName df = plotBarsTopNWith n colName (defaultPlotConfig Bar') df --- Create the ASCII histogram -createHistogram :: HistogramConfig -> [Double] -> String -createHistogram _ [] = [] -createHistogram config values = - let bins = calculateBins values (width config) - maxCount = maximum $ map snd bins - scaleY = fromIntegral maxCount / fromIntegral (height config) +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) - -- Create Y-axis labels - yLabels = [formatNumber (fromIntegral i * scaleY) | i <- [height config, height config-1..0]] - maxYLabelWidth = maximum $ map length yLabels +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) - -- Create X-axis labels - xValues = map fst bins - xLabels = map formatNumber [head xValues, last xValues] +plotValueCounts :: HasCallStack => T.Text -> DataFrame -> IO () +plotValueCounts colName df = plotValueCountsWith colName 10 (defaultPlotConfig Bar') df - -- Create histogram rows - makeRow :: Int -> String - makeRow row = - let threshold = fromIntegral (height config - row) * scaleY - barLine = map (\(_, count) -> - if fromIntegral count >= threshold - then barChar config - else ' ') bins - in printf "%*s |%s" maxYLabelWidth (yLabels !! row) (brightBlue $ L.foldl' (\acc c -> c:'|':acc) "" barLine) +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 - -- Build the complete histogram - histogramRows = map makeRow [0..height config - 1] - xAxis = replicate maxYLabelWidth ' ' ++ " " ++ - L.intercalate (replicate (2 * (width config - length xLabels)) ' ') xLabels +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 - -- Add title if provided - titleLine = case title config of - Just t -> t ++ "\n\n" - Nothing -> "" +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 - in titleLine ++ unlines (histogramRows ++ [xAxis]) +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" 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 @@ -43,21 +43,15 @@ center :: Int -> T.Text -> T.Text center = fillCenter ' ' -showTable :: [T.Text] -> [T.Text] -> [[T.Text]] -> T.Text -showTable header types rows = - let cs = map (\h -> ColDesc center h left) header - widths = [maximum $ map T.length col | col <- transpose $ header : 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] - in T.unlines $ border : fillCols colTitleFill header : separator : fillCols colTitleFill types : separator : map (fillCols colValueFill) rows - -showTableProperMarkdown :: [T.Text] -> [T.Text] -> [[T.Text]] -> T.Text -showTableProperMarkdown header types rows = - let headerWithTypes = zipWith (\h t -> h <> "
" <> t) header types - cs = map (\h -> ColDesc center h left) headerWithTypes - widths = [maximum $ map T.length col | col <- transpose $ headerWithTypes : rows] +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] - in T.unlines $ border : fillCols colTitleFill headerWithTypes : separator : map (fillCols colValueFill) rows + 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/Functions.hs b/src/DataFrame/Functions.hs --- a/src/DataFrame/Functions.hs +++ b/src/DataFrame/Functions.hs @@ -27,6 +27,9 @@ 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 Type.Reflection (typeRep) col :: Columnable a => T.Text -> Expr a col = Col @@ -62,16 +65,13 @@ count (Col name) = GeneralAggregate name "count" VG.length count _ = error "Argument can only be a column reference not an unevaluated expression" -anyValue :: Columnable a => Expr a -> Expr a -anyValue (Col name) = ReductionAggregate name "anyValue" VG.head - minimum :: Columnable a => Expr a -> Expr a -minimum (Col name) = ReductionAggregate name "minimum" VG.minimum +minimum (Col name) = ReductionAggregate name "minimum" min maximum :: Columnable a => Expr a -> Expr a -maximum (Col name) = ReductionAggregate name "maximum" VG.maximum +maximum (Col name) = ReductionAggregate name "maximum" max -sum :: (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 @@ -81,32 +81,99 @@ in total / fromIntegral n in NumericAggregate name "mean" mean' -typeFromString :: String -> Q Type -typeFromString s = do - maybeType <- lookupTypeName s +-- See Section 2.4 of the Haskell Report https://www.haskell.org/definition/haskell2010.pdf +isReservedId :: T.Text -> Bool +isReservedId t = case t of + "case" -> True + "class" -> True + "data" -> True + "default" -> True + "deriving" -> True + "do" -> True + "else" -> True + "foreign" -> True + "if" -> True + "import" -> True + "in" -> True + "infix" -> True + "infixl" -> True + "infixr" -> True + "instance" -> True + "let" -> True + "module" -> True + "newtype" -> True + "of" -> True + "then" -> True + "type" -> True + "where" -> True + _ -> False + +isVarId :: T.Text -> Bool +isVarId t = case T.uncons t of +-- We might want to check c == '_' || Char.isLower c +-- since the haskell report considers '_' a lowercase character +-- However, to prevent an edge case where a user may have a +-- "Name" and an "_Name_" in the same scope, wherein we'd end up +-- with duplicate "_Name_"s, we eschew the check for '_' here. + Just (c, _) -> Char.isLower c && Char.isAlpha c + Nothing -> False + +isHaskellIdentifier :: T.Text -> Bool +isHaskellIdentifier t = not (isVarId t) || isReservedId t + +sanitize :: T.Text -> T.Text +sanitize t + | isValid = t + | isHaskellIdentifier t' = "_" <> t' <> "_" + | otherwise = t' + where + 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 + parentheses c = case c of + '(' -> True + ')' -> True + '{' -> True + '}' -> True + '[' -> True + ']' -> True + _ -> False + +typeFromString :: [String] -> Q Type +typeFromString [] = fail "No type specified" +typeFromString [t] = do + maybeType <- lookupTypeName t case maybeType of Just name -> return (ConT name) - Nothing -> do - if "Maybe " `L.isPrefixOf` s - then do - let innerType = drop 6 s - inner <- typeFromString innerType - return (AppT (ConT ''Maybe) inner) - else if "Either " `L.isPrefixOf` s - then do - let (left: right:_) = tail (words s) - lhs <- typeFromString left - rhs <- typeFromString right - return (AppT (AppT (ConT ''Either) lhs) rhs) - else fail $ "Unsupported type: " ++ s + 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 types = map (columnTypeString . (`unsafeGetColumn` df)) names - specs = zip names types + specs = zipWith (\name type_ -> (sanitize name, type_)) names types in fmap concat $ forM specs $ \(nm, tyStr) -> do - ty <- typeFromString tyStr + 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) |]) [] 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 @@ -24,7 +24,7 @@ import Control.Monad (forM_, zipWithM_, unless, void) import Data.Attoparsec.Text import Data.Char -import DataFrame.Internal.Column (Column(..), freezeColumn', writeColumn, columnLength) +import DataFrame.Internal.Column (Column(..), MutableColumn(..), freezeColumn', writeColumn, columnLength) import DataFrame.Internal.DataFrame (DataFrame(..)) import DataFrame.Internal.Parsing import DataFrame.Operations.Typing @@ -107,13 +107,13 @@ } {-# INLINE readSeparated #-} -getInitialDataVectors :: Int -> VM.IOVector Column -> [T.Text] -> IO () +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" -> MutableUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c) - "Double" -> MutableUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c) - _ -> MutableBoxedColumn <$> ((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 #-} @@ -128,7 +128,7 @@ {-# INLINE inferValueType #-} -- | Reads rows from the handle and stores values in mutable vectors. -fillColumns :: Int -> Char -> VM.IOVector Column -> VM.IOVector [(Int, T.Text)] -> Handle -> IO () +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 @@ -148,7 +148,7 @@ {-# INLINE fillColumns #-} -- | Writes a value into the appropriate column, resizing the vector if necessary. -writeValue :: VM.IOVector Column -> VM.IOVector [(Int, T.Text)] -> Int -> Int -> T.Text -> IO () +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 @@ -157,7 +157,7 @@ {-# INLINE writeValue #-} -- | Freezes a mutable vector into an immutable one, trimming it to the actual row count. -freezeColumn :: VM.IOVector Column -> V.Vector [(Int, T.Text)] -> ReadOptions -> Int -> IO Column +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 diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet.hs @@ -0,0 +1,1560 @@ +{-# 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)) diff --git a/src/DataFrame/IO/Parquet/ColumnStatistics.hs b/src/DataFrame/IO/Parquet/ColumnStatistics.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/ColumnStatistics.hs @@ -0,0 +1,19 @@ +module DataFrame.IO.Parquet.ColumnStatistics where + +import Data.Int +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) + +emptyColumnStatistics :: ColumnStatistics +emptyColumnStatistics = ColumnStatistics [] [] 0 0 [] [] False False diff --git a/src/DataFrame/IO/Parquet/Compression.hs b/src/DataFrame/IO/Parquet/Compression.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Compression.hs @@ -0,0 +1,26 @@ +module DataFrame.IO.Parquet.Compression where + +import Data.Int + +data CompressionCodec + = UNCOMPRESSED + | SNAPPY + | GZIP + | LZO + | BROTLI + | LZ4 + | ZSTD + | LZ4_RAW + | COMPRESSION_CODEC_UNKNOWN + deriving (Show, Eq) + +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/IO/Parquet/Encoding.hs b/src/DataFrame/IO/Parquet/Encoding.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Encoding.hs @@ -0,0 +1,28 @@ +module DataFrame.IO.Parquet.Encoding where + +import Data.Int + +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 diff --git a/src/DataFrame/IO/Parquet/Page.hs b/src/DataFrame/IO/Parquet/Page.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Page.hs @@ -0,0 +1,19 @@ +module DataFrame.IO.Parquet.Page where + +import Data.Int + +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 + diff --git a/src/DataFrame/IO/Parquet/Types.hs b/src/DataFrame/IO/Parquet/Types.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Types.hs @@ -0,0 +1,26 @@ +module DataFrame.IO.Parquet.Types where + +import Data.Int + +data ParquetType + = PBOOLEAN + | PINT32 + | PINT64 + | PINT96 + | PFLOAT + | PDOUBLE + | PBYTE_ARRAY + | PFIXED_LEN_BYTE_ARRAY + | PARQUET_TYPE_UNKNOWN + deriving (Show, Eq) + +parquetTypeFromInt :: Int32 -> ParquetType +parquetTypeFromInt 0 = PBOOLEAN +parquetTypeFromInt 1 = PINT32 +parquetTypeFromInt 2 = PINT64 +parquetTypeFromInt 3 = PINT96 +parquetTypeFromInt 4 = PFLOAT +parquetTypeFromInt 5 = PDOUBLE +parquetTypeFromInt 6 = PBYTE_ARRAY +parquetTypeFromInt 7 = PFIXED_LEN_BYTE_ARRAY +parquetTypeFromInt _ = PARQUET_TYPE_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 @@ -51,13 +51,11 @@ 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 - GroupedBoxedColumn :: Columnable a => VB.Vector (VB.Vector a) -> Column - GroupedUnboxedColumn :: (Columnable a, VU.Unbox a) => VB.Vector (VU.Vector a) -> Column - GroupedOptionalColumn :: (Columnable a) => VB.Vector (VB.Vector (Maybe a)) -> Column - -- These are used purely for I/O, not to store live data. - MutableBoxedColumn :: Columnable a => VBM.IOVector a -> Column - MutableUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> Column +data MutableColumn where + 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. data TypedColumn a where @@ -67,13 +65,6 @@ unwrapTypedColumn :: TypedColumn a -> Column unwrapTypedColumn (TColumn value) = value --- | An internal function that checks if a column --- can be used for aggregation. -isGrouped :: Column -> Bool -isGrouped (GroupedBoxedColumn column) = True -isGrouped (GroupedUnboxedColumn column) = True -isGrouped _ = False - -- | An internal function that checks if a column can -- be used in missing value operations. isOptional :: Column -> Bool @@ -86,10 +77,6 @@ BoxedColumn _ -> "Boxed" UnboxedColumn _ -> "Unboxed" OptionalColumn _ -> "Optional" - GroupedBoxedColumn _ -> "Grouped Boxed" - GroupedUnboxedColumn _ -> "Grouped Unboxed" - GroupedOptionalColumn _ -> "Grouped Optional" - _ -> "Unknown column string" -- | An internal/debugging function to get the type stored in the outermost vector -- of a column. @@ -98,9 +85,6 @@ BoxedColumn (column :: VB.Vector a) -> show (typeRep @a) UnboxedColumn (column :: VU.Vector a) -> show (typeRep @a) OptionalColumn (column :: VB.Vector a) -> show (typeRep @a) - GroupedBoxedColumn (column :: VB.Vector a) -> show (typeRep @a) - GroupedUnboxedColumn (column :: VB.Vector a) -> show (typeRep @a) - GroupedOptionalColumn (column :: VB.Vector a) -> show (typeRep @a) instance (Show a) => Show (TypedColumn a) where show (TColumn col) = show col @@ -110,8 +94,6 @@ show (BoxedColumn column) = show column show (UnboxedColumn column) = show column show (OptionalColumn column) = show column - show (GroupedBoxedColumn column) = show column - show (GroupedUnboxedColumn column) = show column instance Eq Column where (==) :: Column -> Column -> Bool @@ -127,69 +109,15 @@ case testEquality (typeRep @t1) (typeRep @t2) of Nothing -> False Just Refl -> a == b - -- Note: comparing grouped columns is expensive. We do this for stable tests - -- but also you should probably aggregate grouped columns soon after creating them. - (==) (GroupedBoxedColumn (a :: VB.Vector t1)) (GroupedBoxedColumn (b :: VB.Vector t2)) = - case testEquality (typeRep @t1) (typeRep @t2) of - Nothing -> False - Just Refl -> VB.map (L.sort . VG.toList) a == VB.map (L.sort . VG.toList) b - (==) (GroupedUnboxedColumn (a :: VB.Vector t1)) (GroupedUnboxedColumn (b :: VB.Vector t2)) = - case testEquality (typeRep @t1) (typeRep @t2) of - Nothing -> False - Just Refl -> VB.map (L.sort . VG.toList) a == VB.map (L.sort . VG.toList) b - (==) (GroupedOptionalColumn (a :: VB.Vector t1)) (GroupedOptionalColumn (b :: VB.Vector t2)) = - case testEquality (typeRep @t1) (typeRep @t2) of - Nothing -> False - Just Refl -> VB.map (L.sort . VG.toList) a == VB.map (L.sort . VG.toList) b (==) _ _ = False --- | A type with column representations used to select the --- "right" representation when specializing the `toColumn` function. -data Rep - = RBoxed - | RUnboxed - | ROptional - | RGBoxed - | RGUnboxed - | RGOptional - --- | Type-level if statement. -type family If (cond :: Bool) (yes :: k) (no :: k) :: k where - 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 - --- | Compute the column representation tag for any ‘a’. -type family KindOf a :: Rep where - KindOf (Maybe a) = 'ROptional - KindOf (VB.Vector a) = 'RGBoxed - KindOf (VU.Vector a) = 'RGUnboxed - KindOf a = If (Unboxable a) 'RUnboxed 'RBoxed - -- | 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 -- | Constraint synonym for what we can put into columns. -type Columnable a = (Columnable' a, ColumnifyRep (KindOf a) a, UnboxIf a, SBoolI (Unboxable 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 @@ -203,14 +131,6 @@ => ColumnifyRep 'ROptional (Maybe a) where toColumnRep = OptionalColumn -instance Columnable a - => ColumnifyRep 'RGBoxed (VB.Vector a) where - toColumnRep = GroupedBoxedColumn - -instance (Columnable a, VU.Unbox a) - => ColumnifyRep 'RGUnboxed (VU.Vector a) where - toColumnRep = GroupedUnboxedColumn - {- | O(n) Convert a vector to a column. Automatically picks the best representation of a vector to store the underlying data in. __Examples:__ @@ -253,28 +173,6 @@ => [a] -> Column fromList = toColumnRep @(KindOf a) . VB.fromList --- | Type-level boolean for constraint/type comparison. -data SBool (b :: Bool) where - STrue :: SBool 'True - SFalse :: SBool 'False - --- | The runtime witness for our type-level branching. -class SBoolI (b :: Bool) where - sbool :: SBool b - -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 = sbool @(Unboxable a) - -type family When (flag :: Bool) (c :: Constraint) :: Constraint where - When 'True c = c - When 'False c = () -- empty constraint - -type UnboxIf a = When (Unboxable a) (VU.Unbox a) - -- | An internal function to map a function over the values of a column. mapColumn :: forall b c. @@ -299,18 +197,6 @@ STrue -> UnboxedColumn (VU.map f col) SFalse -> fromVector @c (VB.map f (VB.convert col)) | otherwise -> Nothing - GroupedBoxedColumn (col :: VB.Vector (VB.Vector a)) - | Just Refl <- testEquality (typeRep @(VB.Vector a)) (typeRep @b) - -> Just (fromVector @c (VB.map f col)) - | otherwise -> Nothing - GroupedUnboxedColumn (col :: VB.Vector (VU.Vector a)) - | Just Refl <- testEquality (typeRep @(VU.Vector a)) (typeRep @b) - -> Just (fromVector @c (VB.map f col)) - | otherwise -> Nothing - GroupedOptionalColumn (col :: VB.Vector (VB.Vector a)) - | Just Refl <- testEquality (typeRep @(VB.Vector a)) (typeRep @b) - -> Just (fromVector @c (VB.map f col)) - | otherwise -> Nothing -- | O(1) Gets the number of elements in the column. @@ -318,19 +204,20 @@ columnLength (BoxedColumn xs) = VG.length xs columnLength (UnboxedColumn xs) = VG.length xs columnLength (OptionalColumn xs) = VG.length xs -columnLength (GroupedBoxedColumn xs) = VG.length xs -columnLength (GroupedUnboxedColumn xs) = VG.length xs -columnLength (GroupedOptionalColumn xs) = VG.length xs {-# INLINE columnLength #-} +-- | O(n) Gets the number of elements in the column. +numElements :: Column -> Int +numElements (BoxedColumn xs) = VG.length xs +numElements (UnboxedColumn xs) = VG.length xs +numElements (OptionalColumn xs) = VG.foldl' (\acc x -> acc + (fromEnum (isJust x))) 0 xs +{-# INLINE numElements #-} + -- | O(n) Takes the first n values of a column. takeColumn :: Int -> Column -> Column takeColumn n (BoxedColumn xs) = BoxedColumn $ VG.take n xs takeColumn n (UnboxedColumn xs) = UnboxedColumn $ VG.take n xs takeColumn n (OptionalColumn xs) = OptionalColumn $ VG.take n xs -takeColumn n (GroupedBoxedColumn xs) = GroupedBoxedColumn $ VG.take n xs -takeColumn n (GroupedUnboxedColumn xs) = GroupedUnboxedColumn $ VG.take n xs -takeColumn n (GroupedOptionalColumn xs) = GroupedOptionalColumn $ VG.take n xs {-# INLINE takeColumn #-} -- | O(n) Takes the last n values of a column. @@ -343,9 +230,6 @@ sliceColumn start n (BoxedColumn xs) = BoxedColumn $ VG.slice start n xs sliceColumn start n (UnboxedColumn xs) = UnboxedColumn $ VG.slice start n xs sliceColumn start n (OptionalColumn xs) = OptionalColumn $ VG.slice start n xs -sliceColumn start n (GroupedBoxedColumn xs) = GroupedBoxedColumn $ VG.slice start n xs -sliceColumn start n (GroupedUnboxedColumn xs) = GroupedUnboxedColumn $ VG.slice start n xs -sliceColumn start n (GroupedOptionalColumn xs) = GroupedOptionalColumn $ VG.slice start n xs {-# INLINE sliceColumn #-} -- | O(n) Selects the elements at a given set of indices. May change the order. @@ -353,9 +237,6 @@ atIndices indexes (BoxedColumn column) = BoxedColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column atIndices indexes (OptionalColumn column) = OptionalColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column atIndices indexes (UnboxedColumn column) = UnboxedColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column -atIndices indexes (GroupedBoxedColumn column) = GroupedBoxedColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column -atIndices indexes (GroupedUnboxedColumn column) = GroupedUnboxedColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column -atIndices indexes (GroupedOptionalColumn column) = GroupedOptionalColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column {-# INLINE atIndices #-} -- | O(n) Selects the elements at a given set of indices. Does not change the order. @@ -363,9 +244,6 @@ atIndicesStable indexes (BoxedColumn column) = BoxedColumn $ indexes `getIndices` column atIndicesStable indexes (UnboxedColumn column) = UnboxedColumn $ indexes `getIndicesUnboxed` column atIndicesStable indexes (OptionalColumn column) = OptionalColumn $ indexes `getIndices` column -atIndicesStable indexes (GroupedBoxedColumn column) = GroupedBoxedColumn $ indexes `getIndices` column -atIndicesStable indexes (GroupedUnboxedColumn column) = GroupedUnboxedColumn $ indexes `getIndices` column -atIndicesStable indexes (GroupedOptionalColumn column) = GroupedOptionalColumn $ indexes `getIndices` column {-# INLINE atIndicesStable #-} -- | Internal helper to get indices in a boxed vector. @@ -391,12 +269,6 @@ findIndices pred (OptionalColumn (column :: VB.Vector (Maybe b))) = do Refl <- testEquality (typeRep @a) (typeRep @(Maybe b)) pure $ VG.convert (VG.findIndices pred column) -findIndices pred (GroupedBoxedColumn (column :: VB.Vector b)) = do - Refl <- testEquality (typeRep @a) (typeRep @b) - pure $ VG.convert (VG.findIndices pred column) -findIndices pred (GroupedUnboxedColumn (column :: VB.Vector b)) = do - Refl <- testEquality (typeRep @a) (typeRep @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 @@ -415,21 +287,6 @@ 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 (GroupedBoxedColumn 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 (GroupedUnboxedColumn 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 (GroupedOptionalColumn 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. @@ -452,14 +309,6 @@ | Just Refl <- testEquality (typeRep @a) (typeRep @b) -> Just (fromVector @c (VB.imap f col)) | otherwise -> Nothing - GroupedBoxedColumn (col :: VB.Vector a) - | Just Refl <- testEquality (typeRep @a) (typeRep @b) - -> Just (fromVector @c (VB.imap f col)) - | otherwise -> Nothing - GroupedUnboxedColumn (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 @@ -469,12 +318,7 @@ ifilterColumn f c@(UnboxedColumn (column :: VU.Vector b)) = do Refl <- testEquality (typeRep @a) (typeRep @b) return $ UnboxedColumn $ VG.ifilter f column -ifilterColumn f c@(GroupedBoxedColumn (column :: VB.Vector b)) = do - Refl <- testEquality (typeRep @a) (typeRep @b) - return $ GroupedBoxedColumn $ VG.ifilter f column -ifilterColumn f c@(GroupedUnboxedColumn (column :: VB.Vector b)) = do - Refl <- testEquality (typeRep @a) (typeRep @b) - return $ GroupedUnboxedColumn $ 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 @@ -487,12 +331,6 @@ ifoldrColumn f acc c@(UnboxedColumn (column :: VU.Vector d)) = do Refl <- testEquality (typeRep @a) (typeRep @d) return $ VG.ifoldr f acc column -ifoldrColumn f acc c@(GroupedBoxedColumn (column :: VB.Vector d)) = do - Refl <- testEquality (typeRep @a) (typeRep @d) - return $ VG.ifoldr f acc column -ifoldrColumn f acc c@(GroupedUnboxedColumn (column :: VB.Vector d)) = do - 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 @@ -505,13 +343,18 @@ ifoldlColumn f acc c@(UnboxedColumn (column :: VU.Vector d)) = do Refl <- testEquality (typeRep @a) (typeRep @d) return $ VG.ifoldl' f acc column -ifoldlColumn f acc c@(GroupedBoxedColumn (column :: VB.Vector d)) = do - Refl <- testEquality (typeRep @a) (typeRep @d) - return $ VG.ifoldl' f acc column -ifoldlColumn f acc c@(GroupedUnboxedColumn (column :: VB.Vector d)) = do - Refl <- testEquality (typeRep @a) (typeRep @d) - return $ VG.ifoldl' f acc column +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) +headColumn (UnboxedColumn (col :: VU.Vector b)) = do + 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) + -- | Generic reduce function for all Column types. reduceColumn :: forall a b. Columnable a => (a -> b) -> Column -> Maybe b {-# SPECIALIZE reduceColumn :: @@ -526,15 +369,21 @@ reduceColumn f (OptionalColumn (column :: c)) = do Refl <- testEquality (typeRep @c) (typeRep @a) pure $ f column -reduceColumn _ _ = error "UNIMPLEMENTED" {-# INLINE reduceColumn #-} -- | An internal, column version of zip. zipColumns :: Column -> Column -> Column 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. @@ -553,14 +402,14 @@ {-# INLINE zipWithColumns #-} -- Functions for mutable columns (intended for IO). -writeColumn :: Int -> T.Text -> Column -> IO (Either T.Text Bool) -writeColumn i value (MutableBoxedColumn (col :: VBM.IOVector a)) = let +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 then VBM.unsafeWrite col i "" >> return (Left $! value) else VBM.unsafeWrite col i value >> return (Right True)) Nothing -> return (Left value) -writeColumn i value (MutableUnboxedColumn (col :: VUM.IOVector a)) = +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) @@ -572,12 +421,12 @@ Nothing -> VUM.unsafeWrite col i 0 >> return (Left $! value) {-# INLINE writeColumn #-} -freezeColumn' :: [(Int, T.Text)] -> Column -> IO Column -freezeColumn' nulls (MutableBoxedColumn col) +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 -freezeColumn' nulls (MutableUnboxedColumn 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)) @@ -592,12 +441,6 @@ 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 -expandColumn n column@(GroupedBoxedColumn col) - | n > VG.length col = GroupedBoxedColumn $ col <> VB.replicate (n - VG.length col) VB.empty - | otherwise = column -expandColumn n column@(GroupedUnboxedColumn col) - | n > VG.length col = GroupedUnboxedColumn $ col <> VB.replicate (n - VG.length col) VU.empty - | 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 @@ -610,12 +453,6 @@ 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 -leftExpandColumn n column@(GroupedBoxedColumn col) - | n > VG.length col = GroupedBoxedColumn $ VG.replicate (n - VG.length col) VB.empty <> col - | otherwise = column -leftExpandColumn n column@(GroupedUnboxedColumn col) - | n > VG.length col = GroupedUnboxedColumn $ VG.replicate (n - VG.length col) VU.empty <> col - | otherwise = column -- | Concatenates two columns. concatColumns :: Column -> Column -> Maybe Column @@ -628,12 +465,6 @@ concatColumns (UnboxedColumn left) (UnboxedColumn right) = case testEquality (typeOf left) (typeOf right) of Nothing -> Nothing Just Refl -> Just (UnboxedColumn $ left <> right) -concatColumns (GroupedBoxedColumn left) (GroupedBoxedColumn right) = case testEquality (typeOf left) (typeOf right) of - Nothing -> Nothing - Just Refl -> Just (GroupedBoxedColumn $ left <> right) -concatColumns (GroupedUnboxedColumn left) (GroupedUnboxedColumn right) = case testEquality (typeOf left) (typeOf right) of - Nothing -> Nothing - Just Refl -> Just (GroupedUnboxedColumn $ left <> right) concatColumns _ _ = Nothing {- | O(n) Converts a column to a boxed vector. Throws an exception if the wrong type is specified. @@ -652,39 +483,41 @@ Left err -> throw err Right val -> val +{- | O(n) Converts a column to a list. Throws an exception if the wrong type is specified. + +__Examples:__ + +@ +> column = fromList [(1 :: Int), 2, 3, 4] +> toList @Int column +[1,2,3,4] +> toList @Double column +exception: ... +-} +toList :: forall a . Columnable a => Column -> [a] +toList xs = case toVectorSafe @a xs of + Left err -> throw err + Right val -> VB.toList val + -- | A safe version of toVector that returns an Either type. -toVectorSafe :: forall a . Columnable a => Column -> Either DataFrameException (VB.Vector 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 col + 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 col + 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 $ VB.convert col - Nothing -> Left $ TypeMismatchException (MkTypeErrorContext { userType = Right (typeRep @a) - , expectedType = Right (typeRep @b) - , callingFunctionName = Just "toVectorSafe" - , errorColumnName = Nothing}) -toVectorSafe (GroupedBoxedColumn (col :: VB.Vector b)) = - case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> Right col - Nothing -> Left $ TypeMismatchException (MkTypeErrorContext { userType = Right (typeRep @a) - , expectedType = Right (typeRep @b) - , callingFunctionName = Just "toVectorSafe" - , errorColumnName = Nothing}) -toVectorSafe (GroupedUnboxedColumn (col :: VB.Vector b)) = - case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> Right col + Just Refl -> Right $ VG.convert col Nothing -> Left $ TypeMismatchException (MkTypeErrorContext { userType = Right (typeRep @a) , expectedType = Right (typeRep @b) , callingFunctionName = Just "toVectorSafe" 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 @@ -11,14 +11,16 @@ 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 Control.Monad (join) import DataFrame.Display.Terminal.PrettyPrint import DataFrame.Internal.Column import Data.Function (on) -import Data.List (sortBy, transpose) +import Data.List (sortBy, transpose, (\\)) import Data.Maybe (isJust) import Data.Type.Equality (type (:~:)(Refl), TestEquality (testEquality)) +import Text.Printf import Type.Reflection (typeRep) data DataFrame = DataFrame @@ -30,6 +32,22 @@ dataframeDimensions :: (Int, Int) } +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)) + +instance Eq GroupedDataFrame where + (==) (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) && @@ -47,8 +65,6 @@ 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) - getType (GroupedBoxedColumn (column :: V.Vector a)) = T.pack $ show (typeRep @a) - getType (GroupedUnboxedColumn (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 @@ -59,15 +75,13 @@ 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 - get (Just (GroupedBoxedColumn column)) = V.map (T.pack . show) column - get (Just (GroupedUnboxedColumn 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 (if properMarkdown then showTableProperMarkdown else showTable) header ("Int":types) rows + in showTable properMarkdown header ("Int":types) rows -- | O(1) Creates an empty dataframe empty :: DataFrame @@ -85,3 +99,20 @@ null :: DataFrame -> Bool null df = V.null (columns df) + +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)]) + +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 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 @@ -9,6 +9,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE BangPatterns #-} module DataFrame.Internal.Expression where import qualified Data.Map as M @@ -51,10 +52,12 @@ ReductionAggregate :: (Columnable a) => T.Text -- Column name -> T.Text -- Operation name - -> (forall v a. (VG.Vector v a, Columnable a) => v a -> a) + -> (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 @@ -78,32 +81,142 @@ (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 (GeneralAggregate name op (f :: forall v b. (VG.Vector v b, Columnable b) => v b -> c)) = case getColumn name df of +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 +interpret _ expr = error ("Invalid operation for dataframe: " ++ show expr) + +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 (GroupedBoxedColumn col) -> TColumn $ fromVector $ VG.map f col - Just (GroupedUnboxedColumn col) -> TColumn $ fromVector $ VG.map f col - Just (GroupedOptionalColumn col) -> TColumn $ fromVector $ VG.map f col - _ -> error "" -interpret df (ReductionAggregate name op (f :: forall v a. (VG.Vector v a, Columnable a) => v a -> a)) = case getColumn name df of + Just col -> TColumn $ atIndicesStable (VG.map (indices `VG.unsafeIndex`) (VG.init os)) col +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 + (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 +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 (GroupedBoxedColumn col) -> TColumn $ fromVector $ VG.map f col - Just (GroupedUnboxedColumn col) -> TColumn $ fromVector $ VG.map f col - Just (GroupedOptionalColumn col) -> TColumn $ fromVector $ VG.map f col - _ -> error "" -interpret df (NumericAggregate name op (f :: VU.Vector b -> c)) = case getColumn name df of + 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 Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df) - Just (GroupedUnboxedColumn (col :: V.Vector (VU.Vector d))) -> case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> TColumn $ fromVector $ VG.map f col - -- Do the matching trick here. - Nothing -> case testEquality (typeRep @d) (typeRep @Int) of - Just Refl -> case testEquality (typeRep @b) (typeRep @Double) of - Just Refl -> TColumn $ fromVector $ VG.map (f . (VG.map fromIntegral)) col - Nothing -> error $ "Column not a number: " ++ (T.unpack name) - Nothing -> case testEquality (typeRep @d) (typeRep @Double) of - Just Refl -> case testEquality (typeRep @b) (typeRep @Int) of - Just Refl -> TColumn $ fromVector $ VG.map (f . (VG.map round)) col - Nothing -> error $ "Column not a number: " ++ (T.unpack name) - _ -> error "Cannot apply numeric aggregation to boxed column" + 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 +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 + 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))))) + ) + ) + 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 -> 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 (+) :: Expr a -> Expr a -> Expr 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 @@ -3,6 +3,8 @@ {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE ScopedTypeVariables #-} module DataFrame.Internal.Row where import qualified Data.List as L @@ -26,31 +28,51 @@ import Data.Word ( Word8, Word16, Word32, Word64 ) import Type.Reflection (TypeRep, typeOf, typeRep) import Data.Type.Equality (TestEquality(..)) +import Text.ParserCombinators.ReadPrec(ReadPrec) +import Text.Read (Lexeme (Ident), lexP, parens, readListPrec, readListPrecDefault, readPrec) -data RowValue where - Value :: (Columnable' a) => a -> RowValue -instance Eq RowValue where - (==) :: RowValue -> RowValue -> Bool +data Any where + Value :: (Columnable' a) => a -> Any + +instance Eq Any where + (==) :: Any -> Any -> Bool (Value a) == (Value b) = fromMaybe False $ do Refl <- testEquality (typeOf a) (typeOf b) return $ a == b -instance Ord RowValue where - (<=) :: RowValue -> RowValue -> Bool +instance Ord Any where + (<=) :: Any -> Any -> Bool (Value a) <= (Value b) = fromMaybe False $ do Refl <- testEquality (typeOf a) (typeOf b) return $ a <= b -instance Show RowValue where - show :: RowValue -> String - show (Value a) = show a +instance Show Any where + show :: Any -> String + show (Value a) = T.unpack (showValue a) -toRowValue :: forall a . (Columnable' a) => a -> RowValue -toRowValue = Value +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 -type Row = V.Vector RowValue +instance Read Any where + readListPrec :: ReadPrec [Any] + readListPrec = readListPrecDefault + readPrec :: ReadPrec Any + readPrec = parens $ do + Ident "Value" <- lexP + readPrec + + +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 @@ -66,9 +88,9 @@ where get name = case getColumn name df of Nothing -> throw $ ColumnNotFoundException name "[INTERNAL] mkRowFromArgs" (map fst $ M.toList $ columnIndices df) - Just (BoxedColumn column) -> toRowValue (column V.! i) - Just (UnboxedColumn column) -> toRowValue (column VU.! i) - Just (OptionalColumn column) -> toRowValue (column V.! i) + 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)) @@ -82,19 +104,13 @@ ++ show i get name = case getColumn name df of Just (BoxedColumn c) -> case c V.!? i of - Just e -> toRowValue e + Just e -> toAny e Nothing -> throwError name Just (OptionalColumn c) -> case c V.!? i of - Just e -> toRowValue e + Just e -> toAny e Nothing -> throwError name Just (UnboxedColumn c) -> case c VU.!? i of - Just e -> toRowValue e - Nothing -> throwError name - Just (GroupedBoxedColumn c) -> case c V.!? i of - Just e -> toRowValue e - Nothing -> throwError name - Just (GroupedUnboxedColumn c) -> case c V.!? i of - Just e -> toRowValue e + Just e -> toAny e Nothing -> throwError name sortedIndexes' :: Bool -> V.Vector Row -> VU.Vector Int diff --git a/src/DataFrame/Internal/Schema.hs b/src/DataFrame/Internal/Schema.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Internal/Schema.hs @@ -0,0 +1,33 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +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 Type.Reflection (typeRep) + +data SchemaType where + SType :: Columnable a => P.Proxy a -> SchemaType + +instance Show SchemaType where + show (SType (_ :: P.Proxy a)) = show (typeRep @a) + +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 = SType (P.Proxy @a) + +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 @@ -2,20 +2,100 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE ExistentialQuantification #-} {-# LANGUAGE GADTs #-} -{-# LANGUAGE InstanceSigs #-} -{-# LANGUAGE KindSignatures #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} module DataFrame.Internal.Types where import Data.Int ( Int8, Int16, Int32, Int64 ) -import Data.Kind (Type) +import Data.Kind (Type, Constraint) import Data.Maybe (fromMaybe) 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 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. +data Rep + = 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 + +-- | 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 + +-- | 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 + +-- | 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 + +-- | Type-level boolean for constraint/type comparison. +data SBool (b :: Bool) where + STrue :: SBool 'True + SFalse :: SBool 'False + +-- | The runtime witness for our type-level branching. +class SBoolI (b :: Bool) where + sbool :: SBool b + +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 = sbool @(Unboxable 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 + +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 @@ -24,7 +24,7 @@ import Control.Monad (forM_, zipWithM_, unless, when, void, replicateM_) import Data.Attoparsec.Text import Data.Char -import DataFrame.Internal.Column (Column(..), freezeColumn', writeColumn, columnLength) +import DataFrame.Internal.Column (Column(..), MutableColumn(..), freezeColumn', writeColumn, columnLength) import DataFrame.Internal.DataFrame (DataFrame(..)) import DataFrame.Internal.Parsing import DataFrame.Operations.Typing @@ -122,13 +122,13 @@ }, (pos, unconsumed, r + 1)) {-# INLINE readSeparated #-} -getInitialDataVectors :: Int -> VM.IOVector Column -> [T.Text] -> IO () +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" -> MutableUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c) - "Double" -> MutableUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c) - _ -> MutableBoxedColumn <$> ((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 #-} @@ -154,7 +154,7 @@ return (row, unconsumed) -- | Reads rows from the handle and stores values in mutable vectors. -fillColumns :: Int -> Char -> VM.IOVector Column -> VM.IOVector [(Int, T.Text)] -> T.Text -> Handle -> IO (T.Text, Int) +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) @@ -179,7 +179,7 @@ {-# INLINE fillColumns #-} -- | Writes a value into the appropriate column, resizing the vector if necessary. -writeValue :: VM.IOVector Column -> VM.IOVector [(Int, T.Text)] -> Int -> Int -> T.Text -> IO () +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 @@ -188,7 +188,7 @@ {-# INLINE writeValue #-} -- | Freezes a mutable vector into an immutable one, trimming it to the actual row count. -freezeColumn :: VM.IOVector Column -> V.Vector [(Int, T.Text)] -> ReadOptions -> Int -> IO Column +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 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 @@ -27,14 +27,16 @@ import Control.Monad.ST (runST) import DataFrame.Internal.Column (Column(..), fromVector, getIndicesUnboxed, getIndices, + atIndicesStable, Columnable, unwrapTypedColumn, columnVersionString) -import DataFrame.Internal.DataFrame (DataFrame(..), empty, getColumn, unsafeGetColumn) +import DataFrame.Internal.DataFrame (GroupedDataFrame(..), DataFrame(..), 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 @@ -48,43 +50,37 @@ groupBy :: [T.Text] -> DataFrame -> - DataFrame + GroupedDataFrame groupBy names df | any (`notElem` columnNames df) names = throw $ ColumnNotFoundException (T.pack $ show $ names L.\\ columnNames df) "groupBy" (columnNames df) - | otherwise = L.foldl' insertColumns initDf groupingColumns + | 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 = V.fromList $ map (VG.map fst) $ VG.groupBy (\a b -> snd a == snd b) (runST $ do + valueIndices = runST $ do withIndexes <- VG.thaw $ VG.indexed rowRepresentations VA.sortBy (\(a, b) (a', b') -> compare b b') withIndexes - VG.unsafeFreeze withIndexes) + VG.unsafeFreeze withIndexes - -- These are the indexes of the grouping/key rows i.e the minimum elements - -- of the list. - keyIndices = VU.generate (VG.length valueIndices) (\i -> VG.minimum $ valueIndices VG.! i) - -- this will be our main worker function in the fold that takes all - -- indices and replaces each value in a column with a list of - -- the elements with the indices where the grouped row - -- values are the same. - insertColumns = groupColumns valueIndices df - -- Out initial DF will just be all the grouped rows added to an - -- empty dataframe. The entries are dedued and are in their - -- initial order. - initDf = L.foldl' (mkGroupedColumns keyIndices df) empty names - -- All the rest of the columns that we are grouping by. - groupingColumns = columnNames df L.\\ names +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) mkRowRep :: [Int] -> DataFrame -> Int -> Int -mkRowRep groupColumnIndices df i = if length h == 1 then head h else hash h +mkRowRep groupColumnIndices df i = case h of + (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) - getHashedElem _ _ = 0 mkHash j = getHashedElem ((V.!) (columns df) j) i -- | This hash function returns the hash when given a non numeric type but @@ -111,37 +107,21 @@ let vs = indices `getIndicesUnboxed` column in insertUnboxedVector name vs acc -groupColumns :: V.Vector (VU.Vector Int) -> DataFrame -> DataFrame -> T.Text -> DataFrame -groupColumns indices df acc name = - case (V.!) (columns df) (columnIndices df M.! name) of - BoxedColumn column -> - let vs = V.map (`getIndices` column) indices - in insertColumn name (GroupedBoxedColumn vs) acc - OptionalColumn column -> - let vs = V.map (`getIndices` column) indices - in insertColumn name (GroupedBoxedColumn vs) acc - UnboxedColumn column -> - let vs = V.map (`getIndicesUnboxed` column) indices - in insertColumn name (GroupedUnboxedColumn vs) acc - -aggregate :: [(T.Text, UExpr)] -> DataFrame -> DataFrame -aggregate aggs df = let - groupingColumns = Prelude.filter (\c -> not $ T.isPrefixOf "Grouped" (T.pack $ columnVersionString (fromMaybe (error "Unexpected") (getColumn c df)))) (columnNames df) - df' = select groupingColumns df +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 = interpret @a df expr + value = interpretAggregation @a gdf expr in insertColumn name (unwrapTypedColumn value) d in fold f aggs df' -distinct :: DataFrame -> DataFrame -distinct df = groupBy (columnNames df) 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)) + } -distinctBy :: [T.Text] -> DataFrame -> DataFrame -distinctBy names df = let - excluded = (columnNames df) \\ names - distinctColumns = groupBy names df - aggF name = case unsafeGetColumn name distinctColumns of - GroupedBoxedColumn (column :: V.Vector (V.Vector a)) -> (F.anyValue (F.col @a name)) `F.as` name - GroupedUnboxedColumn (column :: V.Vector (VU.Vector a)) -> (F.anyValue (F.col @a name)) `F.as` name - _ -> error $ "Column isn't grouped: " ++ (T.unpack name) - in aggregate (map aggF excluded) distinctColumns +distinct :: DataFrame -> DataFrame +distinct df = selectIndices (VG.map (indices VG.!) (VG.init os)) 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 @@ -107,13 +107,18 @@ values = xs V.++ V.replicate (rows - V.length xs) defaultValue in insertColumn name (fromVector values) d --- TODO: Add existence check in rename. rename :: T.Text -> T.Text -> DataFrame -> DataFrame -rename orig new df = fromMaybe (throw $ ColumnNotFoundException orig "rename" (map fst $ M.toList $ columnIndices df)) $ do +rename orig new df = either throw id (renameSafe orig new df) + +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 df { columnIndices = newAdded } + return (Right df { columnIndices = newAdded }) -- | O(1) Get the number of elements in a given column. columnSize :: T.Text -> DataFrame -> Maybe Int 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 @@ -43,4 +43,5 @@ instance Monoid D.DataFrame where mempty = D.empty - +(|||) :: 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/Statistics.hs b/src/DataFrame/Operations/Statistics.hs --- a/src/DataFrame/Operations/Statistics.hs +++ b/src/DataFrame/Operations/Statistics.hs @@ -8,6 +8,7 @@ {-# LANGUAGE BangPatterns #-} 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 @@ -22,52 +23,32 @@ import Control.Exception (throw) import DataFrame.Errors (DataFrameException(..)) import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame (DataFrame(..), getColumn, empty) +import DataFrame.Internal.DataFrame (DataFrame(..), getColumn, empty, unsafeGetColumn) 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) frequencies :: T.Text -> DataFrame -> DataFrame -frequencies name df = case getColumn name df of - Nothing -> throw $ ColumnNotFoundException name "frequencies" (map fst $ M.toList $ columnIndices df) - Just ((BoxedColumn (column :: V.Vector a))) -> let - counts = valueCounts @a name df - total = P.sum $ map snd counts - vText :: forall a . (Columnable a) => a -> T.Text - vText c' = case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> c' - Nothing -> case testEquality (typeRep @a) (typeRep @String) of - Just Refl -> T.pack c' - Nothing -> (T.pack . show) c' - initDf = empty & insertVector "Statistic" (V.fromList ["Count" :: T.Text, "Percentage (%)"]) - in L.foldl' (\df (col, k) -> insertVector (vText col) (V.fromList [k, k * 100 `div` total]) df) initDf counts - Just ((OptionalColumn (column :: V.Vector a))) -> let - counts = valueCounts @a name df - total = P.sum $ map snd counts - vText :: forall a . (Columnable a) => a -> T.Text - vText c' = case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> c' - Nothing -> case testEquality (typeRep @a) (typeRep @String) of - Just Refl -> T.pack c' - Nothing -> (T.pack . show) c' - initDf = empty & insertVector "Statistic" (V.fromList ["Count" :: T.Text, "Percentage (%)"]) - in L.foldl' (\df (col, k) -> insertVector (vText col) (V.fromList [k, k * 100 `div` total]) df) initDf counts - Just ((UnboxedColumn (column :: VU.Vector a))) -> let - counts = valueCounts @a name df - total = P.sum $ map snd counts - vText :: forall a . (Columnable a) => a -> T.Text - vText c' = case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> c' - Nothing -> case testEquality (typeRep @a) (typeRep @String) of - Just Refl -> T.pack c' - Nothing -> (T.pack . show) c' - initDf = empty & insertVector "Statistic" (V.fromList ["Count" :: T.Text, "Percentage (%)"]) - in L.foldl' (\df (col, k) -> insertVector (vText col) (V.fromList [k, k * 100 `div` total]) df) initDf counts - _ -> error $ "There are ungrouped columns" +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 mean :: T.Text -> DataFrame -> Maybe Double mean = applyStatistic mean' @@ -110,7 +91,7 @@ Nothing -> Nothing applyStatistic :: (VU.Vector Double -> Double) -> T.Text -> DataFrame -> Maybe Double -applyStatistic f name df = case getColumn name df of +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 @@ -123,7 +104,7 @@ _ -> Nothing applyStatistics :: (VU.Vector Double -> VU.Vector Double) -> T.Text -> DataFrame -> Maybe (VU.Vector Double) -applyStatistics f name df = case getColumn name df of +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 @@ -134,9 +115,10 @@ _ -> Nothing summarize :: DataFrame -> DataFrame -summarize df = fold columnStats (columnNames df) (fromNamedColumns [("Statistic", fromList ["Mean" :: T.Text, "Minimum", "25%" ,"Median", "75%", "Max", "StdDev", "IQR", "Skewness"])]) +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 min' = flip (VG.!) 0 <$> quantiles quartile1 = flip (VG.!) 1 <$> quantiles @@ -144,7 +126,8 @@ quartile3 = flip (VG.!) 3 <$> quantiles max' = flip (VG.!) 4 <$> quantiles iqr = (-) <$> quartile3 <*> quartile1 - in [mean name df, + in [count, + mean name df, min', quartile1, median', @@ -153,8 +136,15 @@ standardDeviation name df, iqr, skewness name df] - roundTo :: Int -> Double -> Double - roundTo n x = fromInteger (round $ x * (10^n)) / (10.0^^n) + +-- | Round a @Double@ to Specified Precision +roundTo :: Int -> Double -> Double +roundTo n x = fromInteger (round $ x * (10^n)) / (10.0^^n) + +toPct2dp :: Double -> String +toPct2dp x + | x < 0.00005 = "<0.01%" + | otherwise = printf "%.2f%%" (x * 100) mean' :: VU.Vector Double -> Double mean' samp = let 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 @@ -4,6 +4,8 @@ {-# LANGUAGE TypeApplications #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitNamespaces #-} module DataFrame.Operations.Subset where import qualified Data.List as L @@ -12,6 +14,7 @@ import qualified Data.Text as T import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM import qualified Data.Vector.Generic as VG import qualified Prelude @@ -20,13 +23,15 @@ import DataFrame.Internal.Column import DataFrame.Internal.DataFrame (DataFrame(..), getColumn, empty) import DataFrame.Internal.Expression -import DataFrame.Internal.Row (mkRowFromArgs, RowValue, toRowValue) +import DataFrame.Internal.Row (mkRowFromArgs, Any, toAny) +import Data.Type.Equality (type (:~:)(Refl), TestEquality (..)) 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 -- | O(k * n) Take the first n rows of a DataFrame. take :: Int -> DataFrame -> DataFrame @@ -78,16 +83,34 @@ DataFrame filter filterColumnName condition df = case getColumn filterColumnName df of Nothing -> throw $ ColumnNotFoundException filterColumnName "filter" (map fst $ M.toList $ columnIndices df) - Just column -> case findIndices condition column of - Nothing -> throw $ TypeMismatchException (MkTypeErrorContext - { userType = Right $ typeRep @a - , expectedType = Left (columnTypeString column) :: Either String (TypeRep ()) - , errorColumnName = Just (T.unpack filterColumnName) - , callingFunctionName = Just "filter"}) - Just indexes -> let - c' = snd $ dataframeDimensions df - pick idxs col = atIndicesStable idxs col - in df {columns = V.map (pick indexes) (columns df), dataframeDimensions = (VG.length indexes, c')} + 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 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))} + +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') -- | O(k) a version of filter where the predicate comes first. -- 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 @@ -17,7 +17,7 @@ import DataFrame.Internal.Column (Column(..), columnTypeString, imapColumn, ifoldrColumn, TypedColumn (TColumn), Columnable, mapColumn, unwrapTypedColumn) import DataFrame.Internal.DataFrame (DataFrame(..), getColumn) import DataFrame.Internal.Expression -import DataFrame.Internal.Row (mkRowFromArgs, RowValue, toRowValue) +import DataFrame.Internal.Row (mkRowFromArgs, Any, toAny) import DataFrame.Operations.Core import Data.Maybe import Type.Reflection (typeRep, typeOf, TypeRep) diff --git a/tests/Functions.hs b/tests/Functions.hs new file mode 100644 --- /dev/null +++ b/tests/Functions.hs @@ -0,0 +1,42 @@ +{-# LANGUAGE OverloadedStrings #-} +module Functions where + +import Control.Exception +import Test.HUnit +import DataFrame.Functions (sanitize) +import qualified Data.Text as T + +-- 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 "***") + ] + +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 @@ -16,6 +16,7 @@ import Assertions +import qualified Functions import qualified Operations.Apply import qualified Operations.Derive import qualified Operations.Filter @@ -23,6 +24,7 @@ import qualified Operations.InsertColumn import qualified Operations.Sort import qualified Operations.Take +import qualified Parquet testData :: D.DataFrame testData = D.fromNamedColumns [ ("test1", DI.fromList ([1..26] :: [Int])) @@ -76,6 +78,8 @@ ++ Operations.InsertColumn.tests ++ Operations.Sort.tests ++ Operations.Take.tests + ++ Functions.tests + ++ Parquet.tests ++ parseTests main :: IO () diff --git a/tests/Operations/GroupBy.hs b/tests/Operations/GroupBy.hs --- a/tests/Operations/GroupBy.hs +++ b/tests/Operations/GroupBy.hs @@ -23,40 +23,15 @@ groupBySingleRowWAI :: Test groupBySingleRowWAI = TestCase (assertEqual "Groups by single column" - (D.fromNamedColumns [("test1", DI.fromList [(1::Int)..4]), - -- This just makes rows with [1, 2] for every unique test1 row - ("test2", DI.GroupedUnboxedColumn (V.replicate 4 $ VU.fromList (take 10 $ cycle [1 :: Int, 2]))), - ("test3", DI.GroupedUnboxedColumn (V.generate 4 (\i -> VU.fromList [(i * 10 + 1)..((i + 1) * 10)]))), - ("test4", DI.GroupedUnboxedColumn (V.generate 4 (\i -> VU.fromList [(((3 - i) + 1) * 10),(((3 - i) + 1) * 10 - 1)..((3 - i) * 10 + 1)]))) - ]) - (D.groupBy ["test1"] testData D.|> D.sortBy D.Ascending ["test1"])) + -- 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" - (D.fromNamedColumns [("test1", DI.fromList$ concatMap (replicate 2) [(1::Int)..4]), - ("test2", DI.fromList (take 8 $ cycle [1 :: Int, 2])), - ("test3", DI.GroupedUnboxedColumn (V.fromList [ - VU.fromList [1 :: Int,3..9], - VU.fromList [2,4..10], - VU.fromList [11,13..19], - VU.fromList [12,14..20], - VU.fromList [21,23..29], - VU.fromList [22,24..30], - VU.fromList [31,33..39], - VU.fromList [32,34..40] - ])), - ("test4", DI.GroupedUnboxedColumn (V.fromList $ reverse [ - VU.fromList [1 :: Int,3..9], - VU.fromList [2,4..10], - VU.fromList [11,13..19], - VU.fromList [12,14..20], - VU.fromList [21,23..29], - VU.fromList [22,24..30], - VU.fromList [31,33..39], - VU.fromList [32,34..40] - ])) - ]) - (D.groupBy ["test1", "test2"] testData D.|> D.sortBy D.Ascending ["test1", "test2"])) + -- 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]" diff --git a/tests/Parquet.hs b/tests/Parquet.hs new file mode 100644 --- /dev/null +++ b/tests/Parquet.hs @@ -0,0 +1,17 @@ +{-# LANGUAGE OverloadedStrings #-} +module Parquet where + +import qualified DataFrame as D + +import Assertions +import GHC.IO (unsafePerformIO) +import Test.HUnit + +-- TODO: This currently fails +allTypesPlain :: Test +allTypesPlain = TestCase (assertEqual "allTypesPlain" + (unsafePerformIO (D.readCsv "./data/mtcars.csv")) + (unsafePerformIO (D.readParquet "./data/mtcars.parquet"))) + +tests :: [Test] +tests = []