diff --git a/ChangeLog.md b/ChangeLog.md new file mode 100644 --- /dev/null +++ b/ChangeLog.md @@ -0,0 +1,3 @@ +# Changelog for parquet-hs + +## 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,1 @@ +# parquet-hs diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/parquet-hs.cabal b/parquet-hs.cabal new file mode 100644 --- /dev/null +++ b/parquet-hs.cabal @@ -0,0 +1,108 @@ +cabal-version: 1.12 + +-- This file has been generated from package.yaml by hpack version 0.31.2. +-- +-- see: https://github.com/sol/hpack +-- +-- hash: f82d7918bf922d7a9150a290527c2b0bc63f79b924d159ab0a1af905d6917e17 + +name: parquet-hs +version: 0.1.0.0 +synopsis: Streaming Parquet reader +description: Please see the README on GitHub at +category: Compression +homepage: https://github.com/yigitozkavci/parquet-hs#readme +bug-reports: https://github.com/yigitozkavci/parquet-hs/issues +author: Yigit Ozkavci +maintainer: yigitozkavci8@gmail.com +copyright: 2019 Yigit Ozkavci +license: BSD3 +license-file: LICENSE +build-type: Simple +extra-source-files: + README.md + ChangeLog.md + +source-repository head + type: git + location: https://github.com/yigitozkavci/parquet-hs + +library + exposed-modules: + Main + Parquet.Decoder + Parquet.Monad + Parquet.ParquetObject + Parquet.Reader + Parquet.Stream.Reader + Parquet.ThriftTypes + Parquet.Utils + other-modules: + Paths_parquet_hs + hs-source-dirs: + src + ghc-options: -Wall + build-depends: + aeson + , base >=4.7 && <5 + , binary + , binary-conduit + , bytestring + , conduit + , conduit-extra + , containers + , error-util + , generic-lens + , http-client + , http-conduit + , http-types + , lens + , lifted-async + , monad-logger + , mtl + , pinch + , safe + , serialise + , text + , unordered-containers + default-language: Haskell2010 + +test-suite parquet-hs-integration + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + Paths_parquet_hs + hs-source-dirs: + tests/integration + build-depends: + QuickCheck + , aeson + , base >=4.7 && <5 + , binary + , bytestring + , conduit + , filepath + , hspec + , monad-logger + , mtl + , parquet-hs + , process + , text + default-language: Haskell2010 + +test-suite parquet-hs-unit + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + Parquet.Decoder.Spec + Paths_parquet_hs + hs-source-dirs: + tests/unit + build-depends: + QuickCheck + , base >=4.7 && <5 + , binary + , bytestring + , hspec + , parquet-hs + 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,12 @@ +module Main where + +import Parquet.Reader (readWholeParquetFile) +import Control.Monad.Except +import qualified Conduit as C +import Control.Monad.Logger + +main :: IO () +main = do + let fp = "test.parquet" + void $ runStdoutLoggingT $ C.runResourceT $ runExceptT $ readWholeParquetFile + fp 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,138 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} + +module Parquet.Decoder where + +import Data.Binary.Get +import Data.Binary.Put +import Data.Bits +import qualified Data.ByteString as BS +import Data.Word (Word32, Word8) +import Data.Int (Int32) + +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 + +newtype BitWidth = BitWidth Word8 + deriving (Show, Eq, Ord) + +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/Monad.hs b/src/Parquet/Monad.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/Monad.hs @@ -0,0 +1,14 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} + +module Parquet.Monad where + +import qualified Conduit as C +import Control.Monad.Except +import Control.Monad.Logger +import qualified Data.Text as T + +type PR m + = (C.MonadResource m, MonadLogger m, C.MonadThrow m, MonadError T.Text m) + +type PRS m = (C.MonadResource m, C.MonadThrow 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,40 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE TemplateHaskell #-} + +module Parquet.ParquetObject where + +import Control.Lens +import Data.Binary (Binary(..)) +import Codec.Serialise (Serialise) +import GHC.Generics (Generic) +import qualified Data.ByteString as BS +import qualified Data.HashMap.Strict as HM +import qualified Data.Text as T +import Data.Int (Int64) + +newtype ParquetObject = MkParquetObject (HM.HashMap T.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 (HM.toList hm) + get = MkParquetObject . HM.fromList <$> get + +data ParquetValue = + ParquetObject !ParquetObject + | ParquetInt !Int64 + | ParquetString !BS.ByteString + | ParquetNull + deriving (Eq, Show, Generic, Binary, Serialise) + +makeLenses ''ParquetObject +makePrisms ''ParquetObject + +makeLenses ''ParquetValue +makePrisms ''ParquetValue 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,235 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} + +module Parquet.Reader where + +import Data.Foldable (traverse_) +import qualified Data.Conduit.List as CL +import qualified Data.Conduit.Binary as CB +import Control.Monad.Logger (MonadLogger, runNoLoggingT) +import Control.Monad.Logger.CallStack (logWarn) +import Data.Functor ((<$)) +import Parquet.Stream.Reader + (Value(..), readColumnChunk, ColumnValue(..), decodeConduit) +import Control.Lens +import qualified Data.Map as M +import Control.Arrow ((&&&)) +import qualified Data.Text as T +import qualified Data.HashMap.Strict as HM +import Control.Monad.Except +import qualified Data.List.NonEmpty as NE + +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BS8 +import qualified Conduit as C +import System.IO + (IOMode(ReadMode), openFile, SeekMode(AbsoluteSeek, SeekFromEnd), hSeek) +import Network.HTTP.Types.Status (statusIsSuccessful) +import Network.HTTP.Simple + (getResponseBody, getResponseStatus, httpSource, parseRequest, Header) +import Network.HTTP.Client (Request(requestHeaders)) +import qualified Parquet.ThriftTypes as TT +import Parquet.Utils (failOnExcept) +import qualified Data.Binary.Get as BG +import Parquet.ParquetObject + +newtype ParquetSource m = ParquetSource (Integer -> C.ConduitT () BS.ByteString m ()) + +type Url = String + +readMetadata + :: (MonadError T.Text m, MonadIO 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) => 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 T.Text m + , C.MonadResource m + , MonadLogger m + ) + => String + -> m [ParquetObject] +readWholeParquetFile inputFp = do + metadata <- readMetadata (localParquetFile inputFp) + C.runConduit + $ traverse_ + (sourceRowGroup (localParquetFile inputFp) metadata) + (metadata ^. TT.pinchField @"row_groups") + C..| CL.consume + +type Record = [(ColumnValue, [T.Text])] + +sourceParquet :: FilePath -> C.ConduitT () ParquetObject (C.ResourceT IO) () +sourceParquet fp = runExceptT (readMetadata (localParquetFile fp)) >>= \case + Left err -> fail $ "Could not read metadata: " <> show err + Right metadata -> C.transPipe runNoLoggingT $ traverse_ + (sourceRowGroup (localParquetFile fp) metadata) + (metadata ^. TT.pinchField @"row_groups") + +sourceRowGroupFromRemoteFile + :: (C.MonadResource m, C.MonadIO m, C.MonadThrow m, MonadLogger m) + => String + -> TT.FileMetadata + -> TT.RowGroup + -> C.ConduitT () ParquetObject m () +sourceRowGroupFromRemoteFile url metadata rg = + sourceRowGroup (remoteParquetFile url) metadata rg + +-- | 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) + => ParquetSource m + -> TT.FileMetadata + -> TT.RowGroup + -> C.ConduitT () ParquetObject m () +sourceRowGroup source metadata rg = + C.sequenceSources + (map + (\cc -> sourceColumnChunk source metadata cc + C..| CL.mapMaybe ((<$> mb_path cc) . (,)) + ) + (rg ^. TT.pinchField @"column_chunks") + ) + C..| CL.mapMaybeM parse_record + where + mb_path :: TT.ColumnChunk -> Maybe [T.Text] + mb_path cc = + TT.unField + . TT._ColumnMetaData_path_in_schema + <$> (cc ^. TT.pinchField @"meta_data") + + -- | Given a parquet record (a set of parquet columns, really), converts it to a single JSON-like object. + -- + -- It does it by traversing the list of columns, converting them to objects and combining them. + -- Parsing every column yields a ParquetObject and since ParquetObjects are monoids we combine them. + parse_record :: Record -> m (Maybe ParquetObject) + parse_record = fmap (fmap mconcat) $ traverse $ \(column, paths) -> + case NE.nonEmpty paths of + Nothing -> Nothing <$ logWarn + ( "parse_record: Record with value " + <> T.pack (show (_cvValue column)) + <> " does not have any paths. Record data is corrupted." + ) + Just ne_paths -> parse_column (column, ne_paths) + + -- | Given a parquet column, converts it to a JSON-like object. + -- + -- For a given column: + -- { value = "something", definition_level = 3, path = ["field1", "field2", "field3"] }. + -- + -- We create: + -- { + -- "field1": { + -- "field2": { + -- "field3": "something + -- } + -- } + -- } + parse_column :: (ColumnValue, NE.NonEmpty T.Text) -> m (Maybe ParquetObject) + parse_column (ColumnValue _ 1 _ v, path NE.:| []) = + pure $ Just $ MkParquetObject $ HM.fromList + [(path, value_to_parquet_value v)] + parse_column (ColumnValue _ _ _ _, _ NE.:| []) = do + -- This case means that we are in the last path element but definition level is not 1. + -- For example: + -- + -- { value = "something", definition_level = 5, path = ["field1", "field2", "field3"] }. + -- + -- And we reached to the following point: + -- + -- { value = "something", definition_level = 3, path = ["field3"] }. + -- + -- At this point we should construct the object {"field3": "something"} but this would require + -- a definition level of 1. + -- + -- Hence this record is corrupted. + logWarn + "parse_column: No more paths exist but we still have definition levels. Column data is corrupted." + pure Nothing + parse_column (ColumnValue r d md v, path NE.:| (p : px)) = do + mb_obj <- parse_column (ColumnValue r (d - 1) md v, p NE.:| px) + pure $ mb_obj <&> \obj -> + MkParquetObject $ HM.fromList [(path, ParquetObject obj)] + + value_to_parquet_value :: Value -> ParquetValue + value_to_parquet_value Null = ParquetNull + value_to_parquet_value (ValueInt64 v ) = ParquetInt v + value_to_parquet_value (ValueByteString bs) = ParquetString bs + +sourceColumnChunk + :: (C.MonadIO m, C.MonadResource m, C.MonadThrow m, MonadLogger m) + => ParquetSource m + -> TT.FileMetadata + -> TT.ColumnChunk + -> C.ConduitT () ColumnValue m () +sourceColumnChunk (ParquetSource source) metadata cc = do + let + schema_mapping = + M.fromList + $ map ((^. TT.pinchField @"name") &&& id) + $ metadata + ^. TT.pinchField @"schema" + let offset = cc ^. TT.pinchField @"file_offset" + source (fromIntegral offset) + C..| C.transPipe failOnExcept (readColumnChunk 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,384 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE ViewPatterns #-} + +module Parquet.Stream.Reader where + +import qualified Conduit as C +import Control.Applicative (liftA3) +import Control.Monad.Except +import Control.Monad.Logger +import Control.Monad.Reader +import Data.Bifunctor (first) +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 Data.Int (Int32, Int64) +import qualified Data.List.NonEmpty as NE +import qualified Data.Map as M +import qualified Data.Text as T +import Data.Traversable (for) +import Data.Word (Word32, Word8) +import Control.Lens +import qualified Pinch +import Safe.Exact (zipExactMay) + +import Parquet.Decoder (BitWidth(..), decodeBPBE, decodeRLEBPHybrid) +import Parquet.Monad +import qualified Parquet.ThriftTypes as TT +import Parquet.Utils (()) + +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 + read_page_content 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) 1 + xs -> xs + + read_page_content + :: TT.Encoding + -> [(Word32, Word32)] + -> Int32 + -> Word32 + -> C.ConduitT BS.ByteString ColumnValue m () + read_page_content 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 + 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 + 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) + + -- | Given repetition and definition level data, a dictionary and a set of indexes, + -- constructs values for this dictionary-encoded page. + 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 (NE.head (NE.reverse 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 + path <- asks _pcPath + if NE.length path > 1 + then decodeLevel encoding bit_width num_values + else pure (0, []) + +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 + :: (MonadReader PageCtx m, MonadError T.Text m) => m (Word8, Word8) +calcMaxEncodingLevels = do + schema <- asks _pcSchema + path <- asks _pcPath + filled_path <- for path $ \name -> + M.lookup name schema "Schema Element cannot be found: " <> name + foldM + (\(rep, def) e -> getRepType e >>= \case + (TT.REQUIRED _) -> pure (rep, def) + (TT.OPTIONAL _) -> pure (rep, def + 1) + (TT.REPEATED _) -> pure (rep + 1, def + 1) + ) + (0, 0) + filled_path + +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 + => M.Map T.Text TT.SchemaElement + -> TT.ColumnChunk + -> C.ConduitT BS.ByteString ColumnValue m () +readColumnChunk 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 = metadata ^. TT.pinchField @"path_in_schema" + ne_path <- NE.nonEmpty path "Schema path cannot be empty" + let page_ctx = PageCtx schema ne_path column_ty + C.runReaderC page_ctx $ readPage size Nothing + +readPage + :: (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 + 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/ThriftTypes.hs b/src/Parquet/ThriftTypes.hs new file mode 100644 --- /dev/null +++ b/src/Parquet/ThriftTypes.hs @@ -0,0 +1,426 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +module Parquet.ThriftTypes where + +import Data.ByteString +import qualified Data.Generics.Product.Fields as GL +import qualified Data.Generics.Product.Positions as GL +import Data.Int +import Data.Text +import GHC.Generics +import GHC.TypeLits (AppendSymbol, Symbol) +import Control.Lens +import Pinch +import Data.Binary (Binary) + +-- 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 a => Binary (Field k a) +instance Binary (Enumeration k) + +data StringType = StringType + deriving (Show, Eq, Generic, Binary) +instance Pinchable StringType where + type Tag StringType = TStruct + pinch _ = struct [] + unpinch _ = pure StringType + +data UUIDType = UUIDType + deriving (Show, Eq, Generic, Binary) +instance Pinchable UUIDType where + type Tag UUIDType = TStruct + pinch _ = struct [] + unpinch _ = pure UUIDType + +data MapType = MapType + deriving (Show, Eq, Generic, Binary) +instance Pinchable MapType where + type Tag MapType = TStruct + pinch _ = struct [] + unpinch _ = pure MapType + +data ListType = ListType + deriving (Show, Eq, Generic, Binary) +instance Pinchable ListType where + type Tag ListType = TStruct + pinch _ = struct [] + unpinch _ = pure ListType + +data EnumType = EnumType + deriving (Show, Eq, Generic, Binary) +instance Pinchable EnumType where + type Tag EnumType = TStruct + pinch _ = struct [] + unpinch _ = pure EnumType + +data DateType = DateType + deriving (Show, Eq, Generic, Binary) +instance Pinchable DateType where + type Tag DateType = TStruct + pinch _ = struct [] + unpinch _ = pure DateType + +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 + +data DecimalType = DecimalType + { _DecimalType_scale :: Field 1 Int32 + , _DecimalType_precision :: Field 2 Int32 + } deriving (Show, Eq, Generic, Pinchable, Binary) + +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 MilliSeconds = MilliSeconds + deriving (Show, Eq, Generic, Binary) +instance Pinchable MilliSeconds where + type Tag MilliSeconds = TStruct + pinch _ = struct [] + unpinch _ = pure MilliSeconds + +data MicroSeconds = MicroSeconds + deriving (Show, Eq, Generic, Binary) +instance Pinchable MicroSeconds where + type Tag MicroSeconds = TStruct + pinch _ = struct [] + unpinch _ = pure MicroSeconds + +data NanoSeconds = NanoSeconds + deriving (Show, Eq, Generic, Binary) +instance Pinchable NanoSeconds where + type Tag NanoSeconds = TStruct + pinch _ = struct [] + unpinch _ = pure NanoSeconds + +data TimeUnit + = TimeUnitMILLIS (Field 1 MilliSeconds) + | TimeUnitMICROS (Field 2 MicroSeconds) + | TimeUnitNANOS (Field 3 NanoSeconds) + deriving (Show, Eq, Generic, Binary) +instance Pinchable TimeUnit + +data IntType = IntType + { _IntType_bitWidth :: Field 1 Int8 + , _IntType_isSigned :: Field 2 Bool + } deriving (Show, Eq, Generic, Pinchable, Binary) + +data NullType = NullType + deriving (Show, Eq, Generic, Binary) +instance Pinchable NullType where + type Tag NullType = TStruct + pinch _ = struct [] + unpinch _ = pure NullType + +data JsonType = JsonType + deriving (Show, Eq, Generic, Binary) +instance Pinchable JsonType where + type Tag JsonType = TStruct + pinch _ = struct [] + unpinch _ = pure JsonType + +data BsonType = BsonType + deriving (Show, Eq, Generic, Binary) +instance Pinchable BsonType where + type Tag BsonType = TStruct + pinch _ = struct [] + unpinch _ = pure BsonType + +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 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 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) + +data FieldRepetitionType + = REQUIRED (Enumeration 0) + | OPTIONAL (Enumeration 1) + | REPEATED (Enumeration 2) + 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) + +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 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 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 PageEncodingStats = PageEncodingStats + { _PageEncodingStats_page_type :: Field 1 PageType + , _PageEncodingStats_encoding :: Field 2 Encoding + , _PageEncodingStats_count :: Field 3 Int32 + } 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 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) + +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 EncryptionAlgorithm + = EncryptionAlgorithm_AES_GCM_V1 (Field 1 AesGcmV1) + | EncryptionAlgorithm_AES_GCM_CTR_V1 (Field 2 AesGcmCtrV1) + 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 ColumnOrder + = ColumnOrder_TYPE_ORDER (Field 1 TypeDefinedOrder) + 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 EncryptionWithColumnKey = EncryptionWithColumnKey + { _EncryptionWithColumnKey_path_in_schema :: Field 1 [Text] + , _EncryptionWithColumnKey_key_metadata :: Field 2 (Maybe ByteString) + } deriving (Show, Eq, Generic, Pinchable, Binary) + +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 KeyValue = KeyValue + { _KeyValue_key :: Field 1 Text + , _KeyValue_value :: Field 2 (Maybe Text) + } 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 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 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 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 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) + +data IndexPageHeader = IndexPageHeader + deriving (Show, Eq, Generic, Binary) +instance Pinchable IndexPageHeader where + type Tag IndexPageHeader = TStruct + pinch _ = struct [] + unpinch _ = pure IndexPageHeader + +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 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 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) + +unField :: Field n a -> a +unField (Field a) = a 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,17 @@ +{-# LANGUAGE LambdaCase #-} + +module Parquet.Utils where + +import Control.Monad.Except +import qualified Data.Text as T + +() :: MonadError b m => Maybe a -> b -> m a +() Nothing err = throwError err +() (Just v) _ = pure v + +infixl 4 + +failOnExcept :: Monad m => ExceptT T.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,82 @@ +module Main + ( main + ) +where + +import Test.Hspec +import System.Process +import System.Environment (setEnv, unsetEnv) +import System.FilePath (()) +import Control.Exception (bracket_) +import qualified Data.Aeson as JSON +import qualified Data.Text.IO as TextIO (putStrLn) +import qualified Data.ByteString.Lazy as LByteString (ByteString, toStrict) +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import qualified Data.Text.Lazy as LText (Text) +import qualified Data.Text.Lazy.IO as LTextIO (putStrLn) +import Control.Monad.IO.Class (liftIO) +import Control.Monad.Except (runExceptT) +import Conduit (runResourceT) + +import Control.Monad.Logger +import Parquet.Reader (readWholeParquetFile) + +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 () +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 + performTest $ outParquetFilePath + callProcess "rm" ["-rf", testPath intermediateDir] + -- callProcess "rm" ["-f", outParquetFilePath] + +lazyByteStringToText :: LByteString.ByteString -> T.Text +lazyByteStringToText = T.decodeUtf8 . LByteString.toStrict + +lazyByteStringToString :: LByteString.ByteString -> String +lazyByteStringToString = T.unpack . lazyByteStringToText + +putLazyByteStringLn :: LByteString.ByteString -> IO () +putLazyByteStringLn = TextIO.putStrLn . lazyByteStringToText + +putLazyTextLn :: LText.Text -> IO () +putLazyTextLn = LTextIO.putStrLn + +main :: IO () +main = hspec $ describe "Reader" $ do + it "can read columns" $ do + testParquetFormat "input1.json" $ \parqFile -> do + liftIO . print =<< runResourceT + (runStdoutLoggingT (runExceptT (readWholeParquetFile parqFile))) + 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,69 @@ +module Parquet.Decoder.Spec + ( spec + ) +where + +import Data.Binary.Get (getWord8, lookAhead, runGetOrFail) +import Data.Binary.Put +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as LBS +import Parquet.Decoder +import Test.Hspec +import Test.Hspec.QuickCheck +import Test.QuickCheck + +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,8 @@ +module Main (main) where + +import qualified Parquet.Decoder.Spec +import Test.Hspec + +main :: IO () +main = + hspec Parquet.Decoder.Spec.spec