diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,9 +1,20 @@ # Revision history for dataframe +## 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. + ## 0.3.0.0 * Now supports inner joins -* Aggregations are now expressions allowing for more expressive aggregation logic. +```haskell +ghci> df |> D.innerJoin ["key_1", "key_2"] other +``` +* Aggregations are now expressions allowing for more expressive aggregation logic. Previously: `D.aggregate [("quantity", D.Mean), ("price", D.Sum)] df` now ``D.aggregate [(F.sum (F.col @Double "label") / (F.count (F.col @Double "label")) `F.as` "positive_rate")]`` * In GHCI, you can now create type-safe bindings for each column and use those in expressions. + +```haskell +ghci> :exposeColumns df +ghci> D.aggregate [(F.sum label / F.count label) `F.as` "positive_rate"] +``` * Added pandas and polars benchmarks. * Performance improvements to `groupBy`. * Various bug fixes. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -29,75 +29,16 @@ * Static typing makes code easier to reason about and catches many bugs at compile time—before your code ever runs. * Delivers high performance thanks to Haskell’s optimizing compiler and efficient memory model. * Designed for interactivity: expressive syntax, helpful error messages, and sensible defaults. +* Works seamlessly in both command-line and notebook environments—great for exploration and scripting alike. ## Example usage ### Interactive environment -```haskell -ghci> import qualified DataFrame as D -ghci> import DataFrame ((|>)) -ghci> df <- D.readCsv "./data/housing.csv" -ghci> D.columnInfo df --------------------------------------------------------------------------------------------------------------------- -index | Column Name | # Non-null Values | # Null Values | # Partially parsed | # Unique Values | Type -------|--------------------|-------------------|---------------|--------------------|-----------------|------------- - Int | Text | Int | Int | Int | Int | Text -------|--------------------|-------------------|---------------|--------------------|-----------------|------------- -0 | total_bedrooms | 20433 | 207 | 0 | 1924 | Maybe Double -1 | ocean_proximity | 20640 | 0 | 0 | 5 | Text -2 | median_house_value | 20640 | 0 | 0 | 3842 | Double -3 | median_income | 20640 | 0 | 0 | 12928 | Double -4 | households | 20640 | 0 | 0 | 1815 | Double -5 | population | 20640 | 0 | 0 | 3888 | Double -6 | total_rooms | 20640 | 0 | 0 | 5926 | Double -7 | housing_median_age | 20640 | 0 | 0 | 52 | Double -8 | latitude | 20640 | 0 | 0 | 862 | Double -9 | longitude | 20640 | 0 | 0 | 844 | Double -ghci> :exposeColumns df -ghci> import qualified DataFrame.Functions as F -ghci> df |> D.groupBy ["ocean_proximity"] |> D.aggregate [(F.mean median_house_value) `F.as` "avg_house_value" ] --------------------------------------------- -index | ocean_proximity | avg_house_value -------|-----------------|------------------- - Int | Text | Double -------|-----------------|------------------- -0 | <1H OCEAN | 240084.28546409807 -1 | INLAND | 124805.39200122119 -2 | ISLAND | 380440.0 -3 | NEAR BAY | 259212.31179039303 -4 | NEAR OCEAN | 249433.97742663656 -ghci> df |> D.derive "rooms_per_household" (total_rooms / households) |> D.take 10 --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- -index | longitude | latitude | housing_median_age | total_rooms | total_bedrooms | population | households | median_income | median_house_value | ocean_proximity | rooms_per_household -------|-----------|----------|--------------------|-------------|----------------|------------|------------|--------------------|--------------------|-----------------|-------------------- - Int | Double | Double | Double | Double | Maybe Double | Double | Double | Double | Double | Text | Double -------|-----------|----------|--------------------|-------------|----------------|------------|------------|--------------------|--------------------|-----------------|-------------------- -0 | -122.23 | 37.88 | 41.0 | 880.0 | Just 129.0 | 322.0 | 126.0 | 8.3252 | 452600.0 | NEAR BAY | 6.984126984126984 -1 | -122.22 | 37.86 | 21.0 | 7099.0 | Just 1106.0 | 2401.0 | 1138.0 | 8.3014 | 358500.0 | NEAR BAY | 6.238137082601054 -2 | -122.24 | 37.85 | 52.0 | 1467.0 | Just 190.0 | 496.0 | 177.0 | 7.2574 | 352100.0 | NEAR BAY | 8.288135593220339 -3 | -122.25 | 37.85 | 52.0 | 1274.0 | Just 235.0 | 558.0 | 219.0 | 5.6431000000000004 | 341300.0 | NEAR BAY | 5.8173515981735155 -4 | -122.25 | 37.85 | 52.0 | 1627.0 | Just 280.0 | 565.0 | 259.0 | 3.8462 | 342200.0 | NEAR BAY | 6.281853281853282 -5 | -122.25 | 37.85 | 52.0 | 919.0 | Just 213.0 | 413.0 | 193.0 | 4.0368 | 269700.0 | NEAR BAY | 4.761658031088083 -6 | -122.25 | 37.84 | 52.0 | 2535.0 | Just 489.0 | 1094.0 | 514.0 | 3.6591 | 299200.0 | NEAR BAY | 4.9319066147859925 -7 | -122.25 | 37.84 | 52.0 | 3104.0 | Just 687.0 | 1157.0 | 647.0 | 3.12 | 241400.0 | NEAR BAY | 4.797527047913447 -8 | -122.26 | 37.84 | 42.0 | 2555.0 | Just 665.0 | 1206.0 | 595.0 | 2.0804 | 226700.0 | NEAR BAY | 4.294117647058823 -9 | -122.25 | 37.84 | 52.0 | 3549.0 | Just 707.0 | 1551.0 | 714.0 | 3.6912000000000003 | 261100.0 | NEAR BAY | 4.970588235294118 -ghci> df |> D.derive "nonsense_feature" (latitude + ocean_proximity) |> D.take 10 - -:14:47: error: [GHC-83865] - • Couldn't match type ‘Text’ with ‘Double’ - Expected: Expr Double - Actual: Expr Text - • In the second argument of ‘(+)’, namely ‘ocean_proximity’ - In the second argument of ‘derive’, namely - ‘(latitude + ocean_proximity)’ - In the second argument of ‘(|>)’, namely - ‘derive "nonsense_feature" (latitude + ocean_proximity)’ -``` +![Screencast of usage in GHCI](./static/example.gif) Key features in example: * Intuitive, SQL-like API to get from data to insights. -* Create type-safe references to columns in a dataframe using `:exponseColumns` +* Create typed, completion-ready references to columns in a dataframe using `:exposeColumns` * Type-safe column transformations for faster and safer exploration. * Fluid, chaining API that makes code easy to reason about. @@ -150,22 +91,19 @@ Full example in `./examples` folder using many of the constructs in the API. -### Visual example -![Screencast of usage in GHCI](./static/example.gif) - ## Installing ### Jupyter notebook -* We have a [hosted version of the Jupyter notebook](https://ihaskell-dataframe-crf7g5fvcpahdegz.westus2-01.azurewebsites.net/lab/) on azure sites. -* Use the Dockerfile in the [ihaskell-dataframe](https://github.com/mchav/ihaskell-dataframe) to build and run an image with dataframe integration. +* 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. +* 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. ### CLI -* Install Haskell (ghc + cabal) via [ghcup](https://www.haskell.org/ghcup/install/) selecting all the default options. -* Install snappy (needed for Parquet support) by running: `sudo apt install libsnappy-dev`. -* To install dataframe run `cabal update && cabal install dataframe` -* Open a Haskell repl with dataframe loaded by running `cabal repl --build-depends dataframe`. -* Follow along any one of the tutorials below. +* Run the installation script `curl '=https' --tlsv1.2 -sSf https://raw.githubusercontent.com/mchav/dataframe/refs/heads/main/scripts/install.sh | sh` +* Download the run script with: `curl --output dataframe "https://raw.githubusercontent.com/mchav/dataframe/refs/heads/main/scripts/dataframe.sh"` +* Make the script executable: `chmod +x dataframe` +* Add the script your path: `export PATH=$PATH:./dataframe` +* Run the script with: `dataframe` ## What is exploratory data analysis? 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.0 +version: 0.3.0.1 synopsis: A fast, safe, and intuitive DataFrame library. @@ -46,7 +46,6 @@ DataFrame.Operations.Aggregation, DataFrame.Display.Terminal.Plot, DataFrame.IO.CSV, - DataFrame.IO.Parquet, DataFrame.Lazy.IO.CSV, DataFrame.Lazy.Internal.DataFrame build-depends: base >= 4.17.2.0 && < 4.22, @@ -57,7 +56,6 @@ directory >= 1.3.0.0 && <= 1.3.9.0, filepath >= 1.0.0.0 && <= 1.5.4.0, hashable >= 1.2 && <= 1.5.0.0, - snappy >= 0.2.0.0 && <= 0.2.0.4, statistics >= 0.16.2.1 && <= 0.16.3.0, template-haskell >= 2.0 && <= 2.30, text >= 2.0 && <= 2.1.2, @@ -91,7 +89,6 @@ DataFrame.IO.CSV, DataFrame.Operations.Join, DataFrame.Operations.Merge, - DataFrame.IO.Parquet, DataFrame.Lazy.IO.CSV, DataFrame.Functions build-depends: base >= 4.17.2.0 && < 4.22, @@ -101,7 +98,6 @@ containers >= 0.6.7 && < 0.8, directory >= 1.3.0.0 && <= 1.3.9.0, hashable >= 1.2 && <= 1.5.0.0, - snappy >= 0.2.0.0 && <= 0.2.0.4, statistics >= 0.16.2.1 && <= 0.16.3.0, template-haskell >= 2.0 && <= 2.30, text >= 2.0 && <= 2.1.2, @@ -136,7 +132,6 @@ DataFrame.IO.CSV, DataFrame.Operations.Join, DataFrame.Operations.Merge, - DataFrame.IO.Parquet, DataFrame.Lazy.IO.CSV, DataFrame.Functions build-depends: base >= 4.17.2.0 && < 4.22, @@ -146,7 +141,6 @@ containers >= 0.6.7 && < 0.8, directory >= 1.3.0.0 && <= 1.3.9.0, hashable >= 1.2 && <= 1.5.0.0, - snappy >= 0.2.0.0 && <= 0.2.0.4, statistics >= 0.16.2.1 && <= 0.16.3.0, template-haskell >= 2.0 && <= 2.30, text >= 2.0 && <= 2.1.2, @@ -181,7 +175,6 @@ DataFrame.IO.CSV, DataFrame.Operations.Join, DataFrame.Operations.Merge, - DataFrame.IO.Parquet, DataFrame.Lazy.IO.CSV, DataFrame.Functions build-depends: base >= 4.17.2.0 && < 4.22, @@ -191,8 +184,7 @@ containers >= 0.6.7 && < 0.8, directory >= 1.3.0.0 && <= 1.3.9.0, hashable >= 1.2 && <= 1.5.0.0, - snappy >= 0.2.0.0 && <= 0.2.0.4, - statistics >= 0.16.2.1 && <= 0.16.3.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, @@ -226,7 +218,6 @@ DataFrame.IO.CSV, DataFrame.Operations.Join, DataFrame.Operations.Merge, - DataFrame.IO.Parquet, DataFrame.Lazy.IO.CSV, DataFrame.Functions build-depends: base >= 4.17.2.0 && < 4.22, @@ -237,7 +228,6 @@ directory >= 1.3.0.0 && <= 1.3.9.0, hashable >= 1.2 && <= 1.5.0.0, random >= 1 && <= 1.3.1, - snappy >= 0.2.0.0 && <= 0.2.0.4, statistics >= 0.16.2.1 && <= 0.16.3.0, template-haskell >= 2.0 && <= 2.30, text >= 2.0 && <= 2.1.2, diff --git a/examples/CaliforniaHousing.hs b/examples/CaliforniaHousing.hs --- a/examples/CaliforniaHousing.hs +++ b/examples/CaliforniaHousing.hs @@ -7,7 +7,7 @@ main = do parsed <- D.readCsv "./data/housing.csv" - print $ D.columnInfo parsed + print $ D.describeColumns parsed print $ D.take 5 parsed diff --git a/src/DataFrame.hs b/src/DataFrame.hs --- a/src/DataFrame.hs +++ b/src/DataFrame.hs @@ -24,8 +24,6 @@ import DataFrame.Operations.Aggregation as D import DataFrame.Display.Terminal.Plot as D import DataFrame.IO.CSV as D --- Support for Parquet is still experimental -import DataFrame.IO.Parquet as D import Data.Function import Data.List diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs --- a/src/DataFrame/Functions.hs +++ b/src/DataFrame/Functions.hs @@ -62,6 +62,9 @@ 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 diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs deleted file mode 100644 --- a/src/DataFrame/IO/Parquet.hs +++ /dev/null @@ -1,1554 +0,0 @@ -{-# LANGUAGE ExplicitNamespaces #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE RecordWildCards #-} - -module DataFrame.IO.Parquet ( - readParquet - ) where - -import qualified Codec.Compression.Snappy as Snappy -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 Foreign -import GHC.Float -import GHC.IO (unsafePerformIO) -import System.IO - -footerSize :: Integer -footerSize = 8 - -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 - -data ParquetEncoding - = EPLAIN - | EPLAIN_DICTIONARY - | ERLE - | EBIT_PACKED - | EDELTA_BINARY_PACKED - | EDELTA_LENGTH_BYTE_ARRAY - | EDELTA_BYTE_ARRAY - | ERLE_DICTIONARY - | EBYTE_STREAM_SPLIT - | PARQUET_ENCODING_UNKNOWN - deriving (Show, Eq) - -parquetEncodingFromInt :: Int32 -> ParquetEncoding -parquetEncodingFromInt 0 = EPLAIN -parquetEncodingFromInt 2 = EPLAIN_DICTIONARY -parquetEncodingFromInt 3 = ERLE -parquetEncodingFromInt 4 = EBIT_PACKED -parquetEncodingFromInt 5 = EDELTA_BINARY_PACKED -parquetEncodingFromInt 6 = EDELTA_LENGTH_BYTE_ARRAY -parquetEncodingFromInt 7 = EDELTA_BYTE_ARRAY -parquetEncodingFromInt 8 = ERLE_DICTIONARY -parquetEncodingFromInt 9 = EBYTE_STREAM_SPLIT -parquetEncodingFromInt _ = PARQUET_ENCODING_UNKNOWN - -data CompressionCodec - = UNCOMPRESSED - | SNAPPY - | GZIP - | LZO - | BROTLI - | LZ4 - | ZSTD - | LZ4_RAW - | COMPRESSION_CODEC_UNKNOWN - deriving (Show, Eq) - -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 - -data ColumnStatistics = ColumnStatistics - { columnMin :: [Word8], - columnMax :: [Word8], - columnNullCount :: Int64, - columnDistictCount :: Int64, - columnMinValue :: [Word8], - columnMaxValue :: [Word8], - isColumnMaxValueExact :: Bool, - isColumnMinValueExact :: Bool - } - deriving (Show, Eq) - -emptyColumnStatistics :: ColumnStatistics -emptyColumnStatistics = ColumnStatistics [] [] 0 0 [] [] False False - -data PageType - = DATA_PAGE - | INDEX_PAGE - | DICTIONARY_PAGE - | DATA_PAGE_V2 - | PAGE_TYPE_UNKNOWN - deriving (Show, Eq) - -pageTypeFromInt :: Int32 -> PageType -pageTypeFromInt 0 = DATA_PAGE -pageTypeFromInt 1 = INDEX_PAGE -pageTypeFromInt 2 = DICTIONARY_PAGE -pageTypeFromInt 3 = DATA_PAGE_V2 -pageTypeFromInt _ = PAGE_TYPE_UNKNOWN - -data 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 - - -- Create index decoder - let lvlByteLen = (fromIntegral nbytes + 4) - let rleBytes = drop lvlByteLen (pageBytes p') - let bitWidth = head rleBytes - let indexDecoder = MkRleDecoder (tail rleBytes) (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 (error "PLAIN DATA PAGES NOT SUPPORTED") - Nothing -> pure () - - c' <- readIORef colMap - colNames' <- readIORef colNames - let asscList = map (\name -> (name, c' M.! name)) colNames' - pure $ DI.fromNamedColumns asscList - -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)) .&. 1) [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 (littleEndianWord8 $ 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 -> pure $ Snappy.decompress (BSO.pack compressed) - 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 - -littleEndianWord8 :: [Word8] -> Word8 -littleEndianWord8 bytes - | length bytes == 4 = foldr (\v acc -> acc .|. v) 0 (zipWith (\b i -> b `shiftL` i) bytes [0,8..]) - | otherwise = error "Expected exactly 4 bytes" - -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/Operations/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs --- a/src/DataFrame/Operations/Aggregation.hs +++ b/src/DataFrame/Operations/Aggregation.hs @@ -8,6 +8,7 @@ module DataFrame.Operations.Aggregation where import qualified Data.Set as S +import qualified DataFrame.Functions as F import qualified Data.List as L import qualified Data.Map as M @@ -28,7 +29,7 @@ getIndicesUnboxed, getIndices, Columnable, unwrapTypedColumn, columnVersionString) -import DataFrame.Internal.DataFrame (DataFrame(..), empty, getColumn) +import DataFrame.Internal.DataFrame (DataFrame(..), empty, getColumn, unsafeGetColumn) import DataFrame.Internal.Expression import DataFrame.Internal.Parsing import DataFrame.Internal.Types @@ -37,6 +38,7 @@ import DataFrame.Operations.Subset import Data.Function ((&)) import Data.Hashable +import Data.List ((\\)) import Data.Maybe import Data.Type.Equality (type (:~:)(Refl), TestEquality(..)) import Type.Reflection (typeRep, typeOf) @@ -133,3 +135,13 @@ distinct :: DataFrame -> DataFrame distinct df = groupBy (columnNames df) 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 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 @@ -130,8 +130,8 @@ -- | O(n) Returns the number of non-null columns in the dataframe and the type associated -- with each column. -columnInfo :: DataFrame -> DataFrame -columnInfo df = empty & insertColumn "Column Name" (fromList (map nameOfColumn infos)) +describeColumns :: DataFrame -> DataFrame +describeColumns df = empty & insertColumn "Column Name" (fromList (map nameOfColumn infos)) & insertColumn "# Non-null Values" (fromList (map nonNullValues infos)) & insertColumn "# Null Values" (fromList (map nullValues infos)) & insertColumn "# Partially parsed" (fromList (map partiallyParsedValues infos)) 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 @@ -91,7 +91,7 @@ correlation first second df = do f <- _getColumnAsDouble first df s <- _getColumnAsDouble second df - return $ SS.correlation2 f s + return $ SS.correlation (VG.zip f s) _getColumnAsDouble :: T.Text -> DataFrame -> Maybe (VU.Vector Double) _getColumnAsDouble name df = case getColumn name df of