dataframe-parquet 1.1.0.1 → 1.6.0.0
raw patch · 31 files changed
Files
- benchmark/Writer10GB.hs +49/−0
- dataframe-parquet.cabal +77/−11
- src/DataFrame/IO/Parquet.hs +621/−170
- src/DataFrame/IO/Parquet/Binary.hs +2/−9
- src/DataFrame/IO/Parquet/Dictionary.hs +89/−88
- src/DataFrame/IO/Parquet/Encoding.hs +57/−20
- src/DataFrame/IO/Parquet/Levels.hs +10/−34
- src/DataFrame/IO/Parquet/Page.hs +300/−118
- src/DataFrame/IO/Parquet/Schema.hs +5/−18
- src/DataFrame/IO/Parquet/Seeking.hs +6/−36
- src/DataFrame/IO/Parquet/Thrift.hs +37/−45
- src/DataFrame/IO/Parquet/Time.hs +1/−2
- src/DataFrame/IO/Parquet/Utils.hs +14/−51
- src/DataFrame/IO/Parquet/Writer.hs +438/−0
- src/DataFrame/IO/Parquet/Writer/DefLevels.hs +62/−0
- src/DataFrame/IO/Parquet/Writer/Encoder.hs +316/−0
- src/DataFrame/IO/Parquet/Writer/Metadata.hs +176/−0
- src/DataFrame/IO/Parquet/Writer/Options.hs +34/−0
- src/DataFrame/IO/Utils/RandomAccess.hs +318/−5
- src/DataFrame/Typed/IO/Parquet.hs +65/−0
- stress/DataFrame10GB.hs +137/−0
- stress/StressMain.hs +38/−0
- tests/Main.hs +413/−0
- tests/data/alltypes_dictionary.parquet binary
- tests/data/alltypes_plain.parquet binary
- tests/data/alltypes_plain.snappy.parquet binary
- tests/data/alltypes_tiny_pages.parquet binary
- tests/data/int32_decimal.parquet binary
- tests/data/int64_decimal.parquet binary
- tests/data/mtcars.parquet binary
- tests/data/transactions.parquet binary
+ benchmark/Writer10GB.hs view
@@ -0,0 +1,49 @@+module Main (main) where++import Control.DeepSeq (NFData (rnf))+import Criterion.Main (bench, defaultMain, envWithCleanup, whnfIO)+import DataFrame.IO.Parquet.Writer (writeParquet)+import DataFrame.Internal.DataFrame (DataFrame, forceDataFrame)+import DataFrame10GB (stressDataFrame)+import System.Directory (removeDirectoryRecursive)+import System.FilePath ((</>))+import System.IO.Temp (createTempDirectory, getCanonicalTemporaryDirectory)++data BenchmarkEnvironment = BenchmarkEnvironment+ { benchmarkDataFrame :: DataFrame+ , benchmarkDirectory :: FilePath+ , benchmarkOutput :: FilePath+ }++instance NFData BenchmarkEnvironment where+ rnf environment =+ forceDataFrame (benchmarkDataFrame environment) `seq`+ rnf (benchmarkDirectory environment) `seq`+ rnf (benchmarkOutput environment)++prepareEnvironment :: IO BenchmarkEnvironment+prepareEnvironment = do+ temporary <- getCanonicalTemporaryDirectory+ directory <- createTempDirectory temporary "dataframe-parquet-writer-10gb"+ pure+ BenchmarkEnvironment+ { benchmarkDataFrame = stressDataFrame+ , benchmarkDirectory = directory+ , benchmarkOutput = directory </> "benchmark.parquet"+ }++cleanupEnvironment :: BenchmarkEnvironment -> IO ()+cleanupEnvironment = removeDirectoryRecursive . benchmarkDirectory++main :: IO ()+main =+ defaultMain+ [ envWithCleanup prepareEnvironment cleanupEnvironment $ \environment ->+ -- Memory usage for this benchmark will be north of 20 GB.+ bench "write 10 GiB dataframe" $+ whnfIO+ ( writeParquet+ (benchmarkOutput environment)+ (benchmarkDataFrame environment)+ )+ ]
dataframe-parquet.cabal view
@@ -1,7 +1,6 @@-cabal-version: 2.4+cabal-version: 3.4 name: dataframe-parquet-version: 1.1.0.1-+version: 1.6.0.0 synopsis: Parquet reader and writer for the dataframe ecosystem. description: @DataFrame.IO.Parquet@ — pure-Haskell Parquet 2.0 reader and writer@@ -20,6 +19,7 @@ copyright: (c) 2024-2026 Michael Chavinda category: Data tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2+extra-source-files: tests/data/*.parquet common warnings ghc-options:@@ -29,6 +29,11 @@ -Wunused-local-binds -Wunused-packages +flag stress-tests+ description: Build and run the opt-in 10 GiB Parquet roundtrip stress test.+ default: False+ manual: True+ library import: warnings ghc-options: -O2@@ -45,22 +50,83 @@ DataFrame.IO.Parquet.Thrift DataFrame.IO.Parquet.Time DataFrame.IO.Parquet.Utils+ DataFrame.IO.Parquet.Writer+ DataFrame.IO.Parquet.Writer.DefLevels+ DataFrame.IO.Parquet.Writer.Encoder+ DataFrame.IO.Parquet.Writer.Metadata+ DataFrame.IO.Parquet.Writer.Options DataFrame.IO.Utils.RandomAccess+ DataFrame.Typed.IO.Parquet build-depends: base >= 4 && < 5,- bytestring >= 0.11 && < 0.13,- containers >= 0.6.7 && < 0.9,- dataframe-core ^>= 1.1,- dataframe-operations ^>= 1.1.1,- dataframe-parsing ^>= 1.0.2,+ bytestring >= 0.11 && < 0.14,+ containers >= 0.6.7 && < 0.10,+ dataframe-core >= 2.6 && < 2.7,+ dataframe-operations >= 2.6 && < 2.7,+ dataframe-parsing >= 2.2 && < 2.3,+ primitive >= 0.7 && < 0.11, directory >= 1.3.0.0 && < 2, filepath >= 1.4 && < 2, Glob >= 0.10 && < 1, pinch >= 0.5 && < 1,- snappy-hs ^>= 0.1,+ snappy-hs >= 0.1 && < 0.3, text >= 2.1 && < 3, time >= 1.12 && < 2,- vector ^>= 0.13,+ vector >= 0.13 && < 0.15, zlib >= 0.5 && < 1,- zstd >= 0.1.2.0 && < 0.2+ zstd >= 0.1.2.0 && < 0.3 hs-source-dirs: src default-language: Haskell2010+++test-suite dataframe-parquet-tests+ import: warnings+ type: exitcode-stdio-1.0+ main-is: Main.hs+ hs-source-dirs: tests+ build-depends: base >= 4 && < 5,+ bytestring >= 0.11 && < 0.14,+ dataframe-core >= 2.6 && < 2.7,+ dataframe-operations >= 2.6 && < 2.7,+ dataframe-parquet,+ directory >= 1.3.0.0 && < 2,+ filepath >= 1.4 && < 2,+ temporary >= 1.3 && < 1.5,+ text >= 2.1 && < 3,+ HUnit >= 1.6 && < 1.8+ default-language: Haskell2010++executable dataframe-parquet-10gb-stress+ import: warnings+ main-is: StressMain.hs+ other-modules: DataFrame10GB+ hs-source-dirs: stress+ build-depends: base >= 4 && < 5,+ dataframe-core >= 2.6 && < 2.7,+ dataframe-parquet,+ filepath >= 1.4 && < 2,+ temporary >= 1.3 && < 1.5,+ text >= 2.1 && < 3,+ time >= 1.12 && < 2,+ vector >= 0.13 && < 0.15+ default-language: Haskell2010+ -- ghc-options: -O2 -threaded -rtsopts -with-rtsopts=-N++benchmark dataframe-parquet-writer-10gb+ import: warnings+ type: exitcode-stdio-1.0+ main-is: Writer10GB.hs+ other-modules: DataFrame10GB+ hs-source-dirs: benchmark, stress+ build-depends: base >= 4 && < 5,+ criterion >= 1 && < 2,+ deepseq >= 1.4 && < 2,+ dataframe-core >= 2.6 && < 2.7,+ dataframe-parquet,+ directory >= 1.3 && < 2,+ filepath >= 1.4 && < 2,+ temporary >= 1.3 && < 1.5,+ text >= 2.1 && < 3,+ time >= 1.12 && < 2,+ vector >= 0.13 && < 0.15+ default-language: Haskell2010+ ghc-options: -O2 -threaded -rtsopts -with-rtsopts=-N
src/DataFrame/IO/Parquet.hs view
@@ -1,3 +1,4 @@+{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MonoLocalBinds #-}@@ -7,16 +8,76 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module DataFrame.IO.Parquet where+module DataFrame.IO.Parquet (+ -- * Reading+ readParquet,+ readParquetWithOpts,+ readParquetFiles,+ readParquetFilesWithOpts, + -- * Writing+ writeParquet,+ writeParquetWithOptions,++ -- * Options+ ParquetReadOptions (..),+ defaultParquetReadOptions,+ ParquetWriteOptions (..),+ WriterStrategy (..),+ defaultParquetWriteOptions,++ -- * File metadata+ parseFileMetadata,+ readMetadataFromPath,+ readMetadataFromHandle,+ columnChunksForAll,++ -- * Schema description+ ColumnDescription (..),+ generateColumnDescriptions,+ getColumnNames,++ -- * Decoding+ parseParquetWithOpts,+ parseColumnChunks,+ getNonNullableColumn,+ getNullableColumn,+ getRepeatedColumn,+ applyDescLogicalType,+ applyLogicalType,+ nativeTypeHints,+ restoreNativeType,++ -- * Applying read options to a decoded frame+ applyReadOptions,+ applyPredicate,+ applySelectedColumns,+ applyRowRange,+ applySafeRead,++ -- * Data sources+ RandomAccess (..),+ ReaderIO (runReaderIO),+ FileBufferedOrSeekable,+ ForceNonSeekable,+ withFileBufferedOrSeekable,+ _readParquetWithOpts,+ readParquetWithMetadata,+ FileMetadata,+) where++import Control.Concurrent (getNumCapabilities) import Control.Exception (throw) import Control.Monad import Control.Monad.IO.Class (MonadIO (..))+import Control.Monad.ST (stToIO) import Data.Bits (Bits (shiftL), (.|.)) import qualified Data.ByteString as BS import Data.Either (fromRight) import Data.Functor ((<&>))+import Data.IORef (atomicModifyIORef', newIORef) import Data.Int (Int32, Int64)+import Data.Maybe (fromJust, fromMaybe) #if !MIN_VERSION_base(4,20,0) import Data.List (foldl', transpose) #else@@ -25,35 +86,54 @@ import qualified Data.List as L import qualified Data.Map as Map import qualified Data.Text as T-import Data.Time (UTCTime)-import Data.Time.Clock.POSIX (posixSecondsToUTCTime)+import Data.Time.Calendar (Day (ModifiedJulianDay))+import Data.Time.Clock (UTCTime (UTCTime), picosecondsToDiffTime)+import Data.Type.Equality (testEquality, (:~:) (Refl)) import qualified Data.Vector as Vector+import qualified Data.Vector.Mutable as MVector import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM import DataFrame.Errors (DataFrameException (ColumnsNotFoundException))+import DataFrame.IO.Parquet.Dictionary (+ DictVals,+ dictDouble,+ dictFloat,+ dictInt32,+ dictInt64,+ word32At,+ word64At,+ ) import DataFrame.IO.Parquet.Page ( PageDecoder, UnboxedPageDecoder,+ appendNullableStringPageIO,+ appendStringPageIO, boolDecoder, byteArrayDecoder,+ decodeRequiredNumericInto, doubleDecoder, fixedLenByteArrayDecoder, floatDecoder,+ foldColumnDataPagesM,+ foldColumnPagesM, int32Decoder, int64Decoder, int96Decoder,- foldColumnPagesM, ) import DataFrame.IO.Parquet.Seeking (- FileBufferedOrSeekable,+ FileBufferedOrSeekable (..), ForceNonSeekable, withFileBufferedOrSeekable, ) import DataFrame.IO.Parquet.Thrift ( ColumnChunk (..),+ ColumnMetaData (..), DecimalType (..), FileMetadata (..),+ KeyValue (..), LogicalType (..), RowGroup (..),+ Statistics (..), ThriftType (..), TimeUnit (..), TimestampType (..),@@ -70,62 +150,69 @@ generateColumnDescriptions, getColumnNames, )+import DataFrame.IO.Parquet.Writer (+ ParquetWriteOptions (..),+ WriterStrategy (..),+ defaultParquetWriteOptions,+ nativeTypeKeyPrefix,+ writeParquet,+ writeParquetWithOptions,+ ) import DataFrame.IO.Utils.RandomAccess (+ ChunkBuffer (..), RandomAccess (..),+ Range (Range), ReaderIO (runReaderIO), ) import DataFrame.Internal.Column (Column, Columnable) import qualified DataFrame.Internal.Column as DI+import DataFrame.Internal.Column.Builder (+ concatColumnsIO,+ freezeTextChunk,+ mergeTextChunks,+ newTextBuilder,+ )+import DataFrame.Internal.Control.Concurrent (forkJoin_) import DataFrame.Internal.DataFrame (DataFrame (..))-import DataFrame.Internal.Expression (Expr, getColumns)+import DataFrame.Internal.Expression (+ BinaryOp (binaryName),+ Expr (..),+ getColumns,+ ) import DataFrame.Operations.Merge () import qualified DataFrame.Operations.Subset as DS+import DataFrame.Schema (+ Schema (..),+ SchemaType,+ makeSchema,+ schemaType,+ )+import GHC.Float (castWord32ToFloat, castWord64ToDouble) import qualified Pinch import System.Directory (doesDirectoryExist) import System.FilePath ((</>)) import System.FilePath.Glob (glob) import System.IO (IOMode (ReadMode))---- Options -------------------------------------------------------------------{- | Options for reading Parquet data.--These options are applied in this order:--1. predicate filtering-2. column projection-3. row range-4. safe column promotion+import Type.Reflection (typeRep) -Column selection for @selectedColumns@ uses leaf column names only.--}+-- | Options for reading Parquet data. data ParquetReadOptions = ParquetReadOptions { selectedColumns :: Maybe [T.Text]- {- ^ Columns to keep in the final dataframe. If set, only these columns are returned.- Predicate-referenced columns are read automatically when needed and projected out after filtering.- -}+ -- ^ Columns to keep in the final dataframe. , predicate :: Maybe (Expr Bool) -- ^ Optional row filter expression applied before projection. , rowRange :: Maybe (Int, Int)- -- ^ Optional row slice @(start, end)@ with start-inclusive/end-exclusive semantics.+ -- ^ Optional row slice @(start, end)@, start-inclusive, end-exclusive. , safeColumns :: Bool- -- ^ When True, every column is promoted to OptionalColumn after read, regardless of nullability in the schema.+ -- ^ Promote every column to OptionalColumn after read.+ , widenIntegers :: Bool+ -- ^ Decode INT32 and INT64 columns as 'Int'.+ , parallel :: Bool+ -- ^ Decode row groups concurrently using every RTS capability. } deriving (Show) -{- | Default Parquet read options.--Equivalent to:--@-ParquetReadOptions- { selectedColumns = Nothing- , predicate = Nothing- , rowRange = Nothing- , safeColumns = False- }-@--}+-- | Default Parquet read options. defaultParquetReadOptions :: ParquetReadOptions defaultParquetReadOptions = ParquetReadOptions@@ -133,63 +220,287 @@ , predicate = Nothing , rowRange = Nothing , safeColumns = False+ , widenIntegers = False+ , parallel = True } --- Public API ----------------------------------------------------------------{- | Read a parquet file from path and load it into a dataframe.--==== __Example__-@-ghci> D.readParquet ".\/data\/mtcars.parquet"-@--}+-- | Read a parquet file from path and load it into a dataframe. readParquet :: FilePath -> IO DataFrame readParquet = readParquetWithOpts defaultParquetReadOptions -{- | Read a Parquet file using explicit read options.--==== __Example__-@-ghci> D.readParquetWithOpts-ghci| (D.defaultParquetReadOptions{D.selectedColumns = Just ["id"], D.rowRange = Just (0, 10)})-ghci| "./tests/data/alltypes_plain.parquet"-@--When @selectedColumns@ is set and @predicate@ references other columns, those predicate columns-are auto-included for decoding, then projected back to the requested output columns.--}+-- | Read a Parquet file using explicit read options. readParquetWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame readParquetWithOpts = _readParquetWithOpts Nothing -- | Internal entry point used by tests to force non-seekable mode. _readParquetWithOpts :: ForceNonSeekable -> ParquetReadOptions -> FilePath -> IO DataFrame-_readParquetWithOpts extraConfig opts path =- withFileBufferedOrSeekable extraConfig path ReadMode $ \file ->- runReaderIO (parseParquetWithOpts opts) file+_readParquetWithOpts extraConfig opts path = do+ metadata <-+ withFileBufferedOrSeekable+ extraConfig+ path+ ReadMode+ (runReaderIO parseFileMetadata)+ readParquetWithMetadata extraConfig metadata opts path -{- | Read Parquet files from a directory or glob path.+-- | Read a Parquet file whose footer has already been parsed.+readParquetWithMetadata ::+ ForceNonSeekable ->+ FileMetadata ->+ ParquetReadOptions ->+ FilePath ->+ IO DataFrame+readParquetWithMetadata extraConfig metadata opts path = do+ caps <- getNumCapabilities+ let nRowGroups = length (unField metadata.row_groups)+ if parallel opts && extraConfig /= Just True && caps > 1 && nRowGroups > 1+ then decodeRowGroupsParallel caps opts path metadata+ else withFileBufferedOrSeekable extraConfig path ReadMode $ \file ->+ runReaderIO (parseParquetWithOpts opts) file -This is equivalent to calling 'readParquetFilesWithOpts' with 'defaultParquetReadOptions'.--}-readParquetFiles :: FilePath -> IO DataFrame-readParquetFiles = readParquetFilesWithOpts defaultParquetReadOptions+-- | Decode row groups in parallel, one task per (row group, column).+decodeRowGroupsParallel ::+ Int -> ParquetReadOptions -> FilePath -> FileMetadata -> IO DataFrame+decodeRowGroupsParallel caps opts0 path metadata = do+ let plan = planColumns opts0 metadata+ descs = Vector.fromList (planDescs plan)+ allRowGroupRows =+ map (fromIntegral . unField . rg_num_rows) (unField metadata.row_groups)+ allOffsets = scanl (+) 0 allRowGroupRows+ -- Row groups whose statistics rule out the predicate are skipped.+ statsGroups =+ [ g+ | (g, rg) <- zip [0 ..] (unField metadata.row_groups)+ , maybe True (rowGroupMayMatch metadata rg) (predicate opts0)+ ]+ -- A row range only touches overlapping groups, then is rebased.+ (keptGroups, opts) = case rowRange opts0 of+ Nothing -> (statsGroups, opts0)+ Just (s, e) ->+ let groups =+ [ g+ | (g, off, n) <- zip3 [0 ..] allOffsets allRowGroupRows+ , off < e+ , off + n > s+ ]+ base = case groups of+ [] -> 0+ (g : _) -> allOffsets !! g+ in (groups, opts0{rowRange = Just (s - base, e - base)})+ chunks =+ Vector.fromList+ [Vector.fromList [cs !! g | g <- keptGroups] | cs <- planChunks plan]+ rowGroupRows = Vector.fromList [allRowGroupRows !! g | g <- keptGroups]+ rowOffsets = Vector.prescanl' (+) 0 rowGroupRows+ totalRows = Vector.sum rowGroupRows+ nCols = Vector.length descs+ nRg = Vector.length rowGroupRows+ nTasks = nCols * nRg+ taskChunk t = let (r, c) = t `divMod` nCols in chunks Vector.! c Vector.! r+ sinks <- Vector.mapM (newSink (widenIntegers opts) totalRows) descs+ results <- MVector.new nTasks+ counter <- newIORef (0 :: Int)+ -- Each worker owns a handle, so chunk reads run in parallel too.+ let worker = withFileBufferedOrSeekable Nothing path ReadMode $ \file ->+ let loop = do+ t <- atomicModifyIORef' counter (\n -> (n + 1, n))+ when (t < nTasks) $ do+ let (r, c) = t `divMod` nCols+ desc = descs Vector.! c+ chunk = taskChunk t+ (base, len) = chunkRange chunk+ bytes <- runReaderIO (readBytes (Range base len)) file+ let buffer = ChunkBuffer base bytes+ into ::+ forall a.+ (VU.Unbox a) =>+ (DictVals -> Maybe (VU.Vector a)) ->+ (BS.ByteString -> Int -> a) ->+ Int ->+ VUM.IOVector a ->+ IO ()+ into pick plainAt width out =+ void $+ runReaderIO+ ( decodeRequiredNumericInto+ pick+ plainAt+ width+ desc+ chunk+ out+ (rowOffsets Vector.! r)+ )+ buffer+ case sinks Vector.! c of+ SinkInt out width ->+ if width == 4+ then into (fmap (VU.map fromIntegral) . dictInt32) plainIntFrom32 4 out+ else into (fmap (VU.map fromIntegral) . dictInt64) plainIntFrom64 8 out+ SinkInt32 out -> into dictInt32 plainInt32 4 out+ SinkInt64 out -> into dictInt64 plainInt64 8 out+ SinkFloat out -> into dictFloat plainFloat 4 out+ SinkDouble out -> into dictDouble plainDouble 8 out+ SinkRows -> do+ col <-+ runReaderIO+ (parseColumnChunks (rowGroupRows Vector.! r) [chunk] desc)+ buffer+ MVector.write results (c * nRg + r) col+ loop+ in loop+ forkJoin_ (replicate (max 1 (min caps nTasks)) worker)+ cols <- forM [0 .. nCols - 1] $ \c -> case sinks Vector.! c of+ SinkInt out _ -> DI.UnboxedColumn Nothing <$> VU.unsafeFreeze out+ SinkInt32 out -> DI.UnboxedColumn Nothing <$> VU.unsafeFreeze out+ SinkInt64 out -> DI.UnboxedColumn Nothing <$> VU.unsafeFreeze out+ SinkFloat out -> DI.UnboxedColumn Nothing <$> VU.unsafeFreeze out+ SinkDouble out -> DI.UnboxedColumn Nothing <$> VU.unsafeFreeze out+ SinkRows ->+ concatColumnsIO+ =<< forM [0 .. nRg - 1] (\r -> MVector.read results (c * nRg + r))+ pure (assembleDataFrame opts totalRows metadata plan cols) -{- | Read multiple Parquet files (directory or glob) using explicit options.+-- | Where a column's decoded values go.+data ColumnSink+ = -- | Widened integers; the width is the source's bytes per value.+ SinkInt !(VUM.IOVector Int) !Int+ | SinkInt32 !(VUM.IOVector Int32)+ | SinkInt64 !(VUM.IOVector Int64)+ | SinkFloat !(VUM.IOVector Float)+ | SinkDouble !(VUM.IOVector Double)+ | SinkRows -If @path@ is a directory, all non-directory entries are read.-If @path@ is a glob, matching files are read.+newSink :: Bool -> Int -> ColumnDescription -> IO ColumnSink+newSink widen totalRows description+ | description.maxRepetitionLevel /= 0 || description.maxDefinitionLevel /= 0 =+ pure SinkRows+ | otherwise = case description.colElementType of+ Just (INT32 _)+ | widen -> (`SinkInt` 4) <$> VUM.unsafeNew totalRows+ | otherwise -> SinkInt32 <$> VUM.unsafeNew totalRows+ Just (INT64 _)+ | widen -> (`SinkInt` 8) <$> VUM.unsafeNew totalRows+ | otherwise -> SinkInt64 <$> VUM.unsafeNew totalRows+ Just (FLOAT _) -> SinkFloat <$> VUM.unsafeNew totalRows+ Just (DOUBLE _) -> SinkDouble <$> VUM.unsafeNew totalRows+ _ -> pure SinkRows -For multi-file reads, @rowRange@ is applied once after concatenation (global range semantics).+plainIntFrom32 :: BS.ByteString -> Int -> Int+plainIntFrom32 bs o = fromIntegral (fromIntegral (word32At bs o) :: Int32)+{-# INLINE plainIntFrom32 #-} -==== __Example__-@-ghci> D.readParquetFilesWithOpts-ghci| (D.defaultParquetReadOptions{D.selectedColumns = Just ["id"], D.rowRange = Just (0, 5)})-ghci| "./tests/data/alltypes_plain*.parquet"-@--}+plainIntFrom64 :: BS.ByteString -> Int -> Int+plainIntFrom64 bs o = fromIntegral (word64At bs o)+{-# INLINE plainIntFrom64 #-}++plainInt32 :: BS.ByteString -> Int -> Int32+plainInt32 bs o = fromIntegral (word32At bs o)+{-# INLINE plainInt32 #-}++plainInt64 :: BS.ByteString -> Int -> Int64+plainInt64 bs o = fromIntegral (word64At bs o)+{-# INLINE plainInt64 #-}++plainFloat :: BS.ByteString -> Int -> Float+plainFloat bs o = castWord32ToFloat (word32At bs o)+{-# INLINE plainFloat #-}++plainDouble :: BS.ByteString -> Int -> Double+plainDouble bs o = castWord64ToDouble (word64At bs o)+{-# INLINE plainDouble #-}++-- | Whether a row group can hold a row satisfying the predicate.+rowGroupMayMatch :: FileMetadata -> RowGroup -> Expr Bool -> Bool+rowGroupMayMatch metadata rg = go+ where+ allNames = getColumnNames (drop 1 (unField metadata.schema))+ chunks = unField (rg_columns rg)+ go :: Expr Bool -> Bool+ go (Binary op l r)+ | binaryName op == "and" = case (l, r) of+ (l' :: Expr lt, r' :: Expr rt)+ | Just Refl <- testEquality (typeRep @lt) (typeRep @Bool)+ , Just Refl <- testEquality (typeRep @rt) (typeRep @Bool) ->+ go l' && go r'+ _ -> True+ | otherwise = case (l, r) of+ (Col name :: Expr lt, Lit (v :: rt))+ | Just Refl <- testEquality (typeRep @lt) (typeRep @Int)+ , Just Refl <- testEquality (typeRep @rt) (typeRep @Int) ->+ decide name (binaryName op) v+ (Lit (v :: lt), Col name :: Expr rt)+ | Just Refl <- testEquality (typeRep @lt) (typeRep @Int)+ , Just Refl <- testEquality (typeRep @rt) (typeRep @Int) ->+ decide name (flipped (binaryName op)) v+ _ -> True+ go _ = True+ flipped :: T.Text -> T.Text+ flipped n = case n of+ "lt" -> "gt"+ "gt" -> "lt"+ "leq" -> "geq"+ "geq" -> "leq"+ other -> other+ -- @column <cmp> v@ can hold for some row in [lo, hi].+ decide :: T.Text -> T.Text -> Int -> Bool+ decide name cmp v = case columnRange name of+ Nothing -> True+ Just (lo, hi) -> case cmp of+ "eq" -> lo <= v && v <= hi+ "neq" -> not (lo == v && hi == v)+ "lt" -> lo < v+ "leq" -> lo <= v+ "gt" -> hi > v+ "geq" -> hi >= v+ _ -> True+ columnRange :: T.Text -> Maybe (Int, Int)+ columnRange name = do+ i <- L.findIndex (columnMatches allNames name) allNames+ cc <- chunks `atMay` i+ meta <- unField (cc_meta_data cc)+ stats <- unField (cmd_statistics meta)+ let decodeInt bs = case unField (cmd_type meta) of+ INT32 _+ | BS.length bs == 4 ->+ Just (fromIntegral (fromIntegral (word32At bs 0) :: Int32))+ INT64 _ | BS.length bs == 8 -> Just (fromIntegral (word64At bs 0))+ _ -> Nothing+ lo <-+ decodeInt+ =<< (unField (stats_min_value stats) `orElse` unField (stats_min stats))+ hi <-+ decodeInt+ =<< (unField (stats_max_value stats) `orElse` unField (stats_max stats))+ pure (lo, hi)+ orElse (Just x) _ = Just x+ orElse Nothing y = y+ atMay xs i = if i >= 0 && i < length xs then Just (xs !! i) else Nothing++-- | File offset and byte length of a column chunk (dictionary page included).+chunkRange :: ColumnChunk -> (Integer, Int)+chunkRange columnChunk =+ let meta = fromJust . unField $ columnChunk.cc_meta_data+ dataOffset = fromIntegral . unField $ meta.cmd_data_page_offset+ dictOffset = fromIntegral <$> unField meta.cmd_dictionary_page_offset+ in ( fromMaybe dataOffset dictOffset+ , fromIntegral . unField $ meta.cmd_total_compressed_size+ )++-- | The columns a read decodes, after projection, in schema order.+data ColumnPlan = ColumnPlan+ { planNames :: [T.Text]+ , planDescs :: [ColumnDescription]+ , planChunks :: [[ColumnChunk]]+ -- ^ Per column, one chunk per row group.+ }++-- | Read Parquet files from a directory or glob path.+readParquetFiles :: FilePath -> IO DataFrame+readParquetFiles = readParquetFilesWithOpts defaultParquetReadOptions++-- | Read multiple Parquet files (directory or glob) using options. readParquetFilesWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame readParquetFilesWithOpts opts path = do isDir <- doesDirectoryExist path@@ -198,7 +509,7 @@ matches <- glob pat - files <- filterM (fmap not . doesDirectoryExist) matches+ files <- L.sort <$> filterM (fmap not . doesDirectoryExist) matches case files of [] ->@@ -209,99 +520,131 @@ dfs <- mapM (readParquetWithOpts optsWithoutRowRange) files pure (applyRowRange opts (mconcat dfs)) --- Core parsing pipeline -----------------------------------------------------{- | Parse a Parquet file via the 'RandomAccess' handle, applying all-read options. This is the central parsing entry point used by-'_readParquetWithOpts'.--}+-- | Parse a Parquet file via the 'RandomAccess' handle. parseParquetWithOpts :: (RandomAccess m, MonadIO m) => ParquetReadOptions -> m DataFrame+{-# SPECIALIZE parseParquetWithOpts ::+ ParquetReadOptions -> ReaderIO FileBufferedOrSeekable DataFrame+ #-}+{-# INLINEABLE parseParquetWithOpts #-} parseParquetWithOpts opts = do metadata <- parseFileMetadata+ let plan = planColumns opts metadata+ vectorLength = fileRowCount metadata+ rawCols <-+ zipWithM (parseColumnChunks vectorLength) (planChunks plan) (planDescs plan)+ return (assembleDataFrame opts vectorLength metadata plan rawCols) +-- | Row count from the footer, falling back to the sum of row-group counts.+fileRowCount :: FileMetadata -> Int+fileRowCount metadata =+ let topLevelRows = fromIntegral . unField $ metadata.num_rows :: Int+ rgRows =+ sum $ map (fromIntegral . unField . rg_num_rows) (unField metadata.row_groups) ::+ Int+ in if topLevelRows > 0 then topLevelRows else rgRows++columnMatches :: [T.Text] -> T.Text -> T.Text -> Bool+columnMatches allNames requested name =+ name == requested+ || (leafOf name == requested && requested `notElem` allNames)++-- | The last dot-separated segment of a column path.+leafOf :: T.Text -> T.Text+leafOf = last . T.splitOn "."++-- | Validate the requested columns against the schema and pick what to decode.+planColumns :: ParquetReadOptions -> FileMetadata -> ColumnPlan+planColumns opts metadata = let schemaElems = unField metadata.schema allNames = getColumnNames (drop 1 schemaElems)- leafNames = L.nub (map (last . T.splitOn ".") allNames)+ knownNames = L.nub (allNames ++ map leafOf allNames) predicateColumns = maybe [] (L.nub . getColumns) (predicate opts) selectedColumnsForRead = case selectedColumns opts of Nothing -> Nothing Just selected -> Just (L.nub (selected ++ predicateColumns)) - -- TODO: When rowRange is set, compute cumulative row offsets from- -- rg_num_rows in each RowGroup and skip any group whose row interval does- -- not overlap the requested range, avoiding all decoding for those groups.+ -- TODO: With rowRange set, skip row groups the range does not overlap. - -- TODO: When predicate is set, inspect cmd_statistics min/max values for- -- predicate-referenced columns in each RowGroup and skip groups where- -- statistics prove the predicate cannot be satisfied.+ -- TODO: Skip row groups whose statistics rule out the predicate. - -- Validate selected columns- case selectedColumnsForRead of- Nothing -> pure ()- Just requested ->- let missing = requested L.\\ leafNames- in unless (L.null missing) $- liftIO $- throw- ( ColumnsNotFoundException- missing- "readParquetWithOpts"- leafNames- )+ -- Validate selected columns+ validated = case selectedColumnsForRead of+ Nothing -> ()+ Just requested ->+ let missing = [r | r <- requested, not (any (columnMatches allNames r) allNames)]+ in if L.null missing+ then ()+ else+ throw+ ( ColumnsNotFoundException+ missing+ "readParquetWithOpts"+ knownNames+ ) - let descriptions = generateColumnDescriptions schemaElems+ descriptions = generateColumnDescriptions schemaElems chunks = columnChunksForAll metadata nCols = length chunks nDescs = length descriptions - unless (nCols == nDescs) $- error $- "Column count mismatch: got "- <> show nCols- <> " columns but schema implied "- <> show nDescs- <> " columns"-- -- Some files omit the top-level num_rows field; fall back to summing row-group counts.- let topLevelRows = fromIntegral . unField $ metadata.num_rows :: Int- rgRows =- sum $ map (fromIntegral . unField . rg_num_rows) (unField metadata.row_groups) ::- Int- vectorLength = if topLevelRows > 0 then topLevelRows else rgRows+ counted =+ if nCols == nDescs+ then ()+ else+ error $+ "Column count mismatch: got "+ <> show nCols+ <> " columns but schema implied "+ <> show nDescs+ <> " columns" - -- Column-projection pushdown: decode only the columns needed for the- -- requested output plus any predicate, instead of decoding every column- -- and dropping the unwanted ones afterward. A 'Nothing' selection keeps- -- all columns, so full reads are unchanged.- let keep name = case selectedColumnsForRead of+ -- Decode only the columns needed for the output plus any predicate.+ keep name = case selectedColumnsForRead of Nothing -> True- Just req -> last (T.splitOn "." name) `elem` req+ Just req -> any (\r -> columnMatches allNames r name) req kept = filter (\(n, _, _) -> keep n) (zip3 allNames chunks descriptions)- keptNames = [n | (n, _, _) <- kept]- keptChunks = [c | (_, c, _) <- kept]- keptDescs = [d | (_, _, d) <- kept]-- rawCols <- zipWithM (parseColumnChunks vectorLength) keptChunks keptDescs+ in validated `seq`+ counted `seq`+ ColumnPlan+ { planNames = [n | (n, _, _) <- kept]+ , planDescs = [d | (_, _, d) <- kept]+ , planChunks = [c | (_, c, _) <- kept]+ } - let finalCols = zipWith applyDescLogicalType keptDescs rawCols+-- | Apply logical types and read options to the decoded columns.+assembleDataFrame ::+ ParquetReadOptions -> Int -> FileMetadata -> ColumnPlan -> [Column] -> DataFrame+assembleDataFrame opts totalRows metadata plan rawCols =+ let hints = nativeTypeHints metadata+ keptNames = planNames plan+ widen col+ | not (widenIntegers opts) = col+ | otherwise = case col of+ DI.UnboxedColumn Nothing (v :: VU.Vector t)+ | Just Refl <- testEquality (typeRep @t) (typeRep @Int32) ->+ DI.UnboxedColumn Nothing (VU.map (fromIntegral @Int32 @Int) v)+ | Just Refl <- testEquality (typeRep @t) (typeRep @Int64) ->+ DI.UnboxedColumn Nothing (VU.map (fromIntegral @Int64 @Int) v)+ _ -> col+ finalCols =+ zipWith+ (restoreNativeType hints)+ keptNames+ (map widen (zipWith applyDescLogicalType (planDescs plan) rawCols)) indices = Map.fromList $ zip keptNames [0 ..]- dimensions = (vectorLength, length finalCols)-- let df =+ dimensions = (totalRows, length finalCols)+ df = DataFrame (Vector.fromListN (length finalCols) finalCols) indices dimensions Map.empty-- return (applyReadOptions opts df)+ in applyReadOptions opts df -{- | Parse the file-level Thrift metadata from the Parquet file footer.-Validates the trailing 4-byte magic marker (\"PAR1\") before decoding.--}+-- | Parse the file-level Thrift metadata from the Parquet footer. parseFileMetadata :: (RandomAccess m) => m FileMetadata parseFileMetadata = do footerBytes <- readSuffix 8@@ -344,6 +687,19 @@ [ColumnChunk] -> ColumnDescription -> m Column+{-# SPECIALIZE parseColumnChunks ::+ Int ->+ [ColumnChunk] ->+ ColumnDescription ->+ ReaderIO FileBufferedOrSeekable Column+ #-}+{-# SPECIALIZE parseColumnChunks ::+ Int ->+ [ColumnChunk] ->+ ColumnDescription ->+ ReaderIO ChunkBuffer Column+ #-}+{-# INLINEABLE parseColumnChunks #-} parseColumnChunks totalRows chunks description | description.maxRepetitionLevel == 0 && description.maxDefinitionLevel == 0 = getNonNullableColumn totalRows description chunks@@ -353,6 +709,7 @@ getRepeatedColumn description chunks -- | Decode a required (non-nullable, non-repeated) column.+{-# INLINEABLE getNonNullableColumn #-} getNonNullableColumn :: forall m. (RandomAccess m, MonadIO m) =>@@ -368,7 +725,7 @@ Just (INT96 _) -> go int96Decoder Just (FLOAT _) -> unboxedGo floatDecoder Just (DOUBLE _) -> unboxedGo doubleDecoder- Just (BYTE_ARRAY _) -> go byteArrayDecoder+ Just (BYTE_ARRAY _) -> goPackedText Just (FIXED_LEN_BYTE_ARRAY _) -> case description.typeLength of Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type_length to be set" Just tl -> go (fixedLenByteArrayDecoder (fromIntegral tl))@@ -382,6 +739,22 @@ go decoder = foldNonNullable totalRows (foldColumnPagesM description decoder chunks) + -- Decode non-nullable UTF-8 BYTE_ARRAY straight into one 'PackedText'.+ goPackedText :: m Column+ goPackedText = do+ builder <-+ liftIO $ stToIO (newTextBuilder totalRows (textByteHint totalRows chunks))+ _ <-+ foldColumnDataPagesM+ description+ chunks+ ( \() (dict, enc, nPresent, valBytes, _, _) ->+ liftIO (appendStringPageIO builder dict enc nPresent valBytes)+ )+ ()+ chunk <- liftIO $ stToIO (freezeTextChunk builder)+ pure (mergeTextChunks [chunk])+ unboxedGo :: forall a. (Columnable a, VU.Unbox a) =>@@ -390,7 +763,18 @@ unboxedGo decoder = foldNonNullableUnboxed totalRows (foldColumnPagesM description decoder chunks) +-- | Initial byte capacity for a text column.+textByteHint :: Int -> [ColumnChunk] -> Int+textByteHint totalRows chunks =+ max (totalRows * 8) $+ sum+ [ fromIntegral (unField (cmd_total_uncompressed_size meta))+ | cc <- chunks+ , Just meta <- [unField (cc_meta_data cc)]+ ]+ -- | Decode an optional (nullable) column.+{-# INLINEABLE getNullableColumn #-} getNullableColumn :: forall m. (RandomAccess m, MonadIO m) =>@@ -406,7 +790,7 @@ Just (INT96 _) -> go int96Decoder Just (FLOAT _) -> unboxedGo floatDecoder Just (DOUBLE _) -> unboxedGo doubleDecoder- Just (BYTE_ARRAY _) -> go byteArrayDecoder+ Just (BYTE_ARRAY _) -> goPackedTextNullable Just (FIXED_LEN_BYTE_ARRAY _) -> case description.typeLength of Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type_length to be set" Just tl -> go (fixedLenByteArrayDecoder (fromIntegral tl))@@ -422,15 +806,36 @@ m Column go decoder = foldNullable maxDef totalRows (foldColumnPagesM description decoder chunks)++ -- Nullable UTF-8 BYTE_ARRAY: decode straight into a 'PackedText'.+ goPackedTextNullable :: m Column+ goPackedTextNullable = do+ builder <-+ liftIO $ stToIO (newTextBuilder totalRows (textByteHint totalRows chunks))+ _ <-+ foldColumnDataPagesM+ description+ chunks+ ( \() (dict, enc, nPresent, valBytes, defs, _) ->+ liftIO+ (appendNullableStringPageIO builder maxDef dict enc nPresent valBytes defs)+ )+ ()+ chunk <- liftIO $ stToIO (freezeTextChunk builder)+ pure (mergeTextChunks [chunk]) unboxedGo :: forall a. (Columnable a, VU.Unbox a) => UnboxedPageDecoder a -> m Column unboxedGo decoder =- foldNullableUnboxed maxDef totalRows (foldColumnPagesM description decoder chunks)+ foldNullableUnboxed+ maxDef+ totalRows+ (foldColumnPagesM description decoder chunks) -- | Decode a repeated (list/nested) column.+{-# INLINEABLE getRepeatedColumn #-} getRepeatedColumn :: forall m. (RandomAccess m, MonadIO m) =>@@ -481,8 +886,6 @@ unboxedGo decoder = foldRepeatedUnboxed maxRep maxDef (foldColumnPagesM description decoder chunks) --- Options application ------------------------------------------------------ applyRowRange :: ParquetReadOptions -> DataFrame -> DataFrame applyRowRange opts df = maybe df (`DS.range` df) (rowRange opts)@@ -507,11 +910,7 @@ . applySelectedColumns opts . applyPredicate opts --- Logical type conversion ---------------------------------------------------{- | Apply a column-description's logical type annotation to convert raw-decoded values (e.g. millisecond integers → 'UTCTime').--}+-- | Apply a column's logical type annotation to its decoded values. applyDescLogicalType :: ColumnDescription -> DI.Column -> DI.Column applyDescLogicalType desc = applyLogicalType (colLogicalType desc) @@ -519,34 +918,86 @@ applyLogicalType (Just (LT_TIMESTAMP f)) col = let ts = unField f unit = unField ts.timestamp_unit- divisor = case unit of- MILLIS _ -> 1_000- MICROS _ -> 1_000_000- NANOS _ -> 1_000_000_000- in fromRight col $- DI.mapColumn- (microsecondsToUTCTime . (* (1_000_000 `div` divisor)))- col+ -- Convert at native precision (ticks per second, ps per tick).+ conv = case unit of+ MILLIS _ -> epochToUTCTime 1_000 1_000_000_000+ MICROS _ -> epochToUTCTime 1_000_000 1_000_000+ NANOS _ -> epochToUTCTime 1_000_000_000 1_000+ in fromRight col $ DI.mapColumn conv col applyLogicalType (Just (LT_DECIMAL f)) col = let dt = unField f scale = unField dt.decimal_scale precision = unField dt.decimal_precision in if precision <= 9- then case DI.toVector @Int32 @VU.Vector col of- Right xs ->- DI.fromUnboxedVector $- VU.map (\raw -> fromIntegral @Int32 @Double raw / 10 ^ scale) xs- Left _ -> col+ then+ fromRight col $+ DI.mapColumn (\raw -> fromIntegral @Int32 @Double raw / 10 ^ scale) col else if precision <= 18- then case DI.toVector @Int64 @VU.Vector col of- Right xs ->- DI.fromUnboxedVector $- VU.map (\raw -> fromIntegral @Int64 @Double raw / 10 ^ scale) xs- Left _ -> col+ then+ fromRight col $+ DI.mapColumn (\raw -> fromIntegral @Int64 @Double raw / 10 ^ scale) col else col applyLogicalType _ col = col -microsecondsToUTCTime :: Int64 -> UTCTime-microsecondsToUTCTime us =- posixSecondsToUTCTime (fromIntegral us / 1_000_000)+nativeTypeHints :: FileMetadata -> Schema+nativeTypeHints metadata =+ makeSchema+ [ (name, ty)+ | kv <- concat (unField metadata.key_value_metadata)+ , Just name <- [T.stripPrefix nativeTypeKeyPrefix (unField kv.kv_key)]+ , Just value <- [unField kv.kv_value]+ , Just ty <- [Map.lookup value stampedSchemaTypes]+ ]++stampedSchemaTypes :: Map.Map T.Text SchemaType+stampedSchemaTypes =+ Map.fromList $+ concat+ [ entry @Int "Int"+ , entry @Int32 "Int32"+ , entry @Int64 "Int64"+ , entry @Integer "Integer"+ , entry @Float "Float"+ , entry @Double "Double"+ , entry @Bool "Bool"+ , entry @T.Text "Text"+ , entry @UTCTime "UTCTime"+ ]+ where+ entry ::+ forall a.+ (Columnable a, Read a, Columnable (Maybe a)) =>+ T.Text ->+ [(T.Text, SchemaType)]+ entry name =+ [ (name, schemaType @a)+ , ("Maybe " <> name, schemaType @(Maybe a))+ ]++restoreNativeType :: Schema -> T.Text -> DI.Column -> DI.Column+restoreNativeType hints name col = case lookupHint of+ Just ty+ | stampedAs @Int ty -> narrow (fromIntegral @Int64 @Int)+ | stampedAs @Integer ty -> narrow (fromIntegral @Int64 @Integer)+ _ -> col+ where+ lookupHint =+ case Map.lookup name (elements hints) of+ Just ty -> Just ty+ Nothing -> Map.lookup (leafOf name) (elements hints)+ stampedAs ::+ forall a. (Columnable a, Read a, Columnable (Maybe a)) => SchemaType -> Bool+ stampedAs ty = ty == schemaType @a || ty == schemaType @(Maybe a)+ narrow :: (Columnable a) => (Int64 -> a) -> DI.Column+ narrow f = fromRight col (DI.mapColumn f col)++-- | Convert an epoch timestamp with the given tick rate to 'UTCTime'.+epochToUTCTime :: Int64 -> Integer -> Int64 -> UTCTime+epochToUTCTime ticksPerSecond psPerTick v =+ let (s, subTicks) = v `divMod` ticksPerSecond+ (days, sInDay) = s `divMod` 86_400+ ps = fromIntegral sInDay * 1_000_000_000_000 + fromIntegral subTicks * psPerTick+ in UTCTime+ (ModifiedJulianDay (40_587 + fromIntegral days))+ (picosecondsToDiffTime ps)
src/DataFrame/IO/Parquet/Binary.hs view
@@ -18,12 +18,7 @@ readUVarInt :: BS.ByteString -> (Word64, BS.ByteString) readUVarInt xs = loop xs 0 0 0 where- {-- Each input byte contributes:- - lower 7 payload bits- - The high bit (0x80) is the continuation flag: 1 = more bytes follow, 0 = last byte- Why the magic number 10: For a 64‑bit integer we need at most ceil(64 / 7) = 10 bytes- -}+ -- The high bit flags continuation; 10 bytes cover a 64-bit integer. loop :: BS.ByteString -> Word64 -> Int -> Int -> (Word64, BS.ByteString) loop bs result _ 10 = (result, bs) loop xs' result shiftAmt i = case BS.uncons xs' of@@ -113,9 +108,7 @@ (fromIntegral size) ((`readSingleByte` buf) . fromIntegral) -{- | Allocate a fixed-size buffer, repeat the action on each index.-Fill it into the buffer to get a ByteString.--}+-- | Allocate a fixed-size buffer and fill it by repeating an action. fillByteStringByWord8 :: Int -> (Int -> IO Word8) -> IO BS.ByteString fillByteStringByWord8 size getByte = do bracketOnError
src/DataFrame/IO/Parquet/Dictionary.hs view
@@ -1,6 +1,17 @@ {-# LANGUAGE BangPatterns #-} -module DataFrame.IO.Parquet.Dictionary (DictVals (..), readDictVals, decodeRLEBitPackedHybrid) where+module DataFrame.IO.Parquet.Dictionary (+ DictVals (..),+ readDictVals,+ dictText,+ dictInt32,+ dictInt64,+ dictFloat,+ dictDouble,+ word32At,+ word64At,+ decodeRLEBitPackedHybrid,+) where import Data.Bits import qualified Data.ByteString as BS@@ -10,109 +21,101 @@ import Data.Text.Encoding import Data.Time (UTCTime) import qualified Data.Vector as V+import qualified Data.Vector.Unboxed as VU import Data.Word import DataFrame.IO.Parquet.Binary (readUVarInt) import DataFrame.IO.Parquet.Thrift (ThriftType (..)) import DataFrame.IO.Parquet.Time (int96ToUTCTime)-import DataFrame.Internal.Binary (- littleEndianInt32,- littleEndianWord32,- littleEndianWord64,- )+import DataFrame.Internal.Binary (littleEndianWord32) import GHC.Float +-- | A decoded dictionary page with unboxed numeric entries. data DictVals- = DBool (V.Vector Bool)- | DInt32 (V.Vector Int32)- | DInt64 (V.Vector Int64)- | DInt96 (V.Vector UTCTime)- | DFloat (V.Vector Float)- | DDouble (V.Vector Double)- | DText (V.Vector T.Text)+ = DBool !(VU.Vector Bool)+ | DInt32 !(VU.Vector Int32)+ | DInt64 !(VU.Vector Int64)+ | DInt96 !(V.Vector UTCTime)+ | DFloat !(VU.Vector Float)+ | DDouble !(VU.Vector Double)+ | DText !(VU.Vector Int) !(VU.Vector Int) !BS.ByteString deriving (Show, Eq) -{- | Decode the values from a dictionary page.--The @numVals@ argument is the entry count declared in the dictionary page-header. It is used to limit BOOLEAN decoding (1-bit-per-value encoding has-no natural delimiter).--The @typeLength@ argument is only meaningful for FIXED_LEN_BYTE_ARRAY: it is-the byte-width of each individual dictionary entry, NOT the total number of-entries. Passing @numVals@ here (the old behaviour) would cause it to be-misread as an element size, yielding a dictionary that is far too small.--}+-- | Decode the values from a dictionary page. readDictVals :: ThriftType -> BS.ByteString -> Int32 -> Maybe Int32 -> DictVals-readDictVals (BOOLEAN _) bs count _ = DBool (V.fromList (take (fromIntegral count) $ readPageBool bs))-readDictVals (INT32 _) bs _ _ = DInt32 (V.fromList (readPageInt32 bs))-readDictVals (INT64 _) bs _ _ = DInt64 (V.fromList (readPageInt64 bs))-readDictVals (INT96 _) bs _ _ = DInt96 (V.fromList (readPageInt96Times bs))-readDictVals (FLOAT _) bs _ _ = DFloat (V.fromList (readPageFloat bs))-readDictVals (DOUBLE _) bs _ _ = DDouble (V.fromList (readPageWord64 bs))-readDictVals (BYTE_ARRAY _) bs _ _ = DText (V.fromList (readPageBytes bs))+readDictVals (BOOLEAN _) bs count _ =+ let n = min (fromIntegral count) (8 * BS.length bs)+ in DBool $ VU.generate n $ \i ->+ (BSU.unsafeIndex bs (i `shiftR` 3) `shiftR` (i .&. 7)) .&. 1 == 1+readDictVals (INT32 _) bs _ _ =+ DInt32 $ VU.generate (BS.length bs `div` 4) $ \i -> fromIntegral (word32At bs (4 * i))+readDictVals (INT64 _) bs _ _ =+ DInt64 $ VU.generate (BS.length bs `div` 8) $ \i -> fromIntegral (word64At bs (8 * i))+readDictVals (INT96 _) bs _ _ =+ DInt96 $ V.generate (BS.length bs `div` 12) $ \i ->+ int96ToUTCTime (BSU.unsafeTake 12 (BSU.unsafeDrop (12 * i) bs))+readDictVals (FLOAT _) bs _ _ =+ DFloat $ VU.generate (BS.length bs `div` 4) $ \i -> castWord32ToFloat (word32At bs (4 * i))+readDictVals (DOUBLE _) bs _ _ =+ DDouble $ VU.generate (BS.length bs `div` 8) $ \i -> castWord64ToDouble (word64At bs (8 * i))+readDictVals (BYTE_ARRAY _) bs _ _ = textDict bs readDictVals (FIXED_LEN_BYTE_ARRAY _) bs _ (Just len) =- DText (V.fromList (readPageFixedBytes bs (fromIntegral len)))+ let w = fromIntegral len+ n = if w <= 0 then 0 else BS.length bs `div` w+ in DText (VU.generate n (* w)) (VU.replicate n w) bs readDictVals t _ _ _ = error $ "Unsupported dictionary type: " ++ show t -readPageInt32 :: BS.ByteString -> [Int32]-readPageInt32 xs- | BS.null xs = []- | otherwise = littleEndianInt32 (BS.take 4 xs) : readPageInt32 (BS.drop 4 xs)--readPageWord64 :: BS.ByteString -> [Double]-readPageWord64 xs- | BS.null xs = []- | otherwise =- castWord64ToDouble (littleEndianWord64 (BS.take 8 xs))- : readPageWord64 (BS.drop 8 xs)+-- | Length-prefixed byte-array entries, in two passes.+textDict :: BS.ByteString -> DictVals+textDict bs = DText starts lens bs+ where+ !len = BS.length bs+ count !off !n+ | off + 4 > len = n+ | otherwise = count (off + 4 + fromIntegral (word32At bs off)) (n + 1)+ !n = count 0 0+ positions = VU.unfoldrExactN n step 0+ step off = let l = fromIntegral (word32At bs off) in ((off + 4, l), off + 4 + l)+ (starts, lens) = VU.unzip positions -readPageBytes :: BS.ByteString -> [T.Text]-readPageBytes xs- | BS.null xs = []- | otherwise =- let lenBytes = fromIntegral (littleEndianInt32 $ BS.take 4 xs)- totalBytesRead = lenBytes + 4- in decodeUtf8Lenient (BS.take lenBytes (BS.drop 4 xs))- : readPageBytes (BS.drop totalBytesRead xs)+-- | Dictionary entry @i@ as 'T.Text' (lenient decode, for the boxed paths).+dictText :: DictVals -> Int -> T.Text+dictText (DText starts lens bs) i =+ decodeUtf8Lenient+ ( BSU.unsafeTake+ (VU.unsafeIndex lens i)+ (BSU.unsafeDrop (VU.unsafeIndex starts i) bs)+ )+dictText d _ = error ("dictText: wrong dict type, got " ++ show d) -readPageBool :: BS.ByteString -> [Bool]-readPageBool bs =- concatMap (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7]) (BS.unpack bs)+dictInt32 :: DictVals -> Maybe (VU.Vector Int32)+dictInt32 (DInt32 ds) = Just ds+dictInt32 _ = Nothing -readPageInt64 :: BS.ByteString -> [Int64]-readPageInt64 xs- | BS.null xs = []- | otherwise =- fromIntegral (littleEndianWord64 (BS.take 8 xs)) : readPageInt64 (BS.drop 8 xs)+dictInt64 :: DictVals -> Maybe (VU.Vector Int64)+dictInt64 (DInt64 ds) = Just ds+dictInt64 _ = Nothing -readPageFloat :: BS.ByteString -> [Float]-readPageFloat xs- | BS.null xs = []- | otherwise =- castWord32ToFloat (littleEndianWord32 (BS.take 4 xs))- : readPageFloat (BS.drop 4 xs)+dictFloat :: DictVals -> Maybe (VU.Vector Float)+dictFloat (DFloat ds) = Just ds+dictFloat _ = Nothing -readNInt96Times :: Int -> BS.ByteString -> ([UTCTime], BS.ByteString)-readNInt96Times 0 bs = ([], bs)-readNInt96Times k bs =- let timestamp96 = BS.take 12 bs- utcTime = int96ToUTCTime timestamp96- bs' = BS.drop 12 bs- (times, rest) = readNInt96Times (k - 1) bs'- in (utcTime : times, rest)+dictDouble :: DictVals -> Maybe (VU.Vector Double)+dictDouble (DDouble ds) = Just ds+dictDouble _ = Nothing -readPageInt96Times :: BS.ByteString -> [UTCTime]-readPageInt96Times bs- | BS.null bs = []- | otherwise =- let (times, _) = readNInt96Times (BS.length bs `div` 12) bs- in times+word32At :: BS.ByteString -> Int -> Word32+word32At bs o =+ fromIntegral (BSU.unsafeIndex bs o)+ .|. (fromIntegral (BSU.unsafeIndex bs (o + 1)) `unsafeShiftL` 8)+ .|. (fromIntegral (BSU.unsafeIndex bs (o + 2)) `unsafeShiftL` 16)+ .|. (fromIntegral (BSU.unsafeIndex bs (o + 3)) `unsafeShiftL` 24)+{-# INLINE word32At #-} -readPageFixedBytes :: BS.ByteString -> Int -> [T.Text]-readPageFixedBytes xs len- | BS.null xs = []- | otherwise =- decodeUtf8Lenient (BS.take len xs) : readPageFixedBytes (BS.drop len xs) len+word64At :: BS.ByteString -> Int -> Word64+word64At bs o =+ fromIntegral (word32At bs o)+ .|. (fromIntegral (word32At bs (o + 4)) `unsafeShiftL` 32)+{-# INLINE word64At #-} unpackBitPacked :: Int -> Int -> BS.ByteString -> ([Word32], BS.ByteString) unpackBitPacked bw count bs@@ -124,7 +127,7 @@ rest = BS.drop totalBytes bs in (extractBits bw count chunk, rest) --- | LSB-first bit accumulator: reads each byte once with no intermediate ByteString allocation.+-- | LSB-first bit accumulator that reads each byte once. extractBits :: Int -> Int -> BS.ByteString -> [Word32] extractBits bw count bs = go 0 (0 :: Word64) 0 count where@@ -145,9 +148,7 @@ | bitWidth == 0 = ([0], bs) | BS.null bs = ([], bs) | otherwise =- -- readUVarInt is evaluated here, inside the guard that has already- -- confirmed bs is non-empty. Keeping it in a where clause would cause- -- it to be forced before the BS.null guard under {-# LANGUAGE Strict #-}.+ -- Evaluated inside the guard: a where clause would force it too early. let (hdr64, afterHdr) = readUVarInt bs isPacked = (hdr64 .&. 1) == 1 in if isPacked
src/DataFrame/IO/Parquet/Encoding.hs view
@@ -7,11 +7,13 @@ bitWidthForMaxLevel, -- Vector-based RLE/bit-packed decoder (from new parser) decodeRLEBitPackedHybrid,+ rleHybridFold, extractBitsInto, fillRun, decodeDictIndices, ) where +import Control.Monad (when) import Control.Monad.ST (ST, runST) import Data.Bits import qualified Data.ByteString as BS@@ -22,10 +24,6 @@ import DataFrame.IO.Parquet.Binary (readUVarInt) import DataFrame.Internal.Binary (littleEndianWord32) --- ------------------------------------------------------------------------------ Level-width helpers (used by Levels.hs)--- ---------------------------------------------------------------------------- ceilLog2 :: Int -> Int ceilLog2 x | x <= 1 = 0@@ -34,10 +32,6 @@ bitWidthForMaxLevel :: Int -> Int bitWidthForMaxLevel maxLevel = ceilLog2 (maxLevel + 1) --- ------------------------------------------------------------------------------ Vector-based RLE / bit-packed hybrid decoder--- ---------------------------------------------------------------------------- decodeRLEBitPackedHybrid :: -- | Bit width per value (0 = all zeros, use 'VU.replicate') Int ->@@ -66,8 +60,7 @@ let groups = fromIntegral (hdr64 `shiftR` 1) :: Int totalVals = groups * 8 takeN = min (need - filled) totalVals- -- Consume all the bytes for this group even if we- -- only need a subset of the values.+ -- Consume the whole group even if we need less. bytesN = (bw * totalVals + 7) `div` 8 (chunk, rest) = BS.splitAt bytesN afterHdr extractBitsInto bw takeN chunk mv filled@@ -87,10 +80,7 @@ fillRun mv i end = VUM.set (VUM.unsafeSlice i (end - i) mv) {-# INLINE fillRun #-} -{- | Write @count@ bit-width-@bw@ values from @bs@ into @mv@ starting at-@offset@, reading the byte buffer with a single-pass LSB-first accumulator.-No intermediate list or ByteString allocation.--}+-- | Write bit-packed values into a mutable vector in a single pass. extractBitsInto :: -- | Bit width Int ->@@ -116,12 +106,7 @@ in go (byteIdx + 1) (acc .|. (b `unsafeShiftL` accBits)) (accBits + 8) done {-# INLINE extractBitsInto #-} -{- | Decode @need@ dictionary indices from a DATA_PAGE bit-width-prefixed-stream (the first byte encodes the bit-width of all subsequent RLE\/bitpacked-values).--Returns the index vector (as 'Int') and the unconsumed bytes.--}+-- | Decode dictionary indices from a bit-width-prefixed DATA_PAGE stream. decodeDictIndices :: Int -> BS.ByteString -> (VU.Vector Int, BS.ByteString) decodeDictIndices need bs = case BS.uncons bs of Nothing -> error "decodeDictIndices: empty stream"@@ -130,3 +115,55 @@ (raw, rest1) = decodeRLEBitPackedHybrid bw need rest0 in (VU.map fromIntegral raw, rest1) {-# INLINE decodeDictIndices #-}++-- | Stream RLE\/bit-packed hybrid values to a sink.+rleHybridFold ::+ Int -> Int -> BS.ByteString -> (Int -> Word32 -> ST s ()) -> ST s ()+rleHybridFold bw need bs sink+ | bw == 0 = zeros 0+ | otherwise = go 0 bs+ where+ !mask = if bw == 32 then maxBound else (1 `unsafeShiftL` bw) - 1 :: Word32+ zeros !i = when (i < need) (sink i 0 >> zeros (i + 1))+ go !filled !buf+ | filled >= need = pure ()+ | BS.null buf = pure ()+ | otherwise =+ let (hdr64, afterHdr) = readUVarInt buf+ isPacked = (hdr64 .&. 1) == 1+ in if isPacked+ then do+ let groups = fromIntegral (hdr64 `shiftR` 1) :: Int+ totalVals = groups * 8+ takeN = min (need - filled) totalVals+ bytesN = (bw * totalVals + 7) `div` 8+ (chunk, rest) = BS.splitAt bytesN afterHdr+ extractBitsFold bw takeN chunk (\j w -> sink (filled + j) w)+ go (filled + takeN) rest+ else do+ let runLen = fromIntegral (hdr64 `shiftR` 1) :: Int+ nbytes = (bw + 7) `div` 8+ val = littleEndianWord32 (BS.take 4 afterHdr) .&. mask+ takeN = min (need - filled) runLen+ run !j = when (j < takeN) (sink (filled + j) val >> run (j + 1))+ run 0+ go (filled + takeN) (BS.drop nbytes afterHdr)+{-# INLINE rleHybridFold #-}++-- | 'extractBitsInto' with a sink instead of a destination vector.+extractBitsFold ::+ Int -> Int -> BS.ByteString -> (Int -> Word32 -> ST s ()) -> ST s ()+extractBitsFold bw count bs sink = go 0 (0 :: Word64) 0 0+ where+ !mask = if bw == 32 then maxBound else (1 `unsafeShiftL` bw) - 1 :: Word64+ !len = BS.length bs+ go !byteIdx !acc !accBits !done+ | done >= count = return ()+ | accBits >= bw = do+ sink done (fromIntegral (acc .&. mask))+ go byteIdx (acc `unsafeShiftR` bw) (accBits - bw) (done + 1)+ | byteIdx >= len = return ()+ | otherwise =+ let b = fromIntegral (BSU.unsafeIndex bs byteIdx) :: Word64+ in go (byteIdx + 1) (acc .|. (b `unsafeShiftL` accBits)) (accBits + 8) done+{-# INLINE extractBitsFold #-}
src/DataFrame/IO/Parquet/Levels.hs view
@@ -23,14 +23,7 @@ ) import DataFrame.Internal.Binary (littleEndianWord32) --- ------------------------------------------------------------------------------ Level readers--- -----------------------------------------------------------------------------{- | Convert a 'Word32' level vector to an 'Int' level vector while counting-how many entries equal @maxDef@. Single pass; allocates a single-'VU.Vector Int' of length @VU.length raw@.--}+-- | Convert a 'Word32' level vector to 'Int', counting @maxDef@ entries. convertAndCount :: Int -> VU.Vector Word32 -> (VU.Vector Int, Int) convertAndCount maxDef raw = runST $ do let !n = VU.length raw@@ -65,8 +58,8 @@ (defVec, nPresent, afterDef) = decodeLevelBlock bwDef n afterRep in (defVec, repVec, nPresent, afterDef) where- -- For rep block we don't need nPresent; we still get one cheaply.- decodeLevelBlock 0 n' buf = (VU.replicate n' 0, n' * fromEnum (maxDef == 0), buf)+ -- No levels to read: non-nullable, non-repeated paths never use them.+ decodeLevelBlock 0 n' buf = (VU.empty, n' * fromEnum (maxDef == 0), buf) decodeLevelBlock bw n' buf = let blockLen = fromIntegral (littleEndianWord32 (BS.take 4 buf)) :: Int blockData = BS.take blockLen (BS.drop 4 buf)@@ -94,22 +87,19 @@ bwRep = bitWidthForMaxLevel maxRep bwDef = bitWidthForMaxLevel maxDef repVec- | bwRep == 0 = VU.replicate n 0+ | bwRep == 0 = VU.empty | otherwise = let (raw, _) = decodeRLEBitPackedHybrid bwRep n repBytes (out, _) = convertAndCount maxDef raw in out (defVec, nPresent)- | bwDef == 0 = (VU.replicate n 0, n * fromEnum (maxDef == 0))+ | bwDef == 0 = (VU.empty, n * fromEnum (maxDef == 0)) | otherwise = let (raw, _) = decodeRLEBitPackedHybrid bwDef n defBytes in convertAndCount maxDef raw in (defVec, repVec, nPresent, afterDefBytes) -{- | Stitch a singly-nested list column (@maxRep == 1@) from vector-format-definition and repetition levels plus a compact present-values vector.-Returns one @Maybe [Maybe a]@ per top-level row.--}+-- | Stitch a singly-nested list column from levels and present values. stitchList :: Int -> VU.Vector Int ->@@ -123,9 +113,7 @@ toRow ((_, d, _) : _) | d == 0 = Nothing toRow grp = Just [v | (_, _, v) <- grp] -{- | Stitch a doubly-nested list column (@maxRep == 2@).-@defT1@ is the def threshold at which the depth-1 element is present.--}+-- | Stitch a doubly-nested list column (@maxRep == 2@). stitchList2 :: Int -> Int ->@@ -146,10 +134,7 @@ toLeaf [] = Nothing toLeaf ((_, _, v) : _) = v -{- | Stitch a triply-nested list column (@maxRep == 3@).-@defT1@ and @defT2@ are the def thresholds for depth-1 and depth-2-elements respectively.--}+-- | Stitch a triply-nested list column (@maxRep == 3@). stitchList3 :: Int -> Int ->@@ -174,14 +159,7 @@ toLeaf [] = Nothing toLeaf ((_, _, v) : _) = v --- ------------------------------------------------------------------------------ Internal helpers--- -----------------------------------------------------------------------------{- | Zip rep and def level vectors with a present-values vector, tagging each-position as @Just value@ (when @def == maxDef@) or @Nothing@.-Returns a flat list of @(rep, def, Maybe a)@ triplets for row-splitting.--}+-- | Zip rep and def levels with present values into triplets. pairWithValsV :: Int -> VU.Vector Int ->@@ -200,9 +178,7 @@ then (r, d, Just (VB.unsafeIndex values j)) : go (i + 1) (j + 1) else (r, d, Nothing) : go (i + 1) j -{- | Group a flat triplet list into rows.-A new group begins whenever @rep <= bound@.--}+-- | Group a flat triplet list into rows. splitAtRepBound :: Int -> [(Int, Int, Maybe a)] -> [[(Int, Int, Maybe a)]] splitAtRepBound _ [] = [] splitAtRepBound bound (t : ts) =
src/DataFrame/IO/Parquet/Page.hs view
@@ -17,11 +17,18 @@ fixedLenByteArrayDecoder, -- Page iteration foldColumnPagesM,+ foldColumnDataPagesM,+ RawPage,+ decodeRequiredNumericInto,+ appendStringPageIO,+ appendNullableStringPageIO, ) where import Control.Monad.IO.Class (MonadIO (liftIO))-import Data.Bits (shiftR, (.&.))+import Control.Monad.ST (RealWorld, stToIO)+import Data.Bits (shiftL, shiftR, (.&.), (.|.)) import qualified Data.ByteString as BS+import qualified Data.ByteString.Unsafe as BSU import Data.Int (Int32, Int64) import Data.Maybe (fromJust, fromMaybe) import qualified Data.Text as T@@ -30,12 +37,21 @@ import qualified Data.Vector as VB import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word (Word8) import DataFrame.IO.Parquet.Decompress (decompressData) import DataFrame.IO.Parquet.Dictionary ( DictVals (..),+ dictText, readDictVals,+ word32At,+ word64At, )-import DataFrame.IO.Parquet.Encoding (decodeDictIndices)+import DataFrame.IO.Parquet.Encoding (+ decodeDictIndices,+ decodeRLEBitPackedHybrid,+ rleHybridFold,+ ) import DataFrame.IO.Parquet.Levels (readLevelsV1, readLevelsV2) import DataFrame.IO.Parquet.Thrift ( ColumnChunk (..),@@ -53,62 +69,57 @@ import DataFrame.IO.Parquet.Time (int96ToUTCTime) import DataFrame.IO.Parquet.Utils (ColumnDescription (..)) import DataFrame.IO.Utils.RandomAccess (RandomAccess (..), Range (Range))-import DataFrame.Internal.Binary (- littleEndianInt32,- littleEndianWord32,- littleEndianWord64,+import DataFrame.Internal.Binary (littleEndianInt32)+import DataFrame.Internal.Column.Builder (+ TextBuilder,+ appendNull,+ appendTextSliceFromPtr, )+import Foreign.Ptr (Ptr, castPtr, plusPtr)+import Foreign.Storable (peekByteOff) import GHC.Float (castWord32ToFloat, castWord64ToDouble) import Pinch (decodeWithLeftovers) import qualified Pinch --- ------------------------------------------------------------------------------ Types--- -----------------------------------------------------------------------------{- | A type-specific page decoder.-Given the optional dictionary, the page encoding, the number of present-values, and the decompressed value bytes, returns exactly @nPresent@ values.--}+-- | A type-specific page decoder. type PageDecoder a = Maybe DictVals -> Encoding -> Int -> BS.ByteString -> VB.Vector a type UnboxedPageDecoder a = Maybe DictVals -> Encoding -> Int -> BS.ByteString -> VU.Vector a --- ------------------------------------------------------------------------------ Per-type decoders--- ---------------------------------------------------------------------------- boolDecoder :: UnboxedPageDecoder Bool boolDecoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.fromList (readNBool nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool+ -- PLAIN bools are bit-packed; index the bits directly, avoiding a list.+ PLAIN _ ->+ VU.generate nPresent $ \i ->+ (BSU.unsafeIndex bs (i `shiftR` 3) `shiftR` (i .&. 7)) .&. 1 == 1+ RLE_DICTIONARY _ -> gatherDict "boolDecoder" getBool mDict nPresent bs+ PLAIN_DICTIONARY _ -> gatherDict "boolDecoder" getBool mDict nPresent bs _ -> error ("boolDecoder: unsupported encoding " ++ show enc) where- getBool (DBool ds) i = ds VB.! i- getBool d _ = error ("boolDecoder: wrong dict type, got " ++ show d)+ getBool (DBool ds) = Just ds+ getBool _ = Nothing int32Decoder :: UnboxedPageDecoder Int32 int32Decoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNInt32 nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt32- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt32+ PLAIN _ -> readNInt32 nPresent bs+ RLE_DICTIONARY _ -> gatherDict "int32Decoder" getInt32 mDict nPresent bs+ PLAIN_DICTIONARY _ -> gatherDict "int32Decoder" getInt32 mDict nPresent bs _ -> error ("int32Decoder: unsupported encoding " ++ show enc) where- getInt32 (DInt32 ds) i = ds VB.! i- getInt32 d _ = error ("int32Decoder: wrong dict type, got " ++ show d)+ getInt32 (DInt32 ds) = Just ds+ getInt32 _ = Nothing int64Decoder :: UnboxedPageDecoder Int64 int64Decoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNInt64 nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt64- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getInt64+ PLAIN _ -> readNInt64 nPresent bs+ RLE_DICTIONARY _ -> gatherDict "int64Decoder" getInt64 mDict nPresent bs+ PLAIN_DICTIONARY _ -> gatherDict "int64Decoder" getInt64 mDict nPresent bs _ -> error ("int64Decoder: unsupported encoding " ++ show enc) where- getInt64 (DInt64 ds) i = ds VB.! i- getInt64 d _ = error ("int64Decoder: wrong dict type, got " ++ show d)+ getInt64 (DInt64 ds) = Just ds+ getInt64 _ = Nothing int96Decoder :: PageDecoder UTCTime int96Decoder mDict enc nPresent bs = case enc of@@ -117,28 +128,28 @@ PLAIN_DICTIONARY _ -> lookupDict mDict nPresent bs getInt96 _ -> error ("int96Decoder: unsupported encoding " ++ show enc) where- getInt96 (DInt96 ds) i = ds VB.! i+ getInt96 (DInt96 ds) i = VB.unsafeIndex ds i getInt96 d _ = error ("int96Decoder: wrong dict type, got " ++ show d) floatDecoder :: UnboxedPageDecoder Float floatDecoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNFloat nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getFloat- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getFloat+ PLAIN _ -> readNFloat nPresent bs+ RLE_DICTIONARY _ -> gatherDict "floatDecoder" getFloat mDict nPresent bs+ PLAIN_DICTIONARY _ -> gatherDict "floatDecoder" getFloat mDict nPresent bs _ -> error ("floatDecoder: unsupported encoding " ++ show enc) where- getFloat (DFloat ds) i = ds VB.! i- getFloat d _ = error ("floatDecoder: wrong dict type, got " ++ show d)+ getFloat (DFloat ds) = Just ds+ getFloat _ = Nothing doubleDecoder :: UnboxedPageDecoder Double doubleDecoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.convert (readNDouble nPresent bs)- RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getDouble- PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getDouble+ PLAIN _ -> readNDouble nPresent bs+ RLE_DICTIONARY _ -> gatherDict "doubleDecoder" getDouble mDict nPresent bs+ PLAIN_DICTIONARY _ -> gatherDict "doubleDecoder" getDouble mDict nPresent bs _ -> error ("doubleDecoder: unsupported encoding " ++ show enc) where- getDouble (DDouble ds) i = ds VB.! i- getDouble d _ = error ("doubleDecoder: wrong dict type, got " ++ show d)+ getDouble (DDouble ds) = Just ds+ getDouble _ = Nothing byteArrayDecoder :: PageDecoder T.Text byteArrayDecoder mDict enc nPresent bs = case enc of@@ -147,8 +158,7 @@ PLAIN_DICTIONARY _ -> lookupDict mDict nPresent bs getText _ -> error ("byteArrayDecoder: unsupported encoding " ++ show enc) where- getText (DText ds) i = ds VB.! i- getText d _ = error ("byteArrayDecoder: wrong dict type, got " ++ show d)+ getText = dictText fixedLenByteArrayDecoder :: Int -> PageDecoder T.Text fixedLenByteArrayDecoder len mDict enc nPresent bs = case enc of@@ -157,12 +167,9 @@ PLAIN_DICTIONARY _ -> lookupDict mDict nPresent bs getText _ -> error ("fixedLenByteArrayDecoder: unsupported encoding " ++ show enc) where- getText (DText ds) i = ds VB.! i- getText d _ = error ("fixedLenByteArrayDecoder: wrong dict type, got " ++ show d)+ getText = dictText -{- | Shared dictionary-path helper: decode @nPresent@ RLE/bit-packed indices-and look each one up in the dictionary.--}+-- | Decode RLE\/bit-packed indices and look each up in the dictionary. lookupDict :: Maybe DictVals -> Int ->@@ -175,22 +182,25 @@ let (idxs, _) = decodeDictIndices nPresent bs in VB.generate nPresent (f dict . VU.unsafeIndex idxs) -unboxedLookupDict ::+-- | Decode dictionary indices and gather from an unboxed dictionary.+gatherDict :: (VU.Unbox a) =>+ String ->+ (DictVals -> Maybe (VU.Vector a)) -> Maybe DictVals -> Int -> BS.ByteString ->- (DictVals -> Int -> a) -> VU.Vector a-unboxedLookupDict mDict nPresent bs f = case mDict of- Nothing -> error "Dictionary-encoded page but no dictionary page seen"- Just dict ->- let (idxs, _) = decodeDictIndices nPresent bs- in VU.generate nPresent (f dict . VU.unsafeIndex idxs)---- ------------------------------------------------------------------------------ Core page-iteration loop--- ---------------------------------------------------------------------------+gatherDict who pick mDict nPresent bs = case mDict of+ Nothing -> error (who ++ ": dictionary-encoded page but no dictionary page seen")+ Just dict -> case pick dict of+ Nothing -> error (who ++ ": wrong dict type, got " ++ show dict)+ Just ds -> case BS.uncons bs of+ Nothing -> error (who ++ ": empty index stream")+ Just (w0, rest) ->+ let (idxs, _) = decodeRLEBitPackedHybrid (fromIntegral w0) nPresent rest+ in VU.map (VU.unsafeIndex ds . fromIntegral) idxs+{-# INLINE gatherDict #-} -- | Read the raw (compressed) byte range for a column chunk. readChunkBytes ::@@ -208,37 +218,21 @@ rawBytes <- readBytes (Range offset compLen) return (codec, pType, rawBytes) -{- | Left-fold a monadic step over every DATA page (V1 or V2) of every column-chunk, in order, threading the running dictionary internally.--Replaces the old Streamly @readPages@ 'Unfold' + @unfoldEach@/@fold@ pipeline-with a direct monadic fold: pages are decoded and handed to @step@ one at a-time, so only a single page's decoded values are live at once (the-@unsafeFreeze@-into-preallocated-buffer consumers rely on this — a materialised-page list would roughly double peak memory per column).--Dictionary pages are consumed silently and update the running dictionary that-is threaded through the recursion. @INDEX_PAGE@s are skipped.---- TODO: when a page index is available, use it here to compute which page--- byte ranges to request from the RandomAccess layer instead of reading the--- entire column chunk in one contiguous read.+-- | A decoded DATA page handed to a fold step.+type RawPage =+ (Maybe DictVals, Encoding, Int, BS.ByteString, VU.Vector Int, VU.Vector Int) --- TODO: accept an optional row-range and use the column/offset page index--- (when present in file metadata) to skip pages whose row range does not--- overlap the requested range, avoiding decompression of irrelevant pages--- entirely.--}-foldColumnPagesM ::- forall m v a acc.- (RandomAccess m, MonadIO m, VG.Vector v a) =>+-- | Left-fold a monadic step over every DATA page of every column chunk.+{-# INLINEABLE foldColumnDataPagesM #-}+foldColumnDataPagesM ::+ forall m acc.+ (RandomAccess m, MonadIO m) => ColumnDescription ->- (Maybe DictVals -> Encoding -> Int -> BS.ByteString -> v a) -> [ColumnChunk] ->- (acc -> (v a, VU.Vector Int, VU.Vector Int) -> m acc) ->+ (acc -> RawPage -> m acc) -> acc -> m acc-foldColumnPagesM description decoder chunks step = goChunks chunks+foldColumnDataPagesM description chunks step = goChunks chunks where maxDef = fromIntegral description.maxDefinitionLevel :: Int maxRep = fromIntegral description.maxRepetitionLevel :: Int@@ -252,7 +246,7 @@ goPages dict codec pType bs !acc | BS.null bs = return acc | otherwise = case parsePageHeader bs of- Left e -> error ("foldColumnPagesM: failed to parse page header: " ++ e)+ Left e -> error ("foldColumnDataPagesM: failed to parse page header: " ++ e) Right (rest, hdr) -> do let compSz = fromIntegral . unField $ hdr.ph_compressed_page_size uncmpSz = fromIntegral . unField $ hdr.ph_uncompressed_page_size@@ -277,8 +271,7 @@ decompressed <- liftIO $ decompressData uncmpSz codec pageData let (defLvls, repLvls, nPresent, valBytes) = readLevelsV1 n maxDef maxRep decompressed- triple = (decoder dict enc nPresent valBytes, defLvls, repLvls)- acc' <- step acc triple+ acc' <- step acc (dict, enc, nPresent, valBytes, defLvls, repLvls) goPages dict codec pType rest' acc' DATA_PAGE_V2 _ -> do let dph2 =@@ -289,8 +282,7 @@ enc = unField dph2.dph2_encoding defLen = unField dph2.dph2_definition_levels_byte_length repLen = unField dph2.dph2_repetition_levels_byte_length- -- V2: levels are never compressed; only the value- -- payload is (optionally) compressed.+ -- V2 levels are never compressed, only the values. isCompressed = fromMaybe True (unField dph2.dph2_is_compressed) (defLvls, repLvls, nPresent, compValBytes) = readLevelsV2 n maxDef maxRep repLen defLen pageData@@ -298,49 +290,142 @@ if isCompressed then liftIO $ decompressData uncmpSz codec compValBytes else pure compValBytes- let triple = (decoder dict enc nPresent valBytes, defLvls, repLvls)- acc' <- step acc triple+ acc' <- step acc (dict, enc, nPresent, valBytes, defLvls, repLvls) goPages dict codec pType rest' acc' INDEX_PAGE _ -> goPages dict codec pType rest' acc --- ------------------------------------------------------------------------------ Page header parsing--- ---------------------------------------------------------------------------+-- | Left-fold over per-page value triples with a given decoder.+{-# INLINEABLE foldColumnPagesM #-}+foldColumnPagesM ::+ forall m v a acc.+ (RandomAccess m, MonadIO m, VG.Vector v a) =>+ ColumnDescription ->+ (Maybe DictVals -> Encoding -> Int -> BS.ByteString -> v a) ->+ [ColumnChunk] ->+ (acc -> (v a, VU.Vector Int, VU.Vector Int) -> m acc) ->+ acc ->+ m acc+foldColumnPagesM description decoder chunks step =+ foldColumnDataPagesM description chunks $+ \acc (dict, enc, nPresent, valBytes, defLvls, repLvls) ->+ step acc (decoder dict enc nPresent valBytes, defLvls, repLvls) -parsePageHeader :: BS.ByteString -> Either String (BS.ByteString, PageHeader)-parsePageHeader = decodeWithLeftovers Pinch.compactProtocol+-- | Append a non-nullable BYTE_ARRAY page's strings into a 'TextBuilder'.+appendStringPageIO ::+ TextBuilder RealWorld ->+ Maybe DictVals ->+ Encoding ->+ Int ->+ BS.ByteString ->+ IO ()+appendStringPageIO builder mDict enc nPresent bs = case enc of+ PLAIN _ -> BSU.unsafeUseAsCStringLen bs $ \(cptr, _len) -> do+ let p = castPtr cptr :: Ptr Word8+ go !_ !i | i >= nPresent = pure ()+ go !off !i = do+ len <- readLenLE p off+ stToIO (appendTextSliceFromPtr builder (p `plusPtr` (off + 4)) len)+ go (off + 4 + len) (i + 1)+ go 0 0+ RLE_DICTIONARY _ -> dictAppend+ PLAIN_DICTIONARY _ -> dictAppend+ _ -> error ("appendStringPageIO: unsupported encoding " ++ show enc)+ where+ dictAppend = case mDict of+ Just (DText starts lens dictBytes) ->+ BSU.unsafeUseAsCStringLen dictBytes $ \(cptr, _) -> do+ let p = castPtr cptr :: Ptr Word8+ (idxs, _) = decodeDictIndices nPresent bs+ VU.forM_ idxs $ \i ->+ stToIO+ ( appendTextSliceFromPtr+ builder+ (p `plusPtr` VU.unsafeIndex starts i)+ (VU.unsafeIndex lens i)+ )+ Just d -> error ("appendStringPageIO: wrong dict type, got " ++ show d)+ Nothing -> error "appendStringPageIO: dictionary-encoded page but no dictionary seen" --- ------------------------------------------------------------------------------ Batch value readers--- ---------------------------------------------------------------------------+-- | Read a little-endian 4-byte length prefix from a raw page pointer.+readLenLE :: Ptr Word8 -> Int -> IO Int+readLenLE p off = do+ b0 <- peekByteOff p off :: IO Word8+ b1 <- peekByteOff p (off + 1) :: IO Word8+ b2 <- peekByteOff p (off + 2) :: IO Word8+ b3 <- peekByteOff p (off + 3) :: IO Word8+ pure $+ fromIntegral b0+ .|. (fromIntegral b1 `shiftL` 8)+ .|. (fromIntegral b2 `shiftL` 16)+ .|. (fromIntegral b3 `shiftL` 24)+{-# INLINE readLenLE #-} -readNBool :: Int -> BS.ByteString -> [Bool]-readNBool count bs =- let totalBytes = (count + 7) `div` 8- bits =- concatMap- (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7])- (BS.unpack (BS.take totalBytes bs))- in take count bits+-- | Append a nullable BYTE_ARRAY page into a 'TextBuilder'.+appendNullableStringPageIO ::+ TextBuilder RealWorld ->+ Int ->+ Maybe DictVals ->+ Encoding ->+ Int ->+ BS.ByteString ->+ VU.Vector Int ->+ IO ()+appendNullableStringPageIO builder maxDef mDict enc nPresent bs defs = case enc of+ PLAIN _ -> BSU.unsafeUseAsCStringLen bs $ \(cptr, _len) -> do+ let p = castPtr cptr :: Ptr Word8+ go !_ !i | i >= nDefs = pure ()+ go !off !i+ | VU.unsafeIndex defs i == maxDef = do+ len <- readLenLE p off+ stToIO (appendTextSliceFromPtr builder (p `plusPtr` (off + 4)) len)+ go (off + 4 + len) (i + 1)+ | otherwise = stToIO (appendNull builder) >> go off (i + 1)+ go 0 0+ RLE_DICTIONARY _ -> dictAppend+ PLAIN_DICTIONARY _ -> dictAppend+ _ -> error ("appendNullableStringPageIO: unsupported encoding " ++ show enc)+ where+ nDefs = VU.length defs+ dictAppend = case mDict of+ Just (DText starts lens dictBytes) -> BSU.unsafeUseAsCStringLen dictBytes $ \(cptr, _) -> do+ let p = castPtr cptr :: Ptr Word8+ (idxs, _) = decodeDictIndices nPresent bs+ go !_ !i | i >= nDefs = pure ()+ go !j !i+ | VU.unsafeIndex defs i == maxDef = do+ let d = VU.unsafeIndex idxs j+ stToIO+ ( appendTextSliceFromPtr+ builder+ (p `plusPtr` VU.unsafeIndex starts d)+ (VU.unsafeIndex lens d)+ )+ go (j + 1) (i + 1)+ | otherwise = stToIO (appendNull builder) >> go j (i + 1)+ go 0 0+ Just d -> error ("appendNullableStringPageIO: wrong dict type, got " ++ show d)+ Nothing ->+ error+ "appendNullableStringPageIO: dictionary-encoded page but no dictionary seen" +parsePageHeader :: BS.ByteString -> Either String (BS.ByteString, PageHeader)+parsePageHeader = decodeWithLeftovers Pinch.compactProtocol+ readNInt32 :: Int -> BS.ByteString -> VU.Vector Int32-readNInt32 n bs = VU.generate n $ \i -> littleEndianInt32 (BS.drop (4 * i) bs)+readNInt32 n bs = VU.generate n $ \i -> fromIntegral (word32At bs (4 * i)) readNInt64 :: Int -> BS.ByteString -> VU.Vector Int64-readNInt64 n bs = VU.generate n $ \i ->- fromIntegral (littleEndianWord64 (BS.drop (8 * i) bs))+readNInt64 n bs = VU.generate n $ \i -> fromIntegral (word64At bs (8 * i)) readNInt96 :: Int -> BS.ByteString -> [UTCTime] readNInt96 0 _ = [] readNInt96 n bs = int96ToUTCTime (BS.take 12 bs) : readNInt96 (n - 1) (BS.drop 12 bs) readNFloat :: Int -> BS.ByteString -> VU.Vector Float-readNFloat n bs = VU.generate n $ \i ->- castWord32ToFloat (littleEndianWord32 (BS.drop (4 * i) bs))+readNFloat n bs = VU.generate n $ \i -> castWord32ToFloat (word32At bs (4 * i)) readNDouble :: Int -> BS.ByteString -> VU.Vector Double-readNDouble n bs = VU.generate n $ \i ->- castWord64ToDouble (littleEndianWord64 (BS.drop (8 * i) bs))+readNDouble n bs = VU.generate n $ \i -> castWord64ToDouble (word64At bs (8 * i)) readNTexts :: Int -> BS.ByteString -> [T.Text] readNTexts 0 _ = []@@ -354,3 +439,100 @@ readNFixedTexts len n bs = decodeUtf8Lenient (BS.take len bs) : readNFixedTexts len (n - 1) (BS.drop len bs)++-- | Decode a required numeric column chunk straight into @out[off ..]@.+decodeRequiredNumericInto ::+ forall a m.+ (VU.Unbox a, RandomAccess m, MonadIO m) =>+ (DictVals -> Maybe (VU.Vector a)) ->+ (BS.ByteString -> Int -> a) ->+ Int ->+ ColumnDescription ->+ ColumnChunk ->+ VUM.IOVector a ->+ Int ->+ m Int+decodeRequiredNumericInto pick plainAt width description chunk out off0 = do+ (codec, pType, rawBytes) <- readChunkBytes chunk+ goPages Nothing codec pType rawBytes off0+ where+ goPages ::+ Maybe (VU.Vector a) ->+ CompressionCodec ->+ ThriftType ->+ BS.ByteString ->+ Int ->+ m Int+ goPages dict codec pType bs !off+ | BS.null bs = return off+ | otherwise = case parsePageHeader bs of+ Left e -> error ("decodeRequiredNumericInto: failed to parse page header: " ++ e)+ Right (rest, hdr) -> do+ let compSz = fromIntegral . unField $ hdr.ph_compressed_page_size+ uncmpSz = fromIntegral . unField $ hdr.ph_uncompressed_page_size+ (pageData, rest') = BS.splitAt compSz rest+ case unField hdr.ph_type of+ DICTIONARY_PAGE _ -> do+ let dictHdr =+ fromMaybe+ (error "DICTIONARY_PAGE: missing dictionary page header")+ (unField hdr.ph_dictionary_page_header)+ numVals = unField dictHdr.diph_num_values+ decompressed <- liftIO $ decompressData uncmpSz codec pageData+ let d = readDictVals pType decompressed numVals description.typeLength+ case pick d of+ Nothing -> error "decodeRequiredNumericInto: dictionary page mistyped"+ Just !ds -> goPages (Just ds) codec pType rest' off+ DATA_PAGE _ -> do+ let dph =+ fromMaybe+ (error "DATA_PAGE: missing data page header")+ (unField hdr.ph_data_page_header)+ n = fromIntegral . unField $ dph.dph_num_values+ enc = unField dph.dph_encoding+ decompressed <- liftIO $ decompressData uncmpSz codec pageData+ liftIO (writePage dict enc n decompressed off)+ goPages dict codec pType rest' (off + n)+ DATA_PAGE_V2 _ -> do+ let dph2 =+ fromMaybe+ (error "DATA_PAGE_V2: missing data page header v2")+ (unField hdr.ph_data_page_header_v2)+ n = fromIntegral . unField $ dph2.dph2_num_values+ enc = unField dph2.dph2_encoding+ levelBytes =+ fromIntegral (unField dph2.dph2_definition_levels_byte_length)+ + fromIntegral (unField dph2.dph2_repetition_levels_byte_length)+ isCompressed = fromMaybe True (unField dph2.dph2_is_compressed)+ compValBytes = BS.drop levelBytes pageData+ valBytes <-+ if isCompressed+ then liftIO $ decompressData (uncmpSz - levelBytes) codec compValBytes+ else pure compValBytes+ liftIO (writePage dict enc n valBytes off)+ goPages dict codec pType rest' (off + n)+ INDEX_PAGE _ -> goPages dict codec pType rest' off++ writePage ::+ Maybe (VU.Vector a) -> Encoding -> Int -> BS.ByteString -> Int -> IO ()+ writePage dict enc n vals off = case enc of+ PLAIN _ ->+ let go !i+ | i >= n = pure ()+ | otherwise = do+ VUM.unsafeWrite out (off + i) (plainAt vals (width * i))+ go (i + 1)+ in go 0+ RLE_DICTIONARY _ -> gather+ PLAIN_DICTIONARY _ -> gather+ _ -> error ("decodeRequiredNumericInto: unsupported encoding " ++ show enc)+ where+ gather = case dict of+ Nothing -> error "decodeRequiredNumericInto: dictionary page missing"+ Just ds -> case BS.uncons vals of+ Nothing -> error "decodeRequiredNumericInto: empty index stream"+ Just (w0, rest) ->+ stToIO $+ rleHybridFold (fromIntegral w0) n rest $ \i w ->+ VUM.unsafeWrite out (off + i) (VU.unsafeIndex ds (fromIntegral w))+{-# INLINE decodeRequiredNumericInto #-}
src/DataFrame/IO/Parquet/Schema.hs view
@@ -3,15 +3,7 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeApplications #-} -{- |-Module : DataFrame.IO.Parquet.Schema-License : MIT--Helpers for converting a parquet schema (a list of 'SchemaElement') into an-empty 'DataFrame' whose columns have the right types but no rows. Used by-'DataFrame.TH.declareColumnsFromParquetFile' and by any tooling that wants-to inspect a parquet schema as a 'DataFrame'.--}+-- | Converts a parquet schema into an empty 'DataFrame'. module DataFrame.IO.Parquet.Schema ( schemaToEmptyDataFrame, schemaElemToColumn,@@ -33,21 +25,16 @@ unField, ) import DataFrame.Internal.Column (Column, fromList)-import DataFrame.Internal.DataFrame (DataFrame, fromNamedColumns)+import DataFrame.Internal.DataFrame (DataFrame, fromColumns) -{- | Build an empty 'DataFrame' from a flat list of parquet 'SchemaElement's.-Only leaf elements (those with no children) become columns. Columns whose-name is in @nullableCols@ are typed as @Maybe a@; the rest are typed as @a@.--}+-- | Build an empty 'DataFrame' from a flat list of 'SchemaElement's. schemaToEmptyDataFrame :: S.Set T.Text -> [SchemaElement] -> DataFrame schemaToEmptyDataFrame nullableCols elems = let leafElems = filter (\e -> Maybe.fromMaybe 0 (unField e.num_children) == 0) elems- in fromNamedColumns (map (schemaElemToColumn nullableCols) leafElems)+ in fromColumns (map (schemaElemToColumn nullableCols) leafElems) -{- | Convert a single parquet 'SchemaElement' into a named empty 'Column',-picking a nullable or non-nullable representation based on @nullableCols@.--}+-- | Convert a single 'SchemaElement' into a named empty 'Column'. schemaElemToColumn :: S.Set T.Text -> SchemaElement -> (T.Text, Column) schemaElemToColumn nullableCols element = let colName = unField element.name
src/DataFrame/IO/Parquet/Seeking.hs view
@@ -1,7 +1,4 @@-{- | This module contains low-level utilities around file seeking.--later this module can be renamed / moved to an internal module.--}+-- | Low-level utilities around file seeking. module DataFrame.IO.Parquet.Seeking ( SeekableHandle (getSeekableHandle), SeekMode (..),@@ -22,28 +19,10 @@ import Data.Int import System.IO -{- | This handle carries a proof that it must be seekable.-Note: Handle and SeekableHandle are not thread safe, should not be-shared across threads, beaware when running parallel/concurrent code.--Not seekable:- - stdin / stdout- - pipes / FIFOs--But regular files are always seekable. Parquet fundamentally wants random-access, a non-seekable source will not support effecient access without-buffering the entire file.--}+-- | A handle that carries a proof that it is seekable. newtype SeekableHandle = SeekableHandle {getSeekableHandle :: Handle} -{- | If we truely want to support non-seekable files, we need to also consider the case-to buffer the entire file in memory.--Not thread safe, contains mutable reference (as Handle already is).--If we need concurrent / parallel parsing or something, we need to read into ByteString-first, not sharing the same handle.--}+-- | Handle that is either seekable or buffered in memory. data FileBufferedOrSeekable = FileBuffered !(IORef Int64) !BS.ByteString | FileSeekable !SeekableHandle@@ -57,9 +36,7 @@ -- | For testing only type ForceNonSeekable = Maybe Bool -{- | Smart constructor for FileBufferedOrSeekable, tries to keep in the seekable case-if possible.--}+-- | Smart constructor that stays in the seekable case when possible. mkFileBufferedOrSeekable :: ForceNonSeekable -> Handle -> IO FileBufferedOrSeekable mkFileBufferedOrSeekable forceNonSeek h = do@@ -68,11 +45,7 @@ then FileBuffered <$> newIORef 0 <*> BS.hGetContents h else pure $ FileSeekable $ SeekableHandle h -{- | With / bracket pattern for FileBufferedOrSeekable--Warning: do not return the FileBufferedOrSeekable outside the scope of the action as-it will be closed.--}+-- | Bracket-style access to a 'FileBufferedOrSeekable'. withFileBufferedOrSeekable :: ForceNonSeekable -> FilePath ->@@ -83,10 +56,7 @@ fbos <- mkFileBufferedOrSeekable forceNonSeek h action fbos -{- | Read the last @n@ bytes, useful for reading metadata without loading the-entire file. Uses 'BS.hGet' (not @hGetContents@) so the handle stays open for-the subsequent column-chunk reads.--}+-- | Read the last @n@ bytes, e.g. to get metadata without the whole file. readLastBytes :: Integer -> FileBufferedOrSeekable -> IO BS.ByteString readLastBytes n (FileSeekable sh) = do let h = getSeekableHandle sh
src/DataFrame/IO/Parquet/Thrift.hs view
@@ -12,8 +12,7 @@ import Pinch (Enumeration, Field, Pinchable (..)) import qualified Pinch --- Primitive Parquet Types--- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L32+-- Primitive Parquet Types (parquet.thrift#L32) data ThriftType = BOOLEAN (Enumeration 0) | INT32 (Enumeration 1)@@ -27,7 +26,7 @@ instance Pinchable ThriftType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L183+-- parquet.thrift#L183 data FieldRepetitionType = REQUIRED (Enumeration 0) | OPTIONAL (Enumeration 1)@@ -36,11 +35,11 @@ instance Pinchable FieldRepetitionType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L203+-- parquet.thrift#L203 data Encoding = PLAIN (Enumeration 0) | -- GROUP_VAR_INT Encoding was never used- -- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L578+ -- parquet.thrift#L578 PLAIN_DICTIONARY (Enumeration 2) | RLE (Enumeration 3) | BIT_PACKED (Enumeration 4)@@ -53,7 +52,7 @@ instance Pinchable Encoding --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L244+-- parquet.thrift#L244 data CompressionCodec = UNCOMPRESSED (Enumeration 0) | SNAPPY (Enumeration 1)@@ -67,7 +66,7 @@ instance Pinchable CompressionCodec --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L261+-- parquet.thrift#L261 data PageType = DATA_PAGE (Enumeration 0) | INDEX_PAGE (Enumeration 1)@@ -77,7 +76,7 @@ instance Pinchable PageType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L271+-- parquet.thrift#L271 data BoundaryOrder = UNORDERED (Enumeration 0) | ASCENDING (Enumeration 1)@@ -86,12 +85,9 @@ instance Pinchable BoundaryOrder --- Logical type annotations--- Empty structs can't use deriving Generic with Pinch, so we use a unit-like workaround.--- We represent empty structs as a newtype over () with a manual Pinchable instance.+-- Logical type annotations; empty structs use manual Pinchable instances. --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L283--- struct StringType {}+-- parquet.thrift#L283: struct StringType {} data StringType = StringType deriving (Eq, Show) instance Pinchable StringType where type Tag StringType = Pinch.TStruct@@ -158,7 +154,7 @@ pinch _ = Pinch.struct [] unpinch _ = pure VariantType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L290+-- parquet.thrift#L290 data TimeUnit = MILLIS (Field 1 MilliSeconds) | MICROS (Field 2 MicroSeconds)@@ -185,7 +181,7 @@ pinch _ = Pinch.struct [] unpinch _ = pure NanoSeconds --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L317+-- parquet.thrift#L317 data DecimalType = DecimalType { decimal_scale :: Field 1 Int32@@ -195,7 +191,7 @@ instance Pinchable DecimalType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L328+-- parquet.thrift#L328 data IntType = IntType { int_bitWidth :: Field 1 Int8@@ -205,7 +201,7 @@ instance Pinchable IntType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L338+-- parquet.thrift#L338 data TimeType = TimeType { time_isAdjustedToUTC :: Field 1 Bool@@ -215,7 +211,7 @@ instance Pinchable TimeType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L349+-- parquet.thrift#L349 data TimestampType = TimestampType { timestamp_isAdjustedToUTC :: Field 1 Bool@@ -225,8 +221,7 @@ instance Pinchable TimestampType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L360--- union LogicalType+-- parquet.thrift#L360: union LogicalType data LogicalType = LT_STRING (Field 1 StringType) | LT_MAP (Field 2 MapType)@@ -247,7 +242,7 @@ instance Pinchable LogicalType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L270+-- parquet.thrift#L270 data ConvertedType = UTF8 (Enumeration 0) | MAP (Enumeration 1)@@ -275,10 +270,10 @@ instance Pinchable ConvertedType --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L505+-- parquet.thrift#L505 data SchemaElement = SchemaElement- { schematype :: Field 1 (Maybe ThriftType) -- called just type in parquet.thrift+ { schematype :: Field 1 (Maybe ThriftType) -- named type upstream , type_length :: Field 2 (Maybe Int32) , repetition_type :: Field 3 (Maybe FieldRepetitionType) , name :: Field 4 Text@@ -293,7 +288,7 @@ instance Pinchable SchemaElement --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L560+-- parquet.thrift#L560 data Statistics = Statistics { stats_max :: Field 1 (Maybe ByteString)@@ -309,7 +304,7 @@ instance Pinchable Statistics --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L600+-- parquet.thrift#L600 data PageEncodingStats = PageEncodingStats { pes_page_type :: Field 1 PageType@@ -320,7 +315,7 @@ instance Pinchable PageEncodingStats --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L614+-- parquet.thrift#L614 data ColumnMetaData = ColumnMetaData { cmd_type :: Field 1 ThriftType@@ -343,14 +338,14 @@ instance Pinchable ColumnMetaData --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L875+-- parquet.thrift#L875 data EncryptionWithFooterKey = EncryptionWithFooterKey deriving (Eq, Show) instance Pinchable EncryptionWithFooterKey where type Tag EncryptionWithFooterKey = Pinch.TStruct pinch _ = Pinch.struct [] unpinch _ = pure EncryptionWithFooterKey --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L883+-- parquet.thrift#L883 data EncryptionWithColumnKey = EncryptionWithColumnKey { ewck_path_in_schema :: Field 1 [Text]@@ -360,8 +355,7 @@ instance Pinchable EncryptionWithColumnKey --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L893--- union ColumnCryptoMetaData+-- parquet.thrift#L893: union ColumnCryptoMetaData data ColumnCryptoMetaData = CCM_ENCRYPTION_WITH_FOOTER_KEY (Field 1 EncryptionWithFooterKey) | CCM_ENCRYPTION_WITH_COLUMN_KEY (Field 2 EncryptionWithColumnKey)@@ -369,7 +363,7 @@ instance Pinchable ColumnCryptoMetaData --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L899+-- parquet.thrift#L899 data ColumnChunk = ColumnChunk { cc_file_path :: Field 1 (Maybe Text)@@ -386,7 +380,7 @@ instance Pinchable ColumnChunk --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L940+-- parquet.thrift#L940 data SortingColumn = SortingColumn { sc_column_idx :: Field 1 Int32@@ -397,7 +391,7 @@ instance Pinchable SortingColumn --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L958+-- parquet.thrift#L958 data RowGroup = RowGroup { rg_columns :: Field 1 [ColumnChunk]@@ -412,7 +406,7 @@ instance Pinchable RowGroup --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L980+-- parquet.thrift#L980 data KeyValue = KeyValue { kv_key :: Field 1 Text@@ -422,8 +416,7 @@ instance Pinchable KeyValue --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L990--- union ColumnOrder+-- parquet.thrift#L990: union ColumnOrder newtype ColumnOrder = TYPE_ORDER (Field 1 TypeDefinedOrder) deriving (Eq, Show, Generic)@@ -437,7 +430,7 @@ pinch _ = Pinch.struct [] unpinch _ = pure TypeDefinedOrder --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1094+-- parquet.thrift#L1094 data AesGcmV1 = AesGcmV1 { aes_gcm_v1_aad_prefix :: Field 1 (Maybe ByteString)@@ -448,7 +441,7 @@ instance Pinchable AesGcmV1 --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1107+-- parquet.thrift#L1107 data AesGcmCtrV1 = AesGcmCtrV1 { aes_gcm_ctr_v1_aad_prefix :: Field 1 (Maybe ByteString)@@ -459,8 +452,7 @@ instance Pinchable AesGcmCtrV1 --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1118--- union EncryptionAlgorithm+-- parquet.thrift#L1118: union EncryptionAlgorithm data EncryptionAlgorithm = AES_GCM_V1 (Field 1 AesGcmV1) | AES_GCM_CTR_V1 (Field 2 AesGcmCtrV1)@@ -468,7 +460,7 @@ instance Pinchable EncryptionAlgorithm --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1001+-- parquet.thrift#L1001 data PageLocation = PageLocation { pl_offset :: Field 1 Int64@@ -479,7 +471,7 @@ instance Pinchable PageLocation --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1017+-- parquet.thrift#L1017 data OffsetIndex = OffsetIndex { oi_page_locations :: Field 1 [PageLocation]@@ -489,7 +481,7 @@ instance Pinchable OffsetIndex --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1033+-- parquet.thrift#L1033 data ColumnIndex = ColumnIndex { ci_null_pages :: Field 1 [Bool]@@ -504,7 +496,7 @@ instance Pinchable ColumnIndex --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1248+-- parquet.thrift#L1248 data DataPageHeader = DataPageHeader { dph_num_values :: Field 1 Int32@@ -563,7 +555,7 @@ instance Pinchable PageHeader --- https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift#L1277+-- parquet.thrift#L1277 data FileMetadata = FileMetadata { version :: Field 1 Int32
src/DataFrame/IO/Parquet/Time.hs view
@@ -42,8 +42,7 @@ year = 100 * b + d - 4800 + m `div` 10 in fromGregorian year (fromIntegral month) (fromIntegral day) --- I include this here even though it's unused because we'll likely use--- it for the writer. Since int96 is deprecated this is only included for completeness anyway.+-- Unused for now; kept for the writer since int96 is deprecated. utcTimeToInt96 :: UTCTime -> BS.ByteString utcTimeToInt96 (UTCTime day diffTime) = let julianDay = dayToJulianDay day
src/DataFrame/IO/Parquet/Utils.hs view
@@ -44,15 +44,11 @@ import DataFrame.Internal.Column ( Column (..), Columnable,- buildBitmapFromValid, fromList, )+import DataFrame.Internal.Column.Bitmap (buildBitmapFromValid) -{- | A left-fold driver over a column's per-page triples-@(values, def-levels, rep-levels)@, as produced by-'DataFrame.IO.Parquet.Page.foldColumnPagesM'. Generic in the accumulator so a-single driver serves every fold strategy below.--}+-- | A left-fold driver over a column's per-page triples. type PageFold m v a = forall acc. (acc -> (v a, VU.Vector Int, VU.Vector Int) -> m acc) -> acc -> m acc@@ -161,12 +157,8 @@ childLeaves = go children subPath False in childLeaves ++ go rest path skipThis -{- | Fold a column's value pages into a non-nullable 'Column'.--Pre-allocates a mutable vector of @totalRows@ and fills it page-by-page via a-single streaming left fold ('PageFold'), avoiding any intermediate list or-concatenation allocation. Only one page's values are live at a time.--}+-- | Fold a column's value pages into a non-nullable 'Column'.+{-# INLINEABLE foldNonNullable #-} foldNonNullable :: forall m a. (MonadIO m, Columnable a) =>@@ -186,6 +178,7 @@ v <- liftIO $ VB.unsafeFreeze mv return (BoxedColumn Nothing v) +{-# INLINEABLE foldNonNullableUnboxed #-} foldNonNullableUnboxed :: forall m a. (MonadIO m, Columnable a, VU.Unbox a) =>@@ -198,29 +191,15 @@ runPages ( \off (chunk, _, _) -> liftIO $ do let n = VU.length chunk- go i- | i >= n = return ()- | otherwise = do- VUM.unsafeWrite- mv- (off + i)- (VU.unsafeIndex chunk i)- go (i + 1)- go 0+ VU.unsafeCopy (VUM.unsafeSlice off n mv) chunk return (off + n) ) (0 :: Int) dat <- liftIO $ VU.unsafeFreeze mv return (UnboxedColumn Nothing dat) -{- | Fold a column's (values, def-levels) pages into a nullable 'Column'.--Pre-allocates the output buffer and a valid-mask vector of @totalRows@, then-scatters values inline during a single streaming left fold ('PageFold').--A 'hasNull' flag is accumulated during the scatter so the-'buildBitmapFromValid' call is skipped entirely when all values are present.--}+-- | Fold a column's value and def-level pages into a nullable 'Column'.+{-# INLINEABLE foldNullable #-} foldNullable :: forall m a. (MonadIO m, Columnable a) =>@@ -229,13 +208,7 @@ PageFold m VB.Vector a -> m Column foldNullable maxDef totalRows runPages = do- -- null slots hold an error thunk, guarded by bitmap.- --- -- IMPORTANT: 'VBM.unsafeWrite' for boxed vectors stores a *pointer* to- -- the value without evaluating it, so unsupported-encoding error thunks- -- would be silently swallowed into the column data and only fire lazily- -- when user code reads a cell. The '!v' bang pattern forces each value- -- to WHNF before the write, surfacing decoder errors immediately.+ -- The '!v' bang forces each value so decoder errors surface immediately. mvDat <- liftIO $ VBM.replicate totalRows (error "parquet: null slot accessed") mvValid <- liftIO (VUM.new totalRows :: IO (VUM.IOVector Word8))@@ -264,6 +237,7 @@ else return Nothing return (BoxedColumn maybeBm dat) +{-# INLINEABLE foldNullableUnboxed #-} foldNullableUnboxed :: forall m a. (MonadIO m, Columnable a, VU.Unbox a) =>@@ -299,17 +273,8 @@ else return Nothing return (UnboxedColumn maybeBm dat) -{- | Fold a stream of (values, def-levels, rep-levels) triples into a-repeated (list) 'Column' using Dremel-style level stitching.--The stitching function is selected by @maxRep@:-- * @maxRep == 1@ → 'stitchList' → @[Maybe [Maybe a]]@- * @maxRep == 2@ → 'stitchList2' → @[Maybe [Maybe [Maybe a]]]@- * @maxRep >= 3@ → 'stitchList3' → @[Maybe [Maybe [Maybe [Maybe a]]]]@--Threshold formula: @defT_r = maxDef - 2 * (maxRep - r)@.--}+-- | Fold page triples into a repeated (list) 'Column' by level stitching.+{-# INLINEABLE foldRepeated #-} foldRepeated :: forall m a. ( MonadIO m@@ -333,6 +298,7 @@ fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef allReps allDefs allVals) _ -> fromList (stitchList maxDef allReps allDefs allVals) +{-# INLINEABLE foldRepeatedUnboxed #-} foldRepeatedUnboxed :: forall m a. ( MonadIO m@@ -357,10 +323,7 @@ fromList (stitchList3 (maxDef - 4) (maxDef - 2) maxDef allReps allDefs allVals) _ -> fromList (stitchList maxDef allReps allDefs allVals) -{- | Collect all of a column's page triples into a list, in page order. Used by-the repeated/list folds, where Dremel level-stitching needs the full-concatenated rep/def/value arrays at once (so streaming gives no benefit here).--}+-- | Collect all of a column's page triples into a list, in page order. collectPages :: (Monad m) => PageFold m v a ->
+ src/DataFrame/IO/Parquet/Writer.hs view
@@ -0,0 +1,438 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE OverloadedRecordDot #-}+{-# LANGUAGE OverloadedStrings #-}++module DataFrame.IO.Parquet.Writer (+ writeParquet,+ writeParquetWithOptions,+ ParquetWriteOptions (..),+ WriterStrategy (..),+ defaultParquetWriteOptions,+ nativeTypeKeyPrefix,+ nativeTypeKeyValues,+) where++import Control.Exception (evaluate)+import Control.Monad (forM_, unless, when)+import Control.Monad.ST (RealWorld, ST, stToIO)+import qualified Data.ByteString as BS+import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef)+import Data.Int (Int64)+import Data.Maybe (fromJust)+import Data.Primitive.ByteArray (getSizeofMutableByteArray)+import Data.STRef (STRef, modifySTRef', newSTRef, readSTRef, writeSTRef)+import qualified Data.Text as T+import qualified Data.Vector as VB+import DataFrame.IO.Parquet.Thrift hiding (schema)+import DataFrame.IO.Parquet.Writer.DefLevels (+ DefLevels (..),+ flushDef,+ newDefLevels,+ pushDef,+ )+import DataFrame.IO.Parquet.Writer.Encoder (Encoder (..), buildEncoder)+import DataFrame.IO.Parquet.Writer.Metadata (+ magic,+ mkColumnChunk,+ mkDataPageHeader,+ mkRowGroup,+ mkSchemaElem,+ rootSchemaElement,+ writeFooter,+ )+import DataFrame.IO.Parquet.Writer.Options (+ ParquetWriteOptions (..),+ WriterStrategy (..),+ defaultParquetWriteOptions,+ )+import DataFrame.IO.Utils.RandomAccess (+ MemoryBuffer (..),+ WritableBinaryHandle,+ atomicallyWriteFile,+ bufferResidency,+ ensureCapacity,+ flushBufferToBuffer,+ flushBufferToFile,+ mallocBuffer,+ resetPosition,+ unsafeBufferView,+ withWritableBinaryFile,+ writeByteString,+ writeByteStringToFile,+ writeWord32LE,+ )+import DataFrame.Internal.Column (Column, columnTypeString, hasMissing)+import DataFrame.Internal.DataFrame (+ DataFrame,+ columnNames,+ dataframeDimensions,+ getColumn,+ )+import qualified Pinch+import qualified Snappy+import System.Directory (createDirectoryIfMissing)+import System.FilePath (takeDirectory)+import Text.Printf (printf)++data ParquetWriterState = ParquetWriterState+ { outputFileHandle :: !WritableBinaryHandle+ , columnChunks :: !(VB.Vector (ColumnChunkState RealWorld))+ , currentFileOffsetRef :: !(IORef Int64)+ , scratchBuffer :: !(MemoryBuffer RealWorld)+ , rowGroupMetadataRef :: !(IORef [RowGroup])+ , rowNumberRef :: !(IORef Int)+ }++data ColumnChunkState s = ColumnChunkState+ { columnName :: !T.Text+ , nullable :: !Bool+ , schema :: !SchemaElement+ , encoder :: !(Encoder s)+ , buffer :: !(MemoryBuffer s)+ , uncompressedBufferSize :: !(STRef s Int64)+ , pageState :: !(PageState s)+ }++data PageState s = PageState+ { pageBuffer :: !(MemoryBuffer s)+ , definitionLevels :: !(DefLevels s)+ , currentRowCount :: !(STRef s Int)+ }++writeParquet :: FilePath -> DataFrame -> IO ()+writeParquet = writeParquetWithOptions defaultParquetWriteOptions++writeParquetWithOptions :: ParquetWriteOptions -> FilePath -> DataFrame -> IO ()+writeParquetWithOptions options path df = do+ when (options.strategy == TwoPass) $+ error+ "The Two Pass Strategy for the Parquet Writer has not yet been implemented"+ case options.compressionCodec of+ UNCOMPRESSED _ -> pure ()+ SNAPPY _ -> pure ()+ other -> error ("writeParquet: unsupported codec " <> show other)+ let (totalRows, _) = dataframeDimensions df+ case options.maxRowsPerFile of+ Nothing -> do+ when (isShardPattern path) $+ error+ ( "writeParquet: path "+ <> show path+ <> " contains a '*' placeholder but maxRowsPerFile is not set"+ )+ writeShard options path df 0 totalRows+ Just rowsPerFile -> do+ when (rowsPerFile <= 0) $+ error "writeParquet: maxRowsPerFile must be positive"+ unless (isShardPattern path) $+ error+ ( "writeParquet: maxRowsPerFile requires a path with a '*' placeholder, got "+ <> show path+ )+ let starts = case [0, rowsPerFile .. totalRows - 1] of+ [] -> [0] -- empty frame still produces one (empty) shard+ ss -> ss+ forM_ (zip [0 ..] starts) $ \(shardIndex, start) -> do+ let shardPath = shardPathFor path shardIndex+ createDirectoryIfMissing True (takeDirectory shardPath)+ writeShard options shardPath df start (min totalRows (start + rowsPerFile))++isShardPattern :: FilePath -> Bool+isShardPattern = elem '*'++-- | Replace every @*@ in the pattern with a zero-padded shard index.+shardPathFor :: FilePath -> Int -> FilePath+shardPathFor pattern_ shardIndex =+ concatMap (\c -> if c == '*' then printf "%05d" shardIndex else [c]) pattern_++-- | Write rows @[startRow, endRow)@ of the frame to a single Parquet file.+writeShard ::+ ParquetWriteOptions ->+ FilePath ->+ DataFrame ->+ Int ->+ Int ->+ IO ()+writeShard options path_ df startRow endRow = do+ let names = columnNames df+ shardRows = max 0 (endRow - startRow)+ columnChunks_ <-+ stToIO $+ VB.fromList+ <$> mapM+ ( \columnName_ ->+ initColumnChunkState+ options+ columnName_+ (fromJust (getColumn columnName_ df))+ )+ names+ scratchBuffer_ <- stToIO (mallocBuffer (max 1 options.pageSize))+ atomicallyWriteFile path_ $ \path -> withWritableBinaryFile path $ \output -> do+ writeByteStringToFile output magic+ currentFileOffsetRef_ <- newIORef 4+ rowGroupMetadataRef_ <- newIORef []+ rowNumberRef_ <- newIORef 0+ let writerState =+ ParquetWriterState+ output+ columnChunks_+ currentFileOffsetRef_+ scratchBuffer_+ rowGroupMetadataRef_+ rowNumberRef_+ interval = max 1 options.batchRows+ subBatch = max 1 options.subBatchRows+ writeBatch :: Int -> Int -> IO ()+ writeBatch rowNum batchEnd+ | rowNum >= batchEnd = pure ()+ | otherwise = do+ let count = min subBatch (batchEnd - rowNum)+ VB.forM_ columnChunks_ $ \ccs -> do+ pageFull <- stToIO (writeRows options rowNum count ccs)+ when pageFull (flushPage options scratchBuffer_ ccs)+ modifyIORef' rowNumberRef_ (+ count)+ writeBatch (rowNum + count) batchEnd+ loop :: Int -> IO ()+ loop rowNum+ | rowNum >= endRow = pure ()+ | otherwise = do+ let batchEnd = rowNum + min interval (endRow - rowNum)+ writeBatch rowNum batchEnd+ size <- stToIO (bufferedSize columnChunks_)+ when (size >= options.rowGroupSize) $+ flushRowGroup options writerState+ loop batchEnd+ loop startRow+ flushRowGroup options writerState+ rowGroupMetadata <- reverse <$> readIORef rowGroupMetadataRef_+ let schemaElements =+ rootSchemaElement (VB.length columnChunks_)+ : VB.toList (VB.map schema columnChunks_)+ writeFooter+ output+ schemaElements+ shardRows+ rowGroupMetadata+ (nativeTypeKeyValues names df)++nativeTypeKeyPrefix :: T.Text+nativeTypeKeyPrefix = "dataframe.type."++-- | The type stamp for every column of @df@, as footer key-value pairs.+nativeTypeKeyValues :: [T.Text] -> DataFrame -> [(T.Text, T.Text)]+nativeTypeKeyValues names df =+ [ (nativeTypeKeyPrefix <> name, T.pack (columnTypeString col))+ | name <- names+ , Just col <- [getColumn name df]+ ]++-- | Encode @count@ rows from @firstRow@ into the column's page.+writeRows ::+ ParquetWriteOptions -> Int -> Int -> ColumnChunkState s -> ST s Bool+writeRows options firstRow count ccs = do+ let page = ccs.pageState+ buf = page.pageBuffer+ encode = ccs.encoder.encodeValue+ dl = page.definitionLevels+ end = firstRow + count++ pos0 <- readSTRef buf.positionRef+ let margin = options.pageSize+ arr0 <- ensureCapacity buf (pos0 + max margin (count * 64))+ size0 <- getSizeofMutableByteArray arr0++ let go !size !pos !row+ | row >= end = writeSTRef buf.positionRef pos+ | pos + margin > size = do+ -- Rare: buffer nearly full, grow it+ writeSTRef buf.positionRef pos+ arr' <-+ ensureCapacity+ buf+ (pos + max margin ((end - row) * 64))+ size' <- getSizeofMutableByteArray arr'+ go size' pos row+ | otherwise = do+ (pos', notNull) <- encode buf pos row+ when ccs.nullable $+ pushDef dl (if notNull then 1 else 0)+ go size pos' (row + 1)++ go size0 pos0 firstRow++ -- Batch bookkeeping: once per sub-batch instead of per value+ modifySTRef' page.currentRowCount (+ count)+ flushDef dl+ pageRes <- bufferResidency buf+ defRes <- bufferResidency dl.dlBuf+ pure (pageRes + defRes >= options.pageSize)++flushPage ::+ ParquetWriteOptions ->+ MemoryBuffer RealWorld ->+ ColumnChunkState RealWorld ->+ IO ()+flushPage options scratch columnChunkState = do+ let page = columnChunkState.pageState+ numPageRows <- stToIO (readSTRef page.currentRowCount)+ when (numPageRows > 0) $ do+ body <- stToIO $ do+ pos <- readSTRef page.pageBuffer.positionRef+ pos' <- columnChunkState.encoder.finishValues page.pageBuffer pos+ writeSTRef page.pageBuffer.positionRef pos'+ assemblePageBody scratch columnChunkState+ writeDataPage options.compressionCodec numPageRows body columnChunkState+ stToIO $ do+ resetPosition page.pageBuffer+ resetPosition page.definitionLevels.dlBuf+ resetPosition scratch+ writeSTRef page.currentRowCount 0++assemblePageBody ::+ MemoryBuffer s -> ColumnChunkState s -> ST s (MemoryBuffer s)+assemblePageBody scratch columnChunkState+ | not columnChunkState.nullable = pure columnChunkState.pageState.pageBuffer+ | otherwise = do+ let page = columnChunkState.pageState+ flushDef page.definitionLevels+ resetPosition scratch+ defLevelsSize <- bufferResidency page.definitionLevels.dlBuf+ writeWord32LE scratch (fromIntegral defLevelsSize)+ flushBufferToBuffer page.definitionLevels.dlBuf scratch+ flushBufferToBuffer page.pageBuffer scratch+ pure scratch++writeDataPage ::+ CompressionCodec ->+ Int ->+ MemoryBuffer RealWorld ->+ ColumnChunkState RealWorld ->+ IO ()+writeDataPage codec numPageRows body columnChunkState = do+ uncompressedPageSize <- stToIO (bufferResidency body)+ compressedBody <- case codec of+ UNCOMPRESSED _ -> pure Nothing+ SNAPPY _ -> do+ -- Compress now: the page buffer is reused for the next page.+ page <- unsafeBufferView body+ Just <$> evaluate (Snappy.compress page)+ other -> error ("writeParquet: unsupported codec " <> show other)+ let compressedPageSize = maybe uncompressedPageSize BS.length compressedBody+ headerBytes =+ Pinch.encode+ Pinch.compactProtocol+ (mkDataPageHeader numPageRows uncompressedPageSize compressedPageSize)+ writeByteString columnChunkState.buffer headerBytes+ case compressedBody of+ Nothing -> stToIO (flushBufferToBuffer body columnChunkState.buffer)+ Just bytes -> writeByteString columnChunkState.buffer bytes+ stToIO $+ modifySTRef'+ columnChunkState.uncompressedBufferSize+ (+ fromIntegral (BS.length headerBytes + uncompressedPageSize))++flushRowGroup :: ParquetWriteOptions -> ParquetWriterState -> IO ()+flushRowGroup options writerState = do+ rowNumber <- readIORef writerState.rowNumberRef+ when (rowNumber > 0) $ do+ VB.forM_+ writerState.columnChunks+ (flushPage options writerState.scratchBuffer)+ (reversedColumnChunks, totalCompressed, totalUncompressed) <-+ VB.foldM'+ ( \(acc, totalCompressedSize, totalUncompressedSize) columnChunkState -> do+ offset <- readIORef writerState.currentFileOffsetRef+ (compressedSize, uncompressedSize) <-+ stToIO $+ (,)+ <$> bufferResidency columnChunkState.buffer+ <*> readSTRef columnChunkState.uncompressedBufferSize+ flushBufferToFile+ writerState.outputFileHandle+ columnChunkState.buffer+ writeIORef+ writerState.currentFileOffsetRef+ (offset + fromIntegral compressedSize)+ stToIO (writeSTRef columnChunkState.uncompressedBufferSize 0)+ let columnChunk =+ mkColumnChunk+ options.compressionCodec+ columnChunkState.encoder.encType+ columnChunkState.columnName+ offset+ compressedSize+ uncompressedSize+ rowNumber+ pure+ ( columnChunk : acc+ , totalCompressedSize + fromIntegral compressedSize+ , totalUncompressedSize + uncompressedSize+ )+ )+ ([], 0 :: Int64, 0 :: Int64)+ writerState.columnChunks+ modifyIORef'+ writerState.rowGroupMetadataRef+ ( mkRowGroup+ (reverse reversedColumnChunks)+ totalCompressed+ totalUncompressed+ rowNumber+ :+ )+ writeIORef writerState.rowNumberRef 0++bufferedSize :: VB.Vector (ColumnChunkState s) -> ST s Int+bufferedSize =+ VB.foldM'+ ( \total columnChunkState -> do+ chunkSize <- bufferResidency columnChunkState.buffer+ valuesSize <-+ bufferResidency columnChunkState.pageState.pageBuffer+ defLevelsSize <-+ bufferResidency+ columnChunkState.pageState.definitionLevels.dlBuf+ pure (total + chunkSize + valuesSize + defLevelsSize)+ )+ 0++initColumnChunkState ::+ ParquetWriteOptions -> T.Text -> Column -> ST s (ColumnChunkState s)+initColumnChunkState options columnName_ column = do+ encoder_ <- buildEncoder column+ let nullable_ = hasMissing column+ schema_ =+ mkSchemaElem+ columnName_+ encoder_.encType+ nullable_+ encoder_.convertedType+ encoder_.logicalType+ bufferSize = max 1 options.pageSize+ -- Buffers start at page size and grow, as chunk sizes vary by encoding.+ buffer_ <- mallocBuffer bufferSize+ uncompressedBufferSize_ <- newSTRef 0+ pageState_ <- initPageState bufferSize+ pure+ ColumnChunkState+ { columnName = columnName_+ , nullable = nullable_+ , schema = schema_+ , encoder = encoder_+ , buffer = buffer_+ , uncompressedBufferSize = uncompressedBufferSize_+ , pageState = pageState_+ }++initPageState :: Int -> ST s (PageState s)+initPageState bufferSize = do+ pageBuffer_ <- mallocBuffer bufferSize+ definitionLevels_ <- newDefLevels+ currentRowCount_ <- newSTRef 0+ pure+ PageState+ { pageBuffer = pageBuffer_+ , definitionLevels = definitionLevels_+ , currentRowCount = currentRowCount_+ }
+ src/DataFrame/IO/Parquet/Writer/DefLevels.hs view
@@ -0,0 +1,62 @@+{-# LANGUAGE OverloadedRecordDot #-}++module DataFrame.IO.Parquet.Writer.DefLevels (+ DefLevels (..),+ newDefLevels,+ pushDef,+ flushDef,+) where++import Control.Monad (when)+import Control.Monad.ST (ST)+import Data.Bits (shiftL, shiftR, (.&.), (.|.))+import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef)+import Data.Word (Word64)+import DataFrame.IO.Utils.RandomAccess (MemoryBuffer, mallocBuffer, writeWord8)++data DefLevels s = DefLevels+ { dlBuf :: !(MemoryBuffer s)+ , dlValue :: !(STRef s Int)+ , dlCount :: !(STRef s Int)+ }++newDefLevels :: ST s (DefLevels s)+newDefLevels = DefLevels <$> mallocBuffer 64 <*> newSTRef 0 <*> newSTRef 0++pushDef :: DefLevels s -> Int -> ST s ()+pushDef dl value = do+ count <- readSTRef dl.dlCount+ if count == 0+ then writeSTRef dl.dlValue value >> writeSTRef dl.dlCount 1+ else do+ current <- readSTRef dl.dlValue+ if current == value+ then writeSTRef dl.dlCount (count + 1)+ else do+ writeDefRun dl current count+ writeSTRef dl.dlValue value+ writeSTRef dl.dlCount 1+{-# INLINE pushDef #-}++flushDef :: DefLevels s -> ST s ()+flushDef dl = do+ count <- readSTRef dl.dlCount+ when (count > 0) $ do+ value <- readSTRef dl.dlValue+ writeDefRun dl value count+ writeSTRef dl.dlCount 0+{-# INLINE flushDef #-}++writeDefRun :: DefLevels s -> Int -> Int -> ST s ()+writeDefRun dl value count = do+ writeLeb128 dl.dlBuf (fromIntegral (count `shiftL` 1))+ writeWord8 dl.dlBuf (fromIntegral value)+{-# INLINE writeDefRun #-}++writeLeb128 :: MemoryBuffer s -> Word64 -> ST s ()+writeLeb128 buffer value+ | value < 0x80 = writeWord8 buffer (fromIntegral value)+ | otherwise = do+ writeWord8 buffer (fromIntegral (value .&. 0x7f) .|. 0x80)+ writeLeb128 buffer (value `shiftR` 7)+{-# INLINE writeLeb128 #-}
+ src/DataFrame/IO/Parquet/Writer/Encoder.hs view
@@ -0,0 +1,316 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE OverloadedRecordDot #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module DataFrame.IO.Parquet.Writer.Encoder (+ Encoder (..),+ buildEncoder,+) where++import Control.Monad.ST (ST)+import Data.Bits (shiftL, (.|.))+import Data.Int (Int32, Int64)+import Data.Primitive.ByteArray (withMutableByteArrayContents, writeByteArray)+import Data.STRef (newSTRef, readSTRef, writeSTRef)+import qualified Data.Text as T+import qualified Data.Text.Array as TA+import Data.Text.Internal (Text (Text))+import Data.Time.Calendar (toModifiedJulianDay)+import Data.Time.Clock (UTCTime (UTCTime), diffTimeToPicoseconds)+import Data.Type.Equality (TestEquality (..), (:~:) (Refl))+import qualified Data.Vector as VB+import qualified Data.Vector.Unboxed as VU+import Data.Word (Word8)+import DataFrame.IO.Parquet.Thrift+import DataFrame.IO.Utils.RandomAccess (+ MemoryBuffer (..),+ ensureCapacity,+ writeInteger64At,+ writeWord32At,+ writeWord64At,+ )+import DataFrame.Internal.Column (+ Column (..),+ Columnable,+ columnTypeString,+ hasElemType,+ )+import DataFrame.Internal.Column.Bitmap (+ Bitmap,+ bitmapTestBit,+ )+import DataFrame.Internal.Data.PackedText (+ PackedTextData (..),+ offAt,+ selAt,+ )+import Foreign (plusPtr)+import GHC.Float (castDoubleToWord64, castFloatToWord32)+import Pinch (enum, putField)+import Type.Reflection (typeRep)++data Encoder s = Encoder+ { encType :: !ThriftType+ , convertedType :: !(Maybe ConvertedType)+ , logicalType :: !(Maybe LogicalType)+ , encodeValue :: !(MemoryBuffer s -> Int -> Int -> ST s (Int, Bool))+ , finishValues :: !(MemoryBuffer s -> Int -> ST s Int)+ }++buildEncoder :: Column -> ST s (Encoder s)+buildEncoder col+ | hasElemType @Int32 col =+ pure $+ scalarEncoder+ (INT32 enum)+ Nothing+ Nothing+ ( unboxedColumnWriter @Int32+ col+ (\buffer pos v -> write32 buffer pos (fromIntegral v))+ )+ | hasElemType @Int64 col =+ pure $+ scalarEncoder+ (INT64 enum)+ Nothing+ Nothing+ ( unboxedColumnWriter @Int64+ col+ (\buffer pos v -> write64 buffer pos (fromIntegral v))+ )+ -- Int is 32 or 64 bit by platform, so always write 64-bit integers.+ | hasElemType @Int col =+ pure $+ scalarEncoder+ (INT64 enum)+ Nothing+ Nothing+ ( unboxedColumnWriter @Int+ col+ (\buffer pos v -> write64 buffer pos (fromIntegral v))+ )+ | hasElemType @Integer col =+ pure $+ scalarEncoder+ (INT64 enum)+ Nothing+ Nothing+ (columnWriter @Integer col writeInteger64At)+ | hasElemType @Float col =+ pure $+ scalarEncoder+ (FLOAT enum)+ Nothing+ Nothing+ ( unboxedColumnWriter @Float+ col+ (\buffer pos v -> write32 buffer pos (castFloatToWord32 v))+ )+ | hasElemType @Double col =+ pure $+ scalarEncoder+ (DOUBLE enum)+ Nothing+ Nothing+ ( unboxedColumnWriter @Double+ col+ (\buffer pos v -> write64 buffer pos (castDoubleToWord64 v))+ )+ | hasElemType @Bool col = boolEncoder col+ | hasElemType @T.Text col = pure (textEncoder col)+ | hasElemType @UTCTime col = pure (timestampEncoder col)+ | otherwise =+ error ("writeParquet: unsupported column type " <> columnTypeString col)+ where+ write32 buffer pos value = do+ writeWord32At buffer pos value+ pure (pos + 4)+ write64 buffer pos value = do+ writeWord64At buffer pos value+ pure (pos + 8)++type ColumnWriter s = MemoryBuffer s -> Int -> Int -> ST s (Int, Bool)++scalarEncoder ::+ ThriftType ->+ Maybe ConvertedType ->+ Maybe LogicalType ->+ ColumnWriter s ->+ Encoder s+scalarEncoder tt conv logical encode =+ Encoder tt conv logical encode (\_ pos -> pure pos)++-- | Encode a column whose element type has no 'VU.Unbox' instance.+columnWriter ::+ forall a s.+ (Columnable a) =>+ Column ->+ (MemoryBuffer s -> Int -> a -> ST s Int) ->+ ColumnWriter s+columnWriter = columnWriterWith Nothing+{-# INLINE columnWriter #-}++-- | Encode a column whose element type has a 'VU.Unbox' instance.+unboxedColumnWriter ::+ forall a s.+ (Columnable a, VU.Unbox a) =>+ Column ->+ (MemoryBuffer s -> Int -> a -> ST s Int) ->+ ColumnWriter s+unboxedColumnWriter = columnWriterWith (Just VU.unsafeIndex)+{-# INLINE unboxedColumnWriter #-}++columnWriterWith ::+ forall a s.+ (Columnable a) =>+ Maybe (VU.Vector a -> Int -> a) ->+ Column ->+ (MemoryBuffer s -> Int -> a -> ST s Int) ->+ ColumnWriter s+columnWriterWith unboxedIndex col writePrim = case col of+ BoxedColumn bitmap (values :: VB.Vector b) ->+ case testEquality (typeRep @a) (typeRep @b) of+ Just Refl -> writeFrom bitmap (VB.unsafeIndex values)+ Nothing -> mismatch+ UnboxedColumn bitmap (values :: VU.Vector b) ->+ case testEquality (typeRep @a) (typeRep @b) of+ Just Refl ->+ writeFrom+ bitmap+ (maybe (VU.unsafeIndex values) ($ values) unboxedIndex)+ Nothing -> mismatch+ _ -> mismatch+ where+ writeFrom bitmap at buffer pos row+ | isPresent bitmap row = do+ pos' <- writePrim buffer pos (at row)+ pure (pos', True)+ | otherwise = pure (pos, False)+ mismatch =+ error+ ("writeParquet: incompatible column representation for " <> columnTypeString col)+{-# INLINE columnWriterWith #-}++isPresent :: Maybe Bitmap -> Int -> Bool+isPresent Nothing _ = True+isPresent (Just bitmap) row = bitmapTestBit bitmap row+{-# INLINE isPresent #-}++boolEncoder :: Column -> ST s (Encoder s)+boolEncoder col = do+ bitsRef <- newSTRef (0 :: Word8)+ countRef <- newSTRef (0 :: Int)+ let addBit buffer pos value = do+ bits <- readSTRef bitsRef+ count <- readSTRef countRef+ let bits' = if value then bits .|. ((1 :: Word8) `shiftL` count) else bits+ count' = count + 1+ if count' == 8+ then do+ arr <- readSTRef buffer.arrayRef+ writeByteArray arr pos bits'+ writeSTRef bitsRef 0+ writeSTRef countRef 0+ pure (pos + 1)+ else do+ writeSTRef bitsRef bits'+ writeSTRef countRef count'+ pure pos+ finish buffer pos = do+ count <- readSTRef countRef+ pos' <-+ if count > 0+ then do+ bits <- readSTRef bitsRef+ arr <- readSTRef buffer.arrayRef+ writeByteArray arr pos bits+ pure (pos + 1)+ else pure pos+ writeSTRef bitsRef 0+ writeSTRef countRef 0+ pure pos'+ pure+ ( Encoder+ (BOOLEAN enum)+ Nothing+ Nothing+ (unboxedColumnWriter @Bool col addBit)+ finish+ )++textEncoder :: Column -> Encoder s+textEncoder col =+ Encoder+ (BYTE_ARRAY enum)+ (Just (UTF8 enum))+ (Just (LT_STRING (putField StringType)))+ writePresent+ (\_ pos -> pure pos)+ where+ writePresent = case col of+ BoxedColumn bitmap (values :: VB.Vector a) ->+ case testEquality (typeRep @T.Text) (typeRep @a) of+ Just Refl -> writeBoxed bitmap values+ Nothing -> mismatch+ PackedText bitmap packed -> writePacked bitmap packed+ _ -> mismatch+ writeBoxed bitmap values buffer pos row+ | isPresent bitmap row = do+ let Text bytes offset count = VB.unsafeIndex values row+ pos' <- writeTextSlice buffer pos bytes offset count+ pure (pos', True)+ | otherwise = pure (pos, False)+ writePacked bitmap packed buffer pos row+ | isPresent bitmap row = do+ let baseRow = maybe row (`selAt` row) packed.ptSel+ start = offAt packed.ptOffsets baseRow+ end = offAt packed.ptOffsets (baseRow + 1)+ pos' <- writeTextSlice buffer pos packed.ptBytes start (end - start)+ pure (pos', True)+ | otherwise = pure (pos, False)+ writeTextSlice buffer pos bytes offset count = do+ writeSTRef buffer.positionRef pos+ _ <- ensureCapacity buffer (pos + 4 + count)+ writeWord32At buffer pos (fromIntegral count)+ arr <- readSTRef buffer.arrayRef+ withMutableByteArrayContents arr $ \ptr ->+ TA.copyToPointer bytes offset (ptr `plusPtr` (pos + 4)) count+ pure (pos + 4 + count)+ mismatch =+ error+ ("writeParquet: incompatible text representation for " <> columnTypeString col)++timestampEncoder :: Column -> Encoder s+timestampEncoder col =+ Encoder+ (INT64 enum)+ (Just (TIMESTAMP_MICROS enum))+ (Just timestampLogical)+ (columnWriter @UTCTime col writeMicros)+ (\_ pos -> pure pos)+ where+ writeMicros buffer pos t = do+ writeWord64At buffer pos (fromIntegral (utcToMicros t))+ pure (pos + 8)++timestampLogical :: LogicalType+timestampLogical =+ LT_TIMESTAMP+ ( putField+ TimestampType+ { timestamp_isAdjustedToUTC = putField True+ , timestamp_unit = putField (MICROS (putField MicroSeconds))+ }+ )++utcToMicros :: UTCTime -> Int64+utcToMicros (UTCTime day dt) =+ fromIntegral+ ( (toModifiedJulianDay day - 40587) * 86400 * 1000000+ + diffTimeToPicoseconds dt `div` 1000000+ )+{-# INLINE utcToMicros #-}
+ src/DataFrame/IO/Parquet/Writer/Metadata.hs view
@@ -0,0 +1,176 @@+{-# LANGUAGE OverloadedStrings #-}++module DataFrame.IO.Parquet.Writer.Metadata (+ mkSchemaElem,+ rootSchemaElement,+ mkDataPageHeader,+ mkColumnChunk,+ mkRowGroup,+ writeFooter,+ magic,+) where++import Control.Monad.ST (stToIO)+import qualified Data.ByteString as BS+import Data.Int (Int64)+import qualified Data.Text as T+import DataFrame.IO.Parquet.Thrift+import DataFrame.IO.Utils.RandomAccess (+ WritableBinaryHandle,+ flushBufferToFile,+ mallocBuffer,+ writeByteString,+ writeWord32LE,+ )+import Pinch (enum, putField)+import qualified Pinch++mkDataPageHeader :: Int -> Int -> Int -> PageHeader+mkDataPageHeader rows uncompressedSize compressedSize =+ PageHeader+ { ph_type = putField (DATA_PAGE enum)+ , ph_uncompressed_page_size = putField (fromIntegral uncompressedSize)+ , ph_compressed_page_size = putField (fromIntegral compressedSize)+ , ph_crc = putField Nothing+ , ph_data_page_header = putField (Just dph)+ , ph_index_page_header = putField Nothing+ , ph_dictionary_page_header = putField Nothing+ , ph_data_page_header_v2 = putField Nothing+ }+ where+ dph =+ DataPageHeader+ { dph_num_values = putField (fromIntegral rows)+ , dph_encoding = putField (PLAIN enum)+ , dph_definition_level_encoding = putField (RLE enum)+ , dph_repetition_level_encoding = putField (RLE enum)+ , dph_statistics = putField Nothing+ }++mkSchemaElem ::+ T.Text ->+ ThriftType ->+ Bool ->+ Maybe ConvertedType ->+ Maybe LogicalType ->+ SchemaElement+mkSchemaElem elementName elementType nullable converted logical =+ SchemaElement+ { schematype = putField (Just elementType)+ , type_length = putField Nothing+ , repetition_type =+ putField (Just (if nullable then OPTIONAL enum else REQUIRED enum))+ , name = putField elementName+ , num_children = putField Nothing+ , converted_type = putField converted+ , scale = putField Nothing+ , precision = putField Nothing+ , field_id = putField Nothing+ , logicalType = putField logical+ }++rootSchemaElement :: Int -> SchemaElement+rootSchemaElement count =+ SchemaElement+ { schematype = putField Nothing+ , type_length = putField Nothing+ , repetition_type = putField Nothing+ , name = putField "schema"+ , num_children = putField (Just (fromIntegral count))+ , converted_type = putField Nothing+ , scale = putField Nothing+ , precision = putField Nothing+ , field_id = putField Nothing+ , logicalType = putField Nothing+ }++mkColumnChunk ::+ CompressionCodec ->+ ThriftType ->+ T.Text ->+ Int64 ->+ Int ->+ Int64 ->+ Int ->+ ColumnChunk+mkColumnChunk codec columnType columnName offset compressedSize uncompressedSize rgRows =+ ColumnChunk+ { cc_file_path = putField Nothing+ , cc_file_offset = putField offset+ , cc_meta_data = putField (Just metadata)+ , cc_offset_index_offset = putField Nothing+ , cc_offset_index_length = putField Nothing+ , cc_column_index_offset = putField Nothing+ , cc_column_index_length = putField Nothing+ , cc_crypto_metadata = putField Nothing+ , cc_encrypted_column_metadata = putField Nothing+ }+ where+ metadata =+ ColumnMetaData+ { cmd_type = putField columnType+ , cmd_encodings = putField [PLAIN enum, RLE enum]+ , cmd_path_in_schema = putField [columnName]+ , cmd_codec = putField codec+ , cmd_num_values = putField (fromIntegral rgRows)+ , cmd_total_uncompressed_size = putField uncompressedSize+ , cmd_total_compressed_size = putField (fromIntegral compressedSize)+ , cmd_key_value_metadata = putField Nothing+ , cmd_data_page_offset = putField offset+ , cmd_index_page_offset = putField Nothing+ , cmd_dictionary_page_offset = putField Nothing+ , cmd_statistics = putField Nothing+ , cmd_encoding_stats = putField Nothing+ , cmd_bloom_filter_offset = putField Nothing+ , cmd_bloom_filter_length = putField Nothing+ }++mkRowGroup :: [ColumnChunk] -> Int64 -> Int64 -> Int -> RowGroup+mkRowGroup chunks totalCompressed totalUncompressed rgRows =+ RowGroup+ { rg_columns = putField chunks+ , rg_total_byte_size = putField totalUncompressed+ , rg_num_rows = putField (fromIntegral rgRows)+ , rg_sorting_columns = putField Nothing+ , rg_file_offset = putField Nothing+ , rg_total_compressed_size = putField (Just totalCompressed)+ , rg_ordinal = putField Nothing+ }++writeFooter ::+ WritableBinaryHandle ->+ [SchemaElement] ->+ Int ->+ [RowGroup] ->+ [(T.Text, T.Text)] ->+ IO ()+writeFooter output schemaElements numRows rowGroupMetadata keyValues = do+ let metadata =+ FileMetadata+ { version = putField 1+ , schema = putField schemaElements+ , num_rows = putField (fromIntegral numRows)+ , row_groups = putField rowGroupMetadata+ , key_value_metadata =+ putField $+ if null keyValues+ then Nothing+ else+ Just+ [ KeyValue (putField k) (putField (Just v))+ | (k, v) <- keyValues+ ]+ , created_by = putField (Just "dataframe-parquet")+ , column_orders = putField Nothing+ , encryption_algorithm = putField Nothing+ , footer_signing_key_metadata = putField Nothing+ }+ footer = Pinch.encode Pinch.compactProtocol metadata+ buffer <- stToIO (mallocBuffer (BS.length footer + 8))+ writeByteString buffer footer+ stToIO (writeWord32LE buffer (fromIntegral (BS.length footer)))+ writeByteString buffer magic+ flushBufferToFile output buffer++magic :: BS.ByteString+magic = "PAR1"
+ src/DataFrame/IO/Parquet/Writer/Options.hs view
@@ -0,0 +1,34 @@+module DataFrame.IO.Parquet.Writer.Options (+ ParquetWriteOptions (..),+ WriterStrategy (..),+ defaultParquetWriteOptions,+) where++import DataFrame.IO.Parquet.Thrift+import Pinch (enum)++data WriterStrategy = InMemory | TwoPass+ deriving (Eq, Show)++data ParquetWriteOptions = ParquetWriteOptions+ { pageSize :: !Int+ , rowGroupSize :: !Int+ , batchRows :: !Int+ , subBatchRows :: !Int+ , compressionCodec :: !CompressionCodec+ , strategy :: !WriterStrategy+ , maxRowsPerFile :: !(Maybe Int)+ }+ deriving (Eq, Show)++defaultParquetWriteOptions :: ParquetWriteOptions+defaultParquetWriteOptions =+ ParquetWriteOptions+ { pageSize = 1048576+ , rowGroupSize = 134217728+ , batchRows = 8192+ , subBatchRows = 2048+ , compressionCodec = SNAPPY enum+ , strategy = InMemory+ , maxRowsPerFile = Nothing+ }
src/DataFrame/IO/Utils/RandomAccess.hs view
@@ -1,21 +1,87 @@ {-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE OverloadedRecordDot #-} -module DataFrame.IO.Utils.RandomAccess where+module DataFrame.IO.Utils.RandomAccess (+ uncurry3,+ Range (..),+ RandomAccess (..),+ ReaderIO (runReaderIO),+ LocalFile,+ MMappedFile,+ ChunkBuffer (..),+ unsafeToByteString,+ WritableBinaryHandle,+ openWritableBinaryFile,+ withWritableBinaryFile,+ atomicallyWriteFile,+ MemoryBuffer (..),+ ensureCapacity,+ mallocBuffer,+ writeByteString,+ appendTextArraySlice,+ writeWord8,+ writeWord32LE,+ writeWord64LE,+ writeInteger64,+ writeWord32At,+ writeWord64At,+ writeInteger64At,+ writeFloatLE,+ writeDoubleLE,+ bufferResidency,+ bufferToByteString,+ unsafeBufferView,+ flushBufferToBuffer,+ resetPosition,+ flushBufferToFile,+ writeByteStringToFile,+) where +import Control.Exception (bracket, bracketOnError, finally, throw)+import Control.Monad (when) import Control.Monad.IO.Class (MonadIO (..))-import Data.ByteString (ByteString)-import Data.ByteString.Internal (ByteString (PS))+import Control.Monad.ST (RealWorld, ST, stToIO)+import Data.Bits (shiftR)+import Data.ByteString.Internal (ByteString (BS, PS), create)+import qualified Data.ByteString.Unsafe as BU+import Data.Int (Int64)+import Data.Primitive.ByteArray (+ MutableByteArray (..),+ copyMutableByteArray,+ getSizeofMutableByteArray,+ mutableByteArrayContents,+ newPinnedByteArray,+ withMutableByteArrayContents,+ writeByteArray,+ )+import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef)+import qualified Data.Text.Array as TA import qualified Data.Vector.Storable as VS-import Data.Word (Word8)+import Data.Word (Word32, Word64, Word8) import DataFrame.IO.Parquet.Seeking ( FileBufferedOrSeekable, fGet, fSeek, readLastBytes, )-import Foreign (castForeignPtr)+import Foreign (castForeignPtr, castPtr, copyBytes, plusPtr)+import GHC.Float (castDoubleToWord64, castFloatToWord32)+import GHC.ForeignPtr (ForeignPtr (..), ForeignPtrContents (PlainPtr))+import GHC.Ptr (Ptr (..))+import System.Directory (copyPermissions, doesFileExist, removeFile, renameFile)+import System.FilePath (takeDirectory) import System.IO (+ BufferMode (NoBuffering),+ Handle,+ IOMode (WriteMode), SeekMode (AbsoluteSeek),+ hClose,+ hPutBuf,+ hSetBinaryMode,+ hSetBuffering,+ openBinaryFile,+ openBinaryTempFileWithDefaultPermissions, ) uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d@@ -58,6 +124,14 @@ fGet handle length' readSuffix n = ReaderIO (readLastBytes $ fromIntegral n) +-- | Bytes of one column chunk, already read; @base@ is their file offset.+data ChunkBuffer = ChunkBuffer !Integer !ByteString++instance RandomAccess (ReaderIO ChunkBuffer) where+ readBytes (Range offset' length') = ReaderIO $ \(ChunkBuffer base bs) ->+ pure (BU.unsafeTake length' (BU.unsafeDrop (fromInteger (offset' - base)) bs))+ readSuffix _ = ReaderIO $ \_ -> error "ChunkBuffer: no file footer"+ type MMappedFile = ReaderIO (VS.Vector Word8) -- The instance exists but we don't have the means to mmap the file currently@@ -76,3 +150,242 @@ unsafeToByteString v = PS (castForeignPtr ptr) offset' len where (ptr, offset', len) = VS.unsafeToForeignPtr v++-- Writer buffers are pinned MutableByteArrays so hPutBuf can take a Ptr.++newtype WritableBinaryHandle = WritableBinaryHandle {unHandle :: Handle}++openWritableBinaryFile :: FilePath -> IO WritableBinaryHandle+openWritableBinaryFile filepath = do+ h <- openBinaryFile filepath WriteMode+ hSetBinaryMode h True+ hSetBuffering h NoBuffering+ pure . WritableBinaryHandle $ h++atomicallyWriteFile ::+ FilePath ->+ (FilePath -> IO a) ->+ IO a+atomicallyWriteFile path action =+ bracketOnError+ openAction+ removeFile+ ( \tmpFile -> do+ result <- action tmpFile+ renameFile tmpFile path+ pure result+ )+ where+ openAction =+ bracketOnError+ ( openBinaryTempFileWithDefaultPermissions+ (takeDirectory path)+ "dataframe-parquet.incomplete"+ )+ (\(tmpFile, h) -> hClose h `finally` removeFile tmpFile)+ ( \(tmpFile, h) -> do+ hClose h+ destinationExists <- doesFileExist path+ when destinationExists (copyPermissions path tmpFile)+ pure tmpFile+ )++withWritableBinaryFile ::+ FilePath ->+ (WritableBinaryHandle -> IO a) ->+ IO a+withWritableBinaryFile filepath =+ bracket+ (openWritableBinaryFile filepath)+ (hClose . unHandle)++data MemoryBuffer s = MemoryBuffer+ { arrayRef :: !(STRef s (MutableByteArray s))+ , positionRef :: !(STRef s Int)+ }++mallocBuffer :: Int -> ST s (MemoryBuffer s)+mallocBuffer capacity+ | capacity < 0 = throw (userError "mallocBuffer: negative capacity")+ | otherwise = do+ array <- newPinnedByteArray capacity+ MemoryBuffer <$> newSTRef array <*> newSTRef 0++-- Pinned arrays cannot grow in place, so allocate a new one and copy.+ensureCapacity :: MemoryBuffer s -> Int -> ST s (MutableByteArray s)+ensureCapacity buffer needed = do+ array <- readSTRef buffer.arrayRef+ maxSize <- getSizeofMutableByteArray array+ if needed <= maxSize+ then pure array+ else do+ position <- readSTRef buffer.positionRef+ grown <- newPinnedByteArray (needed + (needed `div` 2))+ copyMutableByteArray grown 0 array 0 position+ writeSTRef buffer.arrayRef grown+ pure grown+{-# INLINE ensureCapacity #-}++writeWord8 :: MemoryBuffer s -> Word8 -> ST s ()+writeWord8 buffer b = do+ position <- readSTRef buffer.positionRef+ array <- ensureCapacity buffer (position + 1)+ writeByteArray array position b+ writeSTRef buffer.positionRef (position + 1)+{-# INLINE writeWord8 #-}++-- Reading the ByteString's pointer needs IO, so this stays out of ST.+writeByteString :: MemoryBuffer RealWorld -> ByteString -> IO ()+writeByteString buffer bs =+ BU.unsafeUseAsCStringLen bs $ \(source, len) -> do+ position <- stToIO (readSTRef buffer.positionRef)+ array <- stToIO (ensureCapacity buffer (position + len))+ withMutableByteArrayContents array $ \dst ->+ copyBytes (dst `plusPtr` position) (castPtr source) len+ stToIO (writeSTRef buffer.positionRef (position + len))+{-# INLINE writeByteString #-}++writeWord32LE :: MemoryBuffer s -> Word32 -> ST s ()+writeWord32LE buffer w = do+ position <- readSTRef buffer.positionRef+ writeWord32At buffer position w+ writeSTRef buffer.positionRef (position + 4)+{-# INLINE writeWord32LE #-}++writeWord64LE :: MemoryBuffer s -> Word64 -> ST s ()+writeWord64LE buffer w = do+ position <- readSTRef buffer.positionRef+ writeWord64At buffer position w+ writeSTRef buffer.positionRef (position + 8)+{-# INLINE writeWord64LE #-}++writeWord32At :: MemoryBuffer s -> Int -> Word32 -> ST s ()+writeWord32At buffer position w = do+ array <- ensureCapacity buffer (position + 4)+ writeByteArray array position (fromIntegral w :: Word8)+ writeByteArray array (position + 1) (fromIntegral (w `shiftR` 8) :: Word8)+ writeByteArray array (position + 2) (fromIntegral (w `shiftR` 16) :: Word8)+ writeByteArray array (position + 3) (fromIntegral (w `shiftR` 24) :: Word8)+{-# INLINE writeWord32At #-}++writeWord64At :: MemoryBuffer s -> Int -> Word64 -> ST s ()+writeWord64At buffer position w = do+ array <- ensureCapacity buffer (position + 8)+ writeByteArray array position (fromIntegral w :: Word8)+ writeByteArray array (position + 1) (fromIntegral (w `shiftR` 8) :: Word8)+ writeByteArray array (position + 2) (fromIntegral (w `shiftR` 16) :: Word8)+ writeByteArray array (position + 3) (fromIntegral (w `shiftR` 24) :: Word8)+ writeByteArray array (position + 4) (fromIntegral (w `shiftR` 32) :: Word8)+ writeByteArray array (position + 5) (fromIntegral (w `shiftR` 40) :: Word8)+ writeByteArray array (position + 6) (fromIntegral (w `shiftR` 48) :: Word8)+ writeByteArray array (position + 7) (fromIntegral (w `shiftR` 56) :: Word8)+{-# INLINE writeWord64At #-}++writeInteger64 :: MemoryBuffer s -> Integer -> ST s ()+writeInteger64 buffer value = do+ position <- readSTRef buffer.positionRef+ newPosition <- writeInteger64At buffer position value+ writeSTRef buffer.positionRef newPosition+{-# INLINE writeInteger64 #-}++writeInteger64At :: MemoryBuffer s -> Int -> Integer -> ST s Int+writeInteger64At buffer position value+ | value < toInteger (minBound :: Int64) = outOfRange+ | value > toInteger (maxBound :: Int64) = outOfRange+ | otherwise = do+ writeWord64At buffer position (fromIntegral value)+ pure (position + 8)+ where+ outOfRange =+ throw (userError "writeParquet: Integer value is outside the INT64 range")+{-# INLINE writeInteger64At #-}++writeFloatLE :: MemoryBuffer s -> Float -> ST s ()+writeFloatLE buffer = writeWord32LE buffer . castFloatToWord32+{-# INLINE writeFloatLE #-}++writeDoubleLE :: MemoryBuffer s -> Double -> ST s ()+writeDoubleLE buffer = writeWord64LE buffer . castDoubleToWord64+{-# INLINE writeDoubleLE #-}++flushBufferToBuffer :: MemoryBuffer s -> MemoryBuffer s -> ST s ()+flushBufferToBuffer source destination+ | source.arrayRef == destination.arrayRef = pure ()+ | otherwise = do+ sourceArray <- readSTRef source.arrayRef+ sourcePosition <- readSTRef source.positionRef+ destinationPosition <- readSTRef destination.positionRef+ destinationArray <-+ ensureCapacity destination (destinationPosition + sourcePosition)+ copyMutableByteArray+ destinationArray+ destinationPosition+ sourceArray+ 0+ sourcePosition+ writeSTRef destination.positionRef (destinationPosition + sourcePosition)+ writeSTRef source.positionRef 0+{-# INLINE flushBufferToBuffer #-}++bufferToByteString :: MemoryBuffer RealWorld -> IO ByteString+bufferToByteString buffer = do+ array <- stToIO (readSTRef buffer.arrayRef)+ position <- stToIO (readSTRef buffer.positionRef)+ create position $ \dst ->+ withMutableByteArrayContents array $ \src ->+ copyBytes dst (castPtr src) position++-- | The buffer's contents as a ByteString, without copying.+unsafeBufferView :: MemoryBuffer RealWorld -> IO ByteString+unsafeBufferView buffer = do+ array@(MutableByteArray array#) <- stToIO (readSTRef buffer.arrayRef)+ position <- stToIO (readSTRef buffer.positionRef)+ let Ptr address# = mutableByteArrayContents array+ pure (BS (ForeignPtr address# (PlainPtr array#)) position)++bufferResidency :: MemoryBuffer s -> ST s Int+bufferResidency buffer = readSTRef buffer.positionRef+{-# INLINE bufferResidency #-}++resetPosition :: MemoryBuffer s -> ST s ()+resetPosition buffer = writeSTRef buffer.positionRef 0+{-# INLINE resetPosition #-}++-- Flush in 256 KiB chunks; larger blocks gave no write speed in dd tests.+flushBufferToFile :: WritableBinaryHandle -> MemoryBuffer RealWorld -> IO ()+flushBufferToFile (WritableBinaryHandle h) buffer = do+ array <- stToIO (readSTRef buffer.arrayRef)+ position <- stToIO (readSTRef buffer.positionRef)+ withMutableByteArrayContents array $ \ptr -> do+ let chunkSize = 262144+ go offset+ | offset >= position = pure ()+ | otherwise = do+ let n = min chunkSize (position - offset)+ hPutBuf h (ptr `plusPtr` offset) n+ go (offset + n)+ go 0+ stToIO (resetPosition buffer)++writeByteStringToFile :: WritableBinaryHandle -> ByteString -> IO ()+writeByteStringToFile (WritableBinaryHandle h) bs =+ BU.unsafeUseAsCStringLen bs $ \(source, len) -> do+ let chunkSize = 262144+ go offset+ | offset >= len = pure ()+ | otherwise = do+ let n = min chunkSize (len - offset)+ hPutBuf h (source `plusPtr` offset) n+ go (offset + n)+ go 0++appendTextArraySlice :: MemoryBuffer s -> TA.Array -> Int -> Int -> ST s ()+appendTextArraySlice buffer source offset count+ | count < 0 = throw (userError "appendTextArraySlice: negative length")+ | otherwise = do+ position <- readSTRef buffer.positionRef+ array <- ensureCapacity buffer (position + count)+ withMutableByteArrayContents array $ \destination ->+ TA.copyToPointer source offset (destination `plusPtr` position) count+ writeSTRef buffer.positionRef (position + count)+{-# INLINE appendTextArraySlice #-}
+ src/DataFrame/Typed/IO/Parquet.hs view
@@ -0,0 +1,65 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Typed Parquet reading, validated against a type-level schema.+module DataFrame.Typed.IO.Parquet (+ readParquet,+ readParquetWithError,+ readParquetWithOpts,+ readParquetFiles,+) where++import Control.Applicative ((<|>))+import Control.Exception (SomeException, try)+import qualified Data.Text as T++import DataFrame.IO.Parquet (ParquetReadOptions (..), defaultParquetReadOptions)+import qualified DataFrame.IO.Parquet as Parquet+import DataFrame.Typed.Freeze (freezeOrThrow, freezeWithError)+import DataFrame.Typed.Schema (KnownSchema, schemaColumnNames)+import DataFrame.Typed.Types (TypedDataFrame)++-- | Read only the columns the schema names, unless the caller chose.+schemaOptions ::+ forall cols. (KnownSchema cols) => ParquetReadOptions -> ParquetReadOptions+schemaOptions opts =+ opts+ { selectedColumns = selectedColumns opts <|> Just (schemaColumnNames @cols)+ }++-- | Read a Parquet file into a typed DataFrame, throwing on mismatch.+readParquet ::+ forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols)+readParquet = readParquetWithOpts @cols defaultParquetReadOptions++-- | Read a Parquet file, returning an error on mismatch.+readParquetWithError ::+ forall cols.+ (KnownSchema cols) =>+ FilePath -> IO (Either T.Text (TypedDataFrame cols))+readParquetWithError path = do+ r <-+ try+ (Parquet.readParquetWithOpts (schemaOptions @cols defaultParquetReadOptions) path)+ pure $ case r of+ Left (e :: SomeException) -> Left (T.pack (show e))+ Right df -> freezeWithError df++-- | Read a Parquet file with custom options, throwing on mismatch.+readParquetWithOpts ::+ forall cols.+ (KnownSchema cols) =>+ ParquetReadOptions -> FilePath -> IO (TypedDataFrame cols)+readParquetWithOpts opts path =+ Parquet.readParquetWithOpts (schemaOptions @cols opts) path+ >>= freezeOrThrow @cols++-- | Read a directory\/glob of Parquet files into a typed DataFrame.+readParquetFiles ::+ forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols)+readParquetFiles path =+ Parquet.readParquetFilesWithOpts+ (schemaOptions @cols defaultParquetReadOptions)+ path+ >>= freezeOrThrow @cols
+ stress/DataFrame10GB.hs view
@@ -0,0 +1,137 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE NumericUnderscores #-}++module DataFrame10GB (+ stressDataFrame,+ stressRows,+ stressColumns,+ stressResidentBytesLowerBound,+) where++import Control.Monad.ST (runST)+import Data.Int (Int32, Int64)+import qualified Data.Text as T+import qualified Data.Text.Array as A+import Data.Time (UTCTime (UTCTime), addDays, fromGregorian, secondsToDiffTime)+import qualified Data.Vector as VB+import qualified Data.Vector.Unboxed as VU+import Data.Word (Word8)+import DataFrame.Internal.Column (Column (..))+import DataFrame.Internal.Column.Bitmap (Bitmap)+import DataFrame.Internal.Data.PackedText (mkPackedContiguous32)+import DataFrame.Internal.DataFrame (DataFrame, fromColumns)++stressRows :: Int+stressRows = 1_000_000++stressGroups :: Int+stressGroups = 16++stressColumns :: Int+stressColumns = stressGroups * 14++textBytesPerRow :: Int+textBytesPerRow = 320++stressResidentBytesLowerBound :: Integer+stressResidentBytesLowerBound =+ fromIntegral stressRows+ * fromIntegral stressGroups+ * fromIntegral (2 * textBytesPerRow + 2 * (4 + 8 + 4 + 8))++stressDataFrame :: DataFrame+stressDataFrame = fromColumns (concatMap columnGroup [0 .. stressGroups - 1])++columnGroup :: Int -> [(T.Text, Column)]+columnGroup group =+ [ named "int32" (UnboxedColumn Nothing (int32Values group))+ , named "int64" (UnboxedColumn Nothing (int64Values group))+ , named "float" (UnboxedColumn Nothing (floatValues group))+ , named "double" (UnboxedColumn Nothing (doubleValues group))+ , named "bool" (UnboxedColumn Nothing (boolValues group))+ , named "timestamp" (BoxedColumn Nothing (timestampValues group))+ , named "text" (textColumn Nothing group)+ , named+ "nullable_int32"+ (UnboxedColumn (Just nullableBitmap) (int32Values (group + stressGroups)))+ , named+ "nullable_int64"+ (UnboxedColumn (Just nullableBitmap) (int64Values (group + stressGroups)))+ , named+ "nullable_float"+ (UnboxedColumn (Just nullableBitmap) (floatValues (group + stressGroups)))+ , named+ "nullable_double"+ (UnboxedColumn (Just nullableBitmap) (doubleValues (group + stressGroups)))+ , named+ "nullable_bool"+ (UnboxedColumn (Just nullableBitmap) (boolValues (group + stressGroups)))+ , named+ "nullable_timestamp"+ (BoxedColumn (Just nullableBitmap) (timestampValues (group + stressGroups)))+ , named "nullable_text" (textColumn (Just nullableBitmap) (group + stressGroups))+ ]+ where+ named suffix column = (T.pack ("group_" <> show group <> "_" <> suffix), column)++nullableBitmap :: Bitmap+nullableBitmap = VU.replicate (stressRows `div` 8) (0xFE :: Word8)++int32Values :: Int -> VU.Vector Int32+int32Values salt =+ VU.generate stressRows $ \row ->+ fromIntegral ((row + salt * 10_007) `mod` 2_000_001 - 1_000_000)++int64Values :: Int -> VU.Vector Int64+int64Values salt =+ VU.generate stressRows $ \row ->+ fromIntegral row * 1_000_003 - fromIntegral salt * 10_000_019++floatValues :: Int -> VU.Vector Float+floatValues salt =+ VU.generate stressRows $ \row ->+ fromIntegral ((row + salt * 101) `mod` 100_003) / 17++doubleValues :: Int -> VU.Vector Double+doubleValues salt =+ VU.generate stressRows $ \row ->+ fromIntegral row / 31.0 - fromIntegral salt * 1_000.25++boolValues :: Int -> VU.Vector Bool+boolValues salt = VU.generate stressRows (\row -> (row + salt) `mod` 3 == 0)++timestampValues :: Int -> VB.Vector UTCTime+timestampValues salt =+ VB.replicate+ stressRows+ ( UTCTime+ (addDays (fromIntegral salt) (fromGregorian 2020 1 1))+ (secondsToDiffTime (fromIntegral (salt * 1_337 `mod` 86_400)))+ )++textColumn :: Maybe Bitmap -> Int -> Column+textColumn bitmap salt = PackedText bitmap $ runST $ do+ target <- A.new (stressRows * textBytesPerRow)+ let template = textTemplate salt+ fill !row+ | row >= stressRows = pure ()+ | otherwise = do+ A.copyI textBytesPerRow target (row * textBytesPerRow) template 0+ fill (row + 1)+ fill 0+ bytes <- A.unsafeFreeze target+ let offsets =+ VU.generate+ (stressRows + 1)+ (\row -> fromIntegral (row * textBytesPerRow) :: Int32)+ pure (mkPackedContiguous32 bytes offsets)++textTemplate :: Int -> A.Array+textTemplate salt = A.run $ do+ bytes <- A.new textBytesPerRow+ let byte = fromIntegral (97 + salt `mod` 26)+ fill !index+ | index >= textBytesPerRow = pure ()+ | otherwise = A.unsafeWrite bytes index byte >> fill (index + 1)+ fill 0+ pure bytes
+ stress/StressMain.hs view
@@ -0,0 +1,38 @@+module Main (main) where++import Control.Exception (evaluate)+import Control.Monad (unless)+import DataFrame.IO.Parquet (readParquet)+import DataFrame.IO.Parquet.Writer (writeParquet)+import DataFrame.Internal.DataFrame (forceDataFrame)+import DataFrame10GB (+ stressColumns,+ stressDataFrame,+ stressResidentBytesLowerBound,+ stressRows,+ )+import System.Exit (exitFailure)+import System.FilePath ((</>))+import System.IO (hPutStrLn, stderr)+import System.IO.Temp (withSystemTempDirectory)++main :: IO ()+main = withSystemTempDirectory "dataframe-parquet-10gb-stress" $ \directory -> do+ expected <- evaluate (forceDataFrame stressDataFrame)+ let output = directory </> "roundtrip.parquet"+ putStrLn+ ( "writing "+ <> show stressRows+ <> " rows x "+ <> show stressColumns+ <> " columns (at least "+ <> show stressResidentBytesLowerBound+ <> " resident payload bytes)"+ )+ writeParquet output expected+ putStrLn "reading the stress dataframe"+ actual <- readParquet output+ putStrLn "checking dataframe equivalence"+ unless (expected == actual) $ do+ hPutStrLn stderr "10 GiB Parquet roundtrip mismatch"+ exitFailure
+ tests/Main.hs view
@@ -0,0 +1,413 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Tests for the writer-buffer logic in "DataFrame.IO.Utils.RandomAccess".+module Main where++import Control.Exception (SomeException, catch, evaluate)+import qualified Data.ByteString as BS+import Data.List (sortOn)+import qualified System.Exit as Exit+import System.FilePath ((</>))+import System.IO.Temp (withSystemTempDirectory)+import Test.HUnit++import Control.Monad (unless)+import Control.Monad.ST (stToIO)+import Data.Int (Int32, Int64)+import Data.Maybe (fromJust)+import qualified Data.Text as T+import DataFrame.IO.Parquet (+ ParquetReadOptions (..),+ defaultParquetReadOptions,+ readParquet,+ readParquetFiles,+ readParquetWithOpts,+ _readParquetWithOpts,+ )+import DataFrame.IO.Parquet.Writer (+ ParquetWriteOptions (..),+ defaultParquetWriteOptions,+ writeParquet,+ writeParquetWithOptions,+ )+import DataFrame.IO.Utils.RandomAccess+import DataFrame.Internal.Column (columnTypeString, fromList)+import DataFrame.Internal.DataFrame (+ DataFrame,+ columnNames,+ fromColumns,+ getColumn,+ )+import DataFrame.Operations.Subset (range, select)+import qualified DataFrame.Typed.IO.Parquet as Typed+import qualified DataFrame.Typed.Operations as DT+import System.Directory (listDirectory)++directWrites :: Test+directWrites = TestCase $ do+ buffer <- stToIO (mallocBuffer 1)+ stToIO (writeWord8 buffer 0xaa)+ stToIO (writeWord32LE buffer 0x78563412)+ stToIO (writeWord64LE buffer 0x0807060504030201)+ stToIO (writeFloatLE buffer 1)+ stToIO (writeDoubleLE buffer 1)+ writeByteString buffer (BS.pack [0xfe, 0xff])+ residency <- stToIO (bufferResidency buffer)+ bytes <- bufferToByteString buffer+ assertEqual "direct write residency" 27 residency+ assertEqual+ "direct write bytes"+ ( BS.pack+ [ 0xaa+ , 0x12+ , 0x34+ , 0x56+ , 0x78+ , 0x01+ , 0x02+ , 0x03+ , 0x04+ , 0x05+ , 0x06+ , 0x07+ , 0x08+ , 0x00+ , 0x00+ , 0x80+ , 0x3f+ , 0x00+ , 0x00+ , 0x00+ , 0x00+ , 0x00+ , 0x00+ , 0xf0+ , 0x3f+ , 0xfe+ , 0xff+ ]+ )+ bytes++directBufferFlush :: Test+directBufferFlush = TestCase $ do+ source <- stToIO (mallocBuffer 0)+ destination <- stToIO (mallocBuffer 0)+ writeByteString destination (BS.pack [1, 2])+ writeByteString source (BS.pack [3, 4, 5])+ stToIO (flushBufferToBuffer source destination)+ sourceResidency <- stToIO (bufferResidency source)+ destinationBytes <- bufferToByteString destination+ assertEqual "source cleared" 0 sourceResidency+ assertEqual "destination appended" (BS.pack [1, 2, 3, 4, 5]) destinationBytes+ stToIO (flushBufferToBuffer destination destination)+ selfFlushedBytes <- bufferToByteString destination+ assertEqual "self flush is a no-op" destinationBytes selfFlushedBytes+ stToIO (resetPosition destination)+ destinationResidency <- stToIO (bufferResidency destination)+ assertEqual "reset position" 0 destinationResidency++directFileFlush :: Test+directFileFlush = TestCase $+ withSystemTempDirectory "dfpq-buffer" $ \dir -> do+ let outPath = dir </> "out.bin"+ payload = BS.pack (take 300000 (cycle [0 .. 255]))+ buffer <- stToIO (mallocBuffer 1)+ writeByteString buffer payload+ withWritableBinaryFile outPath $ \output ->+ flushBufferToFile output buffer+ residency <- stToIO (bufferResidency buffer)+ contents <- BS.readFile outPath+ assertEqual "source cleared after file flush" 0 residency+ assertEqual "large payload round-trips" payload contents++writerRoundTrip :: String -> FilePath -> Test+writerRoundTrip label path = TestCase $+ withSystemTempDirectory "dfpq-writer" $ \dir -> do+ df <- readParquet path+ let out = dir </> "out.parquet"+ writeParquet out df+ df' <- readParquet out+ assertEqual label df df'++writerRoundTripTiny :: String -> FilePath -> Test+writerRoundTripTiny label path = TestCase $+ withSystemTempDirectory "dfpq-writer" $ \dir -> do+ df <- readParquet path+ let out = dir </> "out.parquet"+ writeParquetWithOptions tinyWriteOpts out df+ df' <- readParquet out+ assertEqual label df df'++-- Row groups of 512 rows give the parallel reader several tasks per column.+parallelMatchesSequential :: String -> FilePath -> Test+parallelMatchesSequential label path = TestCase $+ withSystemTempDirectory "dfpq-parallel" $ \dir -> do+ df <- readParquet path+ let out = dir </> "out.parquet"+ writeParquetWithOptions tinyWriteOpts out df+ sequential <-+ readParquetWithOpts defaultParquetReadOptions{parallel = False} out+ parallelRead <-+ readParquetWithOpts defaultParquetReadOptions{parallel = True} out+ assertEqual (label ++ ": sequential") df sequential+ assertEqual (label ++ ": parallel") df parallelRead++writerRoundTripLargeText :: Test+writerRoundTripLargeText = TestCase $+ withSystemTempDirectory "dfpq-writer" $ \dir -> do+ let payload = T.replicate 4096 "abcdefgh"+ df = fromColumns [("text", fromList [payload, "short"])]+ firstOut = dir </> "large-text-1.parquet"+ secondOut = dir </> "large-text-2.parquet"+ writeParquetWithOptions tinyWriteOpts firstOut df+ firstRoundTrip <- readParquet firstOut+ writeParquetWithOptions tinyWriteOpts secondOut firstRoundTrip+ secondRoundTrip <- readParquet secondOut+ assertEqual "large text first round-trip" df firstRoundTrip+ assertEqual "large text second round-trip" df secondRoundTrip++-- | Sharded writes split the frame and reading them back matches.+writerRoundTripSharded :: String -> FilePath -> Int -> Int -> Test+writerRoundTripSharded label path rowsPerFile expectedShards = TestCase $+ withSystemTempDirectory "dfpq-writer" $ \dir -> do+ df <- readParquet path+ let pattern_ = dir </> "shards" </> "part-*.parquet"+ writeParquetWithOptions+ defaultParquetWriteOptions{maxRowsPerFile = Just rowsPerFile}+ pattern_+ df+ shards <- listDirectory (dir </> "shards")+ assertEqual (label <> ": shard count") expectedShards (Prelude.length shards)+ assertEqual+ (label <> ": shard names")+ ["part-" <> pad i <> ".parquet" | i <- [0 .. expectedShards - 1]]+ (sortOn id shards)+ df' <- readParquetFiles pattern_+ assertEqual (label <> ": shards round-trip") df df'+ where+ pad i = let s = show i in replicate (5 - Prelude.length s) '0' <> s++-- | A path without a @*@ placeholder is rejected when sharding is requested.+shardedWriteRequiresPattern :: Test+shardedWriteRequiresPattern = TestCase $+ withSystemTempDirectory "dfpq-writer" $ \dir -> do+ df <- readParquet "tests/data/mtcars.parquet"+ threw <-+ ( False+ <$ writeParquetWithOptions+ defaultParquetWriteOptions{maxRowsPerFile = Just 4}+ (dir </> "out.parquet")+ df+ )+ `catch` (\e -> True <$ evaluate (Prelude.length (show (e :: SomeException))))+ unless threw (assertFailure "expected an error for a path without '*'")++-- | Int, Int64 and Integer columns round-trip with their Haskell type.+writerRoundTripNativeIntTypes :: Test+writerRoundTripNativeIntTypes = TestCase $+ withSystemTempDirectory "dfpq-writer" $ \dir -> do+ let df =+ fromColumns+ [ ("int", fromList [1 :: Int, 2, 3])+ , ("int64", fromList [1 :: Int64, 2, 3])+ , ("int32", fromList [1 :: Int32, 2, 3])+ , ("integer", fromList [1 :: Integer, 2, 3])+ , ("nullableInt", fromList [Just (1 :: Int), Nothing, Just 3])+ , ("nullableInt64", fromList [Just (1 :: Int64), Nothing, Just 3])+ ]+ out = dir </> "int-types.parquet"+ writeParquet out df+ df' <- readParquet out+ assertEqual+ "native int types: column types"+ (columnTypes df)+ (columnTypes df')+ assertEqual "native int types: frame" df df'++columnTypes :: DataFrame -> [(String, String)]+columnTypes df =+ [ (T.unpack name, columnTypeString (fromJust (getColumn name df)))+ | name <- columnNames df+ ]++tinyWriteOpts :: ParquetWriteOptions+tinyWriteOpts =+ defaultParquetWriteOptions+ { pageSize = 64+ , rowGroupSize = 512+ , batchRows = 4+ , subBatchRows = 3+ }++dottedColumnNames :: Test+dottedColumnNames = TestCase $+ withSystemTempDirectory "dfpq-dotted" $ \dir -> do+ let df =+ fromColumns+ [ ("sepal.length", fromList [5.1, 4.9, 4.7 :: Double])+ , ("length", fromList [1, 2, 3 :: Int])+ , ("variety", fromList ["a", "b", "c" :: T.Text])+ ]+ out = dir </> "dotted.parquet"+ writeParquet out df+ whole <- readParquet out+ assertEqual "names survive the round trip" (columnNames df) (columnNames whole)+ projected <-+ readParquetWithOpts+ defaultParquetReadOptions{selectedColumns = Just ["sepal.length", "variety"]}+ out+ assertEqual+ "projection keeps the dotted column"+ ["sepal.length", "variety"]+ (columnNames projected)+ assertEqual+ "projected rows match"+ (select ["sepal.length", "variety"] df)+ projected+ leaf <-+ readParquetWithOpts+ defaultParquetReadOptions{selectedColumns = Just ["length"]}+ out+ assertEqual "a full name wins over a leaf alias" ["length"] (columnNames leaf)+ typed <-+ Typed.readParquet @'[ '("sepal.length", Double), '("variety", T.Text)] out+ assertEqual "the typed reader projects a dotted schema" 3 (DT.nRows typed)++tests :: Test+tests =+ TestList+ [ TestLabel "dotted column names project by full name" dottedColumnNames+ , TestLabel "direct buffer writes" directWrites+ , TestLabel "direct buffer-to-buffer flush" directBufferFlush+ , TestLabel "direct buffer-to-file flush" directFileFlush+ , TestLabel+ "writer roundtrip: alltypes_plain"+ (writerRoundTrip "alltypes_plain" "tests/data/alltypes_plain.parquet")+ , TestLabel+ "writer roundtrip: alltypes_plain.snappy"+ ( writerRoundTrip+ "alltypes_plain.snappy"+ "tests/data/alltypes_plain.snappy.parquet"+ )+ , TestLabel+ "writer roundtrip: alltypes_dictionary"+ (writerRoundTrip "alltypes_dictionary" "tests/data/alltypes_dictionary.parquet")+ , TestLabel+ "writer roundtrip: alltypes_tiny_pages"+ (writerRoundTrip "alltypes_tiny_pages" "tests/data/alltypes_tiny_pages.parquet")+ , TestLabel+ "parallel read matches sequential: transactions"+ (parallelMatchesSequential "transactions" "tests/data/transactions.parquet")+ , TestLabel+ "parallel read matches sequential: mtcars"+ (parallelMatchesSequential "mtcars" "tests/data/mtcars.parquet")+ , TestLabel+ "writer roundtrip: transactions"+ (writerRoundTrip "transactions" "tests/data/transactions.parquet")+ , TestLabel+ "writer roundtrip: mtcars"+ (writerRoundTrip "mtcars" "tests/data/mtcars.parquet")+ , TestLabel+ "writer roundtrip: int32_decimal"+ (writerRoundTrip "int32_decimal" "tests/data/int32_decimal.parquet")+ , TestLabel+ "writer roundtrip: int64_decimal"+ (writerRoundTrip "int64_decimal" "tests/data/int64_decimal.parquet")+ , TestLabel+ "writer roundtrip: sharded mtcars"+ (writerRoundTripSharded "sharded mtcars" "tests/data/mtcars.parquet" 10 4)+ , TestLabel+ "writer roundtrip: sharded exact multiple"+ (writerRoundTripSharded "sharded exact" "tests/data/mtcars.parquet" 32 1)+ , TestLabel+ "sharded write requires a '*' pattern"+ shardedWriteRequiresPattern+ , TestLabel+ "writer roundtrip: Int/Integer keep their Haskell type"+ writerRoundTripNativeIntTypes+ , TestLabel+ "writer roundtrip: alltypes_plain multi-page"+ ( writerRoundTripTiny+ "alltypes_plain multi-page"+ "tests/data/alltypes_plain.parquet"+ )+ , TestLabel "writer roundtrip: large text" writerRoundTripLargeText+ , TestLabel "parallel read: many row groups" parallelReadRowGroups+ , TestLabel "parallel read: row range skips row groups" parallelReadRowRange+ , TestLabel "parallel read: widened integers" parallelReadWidenIntegers+ ]++-- | A frame big enough for 'tinyWriteOpts' to split into many row groups.+manyRowGroupsFrame :: DataFrame+manyRowGroupsFrame =+ fromColumns+ [ ("n", fromList [i * 7 - 3000 | i <- [0 .. 2999 :: Int]])+ , ("d", fromList [fromIntegral i / 8 :: Double | i <- [0 .. 2999 :: Int]])+ ,+ ( "t"+ , fromList+ [T.pack (show i) <> T.replicate (i `mod` 5) "é" | i <- [0 .. 2999 :: Int]]+ )+ ]++withManyRowGroups :: (FilePath -> IO ()) -> IO ()+withManyRowGroups body = withSystemTempDirectory "dfpq-rg" $ \dir -> do+ let out = dir </> "rg.parquet"+ writeParquetWithOptions tinyWriteOpts out manyRowGroupsFrame+ body out++-- | The parallel reader matches the sequential one and the source frame.+parallelReadRowGroups :: Test+parallelReadRowGroups = TestCase $ withManyRowGroups $ \out -> do+ parallel <- readParquet out+ sequential <- _readParquetWithOpts (Just True) defaultParquetReadOptions out+ assertEqual "parallel vs source" manyRowGroupsFrame parallel+ assertEqual "parallel vs sequential" sequential parallel+ projected <-+ readParquetWithOpts+ defaultParquetReadOptions{selectedColumns = Just ["t", "n"]}+ out+ assertEqual "projection" (select ["t", "n"] manyRowGroupsFrame) projected++-- | A row range only decodes the row groups it overlaps, and slices exactly.+parallelReadRowRange :: Test+parallelReadRowRange = TestCase $ withManyRowGroups $ \out -> do+ let ranges = [(0, 5), (7, 1300), (1299, 1301), (2990, 3000), (2995, 4000), (100, 100)]+ mapM_+ ( \r -> do+ got <- readParquetWithOpts defaultParquetReadOptions{rowRange = Just r} out+ assertEqual ("range " ++ show r) (range r manyRowGroupsFrame) got+ )+ ranges++-- | 'widenIntegers' turns Int32 and Int64 columns into Int on both read paths.+parallelReadWidenIntegers :: Test+parallelReadWidenIntegers = TestCase $ withSystemTempDirectory "dfpq-widen" $ \dir -> do+ let n = 3000 :: Int+ df =+ fromColumns+ [ ("i32", fromList [fromIntegral (i - 1500) :: Int32 | i <- [0 .. n - 1]])+ , ("i64", fromList [fromIntegral i * 4000000000 :: Int64 | i <- [0 .. n - 1]])+ ]+ expected =+ fromColumns+ [ ("i32", fromList [i - 1500 | i <- [0 .. n - 1]])+ , ("i64", fromList [i * 4000000000 | i <- [0 .. n - 1]])+ ]+ out = dir </> "ints.parquet"+ opts = defaultParquetReadOptions{widenIntegers = True}+ writeParquetWithOptions tinyWriteOpts out df+ parallel <- readParquetWithOpts opts out+ sequential <- _readParquetWithOpts (Just True) opts out+ assertEqual "parallel" expected parallel+ assertEqual "sequential" expected sequential++main :: IO ()+main = do+ result <- runTestTT tests+ if failures result > 0 || errors result > 0+ then Exit.exitFailure+ else Exit.exitSuccess
+ tests/data/alltypes_dictionary.parquet view
binary file changed (absent → 1698 bytes)
+ tests/data/alltypes_plain.parquet view
binary file changed (absent → 1851 bytes)
+ tests/data/alltypes_plain.snappy.parquet view
binary file changed (absent → 1736 bytes)
+ tests/data/alltypes_tiny_pages.parquet view
binary file changed (absent → 454233 bytes)
+ tests/data/int32_decimal.parquet view
binary file changed (absent → 478 bytes)
+ tests/data/int64_decimal.parquet view
binary file changed (absent → 591 bytes)
+ tests/data/mtcars.parquet view
binary file changed (absent → 4564 bytes)
+ tests/data/transactions.parquet view
binary file changed (absent → 1746 bytes)