parquet-haskell (empty) → 0.1.0.0
raw patch · 25 files changed
+2512/−0 lines, 25 filesdep +QuickCheckdep +aesondep +basesetup-changed
Dependencies added: QuickCheck, aeson, base, binary, binary-conduit, bytestring, conduit, conduit-extra, containers, filepath, generic-lens, hspec, http-client, http-conduit, http-types, lens, lifted-async, monad-logger, mtl, parquet-haskell, pinch, pretty-simple, process, relude, safe, serialise, text, unordered-containers, vector
Files
- ChangeLog.md +3/−0
- LICENSE +30/−0
- README.md +31/−0
- Setup.hs +3/−0
- parquet-haskell.cabal +165/−0
- src/Main.hs +18/−0
- src/Parquet/Decoder.hs +163/−0
- src/Parquet/InstanceOrphanage.hs +17/−0
- src/Parquet/Monad.hs +18/−0
- src/Parquet/ParquetObject.hs +94/−0
- src/Parquet/Pinch.hs +39/−0
- src/Parquet/Prelude.hs +19/−0
- src/Parquet/Reader.hs +738/−0
- src/Parquet/Stream/Reader.hs +393/−0
- src/Parquet/Types.hs +15/−0
- src/Parquet/Types/ColumnChunk.hs +57/−0
- src/Parquet/Types/Enums.hs +129/−0
- src/Parquet/Types/FileMetaData.hs +43/−0
- src/Parquet/Types/PageHeader.hs +73/−0
- src/Parquet/Types/Primitives.hs +241/−0
- src/Parquet/Types/RowGroup.hs +34/−0
- src/Parquet/Utils.hs +38/−0
- tests/integration/Spec.hs +76/−0
- tests/unit/Parquet/Decoder/Spec.hs +66/−0
- tests/unit/Spec.hs +9/−0
+ ChangeLog.md view
@@ -0,0 +1,3 @@+# Changelog for parquet-haskell++## Unreleased changes
+ LICENSE view
@@ -0,0 +1,30 @@+Copyright Author name here (c) 2019++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++ * Redistributions of source code must retain the above copyright+ notice, this list of conditions and the following disclaimer.++ * Redistributions in binary form must reproduce the above+ copyright notice, this list of conditions and the following+ disclaimer in the documentation and/or other materials provided+ with the distribution.++ * Neither the name of Author name here nor the names of other+ contributors may be used to endorse or promote products derived+ from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ README.md view
@@ -0,0 +1,31 @@+# parquet-haskell++A Streaming Parquet reader for Haskell.++[](https://github.com/manishsencha/parquet-haskell/actions/workflows/ci.yml)++## Description++`parquet-haskell` is a library for streaming data from Parquet files in Haskell. It provides a way to read and parse Parquet files efficiently.++## Prerequisites++- [GHC](https://www.haskell.org/ghc/) (Glasgow Haskell Compiler)+- [Cabal](https://www.haskell.org/cabal/) (build system)+- [Python 3](https://www.python.org/) with `pyspark` (for integration tests)++## Building the Project++You can build the library and tests using `cabal`:++```sh+cabal build all+```++## Running Tests++Run the unit and integration tests using:++```sh+cabal test all+```
+ Setup.hs view
@@ -0,0 +1,3 @@+import Distribution.Simple++main = defaultMain
+ parquet-haskell.cabal view
@@ -0,0 +1,165 @@+cabal-version: 1.12++-- This file has been generated from package.yaml by hpack version 0.35.2.+--+-- see: https://github.com/sol/hpack++name: parquet-haskell+version: 0.1.0.0+synopsis: Streaming Parquet reader+description: Please see the README on GitHub at <https://github.com/manishsencha/parquet-haskell#readme>+category: Compression+homepage: https://github.com/manishsencha/parquet-haskell#readme+bug-reports: https://github.com/manishsencha/parquet-haskell/issues+author: Manish Sencha+maintainer: Manish Sencha <manishsencha2001@gmail.com>+copyright: 2019-2021 Yigit Ozkavci, 2026 Manish Sencha+license: BSD3+license-file: LICENSE+build-type: Simple+extra-source-files:+ ChangeLog.md+ README.md++source-repository head+ type: git+ location: https://github.com/manishsencha/parquet-haskell++library+ exposed-modules:+ Main+ Parquet.Decoder+ Parquet.InstanceOrphanage+ Parquet.Monad+ Parquet.ParquetObject+ Parquet.Pinch+ Parquet.Prelude+ Parquet.Reader+ Parquet.Stream.Reader+ Parquet.Types+ Parquet.Types.ColumnChunk+ Parquet.Types.Enums+ Parquet.Types.FileMetaData+ Parquet.Types.PageHeader+ Parquet.Types.Primitives+ Parquet.Types.RowGroup+ Parquet.Utils+ other-modules:+ Paths_parquet_haskell+ hs-source-dirs:+ src+ default-extensions:+ DataKinds+ DeriveAnyClass+ DeriveGeneric+ FlexibleContexts+ OverloadedStrings+ LambdaCase+ MultiWayIf+ NamedFieldPuns+ NoImplicitPrelude+ ScopedTypeVariables+ TypeApplications+ TypeOperators+ TypeFamilies+ ViewPatterns+ ghc-options: -Wall -fspecialise-aggressively -Wall+ build-depends:+ aeson >=2.2.3 && <2.3+ , base >=4.7 && <5+ , binary >=0.8.9 && <0.9+ , binary-conduit >=1.3.1 && <1.4+ , bytestring >=0.11.5 && <0.12+ , conduit >=1.3.6 && <1.4+ , conduit-extra >=1.3.8 && <1.4+ , containers >=0.6.7 && <0.7+ , generic-lens >=2.3.0 && <2.4+ , http-client >=0.7.19 && <0.8+ , http-conduit >=2.3.9 && <2.4+ , http-types >=0.12.4 && <0.13+ , lens >=5.3.6 && <5.4+ , lifted-async >=0.11.0 && <0.12+ , monad-logger >=0.3.42 && <0.4+ , mtl >=2.3.1 && <2.4+ , pinch >=0.5.2 && <0.6+ , pretty-simple >=4.1.4 && <4.2+ , relude >=1.2.2 && <1.3+ , safe >=0.3.21 && <0.4+ , serialise >=0.2.6 && <0.3+ , text >=2.0.2 && <2.1+ , unordered-containers >=0.2.21 && <0.3+ , vector >=0.13.2 && <0.14+ default-language: Haskell2010++test-suite parquet-haskell-integration+ type: exitcode-stdio-1.0+ main-is: Spec.hs+ other-modules:+ Paths_parquet_haskell+ hs-source-dirs:+ tests/integration+ default-extensions:+ DataKinds+ DeriveAnyClass+ DeriveGeneric+ FlexibleContexts+ OverloadedStrings+ LambdaCase+ MultiWayIf+ NamedFieldPuns+ NoImplicitPrelude+ ScopedTypeVariables+ TypeApplications+ TypeOperators+ TypeFamilies+ ViewPatterns+ ghc-options: -Wall -fspecialise-aggressively -threaded -rtsopts -with-rtsopts=-N+ build-depends:+ QuickCheck >=2.14 && <3.0+ , aeson >=2.2.3 && <2.3+ , base >=4.7 && <5+ , binary >=0.8.9 && <0.9+ , bytestring >=0.11.5 && <0.12+ , conduit >=1.3.6 && <1.4+ , filepath ==1.4.*+ , hspec >=2.10 && <3.0+ , monad-logger >=0.3.42 && <0.4+ , mtl >=2.3.1 && <2.4+ , parquet-haskell+ , pretty-simple >=4.1.4 && <4.2+ , process ==1.6.*+ , text >=2.0.2 && <2.1+ default-language: Haskell2010++test-suite parquet-haskell-unit+ type: exitcode-stdio-1.0+ main-is: Spec.hs+ other-modules:+ Parquet.Decoder.Spec+ Paths_parquet_haskell+ hs-source-dirs:+ tests/unit+ default-extensions:+ DataKinds+ DeriveAnyClass+ DeriveGeneric+ FlexibleContexts+ OverloadedStrings+ LambdaCase+ MultiWayIf+ NamedFieldPuns+ NoImplicitPrelude+ ScopedTypeVariables+ TypeApplications+ TypeOperators+ TypeFamilies+ ViewPatterns+ ghc-options: -Wall -fspecialise-aggressively -threaded -rtsopts -with-rtsopts=-N+ build-depends:+ QuickCheck >=2.14 && <3.0+ , base >=4.7 && <5+ , binary >=0.8.9 && <0.9+ , bytestring >=0.11.5 && <0.12+ , hspec >=2.10 && <3.0+ , parquet-haskell+ default-language: Haskell2010
+ src/Main.hs view
@@ -0,0 +1,18 @@+-- |+module Main+ ( main,+ )+where++------------------------------------------------------------------------------++import Conduit (runResourceT)+import Control.Monad.Logger+import Parquet.Prelude+import Parquet.Reader (readWholeParquetFile)++------------------------------------------------------------------------------+main :: IO ()+main = void . runT $ readWholeParquetFile "test.parquet"+ where+ runT = runStdoutLoggingT . runResourceT . runExceptT
+ src/Parquet/Decoder.hs view
@@ -0,0 +1,163 @@+{-# LANGUAGE BangPatterns #-}++-- |+module Parquet.Decoder+ ( -- * Type definitions+ BitWidth (..),++ -- * Decoder functions+ decodeBPBE,+ decodeBPLE,+ decodeRLE,+ decodeRLEBPHybrid,+ decodeVarint,+ encodeVarint,+ takeBytesLe,+ )+where++------------------------------------------------------------------------------++import Data.Binary.Get (Get, getByteString, getWord8)+import Data.Binary.Put (Put, putWord8)+import qualified Data.ByteString as BS+import Parquet.Prelude++------------------------------------------------------------------------------+newtype BitWidth = BitWidth Word8+ deriving (Show, Eq, Ord)++------------------------------------------------------------------------------+cLeb128ByteLimit :: Int+cLeb128ByteLimit = 32++------------------------------------------------------------------------------+takeBytesLe :: Word8 -> Get Integer+takeBytesLe 0 = pure 0+takeBytesLe n = do+ v <- getWord8+ rest <- takeBytesLe (n - 1)+ pure $ (rest `shiftL` 8) .|. fromIntegral v++------------------------------------------------------------------------------+takeBytesBe :: Int -> Word8 -> Get Integer+takeBytesBe = go+ where+ go :: Int -> Word8 -> Get Integer+ go _ 0 = pure 0+ go sh n = do+ v <- getWord8+ rest <- go (sh - 1) (n - 1)+ pure $ (fromIntegral v `shiftL` (8 * (sh - 1))) .|. rest++------------------------------------------------------------------------------+decodeBPBE :: BitWidth -> Get [Word32]+decodeBPBE (BitWidth bit_width) = do+ header <- decodeVarint+ let run_len = header `shiftR` 1+ run $ fromInteger run_len+ where+ run :: Word32 -> Get [Word32]+ run 0 = pure []+ run scaled_run_len = do+ v <- takeBytesBe (fromIntegral bit_width) bit_width+ !batch_bytes <- go 8 v+ (batch_bytes <>) <$> run (scaled_run_len - 1)++ go :: Int -> Integer -> Get [Word32]+ go 0 _ = pure []+ go rem_vals data_bytes = do+ let mask :: Integer+ mask = ((2 ^ bit_width) - 1) `shiftL` (fromIntegral bit_width * 7)+ -- Unsafe fromInteger justification:+ -- Max bit_width = 32 and masking any value with+ -- (2 ^ 32) - 1 is in unsigned 32-bit bound.+ let val =+ fromInteger $+ (data_bytes .&. mask)+ `shiftR` (fromIntegral bit_width * 7)+ rest <- go (rem_vals - 1) (data_bytes `shiftL` fromIntegral bit_width)+ pure $ val : rest++------------------------------------------------------------------------------+decodeBPLE :: BitWidth -> Word32 -> Get [Word32]+decodeBPLE _ 0 = pure []+decodeBPLE bw@(BitWidth bit_width) scaled_run_len = do+ !v <- takeBytesLe bit_width+ batch_bytes <- go 8 v+ (batch_bytes <>) <$> decodeBPLE bw (scaled_run_len - 1)+ where+ go :: Int -> Integer -> Get [Word32]+ go 0 _ = pure []+ go rem_vals data_bytes = do+ let mask = (2 ^ bit_width) - 1+ -- Unsafe fromInteger justification:+ -- Max bit_width = 32 and masking any value with+ -- (2 ^ 32) - 1 is in unsigned 32-bit bound.+ let val = fromInteger $ data_bytes .&. mask+ rest <- go (rem_vals - 1) (data_bytes `shiftR` fromIntegral bit_width)+ pure $ val : rest++------------------------------------------------------------------------------+decodeRLE :: BitWidth -> Word32 -> Get [Word32]+decodeRLE (BitWidth bit_width) run_len = do+ !result <-+ unsafe_bs_to_w32 . BS.unpack+ <$> getByteString+ (fromIntegral fixed_width)+ pure (replicate (fromIntegral run_len) result)+ where+ fixed_width :: Word8+ fixed_width = if bit_width == 0 then 0 else ((bit_width - 1) `div` 8) + 1++ -- TODO(yigitozkavci): We can do a safety check here. In+ -- case of overflow we get 0 as an answer.+ unsafe_bs_to_w32 :: [Word8] -> Word32+ unsafe_bs_to_w32 = foldr (\x -> (fromIntegral x .|.) . (`shiftL` 8)) 0++------------------------------------------------------------------------------+decodeRLEBPHybrid :: BitWidth -> Int32 -> Get [Word32]+decodeRLEBPHybrid bit_width num_values = do+ header <- decodeVarint+ let encoding_ty = header .&. 0x01+ let !run_len = header `shiftR` 1++ -- Unsafe fromInteger justification:+ -- run_len value being in range [1, 2^31-1]+ -- is guaranteed by the protocol.+ case encoding_ty of+ 0x00 -> decodeRLE bit_width $ fromInteger run_len+ 0x01 ->+ take (fromIntegral num_values)+ <$> decodeBPLE bit_width (fromInteger run_len)+ _ ->+ fail+ "Impossible happened! 0x01 .&. _ resulted in a value larger than 0x01"++------------------------------------------------------------------------------+decodeVarint :: Get Integer+decodeVarint = go cLeb128ByteLimit 0 0+ where+ go :: Int -> Integer -> Int -> Get Integer+ go 0 _ _ =+ fail $+ "Could not find a LEB128-encoded value in "+ <> show cLeb128ByteLimit+ <> "bytes"+ go rem_limit acc sh = do+ byte <- getWord8+ let high = byte .&. 0x80+ let low = fromIntegral $ byte .&. 0x7F+ let res = (low `shiftL` sh) .|. acc+ if high == 0x80 then go (rem_limit - 1) res (sh + 7) else pure res++------------------------------------------------------------------------------+encodeVarint :: Integer -> Put+encodeVarint 0 = pure ()+encodeVarint val = do+ let low = fromIntegral $ val .&. 0x7F+ if val `shiftR` 7 == 0+ then putWord8 low+ else do+ putWord8 $ low .|. 0x80+ encodeVarint (val `shiftR` 7)
+ src/Parquet/InstanceOrphanage.hs view
@@ -0,0 +1,17 @@+{-# OPTIONS_GHC -fno-warn-orphans #-}++-- |+module Parquet.InstanceOrphanage where++------------------------------------------------------------------------------++import Data.Binary (Binary)+import Pinch++------------------------------------------------------------------------------++-- TODO(yigitozkavci): Move these orphan instances to the +pinch library code.+-- This will require opening a PR to https://github.com/abhinav/pinch.+instance Binary (Enumeration k)++instance Binary a => Binary (Field k a)
+ src/Parquet/Monad.hs view
@@ -0,0 +1,18 @@+{-# LANGUAGE ConstraintKinds #-}++-- |+module Parquet.Monad+ ( -- *+ PR,+ )+where++------------------------------------------------------------------------------++import Conduit (MonadResource, MonadThrow)+import Control.Monad.Except (MonadError)+import Control.Monad.Logger (MonadLogger)+import Parquet.Prelude++------------------------------------------------------------------------------+type PR m = (MonadResource m, MonadLogger m, MonadThrow m, MonadError Text m)
+ src/Parquet/ParquetObject.hs view
@@ -0,0 +1,94 @@+{-# LANGUAGE TemplateHaskell #-}++module Parquet.ParquetObject+ ( -- * Type definitions+ ParquetValue (..),+ ParquetObject (..),+ ParquetList (..),+ )+where++------------------------------------------------------------------------------++import Control.Lens (makeLenses, makePrisms)+import qualified Data.Aeson as JSON+import qualified Data.Aeson.KeyMap as JSONKeyMap+import Data.Binary (Binary (get, put))+import Parquet.Prelude hiding (get, put)++------------------------------------------------------------------------------+newtype ParquetObject = MkParquetObject (HashMap Text ParquetValue)+ deriving (Eq, Show, Generic, Serialise)++instance Semigroup ParquetObject where+ MkParquetObject hm1 <> MkParquetObject hm2 = MkParquetObject (hm1 <> hm2)++instance Monoid ParquetObject where+ mempty = MkParquetObject mempty++instance Binary ParquetObject where+ put (MkParquetObject hm) = put (toList hm)+ get = MkParquetObject . fromList <$> get++instance ToJSON ParquetObject where+ toJSON (MkParquetObject obj) = toJSON obj++------------------------------------------------------------------------------+newtype ParquetList = MkParquetList [ParquetValue]+ deriving (Eq, Show, Generic, Serialise)++instance Semigroup ParquetList where+ MkParquetList l1 <> MkParquetList l2 = MkParquetList (l1 <> l2)++instance Monoid ParquetList where+ mempty = MkParquetList mempty++instance Binary ParquetList where+ put (MkParquetList l) = put l+ get = MkParquetList <$> get++instance ToJSON ParquetList where+ toJSON (MkParquetList l) = toJSON l++------------------------------------------------------------------------------+data ParquetValue+ = ParquetObject !ParquetObject+ | ParquetList !ParquetList+ | ParquetInt !Int64+ | ParquetString !ByteString+ | ParquetBool !Bool+ | ParquetNull+ | EmptyValue+ deriving (Eq, Show, Generic, Binary, Serialise)++instance FromJSON ParquetValue where+ parseJSON = \case+ JSON.Object obj -> do+ ParquetObject . MkParquetObject . JSONKeyMap.toHashMapText <$> traverse parseJSON obj+ JSON.Array vec -> do+ ParquetList . MkParquetList . toList <$> traverse parseJSON vec+ JSON.Number sci ->+ pure $ ParquetInt $ fromInteger $ numerator $ toRational sci+ JSON.String s ->+ pure $ ParquetString $ encodeUtf8 s+ JSON.Bool b -> pure $ ParquetBool b+ JSON.Null -> pure ParquetNull++instance ToJSON ParquetValue where+ toJSON = \case+ ParquetObject obj -> toJSON obj+ ParquetList l -> toJSON l+ ParquetInt i64 -> JSON.Number (fromIntegral i64)+ ParquetString bs -> case decodeUtf8' bs of+ Right t -> JSON.String t+ Left _ -> JSON.String "<non-utf8-string>"+ ParquetBool b -> JSON.Bool b+ ParquetNull -> JSON.Null+ EmptyValue -> JSON.Null++------------------------------------------------------------------------------+makeLenses ''ParquetObject+makeLenses ''ParquetValue++makePrisms ''ParquetObject+makePrisms ''ParquetValue
+ src/Parquet/Pinch.hs view
@@ -0,0 +1,39 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UndecidableInstances #-}++module Parquet.Pinch where++------------------------------------------------------------------------------++import Control.Lens+import qualified Data.Generics.Product.Fields as GL+import qualified Data.Generics.Product.Positions as GL+import GHC.Generics (D, M1, Meta (..), Rep)+import GHC.TypeLits (AppendSymbol, Symbol)+import Parquet.Prelude+import Pinch++------------------------------------------------------------------------------+type family TypeName a :: Symbol where+ TypeName (M1 D ('MetaData name _ _ _) f ()) = name+ TypeName a = TypeName (Rep a ())++pinchPos ::+ forall pos s t a1 b1 a2 b2.+ (GL.HasPosition 1 a1 b1 a2 b2, GL.HasPosition pos s t a1 b1) =>+ Lens s t a2 b2+pinchPos = GL.position @pos . GL.position @1++pinchField ::+ forall field s i r field_name.+ ( field_name ~ ("_" `AppendSymbol` TypeName s `AppendSymbol` "_" `AppendSymbol` field),+ GL.HasPosition 1 i i r r,+ GL.HasField field_name s s i i+ ) =>+ Lens s s r r+pinchField = GL.field @field_name . GL.position @1++------------------------------------------------------------------------------+unField :: Field n a -> a+unField (Field a) = a
+ src/Parquet/Prelude.hs view
@@ -0,0 +1,19 @@+-- |+module Parquet.Prelude+ ( module X,+ )+where++------------------------------------------------------------------------------++import Codec.Serialise as X (Serialise)+import Data.Aeson as X (FromJSON (..), ToJSON (..))+import Data.Binary as X (Binary)+import Data.Bits as X (shiftL, shiftR, (.&.), (.|.))+import Data.Traversable as X (for)+import Parquet.InstanceOrphanage ()+import Relude as X hiding (Type)+import Safe as X (headMay)+import Safe.Exact as X (zipExactMay)++------------------------------------------------------------------------------
+ src/Parquet/Reader.hs view
@@ -0,0 +1,738 @@+{-# LANGUAGE TypeApplications #-}++{-+An example schema:++[root] spark_schema: {+ optional f1: {+ repeated group list: {+ optional element: {+ repeated group list: {+ optional int64 element = 1;+ }+ }+ }+ }+ optional f2: {+ repeated group list: {+ optional int64 element = 2;+ }+ }+ optional f3: {+ repeated group list: {+ optional int64 element = 3;+ }+ }+ optional f4: {+ repeated group list: {+ optional int64 element = 4;+ }+ }+ optional f5: {+ repeated group list: {+ optional int64 element = 5;+ }+ }+ optional f6: {+ repeated group list: {+ optional int64 element = 6;+ }+ }+}++Then, the following column values:+____________________________________________________________________++\| rep_level | def_level | path | value |+\|___________|___________|__________________________________|_______|+\| 0 | 5 | f1, list, element, list, element | 1 |+\| 2 | 5 | f1, list, element, list, element | 2 |+\| 1 | 5 | f1, list, element, list, element | 3 |+\| 2 | 5 | f1, list, element, list, element | 4 |+\| 1 | 5 | f1, list, element, list, element | 5 |+\| 0 | 0 | f1, list, element, list, element | 1 |+\| 0 | 0 | f1, list, element, list, element | 2 |+\| 0 | 0 | f1, list, element, list, element | 3 |+\| 0 | 0 | f1, list, element, list, element | 2 |+\| 0 | 0 | f1, list, element, list, element | 3 |+\|___________|___________|__________________________________|_______|++should produce the following data:+Note: Values between "_" characters describe the accumulator we use during the recursion.++== Value 1 ==+_{}_+> (r: 0, d: 5, v: 1, p: [f1, list, element, list, element])+instruction: IObjectField "f1"+{f1: _?_} (f1's type is OPTIONAL)+> (r: 0, d: 4, v: 1, p: [list, element, list, element])+instruction: NewList+{f1: [_?_]} (list's type is REPEATED)+> (r: 0, d: 3, v: 1, p: [element, list, element])+instruction: IObjectField "element"+{f1: [{ element: _?_ }]} (element's type is OPTIONAL)+> (r: 0, d: 2, v: 1, p: [list, element])+instruction: NewList+{f1: [{ element: [_?_] }]} (list's type is REPEATED)+> (r: 0, d: 1, v: 1, p: [element])+instruction: IObjectField "element"+{f1: [{ element: [{ element: _?_ }] }]} (element's type is OPTIONAL)+> (r: 0, d: 0, v: 1, p: [])+instruction: IValue 1+{f1: [{ element: [{ element: 1}] }]}++== Value 2 ==+_{f1: [{ element: [{ element: 1 }] }]}_+> (r: 2, d: 5, v: 2, p: [f1, list, element, list, element])+instruction: IObjectField "f1"+(Note: f1 exists in the accumulator, so we use it.)+{f1: _[{ element: [{ element: 1 }] }]_}+> (r: 2, d: 4, v: 2, p: [list, element, list, element])+instruction: IListElement+(Note: since repetition level is non-zero, we use the last element in the list.)+{f1: [_{ element: [{ element: 1 }] }_]}+> (r: 1, d: 3, v: 2, p: [element, list, element])+instruction: IObjectField "element"+{f1: [{ element: _[{ element: 1 }]_ }]}+> (r: 1, d: 2, v: 2, p: [list, element])+instruction: INewListElement+(Note: repetition level is 1 and we see a REPEATED type. Create a new element.)+(NOTE(yigitozkavci): Is this an edge case or am I not smart enough? Probably the latter.)+{f1: [{ element: [{ element: 1 }, _{}_] }]}+> (r: 0, d: 1, v: 2, p: [element])+instruction: IObjectField "element"+{f1: [{ element: [{ element: 1 }, { element: _?_ }] }]}+> (r: 0, d: 0, v: 2, p: [])+instruction: IValue 2+{f1: [{ element: [{ element: 1 }, { element: 2 }] }]}++== Value 3 ==+_{f1: [{ element: [{ element: 1 }, { element: 2 }] }]}_+> (r: 1, d: 5, v: 3, p: [f1, list, element, list, element])+instruction: IObjectField "f1"+{f1: _[{ element: [{ element: 1 }, { element: 2 }] }]_}+> (r: 1, d: 4, v: 3, p: [list, element, list, element])+instruction: INewListElement+{f1: [{ element: [{ element: 1 }, { element: 2 }] }, _{}_]}+> (r: 0, d: 3, v: 3, p: [element, list, element])+instruction: IObjectField "element"+{f1: [{ element: [{ element: 1 }, { element: 2 }] }, { element: _?_}]}+> (r: 0, d: 2, v: 3, p: [list, element])+instruction: INewList+(Note: repetition level is 0 and the type is REPEATED, so create a new list.)+{f1: [{ element: [{ element: 1 }, { element: 2 }] }, { element: [_?_]}]}+> (r: 0, d: 1, v: 3, p: [element])+instruction: IObjectField "element"+{f1: [{ element: [{ element: 1 }, { element: 2 }] }, { element: [{ element: _?_ }]}]}+> (r: 0, d: 0, v: 3, p: [])+instruction: IValue 3+{f1: [{ element: [{ element: 1 }, { element: 2 }] }, { element: [{ element: 3 }]}]}++-}+module Parquet.Reader where++------------------------------------------------------------------------------++import qualified Conduit as C+import Control.Lens hiding (ix)+import Control.Monad (foldM)+import Control.Monad.Except+import Control.Monad.Logger (MonadLogger, runNoLoggingT)+import Control.Monad.Logger.CallStack (logError, logInfo, logWarn)+import qualified Data.Binary.Get as BG+import qualified Data.ByteString.Char8 as BS8+import qualified Data.Conduit.Binary as CB+import qualified Data.Conduit.List as CL+import qualified Data.HashMap.Strict as HM+import qualified Data.Map as M+import qualified Data.Sequence as Seq+import qualified Data.Text as T+import qualified Data.Text.Lazy as LT+import Network.HTTP.Client (Request (requestHeaders))+import Network.HTTP.Simple+ ( Header,+ getResponseBody,+ getResponseStatus,+ httpSource,+ parseRequest,+ )+import Network.HTTP.Types.Status (statusIsSuccessful)+import Parquet.ParquetObject+import qualified Parquet.Pinch as TT+import Parquet.Prelude+import Parquet.Stream.Reader+ ( ColumnValue (..),+ Value (..),+ decodeConduit,+ readColumnChunk,+ )+import qualified Parquet.Types as TT+import Parquet.Utils (failOnExcept, failOnMay)+import System.IO+ ( SeekMode (AbsoluteSeek, SeekFromEnd),+ hSeek,+ openFile,+ )+import Text.Pretty.Simple (pString)++------------------------------------------------------------------------------+type Url = String++------------------------------------------------------------------------------+newtype ParquetSource m = ParquetSource (Integer -> C.ConduitT () ByteString m ())++------------------------------------------------------------------------------+readFieldTypeMapping ::+ (MonadError Text m) => TT.FileMetaData -> m (HashMap Text TT.Type)+readFieldTypeMapping fm =+ let schemaElements = fm ^. TT.pinchField @"schema"+ in fmap fromList+ $ for schemaElements+ $ \se -> do+ let name = se ^. TT.pinchField @"name"+ case se ^. TT.pinchField @"type" of+ Nothing ->+ throwError $ "Type info for field " <> name <> " doesn't exist"+ Just ty -> pure (name, ty)++------------------------------------------------------------------------------+readMetadata ::+ ( MonadError Text m,+ MonadIO m,+ MonadFail m+ ) =>+ ParquetSource m ->+ m TT.FileMetaData+readMetadata (ParquetSource source) = do+ bs <- C.runConduit (source (-8) C..| CB.take 8)+ case BG.runGetOrFail BG.getWord32le bs of+ Left err -> fail $ "Could not fetch metadata size: " <> show err+ Right (_, _, metadataSize) ->+ fmap (snd . fst)+ $ C.runConduit+ $ source (-(8 + fromIntegral metadataSize))+ C..| decodeConduit metadataSize+ `C.fuseBoth` pure ()++------------------------------------------------------------------------------+localParquetFile :: (C.MonadResource m) => FilePath -> ParquetSource m+localParquetFile fp = ParquetSource $ \pos -> C.sourceIOHandle $ do+ h <- openFile fp ReadMode+ if pos > 0 then hSeek h AbsoluteSeek pos else hSeek h SeekFromEnd pos+ pure h++------------------------------------------------------------------------------+remoteParquetFile ::+ ( C.MonadResource m,+ C.MonadThrow m,+ C.MonadIO m,+ MonadFail m+ ) =>+ Url ->+ ParquetSource m+remoteParquetFile url = ParquetSource $ \pos -> do+ req <- parseRequest url+ let rangedReq = req {requestHeaders = mkRangeHeader pos : requestHeaders req}+ httpSource rangedReq call+ where+ mkRangeHeader :: Integer -> Header+ mkRangeHeader pos =+ let rangeVal = if pos > 0 then show pos <> "-" else show pos+ in ("Range", "bytes=" <> BS8.pack rangeVal)++ call req =+ let status = getResponseStatus req+ in if statusIsSuccessful status+ then getResponseBody req+ else+ fail+ $ "Non-success response code from remoteParquetFile call: "+ ++ show status++------------------------------------------------------------------------------+readWholeParquetFile ::+ ( C.MonadThrow m,+ MonadIO m,+ MonadError Text m,+ C.MonadResource m,+ MonadLogger m,+ MonadFail m+ ) =>+ String ->+ m [ParquetValue]+readWholeParquetFile inputFp = do+ metadata <- readMetadata (localParquetFile inputFp)+ (`runReaderT` metadata)+ $ C.runConduit+ $ traverse_+ (sourceRowGroup (localParquetFile inputFp))+ (metadata ^. TT.pinchField @"row_groups")+ C..| CL.map convertLiteralJsonLists+ C..| CL.consume++------------------------------------------------------------------------------+convertLiteralJsonLists :: ParquetValue -> ParquetValue+convertLiteralJsonLists (ParquetObject (MkParquetObject (HM.toList -> [("element", parquetValue)]))) =+ convertLiteralJsonLists parquetValue+convertLiteralJsonLists (ParquetObject (MkParquetObject kvMapping)) =+ ParquetObject $ MkParquetObject $ convertLiteralJsonLists <$> kvMapping+convertLiteralJsonLists (ParquetList (MkParquetList xs)) =+ ParquetList $ MkParquetList $ map convertLiteralJsonLists xs+convertLiteralJsonLists v = v++------------------------------------------------------------------------------+sourceParquet :: FilePath -> C.ConduitT () ParquetValue (C.ResourceT IO) ()+sourceParquet fp =+ runExceptT (readMetadata (localParquetFile fp)) >>= \case+ Left err -> fail $ "Could not read metadata: " <> show err+ Right metadata ->+ C.transPipe (runNoLoggingT . (`runReaderT` metadata))+ $ traverse_+ (sourceRowGroup (localParquetFile fp))+ (metadata ^. TT.pinchField @"row_groups")++------------------------------------------------------------------------------+foldMaybeM ::+ (Foldable t, Monad m) => (b -> a -> m (Maybe b)) -> b -> t a -> m b+foldMaybeM action = foldM $ \b a ->+ action b a >>= \case+ Nothing -> pure b+ Just newB -> pure newB++------------------------------------------------------------------------------+sourceRowGroupFromRemoteFile ::+ ( C.MonadResource m,+ C.MonadIO m,+ C.MonadThrow m,+ MonadLogger m,+ MonadReader TT.FileMetaData m,+ MonadFail m+ ) =>+ String ->+ TT.RowGroup ->+ C.ConduitT () ParquetValue m ()+sourceRowGroupFromRemoteFile url = sourceRowGroup (remoteParquetFile url)++------------------------------------------------------------------------------+throwOnNothing :: (MonadError err m) => err -> Maybe a -> m a+throwOnNothing err Nothing = throwError err+throwOnNothing _ (Just v) = pure v++------------------------------------------------------------------------------+initColumnState :: ParquetValue+initColumnState = EmptyValue++-- |+-- Instruction generator for a single column.+--+-- In a parquet column, a repetition level of 0 denotes start of a new record.+-- Example:+--+-- For the following json:+--+-- [+-- { "f1": [[1, 2], [3, 4], [5]]+-- },+-- { "f2": [1, 2, 3]+-- },+-- { "f3": [1, 2, 3]+-- },+-- { "f4": [1, 2, 3]+-- },+-- { "f5": [1, 2, 3]+-- },+-- { "f6": [1, 2, 3]+-- }+-- ]+--+--+-- Values look like the following:+-- ____________________________________________________________________++-- | rep_level | def_level | path | value |+-- |___________|___________|__________________________________|_______|+-- | 0 | 5 | f1, list, element, list, element | 1 |+-- | 2 | 5 | f1, list, element, list, element | 2 |+-- | 1 | 5 | f1, list, element, list, element | 3 |+-- | 2 | 5 | f1, list, element, list, element | 4 |+-- | 1 | 5 | f1, list, element, list, element | 5 |+-- | 0 | 0 | f1, list, element, list, element | 1 |+-- | 0 | 0 | f1, list, element, list, element | 2 |+-- | 0 | 0 | f1, list, element, list, element | 3 |+-- | 0 | 0 | f1, list, element, list, element | 2 |+-- | 0 | 0 | f1, list, element, list, element | 3 |+-- |___________|___________|__________________________________|_______|+--+-- Consuming the stream above will yield the following @ColumnConstructor@s:+--+-- ___________________________________+-- | { "f1": [[1, 2], [3, 4], [5]] } |+-- |_________________________________|+generateInstructions ::+ forall m.+ ( C.MonadResource m,+ C.MonadIO m,+ C.MonadThrow m,+ MonadLogger m,+ MonadFail m,+ MonadReader TT.FileMetaData m+ ) =>+ C.ConduitT (ColumnValue, [Text]) ColumnConstructor m ()+generateInstructions = loop Seq.empty+ where+ loop ::+ Seq.Seq InstructionSet ->+ C.ConduitT (ColumnValue, [Text]) ColumnConstructor m ()+ loop instructions =+ C.await >>= \case+ Nothing ->+ unless (Seq.null instructions) $ C.yield $ ColumnConstructor instructions+ Just cv@(ColumnValue {_cvRepetitionLevel}, _) ->+ case (_cvRepetitionLevel, instructions) of+ (0, Seq.Empty) -> go Seq.empty cv+ (0, _) -> do+ C.yield $ ColumnConstructor instructions+ C.leftover cv+ loop Seq.empty+ (_, Seq.Empty) ->+ logError+ "generateInstructions: Repetition level is nonzero but we don't have any accumulated instructions. This either means there is a bug in this logic or record is corrupted."+ _ -> go instructions cv++ go ::+ Seq.Seq InstructionSet ->+ (ColumnValue, [Text]) ->+ C.ConduitT (ColumnValue, [Text]) ColumnConstructor m ()+ go ix cv =+ mkInstructions cv >>= \case+ Nothing -> logError "Could not create instructions: "+ Just is -> loop (ix Seq.|> is)++------------------------------------------------------------------------------+readSchemaMapping ::+ (MonadReader TT.FileMetaData m) =>+ m (M.Map Text TT.SchemaElement)+readSchemaMapping = do+ metadata <- ask+ pure $ mk_schema_mapping $ metadata ^. TT.pinchField @"schema"+ where+ mk_schema_mapping :: [TT.SchemaElement] -> M.Map Text TT.SchemaElement+ mk_schema_mapping schema = snd $ execState (go mempty) (schema, mempty)++ go ::+ (MonadState ([TT.SchemaElement], M.Map Text TT.SchemaElement) m) =>+ Text ->+ m ()+ go prefix =+ get >>= \case+ ([], _) -> pure ()+ (schema_element : rest, schema_mapping) -> do+ let mb_num_children = schema_element ^. TT.pinchField @"num_children"+ let name = schema_element ^. TT.pinchField @"name"+ case mb_num_children of+ Nothing -> do+ put (rest, M.insert (prefix <> name) schema_element schema_mapping)+ Just num_children -> do+ put (rest, M.insert (prefix <> name) schema_element schema_mapping)+ replicateM_ (fromIntegral num_children) (go (prefix <> name <> "."))++------------------------------------------------------------------------------+findSchemaElement ::+ (MonadReader TT.FileMetaData m, MonadFail m) =>+ Text ->+ m (Maybe TT.SchemaElement)+findSchemaElement path = do+ schema_mapping <- readSchemaMapping+ schema_root <- readSchemaRoot+ pure $ M.lookup (schema_root ^. TT.pinchField @"name" <> "." <> path) schema_mapping++------------------------------------------------------------------------------++-- | Given a single column, generates instructions for how to build an object with that column.+--+-- For example, for the following column:+-- (ColumnValue {_cvRepetitionLevel = 0, _cvDefinitionLevel = 5, _cvMaxDefinitionLevel = 5, _cvValue = ValueInt64 1},["arr2","list","element","list","element"])+-- [ IObjectField "arr2"+-- , INewList+-- , INewList+-- , IValue (ValueInt64 1)+-- ]+--+-- See the following blog post to understand:+-- https://blog.twitter.com/engineering/en_us/a/2013/dremel-made-simple-with-parquet.html+mkInstructions ::+ forall m.+ ( C.MonadIO m,+ C.MonadThrow m,+ MonadLogger m,+ MonadReader TT.FileMetaData m,+ MonadFail m+ ) =>+ (ColumnValue, [Text]) ->+ m (Maybe InstructionSet)+mkInstructions (c, path) = do+ logInfo $ "Creating instruction for column value: " <> T.pack (show c) <> " and path " <> T.intercalate "." path+ result <- go [] (c, path)+ logInfo $ "Instruction set: " <> T.pack (show result)+ pure result+ where+ go :: [Text] -> (ColumnValue, [Text]) -> m (Maybe InstructionSet)+ go pathSoFar columnValue = do+ schema_mapping <- readSchemaMapping+ schema_root <- readSchemaRoot+ case columnValue of+ (ColumnValue r d md v, fieldName : restPath) -> do+ let fullPathSoFar = T.intercalate "." $ pathSoFar <> [fieldName]+ case M.lookup (schema_root ^. TT.pinchField @"name" <> "." <> fullPathSoFar) schema_mapping of+ Nothing -> Nothing <$ logWarn ("Couldn't find the schema element for path " <> fullPathSoFar)+ Just schema_element ->+ case schema_element ^. TT.pinchField @"repetition_type" of+ Nothing ->+ Nothing <$ logError ("Path doesn't have a repetition type: " <> fullPathSoFar)+ Just (TT.REQUIRED _)+ | d == 0 ->+ Nothing <$ logError ("Found a REQUIRED schema element while definition level is 0: " <> fullPathSoFar)+ | otherwise -> do+ mb_rest_instructions <-+ go+ (pathSoFar <> [fieldName])+ (ColumnValue r d md v, restPath)+ pure $ (IObjectField fieldName Seq.<|) <$> mb_rest_instructions+ Just (TT.OPTIONAL _)+ | d == 0 ->+ case pathSoFar of+ [] ->+ pure $ Just $ Seq.singleton INullOpt+ _ ->+ pure $ Just $ Seq.singleton (IValue Null)+ | otherwise -> do+ mb_rest_instructions <-+ go+ (pathSoFar <> [fieldName])+ (ColumnValue r (d - 1) md v, restPath)+ pure $ (IObjectField fieldName Seq.<|) <$> mb_rest_instructions+ Just (TT.REPEATED _)+ | d == 0 ->+ Nothing <$ logError ("Found a REPEATED schema element while definition level is 0: " <> fullPathSoFar)+ | r == 0 -> do+ mb_rest_instructions <-+ go+ (pathSoFar <> [fieldName])+ (ColumnValue r (d - 1) md v, restPath)+ pure $ (INewList Seq.<|) <$> mb_rest_instructions+ | r == 1 -> do+ mb_rest_instructions <-+ go+ (pathSoFar <> [fieldName])+ (ColumnValue (r - 1) (d - 1) md v, restPath)+ pure $ (INewListElement Seq.<|) <$> mb_rest_instructions+ | otherwise -> do+ mb_rest_instructions <-+ go+ (pathSoFar <> [fieldName])+ (ColumnValue (r - 1) (d - 1) md v, restPath)+ pure $ (IListElement Seq.<|) <$> mb_rest_instructions+ (ColumnValue _ 0 _ v, []) -> pure $ Just $ Seq.singleton $ IValue v+ (ColumnValue {}, []) ->+ Nothing+ <$ logWarn "Saw column with nonzero rep/def levels and empty path."++------------------------------------------------------------------------------+newtype ColumnConstructor = ColumnConstructor+ { ccInstrSet :: Seq.Seq InstructionSet+ }+ deriving (Eq, Show)++------------------------------------------------------------------------------++-- | Streams the values for every column chunk and zips them into records.+--+-- Illustration:+--+-- _____________________+-- | col1 | col2 | col3 |+-- | 1 | a | x |+-- | 2 | b | y |+-- | 3 | c | z |+-- |______|______|______|+--+-- @sourceRowGroup@ yields the following values in a stream:+--+-- (1, a, x)+-- (2, b, y)+-- (3, c, z)+sourceRowGroup ::+ forall m.+ ( C.MonadResource m,+ C.MonadIO m,+ C.MonadThrow m,+ MonadLogger m,+ MonadReader TT.FileMetaData m,+ MonadFail m+ ) =>+ ParquetSource m ->+ TT.RowGroup ->+ C.ConduitT () ParquetValue m ()+sourceRowGroup source rg = do+ logInfo $ "Parsing new row group. Metadata: " <> LT.toStrict (pString (show rg))+ C.sequenceSources+ ( map+ ( \cc ->+ sourceColumnChunk source cc+ C..| CL.mapMaybe ((<$> mb_path cc) . (,))+ C..| generateInstructions+ )+ (rg ^. TT.pinchField @"column_chunks")+ )+ C..| CL.mapM (construct_record initColumnState)+ where+ mb_path :: TT.ColumnChunk -> Maybe [Text]+ mb_path cc =+ TT.unField+ . TT._ColumnMetaData_path_in_schema+ <$> cc+ ^. TT.pinchField @"meta_data"++ construct_record :: ParquetValue -> [ColumnConstructor] -> m ParquetValue+ construct_record = foldM construct_column++ construct_column :: ParquetValue -> ColumnConstructor -> m ParquetValue+ construct_column pv = foldM apply_instructions pv . ccInstrSet++ apply_instructions :: ParquetValue -> InstructionSet -> m ParquetValue+ apply_instructions val instrSet =+ runExceptT (interpretInstructions val instrSet) >>= \case+ Left err ->+ ParquetNull+ <$ logError ("Error while interpreting instructions: " <> err)+ Right newVal -> pure newVal++------------------------------------------------------------------------------+valueToParquetValue :: Value -> ParquetValue+valueToParquetValue Null = ParquetNull+valueToParquetValue (ValueInt64 v) = ParquetInt v+valueToParquetValue (ValueByteString bs) = ParquetString bs++------------------------------------------------------------------------------+type InstructionSet = Seq.Seq Instruction++------------------------------------------------------------------------------+data Instruction+ = IValue Value+ | IListElement+ | INewList+ | INullOpt+ | IObjectField Text+ | INewListElement+ deriving (Eq, Show)++------------------------------------------------------------------------------++-- | Traverses through given instruction list and changes the given ParquetValue accordingly.+--+-- Given;+-- Value: {}+-- Instruction Set: [IObjectField "f1",INewList,INewList,IValue (ValueInt64 1)]+--+-- Returns;+-- { "f1": [[1]] }+interpretInstructions ::+ (MonadLogger m, MonadError Text m) =>+ ParquetValue ->+ InstructionSet ->+ m ParquetValue+interpretInstructions parquetVal is = do+ logInfo+ $ "Interpreting instructions: "+ <> T.pack (show parquetVal)+ <> ", "+ <> T.pack (show is)+ case (parquetVal, is) of+ (EmptyValue, Seq.Empty) -> pure parquetVal+ (ParquetNull, _) -> pure ParquetNull+ (pv, Seq.Empty) -> pure pv+ (pv, i Seq.:<| ix) -> case i of+ IValue val -> pure $ valueToParquetValue val+ INewListElement -> case pv of+ ParquetList (MkParquetList xs) -> do+ newValueIntoTheList <- interpretInstructions EmptyValue ix+ pure $ ParquetList $ MkParquetList $ xs <> [newValueIntoTheList]+ v ->+ throwError+ $ "Wrong parquet value "+ <> T.pack (show v)+ <> " type for instruction IListElement"+ IListElement -> case pv of+ ParquetList (MkParquetList xs) -> case reverse xs of+ (revX : revXs) -> do+ newRevX <- interpretInstructions revX ix+ pure $ ParquetList $ MkParquetList $ reverse $ newRevX : revXs+ _ -> throwError "List is empty for NestedListElement instruction"+ v ->+ throwError+ $ "Wrong parquet value "+ <> T.pack (show v)+ <> " type for instruction IListElement"+ INewList -> case pv of+ EmptyValue -> do+ newX <- interpretInstructions EmptyValue ix+ pure $ ParquetList $ MkParquetList [newX]+ ParquetList (MkParquetList xs) -> do+ newX <- interpretInstructions EmptyValue ix+ pure $ ParquetList $ MkParquetList $ xs <> [newX]+ v ->+ throwError+ $ "Wrong parquet value "+ <> T.pack (show v)+ <> " type for instruction INewList"+ INullOpt -> interpretInstructions pv ix+ IObjectField fieldName -> case pv of+ EmptyValue -> do+ val <- interpretInstructions EmptyValue ix+ pure+ $ ParquetObject+ $ MkParquetObject+ $ fromList+ [(fieldName, val)]+ ParquetObject (MkParquetObject hm) -> do+ newObj <- flip (at fieldName) hm $ \mbExistingParquetVal ->+ Just+ <$> interpretInstructions+ (fromMaybe EmptyValue mbExistingParquetVal)+ ix+ pure $ ParquetObject $ MkParquetObject newObj+ v ->+ throwError+ $ "Cannot apply IObjectField instruction on parquet value "+ <> T.pack (show v)++------------------------------------------------------------------------------+readSchemaRoot :: (MonadReader TT.FileMetaData m, MonadFail m) => m TT.SchemaElement+readSchemaRoot = do+ metadata <- ask+ headMay (metadata ^. TT.pinchField @"schema") `failOnMay` "Schema cannot be empty"++------------------------------------------------------------------------------+sourceColumnChunk ::+ ( MonadReader TT.FileMetaData m,+ C.MonadIO m,+ C.MonadResource m,+ C.MonadThrow m,+ MonadLogger m,+ MonadFail m+ ) =>+ ParquetSource m ->+ TT.ColumnChunk ->+ C.ConduitT () ColumnValue m ()+sourceColumnChunk (ParquetSource source) cc = do+ metadata <- ask+ schema_mapping <- readSchemaMapping+ col_metadata <- (cc ^. TT.pinchField @"meta_data") `failOnMay` "Metadata could not be found"+ let data_offset = col_metadata ^. TT.pinchField @"data_page_offset"+ dict_offset = col_metadata ^. TT.pinchField @"dictionary_page_offset"+ offset = maybe data_offset (min data_offset) dict_offset+ logInfo $ "Schema 1: " <> LT.toStrict (pString $ show schema_mapping)+ logInfo $ "Schema 2: " <> LT.toStrict (pString $ show (metadata ^. TT.pinchField @"schema"))+ root <- readSchemaRoot+ source (fromIntegral offset)+ C..| C.transPipe failOnExcept (readColumnChunk root schema_mapping cc)
+ src/Parquet/Stream/Reader.hs view
@@ -0,0 +1,393 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MultiWayIf #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE ViewPatterns #-}++module Parquet.Stream.Reader where++import qualified Conduit as C+import Control.Lens+import Control.Monad.Except+import Control.Monad.Logger (MonadLogger)+import Control.Monad.Logger.CallStack (logInfo)+import Control.Monad.Reader+import qualified Data.Binary.Get as BG+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as LBS+import qualified Data.Conduit.Binary as CB+import qualified Data.Conduit.Serialization.Binary as CB+import qualified Data.List.NonEmpty as NE+import qualified Data.Map as M+import qualified Data.Text as T+import qualified Data.Text.Lazy as LT+import Parquet.Decoder (BitWidth (..), decodeBPBE, decodeRLEBPHybrid)+import Parquet.Monad+import qualified Parquet.Pinch as TT+import Parquet.Prelude+import qualified Parquet.Types as TT+import Parquet.Utils ((<??>))+import qualified Pinch+import Text.Pretty.Simple (pString)++data ColumnValue = ColumnValue+ { _cvRepetitionLevel :: Word32,+ _cvDefinitionLevel :: Word32,+ _cvMaxDefinitionLevel :: Word32,+ _cvValue :: Value+ }+ deriving (Eq, Show)++-- | TODO: This is so unoptimized that my eyes bleed.+replicateMSized :: (Monad m) => Int -> m (Int64, result) -> m (Int64, [result])+replicateMSized 0 _ = pure (0, [])+replicateMSized n comp = do+ (consumed, result) <- comp+ (rest_consumed, rest_result) <- replicateMSized (n - 1) comp+ pure (consumed + rest_consumed, result : rest_result)++forSized :: (Monad m) => [a] -> (a -> m (Int64, result)) -> m (Int64, [result])+forSized = flip traverseSized++traverseSized ::+ (Monad m) => (a -> m (Int64, result)) -> [a] -> m (Int64, [result])+traverseSized _ [] = pure (0, [])+traverseSized comp (x : xs) = do+ (consumed, result) <- comp x+ (rest_consumed, rest_result) <- traverseSized comp xs+ pure (consumed + rest_consumed, result : rest_result)++maxLevelToBitWidth :: Word8 -> BitWidth+maxLevelToBitWidth 0 = BitWidth 0+maxLevelToBitWidth max_level =+ BitWidth $ floor (logBase 2 (fromIntegral max_level) :: Double) + 1++dataPageReader ::+ forall m.+ (PR m, MonadReader PageCtx m) =>+ TT.DataPageHeader ->+ Maybe [Value] ->+ C.ConduitT BS.ByteString ColumnValue m ()+dataPageReader header mb_dict = do+ let num_values = header ^. TT.pinchField @"num_values"+ let def_level_encoding = header ^. TT.pinchField @"definition_level_encoding"+ let rep_level_encoding = header ^. TT.pinchField @"repetition_level_encoding"+ let encoding = header ^. TT.pinchField @"encoding"+ (max_rep_level, max_def_level) <- calcMaxEncodingLevels+ (_rep_consumed, fill_level_default num_values -> rep_data) <-+ readRepetitionLevel+ rep_level_encoding+ (maxLevelToBitWidth max_rep_level)+ num_values+ (_def_consumed, fill_level_default num_values -> def_data) <-+ readDefinitionLevel+ def_level_encoding+ (maxLevelToBitWidth max_def_level)+ num_values+ level_data <- zip_level_data rep_data def_data+ path <- asks _pcPath+ logInfo $ LT.toStrict $ pString $ show path+ logInfo $ LT.toStrict $ pString $ show level_data+ read_page_content path encoding level_data num_values (fromIntegral max_def_level)+ pure ()+ where+ find_from_dict ::+ forall m0. MonadError T.Text m0 => [Value] -> Word32 -> m0 Value+ find_from_dict dict (fromIntegral -> d_index) = case dict ^? ix d_index of+ Nothing ->+ throwError $+ "A dictionary value couldn't be found in index "+ <> T.pack+ (show d_index)+ Just val -> pure val++ zip_level_data ::+ forall m0 a b. MonadError T.Text m0 => [a] -> [b] -> m0 [(a, b)]+ zip_level_data rep_data def_data =+ zipExactMay rep_data def_data+ <??> ( "Repetition and Definition data sizes differ: page has "+ <> T.pack (show (length rep_data))+ <> " repetition values and "+ <> T.pack (show (length def_data))+ <> " definition values."+ )++ fill_level_default :: Int32 -> [Word32] -> [Word32]+ fill_level_default num_values = \case+ [] -> replicate (fromIntegral num_values) 0+ xs -> xs++ read_page_content ::+ NE.NonEmpty T.Text ->+ TT.Encoding ->+ [(Word32, Word32)] ->+ Int32 ->+ Word32 ->+ C.ConduitT BS.ByteString ColumnValue m ()+ read_page_content path encoding level_data num_values max_def_level =+ case (mb_dict, encoding) of+ (Nothing, TT.PLAIN _) -> do+ vals <- for level_data $ \(r, d) ->+ if+ | d == max_def_level -> do+ (_, val) <- decodeValue+ pure (ColumnValue r d max_def_level val)+ | otherwise -> pure $ ColumnValue r d max_def_level Null+ logInfo $ "Values (" <> T.pack (show path) <> "): " <> (LT.toStrict $ pString $ show vals)+ C.yieldMany vals+ (Just _, TT.PLAIN _) ->+ throwError+ "We shouldn't have PLAIN-encoded data pages with a dictionary."+ (Just dict, TT.PLAIN_DICTIONARY _) -> do+ !bit_width <- CB.sinkGet BG.getWord8+ val_indexes <-+ CB.sinkGet $+ decodeRLEBPHybrid (BitWidth bit_width) num_values+ vals <- construct_dict_values max_def_level dict level_data val_indexes+ logInfo $ "Values: " <> (LT.toStrict $ pString $ show vals)+ C.yieldMany vals+ (Nothing, TT.PLAIN_DICTIONARY _) ->+ throwError+ "Data page has PLAIN_DICTIONARY encoding but we don't have a dictionary yet."+ other ->+ throwError $+ "Don't know how to encode data pages with encoding: "+ <> T.pack (show other)+ construct_dict_values ::+ forall m0.+ (MonadError T.Text m0) =>+ Word32 ->+ [Value] ->+ [(Word32, Word32)] ->+ [Word32] ->+ m0 [ColumnValue]+ construct_dict_values _ _ [] _ = pure []+ construct_dict_values _ _ _ [] =+ throwError+ "There are not enough level data for given amount of dictionary indexes."+ construct_dict_values max_def_level dict ((r, d) : lx) (v : vx)+ | d == max_def_level = do+ val <- find_from_dict dict v+ (ColumnValue r d max_def_level val :)+ <$> construct_dict_values max_def_level dict lx vx+ | otherwise = do+ (ColumnValue r d max_def_level Null :)+ <$> construct_dict_values max_def_level dict lx (v : vx)++data Value+ = ValueInt64 Int64+ | ValueByteString BS.ByteString+ | Null+ deriving (Show, Eq)++decodeValue ::+ (PR m, MonadReader PageCtx m) =>+ C.ConduitT BS.ByteString o m (Int64, Value)+decodeValue =+ asks _pcColumnTy >>= \case+ (TT.BYTE_ARRAY _) -> do+ !len <- CB.sinkGet BG.getWord32le+ (consumed, result) <-+ replicateMSized+ (fromIntegral len)+ (CB.sinkGet (sizedGet BG.getWord8))+ pure (consumed + 4, ValueByteString (BS.pack result))+ (TT.INT64 _) -> do+ (consumed, result) <- CB.sinkGet (sizedGet BG.getInt64le)+ pure (consumed, ValueInt64 result)+ ty ->+ throwError $+ "Don't know how to decode value of type "+ <> T.pack (show ty)+ <> " yet."++dictPageReader ::+ (MonadReader PageCtx m, PR m) =>+ TT.DictionaryPageHeader ->+ C.ConduitT BS.ByteString o m (Int64, [Value])+dictPageReader header = do+ let num_values = header ^. TT.pinchField @"num_values"+ let _encoding = header ^. TT.pinchField @"encoding"+ let _is_sorted = header ^. TT.pinchField @"is_sorted"+ (consumed, vals) <- replicateMSized (fromIntegral num_values) decodeValue+ pure (consumed, vals)++data PageCtx = PageCtx+ { _pcSchema :: M.Map T.Text TT.SchemaElement,+ _pcPath :: NE.NonEmpty T.Text,+ _pcColumnTy :: TT.Type+ }+ deriving (Show, Eq)++getLastSchemaElement ::+ (MonadError T.Text m, MonadReader PageCtx m) => m TT.SchemaElement+getLastSchemaElement = do+ path <- asks _pcPath+ schema <- asks _pcSchema+ M.lookup (T.intercalate "." (NE.toList path)) schema+ <??> "Schema element could not be found"++readDefinitionLevel ::+ (PR m, MonadReader PageCtx m) =>+ TT.Encoding ->+ BitWidth ->+ Int32 ->+ C.ConduitT BS.ByteString a m (Int64, [Word32])+readDefinitionLevel _ (BitWidth 0) _ = pure (0, [])+readDefinitionLevel encoding bit_width num_values =+ getLastSchemaElement >>= getRepType >>= \case+ TT.OPTIONAL _ -> decodeLevel encoding bit_width num_values+ TT.REPEATED _ -> decodeLevel encoding bit_width num_values+ TT.REQUIRED _ -> pure (0, [])++readRepetitionLevel ::+ (C.MonadThrow m, MonadError T.Text m, MonadReader PageCtx m, MonadLogger m) =>+ TT.Encoding ->+ BitWidth ->+ Int32 ->+ C.ConduitT BS.ByteString a m (Int64, [Word32])+readRepetitionLevel encoding bit_width num_values = do+ decodeLevel encoding bit_width num_values++sizedGet :: BG.Get result -> BG.Get (Int64, result)+sizedGet g = do+ (before, result, after) <- liftA3 (,,) BG.bytesRead g BG.bytesRead+ pure (after - before, result)++decodeLevel ::+ (C.MonadThrow m, MonadError T.Text m, MonadLogger m, MonadReader PageCtx m) =>+ TT.Encoding ->+ BitWidth ->+ Int32 ->+ C.ConduitT BS.ByteString a m (Int64, [Word32])+decodeLevel _ (BitWidth 0) _ = pure (0, [])+decodeLevel encoding bit_width (fromIntegral -> num_values) = case encoding of+ TT.RLE _ ->+ CB.sinkGet $+ sizedGet $+ BG.getWord32le+ *> decodeRLEBPHybrid bit_width num_values+ TT.BIT_PACKED _ ->+ CB.sinkGet $+ sizedGet $+ BG.getWord32le+ *> (take (fromIntegral num_values) <$> decodeBPBE bit_width)+ _ -> throwError "Only RLE and BIT_PACKED encodings are supported for levels"++-- | Algorithm:+-- https://blog.twitter.com/engineering/en_us/a/2013/dremel-made-simple-with-parquet.html+calcMaxEncodingLevels ::+ forall m.+ (MonadReader PageCtx m, MonadError T.Text m) =>+ m (Word8, Word8)+calcMaxEncodingLevels = do+ schema <- asks _pcSchema+ path <- asks _pcPath+ calc_max_encoding_levels schema path+ where+ calc_max_encoding_levels :: M.Map T.Text TT.SchemaElement -> NE.NonEmpty T.Text -> m (Word8, Word8)+ calc_max_encoding_levels schema path = do+ let pathText = T.intercalate "." (NE.toList path)+ schema_element <- M.lookup pathText schema <??> ("Schema Element cannot be found: " <> pathText)+ (repVal, defVal) <-+ getRepType schema_element >>= \case+ (TT.REQUIRED _) -> pure (0, 0)+ (TT.OPTIONAL _) -> pure (0, 1)+ (TT.REPEATED _) -> pure (1, 1)+ case NE.init path of+ (root : v : vx) -> bimap (+ repVal) (+ defVal) <$> calc_max_encoding_levels schema (root NE.:| v : vx)+ _ -> pure (repVal, defVal)++getRepType ::+ MonadError T.Text m => TT.SchemaElement -> m TT.FieldRepetitionType+getRepType e =+ e+ ^. TT.pinchField @"repetition_type"+ <??> "Repetition type could not be found for elem "+ <> T.pack (show e)++validateCompression :: MonadError T.Text m => TT.ColumnMetaData -> m ()+validateCompression metadata =+ let compression = metadata ^. TT.pinchField @"codec"+ in case compression of+ TT.UNCOMPRESSED _ -> pure ()+ _ ->+ throwError "This library doesn't support compression algorithms yet."++readColumnChunk ::+ PR m =>+ TT.SchemaElement ->+ M.Map T.Text TT.SchemaElement ->+ TT.ColumnChunk ->+ C.ConduitT BS.ByteString ColumnValue m ()+readColumnChunk root schema cc = do+ let mb_metadata = cc ^. TT.pinchField @"meta_data"+ metadata <- mb_metadata <??> "Metadata could not be found"+ validateCompression metadata+ let size = metadata ^. TT.pinchField @"total_compressed_size"+ let column_ty = metadata ^. TT.pinchField @"type"+ let path = root ^. TT.pinchField @"name" NE.:| metadata ^. TT.pinchField @"path_in_schema"+ let page_ctx = PageCtx schema path column_ty+ C.runReaderC page_ctx $ readPage size Nothing++readPage ::+ (MonadLogger m, MonadReader PageCtx m, PR m) =>+ Int64 ->+ Maybe [Value] ->+ C.ConduitT BS.ByteString ColumnValue m ()+readPage 0 _ = pure ()+readPage remaining mb_dict = do+ (page_header_size, page_header :: TT.PageHeader) <- decodeConduit remaining+ logInfo $ LT.toStrict $ pString $ show page_header+ let page_content_size = page_header ^. TT.pinchField @"uncompressed_page_size"+ let validate_consumed_page_bytes consumed =+ unless (fromIntegral page_content_size == consumed) $+ throwError "Reader did not consume the whole page!"+ let page_size = fromIntegral page_header_size + page_content_size+ case ( page_header ^. TT.pinchField @"dictionary_page_header",+ page_header ^. TT.pinchField @"data_page_header",+ mb_dict+ ) of+ (Just dict_page_header, Nothing, Nothing) -> do+ (page_consumed, dict) <- dictPageReader dict_page_header+ validate_consumed_page_bytes page_consumed+ readPage (remaining - fromIntegral page_size) (Just dict)+ (Just _dict_page_header, Nothing, Just _dict) ->+ throwError "Found dictionary page while we already had a dictionary."+ (Nothing, Just dp_header, Nothing) -> do+ dataPageReader dp_header Nothing+ (Nothing, Just dp_header, Just dict) -> do+ dataPageReader dp_header (Just dict)+ (Nothing, Nothing, _) ->+ throwError+ "Page doesn't have any of the dictionary or data page header."+ (Just _, Just _, _) ->+ throwError "Page has both dictionary and data page headers."++failOnError :: Show err => IO (Either err b) -> IO b+failOnError v =+ v >>= \case+ Left err -> fail $ show err+ Right val -> pure val++decodeConduit ::+ forall a size m o.+ (MonadError T.Text m, MonadIO m, Integral size, Pinch.Pinchable a) =>+ size ->+ C.ConduitT BS.ByteString o m (Int, a)+decodeConduit (fromIntegral -> size) = do+ (left, val) <-+ liftEither+ . first T.pack+ . Pinch.decodeWithLeftovers Pinch.compactProtocol+ . LBS.toStrict+ =<< CB.take size+ C.leftover left+ pure (size - BS.length left, val)
+ src/Parquet/Types.hs view
@@ -0,0 +1,15 @@+module Parquet.Types+ ( module X,+ )+where++------------------------------------------------------------------------------++import Parquet.Types.ColumnChunk as X+import Parquet.Types.Enums as X+import Parquet.Types.FileMetaData as X+import Parquet.Types.PageHeader as X+import Parquet.Types.Primitives as X+import Parquet.Types.RowGroup as X++------------------------------------------------------------------------------
+ src/Parquet/Types/ColumnChunk.hs view
@@ -0,0 +1,57 @@+module Parquet.Types.ColumnChunk+ ( -- * Type definitions+ ColumnChunk (..),+ ColumnMetaData (..),++ -- * Internal type definitions+ PageEncodingStats (..),+ )+where++------------------------------------------------------------------------------++import Parquet.Prelude+import Parquet.Types.Enums+import Parquet.Types.Primitives+import Pinch++------------------------------------------------------------------------------+data ColumnChunk = ColumnChunk+ { _ColumnChunk_file_path :: Field 1 (Maybe Text),+ _ColumnChunk_file_offset :: Field 2 Int64,+ _ColumnChunk_meta_data :: Field 3 (Maybe ColumnMetaData),+ _ColumnChunk_offset_index_offset :: Field 4 (Maybe Int64),+ _ColumnChunk_offset_index_length :: Field 5 (Maybe Int32),+ _ColumnChunk_column_index_offset :: Field 6 (Maybe Int64),+ _ColumnChunk_column_index_length :: Field 7 (Maybe Int32),+ _ColumnChunk_crypto_metadata :: Field 8 (Maybe ColumnCryptoMetaData),+ _ColumnChunk_encrypted_column_metadata :: Field 9 (Maybe ByteString)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data ColumnMetaData = ColumnMetaData+ { _ColumnMetaData_type :: Field 1 Type,+ _ColumnMetaData_encodings :: Field 2 [Encoding],+ _ColumnMetaData_path_in_schema :: Field 3 [Text],+ _ColumnMetaData_codec :: Field 4 CompressionCodec,+ _ColumnMetaData_num_values :: Field 5 Int64,+ _ColumnMetaData_total_uncompressed_size :: Field 6 Int64,+ _ColumnMetaData_total_compressed_size :: Field 7 Int64,+ _ColumnMetaData_key_value_metadata :: Field 8 (Maybe [KeyValue]),+ _ColumnMetaData_data_page_offset :: Field 9 Int64,+ _ColumnMetaData_index_page_offset :: Field 10 (Maybe Int64),+ _ColumnMetaData_dictionary_page_offset :: Field 11 (Maybe Int64),+ _ColumnMetaData_statistics :: Field 12 (Maybe Statistics),+ _ColumnMetaData_encoding_stats :: Field 13 (Maybe [PageEncodingStats]),+ _ColumnMetaData_bloom_filter_offset :: Field 14 (Maybe Int64)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data PageEncodingStats = PageEncodingStats+ { _PageEncodingStats_page_type :: Field 1 PageType,+ _PageEncodingStats_encoding :: Field 2 Encoding,+ _PageEncodingStats_count :: Field 3 Int32+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)
+ src/Parquet/Types/Enums.hs view
@@ -0,0 +1,129 @@+module Parquet.Types.Enums+ ( -- * Type definitions+ ColumnCryptoMetaData (..),+ ColumnOrder (..),+ CompressionCodec (..),+ ConvertedType (..),+ Encoding (..),+ EncryptionAlgorithm (..),+ FieldRepetitionType (..),+ LogicalType (..),+ PageType (..),+ Type (..),+ )+where++------------------------------------------------------------------------------++import Parquet.Prelude+import Parquet.Types.Primitives+import Pinch++------------------------------------------------------------------------------+data ColumnCryptoMetaData+ = ColumnCryptoMetaData_ENCRYPTION_WITH_FOOTER_KEY (Field 1 EncryptionWithFooterKey)+ | ColumnCryptoMetaData_ENCRYPTION_WITH_COLUMN_KEY (Field 2 EncryptionWithColumnKey)+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data ColumnOrder+ = ColumnOrder_TYPE_ORDER (Field 1 TypeDefinedOrder)+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data CompressionCodec+ = UNCOMPRESSED (Enumeration 0)+ | SNAPPY (Enumeration 1)+ | GZIP (Enumeration 2)+ | LZO (Enumeration 3)+ | BROTLI (Enumeration 4)+ | LZ4 (Enumeration 5)+ | ZSTD (Enumeration 6)+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data ConvertedType+ = UTF8 (Enumeration 0)+ | MAP (Enumeration 1)+ | MAP_KEY_VALUE (Enumeration 2)+ | LIST (Enumeration 3)+ | DECIMAL (Enumeration 5)+ | DATE (Enumeration 6)+ | TIME_MILLIS (Enumeration 7)+ | TIME_MICROS (Enumeration 8)+ | TIMESTAMP_MILLIS (Enumeration 9)+ | TIMESTAMP_MICROS (Enumeration 10)+ | UINT_8 (Enumeration 11)+ | UINT_16 (Enumeration 12)+ | UINT_32 (Enumeration 13)+ | UINT_64 (Enumeration 14)+ | INT_8 (Enumeration 15)+ | INT_16 (Enumeration 16)+ | INT_32 (Enumeration 17)+ | INT_64 (Enumeration 18)+ | JSON (Enumeration 19)+ | BSON (Enumeration 20)+ | INTERVAL (Enumeration 21)+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data Encoding+ = PLAIN (Enumeration 0)+ | PLAIN_DICTIONARY (Enumeration 2)+ | RLE (Enumeration 3)+ | BIT_PACKED (Enumeration 4)+ | DELTA_BINARY_PACKED (Enumeration 5)+ | DELTA_LENGTH_BYTE_ARRAY (Enumeration 6)+ | DELTA_BYTE_ARRAY (Enumeration 7)+ | RLE_DICTIONARY (Enumeration 8)+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data EncryptionAlgorithm+ = EncryptionAlgorithm_AES_GCM_V1 (Field 1 AesGcmV1)+ | EncryptionAlgorithm_AES_GCM_CTR_V1 (Field 2 AesGcmCtrV1)+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data FieldRepetitionType+ = REQUIRED (Enumeration 0)+ | OPTIONAL (Enumeration 1)+ | REPEATED (Enumeration 2)+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data LogicalType+ = LogicalTypeSTRING (Field 1 StringType)+ | LogicalTypeMAP (Field 2 MapType)+ | LogicalTypeLIST (Field 3 ListType)+ | LogicalTypeENUM (Field 4 EnumType)+ | LogicalTypeDECIMAL (Field 5 DecimalType)+ | LogicalTypeDATE (Field 6 DateType)+ | LogicalTypeTIME (Field 7 TimeType)+ | LogicalTypeTIMESTAMP (Field 8 TimestampType)+ | LogicalTypeINTEGER (Field 10 IntType)+ | LogicalTypeUNKNOWN (Field 11 NullType)+ | LogicalTypeJSON (Field 12 JsonType)+ | LogicalTypeBSON (Field 13 BsonType)+ | LogicalTypeUUID (Field 14 UUIDType)+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data PageType+ = DATA_PAGE (Enumeration 0)+ | INDEX_PAGE (Enumeration 1)+ | DICTIONARY_PAGE (Enumeration 2)+ | DATA_PAGE_V2 (Enumeration 3)+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data Type+ = BOOLEAN (Enumeration 0)+ | INT32 (Enumeration 1)+ | INT64 (Enumeration 2)+ | INT96 (Enumeration 3)+ | FLOAT (Enumeration 4)+ | DOUBLE (Enumeration 5)+ | BYTE_ARRAY (Enumeration 6)+ | FIXED_LEN_BYTE_ARRAY (Enumeration 7)+ deriving (Show, Eq, Generic, Pinchable, Binary)
+ src/Parquet/Types/FileMetaData.hs view
@@ -0,0 +1,43 @@+module Parquet.Types.FileMetaData+ ( -- * Type definitions+ FileMetaData (..),+ SchemaElement (..),+ )+where++------------------------------------------------------------------------------++import Parquet.Prelude+import Parquet.Types.Enums+import Parquet.Types.Primitives+import Parquet.Types.RowGroup+import Pinch++------------------------------------------------------------------------------+data FileMetaData = FileMetaData+ { _FileMetaData_version :: Field 1 Int32,+ _FileMetaData_schema :: Field 2 [SchemaElement],+ _FileMetaData_num_rows :: Field 3 Int64,+ _FileMetaData_row_groups :: Field 4 [RowGroup],+ _FileMetaData_key_value_metadata :: Field 5 (Maybe [KeyValue]),+ _FileMetaData_created_by :: Field 6 (Maybe Text),+ _FileMetaData_column_orders :: Field 7 (Maybe [ColumnOrder]),+ _FileMetaData_encryption_algorithm :: Field 8 (Maybe EncryptionAlgorithm),+ _FileMetaData_footer_signing_key_metadata :: Field 9 (Maybe ByteString)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data SchemaElement = SchemaElement+ { _SchemaElement_type :: Field 1 (Maybe Type),+ _SchemaElement_type_length :: Field 2 (Maybe Int32),+ _SchemaElement_repetition_type :: Field 3 (Maybe FieldRepetitionType),+ _SchemaElement_name :: Field 4 Text,+ _SchemaElement_num_children :: Field 5 (Maybe Int32),+ _SchemaElement_converted_type :: Field 6 (Maybe ConvertedType),+ _SchemaElement_scale :: Field 7 (Maybe Int32),+ _SchemaElement_precision :: Field 8 (Maybe Int32),+ _SchemaElement_field_id :: Field 9 (Maybe Int32),+ _SchemaElement_logicalType :: Field 10 (Maybe LogicalType)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)
+ src/Parquet/Types/PageHeader.hs view
@@ -0,0 +1,73 @@+module Parquet.Types.PageHeader+ ( -- * Type definitions+ DataPageHeader (..),+ DictionaryPageHeader (..),+ PageHeader (..),++ -- * Internal type definitions+ IndexPageHeader (..),++ -- * Unused type definitions+ DataPageHeaderV2 (..),+ )+where++------------------------------------------------------------------------------++import Parquet.Prelude+import Parquet.Types.Enums+import Parquet.Types.Primitives+import Pinch++------------------------------------------------------------------------------+data DataPageHeader = DataPageHeader+ { _DataPageHeader_num_values :: Field 1 Int32,+ _DataPageHeader_encoding :: Field 2 Encoding,+ _DataPageHeader_definition_level_encoding :: Field 3 Encoding,+ _DataPageHeader_repetition_level_encoding :: Field 4 Encoding,+ _DataPageHeader_statistics :: Field 5 (Maybe Statistics)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data DataPageHeaderV2 = DataPageHeaderV2+ { _DataPageHeaderV2_num_values :: Field 1 Int32,+ _DataPageHeaderV2_num_nulls :: Field 2 Int32,+ _DataPageHeaderV2_num_rows :: Field 3 Int32,+ _DataPageHeaderV2_encoding :: Field 4 Encoding,+ _DataPageHeaderV2_definition_levels_byte_length :: Field 5 Int32,+ _DataPageHeaderV2_repetition_levels_byte_length :: Field 6 Int32,+ _DataPageHeaderV2_is_compressed :: Field 7 (Maybe Bool),+ _DataPageHeaderV2_statistics :: Field 8 (Maybe Statistics)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data DictionaryPageHeader = DictionaryPageHeader+ { _DictionaryPageHeader_num_values :: Field 1 Int32,+ _DictionaryPageHeader_encoding :: Field 2 Encoding,+ _DictionaryPageHeader_is_sorted :: Field 3 (Maybe Bool)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data IndexPageHeader = IndexPageHeader+ deriving (Show, Eq, Generic, Binary)++instance Pinchable IndexPageHeader where+ type Tag IndexPageHeader = TStruct+ pinch _ = struct []+ unpinch _ = pure IndexPageHeader++------------------------------------------------------------------------------+data PageHeader = PageHeader+ { _PageHeader_type :: Field 1 PageType,+ _PageHeader_uncompressed_page_size :: Field 2 Int32,+ _PageHeader_compressed_page_size :: Field 3 Int32,+ _PageHeader_crc :: Field 4 (Maybe Int32),+ _PageHeader_data_page_header :: Field 5 (Maybe DataPageHeader),+ _PageHeader_index_page_header :: Field 6 (Maybe IndexPageHeader),+ _PageHeader_dictionary_page_header :: Field 7 (Maybe DictionaryPageHeader),+ _PageHeader_data_page_header_v2 :: Field 8 (Maybe DataPageHeaderV2)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)
+ src/Parquet/Types/Primitives.hs view
@@ -0,0 +1,241 @@+module Parquet.Types.Primitives+ ( -- * Type definitions+ AesGcmV1 (..),+ AesGcmCtrV1 (..),+ EncryptionWithColumnKey (..),+ EncryptionWithFooterKey (..),+ KeyValue (..),+ Statistics (..),+ TypeDefinedOrder (..),++ -- * Primitive wrappers+ BsonType (..),+ DateType (..),+ DecimalType (..),+ EnumType (..),+ IntType (..),+ JsonType (..),+ ListType (..),+ MapType (..),+ NullType (..),+ StringType (..),+ TimestampType (..),+ TimeType (..),+ UUIDType (..),++ -- * Internal type definitions+ TimeUnit (..),+ MicroSeconds (..),+ MilliSeconds (..),+ NanoSeconds (..),+ )+where++------------------------------------------------------------------------------++import Parquet.Prelude+import Pinch++------------------------------------------------------------------------------+data AesGcmV1 = AesGcmV1+ { _AesGcmV1_aad_prefix :: Field 1 (Maybe ByteString),+ _AesGcmV1_aad_file_unique :: Field 2 (Maybe ByteString),+ _AesGcmV1_supply_aad_prefix :: Field 3 (Maybe Bool)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data AesGcmCtrV1 = AesGcmCtrV1+ { _AesGcmCtrV1_aad_prefix :: Field 1 (Maybe ByteString),+ _AesGcmCtrV1_aad_file_unique :: Field 2 (Maybe ByteString),+ _AesGcmCtrV1_supply_aad_prefix :: Field 3 (Maybe Bool)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data EncryptionWithColumnKey = EncryptionWithColumnKey+ { _EncryptionWithColumnKey_path_in_schema :: Field 1 [Text],+ _EncryptionWithColumnKey_key_metadata :: Field 2 (Maybe ByteString)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data EncryptionWithFooterKey = EncryptionWithFooterKey+ deriving (Show, Eq, Generic, Binary)++instance Pinchable EncryptionWithFooterKey where+ type Tag EncryptionWithFooterKey = TStruct+ pinch _ = struct []+ unpinch _ = pure EncryptionWithFooterKey++------------------------------------------------------------------------------+data KeyValue = KeyValue+ { _KeyValue_key :: Field 1 Text,+ _KeyValue_value :: Field 2 (Maybe Text)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data Statistics = Statistics+ { _Statistics_max :: Field 1 (Maybe ByteString),+ _Statistics_min :: Field 2 (Maybe ByteString),+ _Statistics_null_count :: Field 3 (Maybe Int64),+ _Statistics_distinct_count :: Field 4 (Maybe Int64),+ _Statistics_max_value :: Field 5 (Maybe ByteString),+ _Statistics_min_value :: Field 6 (Maybe ByteString)+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data TypeDefinedOrder = TypeDefinedOrder+ deriving (Show, Eq, Generic, Binary)++instance Pinchable TypeDefinedOrder where+ type Tag TypeDefinedOrder = TStruct+ pinch _ = struct []+ unpinch _ = pure TypeDefinedOrder++------------------------------------------------------------------------------+data BsonType = BsonType+ deriving (Show, Eq, Generic, Binary)++instance Pinchable BsonType where+ type Tag BsonType = TStruct+ pinch _ = struct []+ unpinch _ = pure BsonType++------------------------------------------------------------------------------+data DateType = DateType+ deriving (Show, Eq, Generic, Binary)++instance Pinchable DateType where+ type Tag DateType = TStruct+ pinch _ = struct []+ unpinch _ = pure DateType++------------------------------------------------------------------------------+data DecimalType = DecimalType+ { _DecimalType_scale :: Field 1 Int32,+ _DecimalType_precision :: Field 2 Int32+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data EnumType = EnumType+ deriving (Show, Eq, Generic, Binary)++instance Pinchable EnumType where+ type Tag EnumType = TStruct+ pinch _ = struct []+ unpinch _ = pure EnumType++------------------------------------------------------------------------------+data IntType = IntType+ { _IntType_bitWidth :: Field 1 Int8,+ _IntType_isSigned :: Field 2 Bool+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data JsonType = JsonType+ deriving (Show, Eq, Generic, Binary)++instance Pinchable JsonType where+ type Tag JsonType = TStruct+ pinch _ = struct []+ unpinch _ = pure JsonType++------------------------------------------------------------------------------+data ListType = ListType+ deriving (Show, Eq, Generic, Binary)++instance Pinchable ListType where+ type Tag ListType = TStruct+ pinch _ = struct []+ unpinch _ = pure ListType++------------------------------------------------------------------------------+data MapType = MapType+ deriving (Show, Eq, Generic, Binary)++instance Pinchable MapType where+ type Tag MapType = TStruct+ pinch _ = struct []+ unpinch _ = pure MapType++------------------------------------------------------------------------------+data MicroSeconds = MicroSeconds+ deriving (Show, Eq, Generic, Binary)++instance Pinchable MicroSeconds where+ type Tag MicroSeconds = TStruct+ pinch _ = struct []+ unpinch _ = pure MicroSeconds++------------------------------------------------------------------------------+data MilliSeconds = MilliSeconds+ deriving (Show, Eq, Generic, Binary)++instance Pinchable MilliSeconds where+ type Tag MilliSeconds = TStruct+ pinch _ = struct []+ unpinch _ = pure MilliSeconds++------------------------------------------------------------------------------+data NanoSeconds = NanoSeconds+ deriving (Show, Eq, Generic, Binary)++instance Pinchable NanoSeconds where+ type Tag NanoSeconds = TStruct+ pinch _ = struct []+ unpinch _ = pure NanoSeconds++------------------------------------------------------------------------------+data NullType = NullType+ deriving (Show, Eq, Generic, Binary)++instance Pinchable NullType where+ type Tag NullType = TStruct+ pinch _ = struct []+ unpinch _ = pure NullType++------------------------------------------------------------------------------+data StringType = StringType+ deriving (Show, Eq, Generic, Binary)++instance Pinchable StringType where+ type Tag StringType = TStruct+ pinch _ = struct []+ unpinch _ = pure StringType++------------------------------------------------------------------------------+data TimestampType = TimestampType+ { _TimestampType_isAdjustedToUTC :: Field 1 Bool,+ _TimestampType_unit :: Field 2 TimeUnit+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data TimeType = TimeType+ { _TimeType_isAdjustedToUTC :: Field 1 Bool,+ _TimeType_unit :: Field 2 TimeUnit+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)++------------------------------------------------------------------------------+data TimeUnit+ = TimeUnitMILLIS (Field 1 MilliSeconds)+ | TimeUnitMICROS (Field 2 MicroSeconds)+ | TimeUnitNANOS (Field 3 NanoSeconds)+ deriving (Show, Eq, Generic, Binary)++instance Pinchable TimeUnit++------------------------------------------------------------------------------+data UUIDType = UUIDType+ deriving (Show, Eq, Generic, Binary)++instance Pinchable UUIDType where+ type Tag UUIDType = TStruct+ pinch _ = struct []+ unpinch _ = pure UUIDType
+ src/Parquet/Types/RowGroup.hs view
@@ -0,0 +1,34 @@+module Parquet.Types.RowGroup+ ( -- * Type definitions+ RowGroup (..),++ -- * Internal type definitions+ SortingColumn (..),+ )+where++------------------------------------------------------------------------------++import Parquet.Prelude+import Parquet.Types.ColumnChunk+import Pinch++------------------------------------------------------------------------------+data RowGroup = RowGroup+ { _RowGroup_column_chunks :: Field 1 [ColumnChunk],+ _RowGroup_total_byte_size :: Field 2 Int64,+ _RowGroup_num_rows :: Field 3 Int64,+ _RowGroup_sorting_columns :: Field 4 (Maybe [SortingColumn]),+ _RowGroup_file_offset :: Field 5 (Maybe Int64),+ _RowGroup_total_compressed_size :: Field 6 (Maybe Int64),+ _RowGroup_ordinal :: Field 7 (Maybe Int16)+ }+ deriving (Show, Eq, Generic, Binary, Pinchable)++------------------------------------------------------------------------------+data SortingColumn = SortingColumn+ { _SortingColumn_column_idx :: Field 1 Int32,+ _SortingColumn_descending :: Field 2 Bool,+ _SortingColumn_nulls_first :: Field 3 Bool+ }+ deriving (Show, Eq, Generic, Pinchable, Binary)
+ src/Parquet/Utils.hs view
@@ -0,0 +1,38 @@+-- |+module Parquet.Utils where++------------------------------------------------------------------------------++import Control.Monad.Except+import qualified Data.Text as T+import Parquet.Prelude++------------------------------------------------------------------------------+(<??>) :: MonadError b m => Maybe a -> b -> m a+(<??>) Nothing err = throwError err+(<??>) (Just v) _ = pure v++infixl 4 <??>++------------------------------------------------------------------------------+failOnMay ::+ ( Monad m,+ MonadFail m+ ) =>+ Maybe a ->+ String ->+ m a+failOnMay Nothing s = fail s+failOnMay (Just a) _ = pure a++------------------------------------------------------------------------------+failOnExcept ::+ ( Monad m,+ MonadFail m+ ) =>+ ExceptT Text m a ->+ m a+failOnExcept =+ runExceptT >=> \case+ Left err -> fail (T.unpack err)+ Right v -> pure v
+ tests/integration/Spec.hs view
@@ -0,0 +1,76 @@+{-# LANGUAGE ScopedTypeVariables #-}++module Main+ ( main,+ )+where++import Conduit (runResourceT)+import Control.Exception (bracket_)+import Control.Monad.Except (runExceptT)+import Control.Monad.Logger+import qualified Data.Aeson as JSON+import qualified Data.ByteString.Lazy as LBS+import Parquet.Reader (readWholeParquetFile)+import System.Environment (setEnv, unsetEnv)+import System.FilePath ((</>))+import System.Process+import Test.Hspec+import Prelude++testPath :: String+testPath = "tests" </> "integration"++testDataPath :: String+testDataPath = testPath </> "testdata"++intermediateDir :: String+intermediateDir = "nested.parquet"++encoderScriptPath :: String+encoderScriptPath = "gen_parquet.py"++outParquetFilePath :: String+outParquetFilePath = testPath </> "test.parquet"++pysparkPythonEnvName :: String+pysparkPythonEnvName = "PYSPARK_PYTHON"++testParquetFormat :: String -> (String -> IO () -> IO ()) -> IO ()+testParquetFormat _inputFile performTest =+ bracket_+ (setEnv pysparkPythonEnvName "/usr/bin/python3")+ (unsetEnv pysparkPythonEnvName)+ $ do+ callProcess+ "python3"+ [ testPath </> encoderScriptPath,+ testDataPath </> "input1.json",+ testPath </> intermediateDir+ ]+ callCommand $+ "cp "+ <> testPath+ </> intermediateDir+ </> "*.parquet "+ <> outParquetFilePath++ let close = callProcess "rm" ["-rf", testPath </> intermediateDir]+ performTest outParquetFilePath close+ close++main :: IO ()+main = hspec $+ describe "Reader" $ do+ it "can read columns" $ do+ testParquetFormat "input1.json" $ \parqFile closePrematurely -> do+ result <-+ runResourceT+ (runStdoutLoggingT (runExceptT (readWholeParquetFile parqFile)))+ case result of+ Left err -> fail $ show err+ Right v -> do+ origJson :: Maybe JSON.Value <- JSON.decode <$> LBS.readFile (testDataPath </> "input1.json")+ closePrematurely+ Just (JSON.encode v) `shouldBe` (JSON.encode <$> origJson)+ pure ()
+ tests/unit/Parquet/Decoder/Spec.hs view
@@ -0,0 +1,66 @@+module Parquet.Decoder.Spec+ ( spec,+ )+where++import Data.Binary.Put+import Data.Binary.Get (runGetOrFail)+import qualified Data.ByteString.Lazy as LBS+import Parquet.Decoder+import Test.Hspec+import Prelude++spec :: Spec+spec = describe "Decoder" $ do+ it "can decode from little endian bit packing" $+ runGetOrFail (decodeBPLE (BitWidth 3) 1) (LBS.pack [136, 198, 250])+ `shouldBe` Right (LBS.empty, 3, [0, 1, 2, 3, 4, 5, 6, 7])++ it "can decode from little endian bit packing with 0 padding" $+ -- There are 5 values: [0, 5)+ -- Bit packing is 3+ -- Scaled run len is: 1 (and run len is 8) because we keep 5 values+ --+ -- Result type is (Either Error (Unconsumed bytes, consumed byte len, result))+ -- In this case we decode 5 values but due to parquet only bit-packs+ -- a multiple of 8 values at once, we still have to consume them.+ -- Docs:+ -- https://github.com/apache/parquet-format/blob/master/Encodings.md#run-length-encoding--bit-packing-hybrid-rle--3+ runGetOrFail+ (decodeBPLE (BitWidth 3) 1)+ (LBS.pack [0x88, 0x46, 0x00, 0x00])+ `shouldBe` Right (LBS.pack [0], 3, [0, 1, 2, 3, 4, 0, 0, 0])++ it "can decode from big endian bit packing" $+ runGetOrFail+ (decodeBPBE (BitWidth 3))+ (runPut (encodeVarint 3) <> LBS.pack [5, 57, 119])+ `shouldBe` Right (LBS.empty, 4, [0, 1, 2, 3, 4, 5, 6, 7])++ it "can do run length encoding (RLE)" $+ runGetOrFail (decodeRLE (BitWidth 3) 4) (LBS.pack [1, 2, 3, 4, 5])+ `shouldBe` Right (LBS.pack [2, 3, 4, 5], 1, [1, 1, 1, 1])++ it "can takeBytesLe" $+ runGetOrFail (takeBytesLe 3) (LBS.pack [136, 198, 250])+ `shouldBe` Right (LBS.empty, 3, 16434824)++ it "can takeBytesLe with leftovers" $+ runGetOrFail (takeBytesLe 3) (LBS.pack [136, 198, 250, 1, 2])+ `shouldBe` Right (LBS.pack [1, 2], 3, 16434824)++ it "can encode varint (LEB128) exactly" $+ runPut (encodeVarint 624485)+ `shouldBe` LBS.pack [0xE5, 0x8E, 0x26]++ it "can decode from varint (LEB128) exactly" $+ runGetOrFail decodeVarint (LBS.pack [0xE5, 0x8E, 0x26])+ `shouldBe` Right (LBS.empty, 3, 624485)++ it "can decode from varint (LEB128) with leftovers" $+ runGetOrFail decodeVarint (LBS.pack [0xE5, 0x8E, 0x26, 1, 2])+ `shouldBe` Right (LBS.pack [1, 2], 3, 624485)++ it "can decode from varint (LEB128) with leftovers and lookahead" $+ runGetOrFail decodeVarint (LBS.pack [0xE5, 0x8E, 0x26, 1, 2])+ `shouldBe` Right (LBS.pack [1, 2], 3, 624485)
+ tests/unit/Spec.hs view
@@ -0,0 +1,9 @@+module Main (main) where++import qualified Parquet.Decoder.Spec+import Test.Hspec+import Prelude++main :: IO ()+main =+ hspec Parquet.Decoder.Spec.spec