diff --git a/ChangeLog.md b/ChangeLog.md new file mode 100644 --- /dev/null +++ b/ChangeLog.md @@ -0,0 +1,3 @@ +# Changelog for parquet-haskell + +## Unreleased changes diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,31 @@ +# parquet-haskell + +A Streaming Parquet reader for Haskell. + +[![Haskell CI](https://github.com/manishsencha/parquet-haskell/actions/workflows/ci.yml/badge.svg)](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 +``` diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,3 @@ +import Distribution.Simple + +main = defaultMain diff --git a/parquet-haskell.cabal b/parquet-haskell.cabal new file mode 100644 --- /dev/null +++ b/parquet-haskell.cabal @@ -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 +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 +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 diff --git a/src/Main.hs b/src/Main.hs new file mode 100644 --- /dev/null +++ b/src/Main.hs @@ -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 diff --git a/src/Parquet/Decoder.hs b/src/Parquet/Decoder.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Decoder.hs @@ -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) diff --git a/src/Parquet/InstanceOrphanage.hs b/src/Parquet/InstanceOrphanage.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/InstanceOrphanage.hs @@ -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) diff --git a/src/Parquet/Monad.hs b/src/Parquet/Monad.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Monad.hs @@ -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) diff --git a/src/Parquet/ParquetObject.hs b/src/Parquet/ParquetObject.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/ParquetObject.hs @@ -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 "" + ParquetBool b -> JSON.Bool b + ParquetNull -> JSON.Null + EmptyValue -> JSON.Null + +------------------------------------------------------------------------------ +makeLenses ''ParquetObject +makeLenses ''ParquetValue + +makePrisms ''ParquetObject +makePrisms ''ParquetValue diff --git a/src/Parquet/Pinch.hs b/src/Parquet/Pinch.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Pinch.hs @@ -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 diff --git a/src/Parquet/Prelude.hs b/src/Parquet/Prelude.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Prelude.hs @@ -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) + +------------------------------------------------------------------------------ diff --git a/src/Parquet/Reader.hs b/src/Parquet/Reader.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Reader.hs @@ -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) diff --git a/src/Parquet/Stream/Reader.hs b/src/Parquet/Stream/Reader.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Stream/Reader.hs @@ -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) diff --git a/src/Parquet/Types.hs b/src/Parquet/Types.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Types.hs @@ -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 + +------------------------------------------------------------------------------ diff --git a/src/Parquet/Types/ColumnChunk.hs b/src/Parquet/Types/ColumnChunk.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Types/ColumnChunk.hs @@ -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) diff --git a/src/Parquet/Types/Enums.hs b/src/Parquet/Types/Enums.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Types/Enums.hs @@ -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) diff --git a/src/Parquet/Types/FileMetaData.hs b/src/Parquet/Types/FileMetaData.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Types/FileMetaData.hs @@ -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) diff --git a/src/Parquet/Types/PageHeader.hs b/src/Parquet/Types/PageHeader.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Types/PageHeader.hs @@ -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) diff --git a/src/Parquet/Types/Primitives.hs b/src/Parquet/Types/Primitives.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Types/Primitives.hs @@ -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 diff --git a/src/Parquet/Types/RowGroup.hs b/src/Parquet/Types/RowGroup.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Types/RowGroup.hs @@ -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) diff --git a/src/Parquet/Utils.hs b/src/Parquet/Utils.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Utils.hs @@ -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 diff --git a/tests/integration/Spec.hs b/tests/integration/Spec.hs new file mode 100644 --- /dev/null +++ b/tests/integration/Spec.hs @@ -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 () diff --git a/tests/unit/Parquet/Decoder/Spec.hs b/tests/unit/Parquet/Decoder/Spec.hs new file mode 100644 --- /dev/null +++ b/tests/unit/Parquet/Decoder/Spec.hs @@ -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) diff --git a/tests/unit/Spec.hs b/tests/unit/Spec.hs new file mode 100644 --- /dev/null +++ b/tests/unit/Spec.hs @@ -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