dataframe-parquet 1.1.0.1 → 1.5.1.0
raw patch · 23 files changed
Files
- benchmark/Writer10GB.hs +49/−0
- dataframe-parquet.cabal +76/−11
- src/DataFrame/IO/Parquet.hs +226/−21
- src/DataFrame/IO/Parquet/Page.hs +160/−31
- src/DataFrame/IO/Parquet/Utils.hs +7/−1
- src/DataFrame/IO/Parquet/Writer.hs +413/−0
- src/DataFrame/IO/Parquet/Writer/DefLevels.hs +61/−0
- src/DataFrame/IO/Parquet/Writer/Encoder.hs +394/−0
- src/DataFrame/IO/Parquet/Writer/Metadata.hs +175/−0
- src/DataFrame/IO/Parquet/Writer/Options.hs +34/−0
- src/DataFrame/IO/Utils/RandomAccess.hs +338/−5
- src/DataFrame/Typed/IO/Parquet.hs +111/−0
- stress/DataFrame10GB.hs +137/−0
- stress/StressMain.hs +38/−0
- tests/Main.hs +283/−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.5.1.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,82 @@ 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.5 && < 2.6,+ dataframe-operations >= 2.5 && < 2.6,+ 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.5 && < 2.6,+ 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.5 && < 2.6,+ 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.5 && < 2.6,+ 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,11 +8,66 @@ {-# 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,+) where+ 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)@@ -25,26 +81,29 @@ 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 qualified Data.Vector as Vector import qualified Data.Vector.Unboxed as VU import DataFrame.Errors (DataFrameException (ColumnsNotFoundException)) import DataFrame.IO.Parquet.Page ( PageDecoder, UnboxedPageDecoder,+ appendNullableStringPageIO,+ appendStringPageIO, boolDecoder, byteArrayDecoder, doubleDecoder, fixedLenByteArrayDecoder, floatDecoder,+ foldColumnDataPagesM,+ foldColumnPagesM, int32Decoder, int64Decoder, int96Decoder,- foldColumnPagesM, ) import DataFrame.IO.Parquet.Seeking (- FileBufferedOrSeekable,+ FileBufferedOrSeekable (..), ForceNonSeekable, withFileBufferedOrSeekable, )@@ -52,6 +111,7 @@ ColumnChunk (..), DecimalType (..), FileMetadata (..),+ KeyValue (..), LogicalType (..), RowGroup (..), ThriftType (..),@@ -70,16 +130,35 @@ generateColumnDescriptions, getColumnNames, )+import DataFrame.IO.Parquet.Writer (+ ParquetWriteOptions (..),+ WriterStrategy (..),+ defaultParquetWriteOptions,+ nativeTypeKeyPrefix,+ writeParquet,+ writeParquetWithOptions,+ ) import DataFrame.IO.Utils.RandomAccess ( RandomAccess (..), ReaderIO (runReaderIO), ) import DataFrame.Internal.Column (Column, Columnable) import qualified DataFrame.Internal.Column as DI+import DataFrame.Internal.Column.Builder (+ freezeTextChunk,+ mergeTextChunks,+ newTextBuilder,+ ) import DataFrame.Internal.DataFrame (DataFrame (..)) import DataFrame.Internal.Expression (Expr, getColumns) import DataFrame.Operations.Merge () import qualified DataFrame.Operations.Subset as DS+import DataFrame.Schema (+ Schema (..),+ SchemaType,+ makeSchema,+ schemaType,+ ) import qualified Pinch import System.Directory (doesDirectoryExist) import System.FilePath ((</>))@@ -198,7 +277,7 @@ matches <- glob pat - files <- filterM (fmap not . doesDirectoryExist) matches+ files <- L.sort <$> filterM (fmap not . doesDirectoryExist) matches case files of [] ->@@ -219,6 +298,10 @@ (RandomAccess m, MonadIO m) => ParquetReadOptions -> m DataFrame+{-# SPECIALIZE parseParquetWithOpts ::+ ParquetReadOptions -> ReaderIO FileBufferedOrSeekable DataFrame+ #-}+{-# INLINEABLE parseParquetWithOpts #-} parseParquetWithOpts opts = do metadata <- parseFileMetadata @@ -286,7 +369,12 @@ rawCols <- zipWithM (parseColumnChunks vectorLength) keptChunks keptDescs - let finalCols = zipWith applyDescLogicalType keptDescs rawCols+ let hints = nativeTypeHints metadata+ finalCols =+ zipWith+ (restoreNativeType hints)+ keptNames+ (zipWith applyDescLogicalType keptDescs rawCols) indices = Map.fromList $ zip keptNames [0 ..] dimensions = (vectorLength, length finalCols) @@ -344,6 +432,13 @@ [ColumnChunk] -> ColumnDescription -> m Column+{-# SPECIALIZE parseColumnChunks ::+ Int ->+ [ColumnChunk] ->+ ColumnDescription ->+ ReaderIO FileBufferedOrSeekable Column+ #-}+{-# INLINEABLE parseColumnChunks #-} parseColumnChunks totalRows chunks description | description.maxRepetitionLevel == 0 && description.maxDefinitionLevel == 0 = getNonNullableColumn totalRows description chunks@@ -353,6 +448,7 @@ getRepeatedColumn description chunks -- | Decode a required (non-nullable, non-repeated) column.+{-# INLINEABLE getNonNullableColumn #-} getNonNullableColumn :: forall m. (RandomAccess m, MonadIO m) =>@@ -368,7 +464,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 +478,27 @@ go decoder = foldNonNullable totalRows (foldColumnPagesM description decoder chunks) + -- Decode a non-nullable BYTE_ARRAY (UTF-8) column straight into a single+ -- shared byte buffer + offsets ('PackedText'), instead of a boxed vector+ -- of per-row 'Text'. Each page's decoded 'Text' values (which share the+ -- chunk dictionary for dictionary-encoded pages) are appended by memcpy+ -- into one builder across all pages/chunks, then frozen once. This is the+ -- same representation the fast CSV reader uses and matches Arrow's string+ -- layout: no retained per-row 'Text' headers, no eager UTF-8 validation.+ goPackedText :: m Column+ goPackedText = do+ builder <- liftIO $ stToIO (newTextBuilder totalRows (totalRows * 8))+ _ <-+ 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) =>@@ -391,6 +508,7 @@ foldNonNullableUnboxed totalRows (foldColumnPagesM description decoder chunks) -- | Decode an optional (nullable) column.+{-# INLINEABLE getNullableColumn #-} getNullableColumn :: forall m. (RandomAccess m, MonadIO m) =>@@ -406,7 +524,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 +540,38 @@ m Column go decoder = foldNullable maxDef totalRows (foldColumnPagesM description decoder chunks)++ -- Nullable BYTE_ARRAY (UTF-8): decode straight into a 'PackedText' (shared+ -- byte buffer + offsets + validity bitmap) via the text builder, walking+ -- def-levels to interleave nulls. Avoids the boxed @Vector Text@ the+ -- generic 'foldNullable' path would build.+ goPackedTextNullable :: m Column+ goPackedTextNullable = do+ builder <- liftIO $ stToIO (newTextBuilder totalRows (totalRows * 8))+ _ <-+ 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) =>@@ -519,14 +660,15 @@ 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+ -- (ticks per second, picoseconds per tick) for each unit. Convert at+ -- native precision: a millisecond grid keeps ms, a nanosecond grid+ -- keeps ns. (The old code multiplied by @1_000_000 `div` divisor@,+ -- which truncated NANOS to 0 and collapsed every value to the epoch.)+ 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@@ -547,6 +689,69 @@ 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 Map.lookup leaf (elements hints) of+ Just ty+ | stampedAs @Int ty -> narrow (fromIntegral @Int64 @Int)+ | stampedAs @Integer ty -> narrow (fromIntegral @Int64 @Integer)+ _ -> col+ where+ leaf = last (T.splitOn "." name)+ 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 expressed as @ticksPerSecond@ ticks/second+(each tick = @psPerTick@ picoseconds) to 'UTCTime', at full precision.++Splits into days + within-day picoseconds with integer 'divMod' (which floors,+so the split is correct for pre-epoch negative values too), and never forms+picoseconds-since-epoch — that would overflow 'Int64' for modern dates — only+the bounded within-day picosecond count. 40587 is the Modified Julian Day of+the Unix epoch (1970-01-01).+-}+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/Page.hs view
@@ -17,11 +17,17 @@ fixedLenByteArrayDecoder, -- Page iteration foldColumnPagesM,+ foldColumnDataPagesM,+ RawPage,+ 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,6 +36,7 @@ import qualified Data.Vector as VB import qualified Data.Vector.Generic as VG import qualified Data.Vector.Unboxed as VU+import Data.Word (Word8) import DataFrame.IO.Parquet.Decompress (decompressData) import DataFrame.IO.Parquet.Dictionary ( DictVals (..),@@ -58,6 +65,14 @@ littleEndianWord32, littleEndianWord64, )+import DataFrame.Internal.Column.Builder (+ TextBuilder,+ appendNull,+ appendText,+ appendTextSliceFromPtr,+ )+import Foreign.Ptr (Ptr, castPtr, plusPtr)+import Foreign.Storable (peekByteOff) import GHC.Float (castWord32ToFloat, castWord64ToDouble) import Pinch (decodeWithLeftovers) import qualified Pinch@@ -82,7 +97,12 @@ boolDecoder :: UnboxedPageDecoder Bool boolDecoder mDict enc nPresent bs = case enc of- PLAIN _ -> VU.fromList (readNBool nPresent bs)+ -- PLAIN bools are bit-packed (1 bit/value, LSB-first). Generate the+ -- unboxed vector directly by indexing the bit for each row, avoiding an+ -- intermediate @[Bool]@ list.+ PLAIN _ ->+ VU.generate nPresent $ \i ->+ (BSU.unsafeIndex bs (i `shiftR` 3) `shiftR` (i .&. 7)) .&. 1 == 1 RLE_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool PLAIN_DICTIONARY _ -> unboxedLookupDict mDict nPresent bs getBool _ -> error ("boolDecoder: unsupported encoding " ++ show enc)@@ -208,17 +228,20 @@ 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.+{- | A decoded DATA page handed to a fold step: the running dictionary, the+page encoding, the present-value count, the (decompressed) value bytes, and the+definition/repetition level vectors.+-}+type RawPage =+ (Maybe DictVals, Encoding, Int, BS.ByteString, VU.Vector Int, VU.Vector Int) -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).+{- | Left-fold a monadic step over every DATA page (V1 or V2) of every column+chunk, in order, threading the running dictionary internally and handing each+page to @step@ as a 'RawPage' (no value decoding — the step decides how to+materialize, e.g. into a typed vector or straight into a text buffer). -Dictionary pages are consumed silently and update the running dictionary that-is threaded through the recursion. @INDEX_PAGE@s are skipped.+Dictionary pages update the running dictionary; @INDEX_PAGE@s are skipped. Only+one page's bytes are live at a time. -- 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@@ -229,16 +252,16 @@ -- overlap the requested range, avoiding decompression of irrelevant pages -- entirely. -}-foldColumnPagesM ::- forall m v a acc.- (RandomAccess m, MonadIO m, VG.Vector v a) =>+{-# 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 +275,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 +300,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 =@@ -298,11 +320,127 @@ 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 +{- | Left-fold over per-page value triples, decoding each page with @decoder@.+A thin wrapper over 'foldColumnDataPagesM' for the typed (numeric/boxed) column+paths.+-}+{-# 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)++{- | Append a non-nullable BYTE_ARRAY page's strings straight into a+'TextBuilder', no intermediate boxed 'Text' vector.++PLAIN pages copy each value's UTF-8 bytes directly from the page buffer pointer+('appendTextSliceFromPtr'), skipping the per-value 'decodeUtf8Lenient' ++'ByteString' slicing + 'Text' allocation that @readNTexts@ would do.+Dictionary pages append the shared dictionary 'Text's by reference.+-}+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 ds) ->+ let (idxs, _) = decodeDictIndices nPresent bs+ in stToIO (VU.mapM_ (\i -> appendText builder (ds VB.! i)) idxs)+ Just d -> error ("appendStringPageIO: wrong dict type, got " ++ show d)+ Nothing -> error "appendStringPageIO: dictionary-encoded page but no dictionary seen"++{- | Read a little-endian 4-byte length prefix at byte @off@ 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 #-}++{- | Append a NULLABLE BYTE_ARRAY page into a 'TextBuilder', interleaving nulls+by walking the definition levels: a present row (@def == maxDef@) consumes the+next stored value (dictionary entry by reference, or the next length-prefixed+PLAIN slice by memcpy); a null row ('appendNull') writes an offset and clears a+validity bit. Builds 'PackedText' + validity bitmap with no per-row 'Text' or+@[Maybe a]@ allocation. Only present values are stored in the page payload.+-}+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 ds) -> do+ let (idxs, _) = decodeDictIndices nPresent bs+ go !_ !i | i >= nDefs = pure ()+ go !j !i+ | VU.unsafeIndex defs i == maxDef =+ stToIO (appendText builder (ds VB.! VU.unsafeIndex idxs j))+ >> 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"+ -- --------------------------------------------------------------------------- -- Page header parsing -- ---------------------------------------------------------------------------@@ -313,15 +451,6 @@ -- --------------------------------------------------------------------------- -- Batch value readers -- -----------------------------------------------------------------------------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 readNInt32 :: Int -> BS.ByteString -> VU.Vector Int32 readNInt32 n bs = VU.generate n $ \i -> littleEndianInt32 (BS.drop (4 * i) bs)
src/DataFrame/IO/Parquet/Utils.hs view
@@ -44,9 +44,9 @@ 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@@ -167,6 +167,7 @@ single streaming left fold ('PageFold'), avoiding any intermediate list or concatenation allocation. Only one page's values are live at a time. -}+{-# INLINEABLE foldNonNullable #-} foldNonNullable :: forall m a. (MonadIO m, Columnable a) =>@@ -186,6 +187,7 @@ v <- liftIO $ VB.unsafeFreeze mv return (BoxedColumn Nothing v) +{-# INLINEABLE foldNonNullableUnboxed #-} foldNonNullableUnboxed :: forall m a. (MonadIO m, Columnable a, VU.Unbox a) =>@@ -221,6 +223,7 @@ A 'hasNull' flag is accumulated during the scatter so the 'buildBitmapFromValid' call is skipped entirely when all values are present. -}+{-# INLINEABLE foldNullable #-} foldNullable :: forall m a. (MonadIO m, Columnable a) =>@@ -264,6 +267,7 @@ else return Nothing return (BoxedColumn maybeBm dat) +{-# INLINEABLE foldNullableUnboxed #-} foldNullableUnboxed :: forall m a. (MonadIO m, Columnable a, VU.Unbox a) =>@@ -310,6 +314,7 @@ Threshold formula: @defT_r = maxDef - 2 * (maxRep - r)@. -}+{-# INLINEABLE foldRepeated #-} foldRepeated :: forall m a. ( MonadIO m@@ -333,6 +338,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
+ src/DataFrame/IO/Parquet/Writer.hs view
@@ -0,0 +1,413 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE OverloadedRecordDot #-}+{-# LANGUAGE OverloadedStrings #-}++module DataFrame.IO.Parquet.Writer (+ writeParquet,+ writeParquetWithOptions,+ ParquetWriteOptions (..),+ WriterStrategy (..),+ defaultParquetWriteOptions,+ nativeTypeKeyPrefix,+ nativeTypeKeyValues,+) where++import Control.Monad (forM_, unless, when)+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 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,+ bufferToByteString,+ ensureCapacity,+ flushBufferToBuffer,+ flushBufferToFile,+ mallocBuffer,+ resetPosition,+ 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)+ , currentFileOffsetRef :: !(IORef Int64)+ , scratchBuffer :: !MemoryBuffer+ , rowGroupMetadataRef :: !(IORef [RowGroup])+ , rowNumberRef :: !(IORef Int)+ }++data ColumnChunkState = ColumnChunkState+ { columnName :: !T.Text+ , nullable :: !Bool+ , schema :: !SchemaElement+ , encoder :: !Encoder+ , buffer :: !MemoryBuffer+ , uncompressedBufferSize :: !(IORef Int64)+ , pageState :: !PageState+ }++data PageState = PageState+ { pageBuffer :: !MemoryBuffer+ , definitionLevels :: !DefLevels+ , currentRowCount :: !(IORef 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_ <-+ VB.fromList+ <$> mapM+ ( \columnName_ ->+ initColumnChunkState+ options+ columnName_+ (fromJust (getColumn columnName_ df))+ )+ names+ scratchBuffer_ <- 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_ (writeRows options scratchBuffer_ rowNum count)+ 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 <- 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]+ ]++writeRows ::+ ParquetWriteOptions -> MemoryBuffer -> Int -> Int -> ColumnChunkState -> IO ()+writeRows options scratch firstRow count ccs = do+ let page = ccs.pageState+ buf = page.pageBuffer+ encode = ccs.encoder.encodeValue+ dl = page.definitionLevels+ end = firstRow + count++ pos0 <- readIORef buf.positionRef+ let margin = options.pageSize+ arr0 <- ensureCapacity buf (pos0 + max margin (count * 64))+ size0 <- getSizeofMutableByteArray arr0++ let go !size !pos !row+ | row >= end = writeIORef buf.positionRef pos+ | pos + margin > size = do+ -- Rare: buffer nearly full, grow it+ writeIORef 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+ modifyIORef' page.currentRowCount (+ count)+ flushDef dl+ pageRes <- bufferResidency buf+ defRes <- bufferResidency dl.dlBuf+ when+ (pageRes + defRes >= options.pageSize)+ (flushPage options scratch ccs)++flushPage :: ParquetWriteOptions -> MemoryBuffer -> ColumnChunkState -> IO ()+flushPage options scratch columnChunkState = do+ let page = columnChunkState.pageState+ numPageRows <- readIORef page.currentRowCount+ when (numPageRows > 0) $ do+ pos <- readIORef page.pageBuffer.positionRef+ pos' <- columnChunkState.encoder.finishValues page.pageBuffer pos+ writeIORef page.pageBuffer.positionRef pos'+ body <- assemblePageBody scratch columnChunkState+ writeDataPage options.compressionCodec numPageRows body columnChunkState+ resetPosition page.pageBuffer+ resetPosition page.definitionLevels.dlBuf+ resetPosition scratch+ writeIORef page.currentRowCount 0++assemblePageBody :: MemoryBuffer -> ColumnChunkState -> IO MemoryBuffer+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 -> ColumnChunkState -> IO ()+writeDataPage codec numPageRows body columnChunkState = do+ uncompressedPageSize <- bufferResidency body+ compressedBody <- case codec of+ UNCOMPRESSED _ -> pure Nothing+ SNAPPY _ -> Just . Snappy.compress <$> bufferToByteString body+ 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 -> flushBufferToBuffer body columnChunkState.buffer+ Just bytes -> writeByteString columnChunkState.buffer bytes+ modifyIORef'+ 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 <- bufferResidency columnChunkState.buffer+ uncompressedSize <- readIORef columnChunkState.uncompressedBufferSize+ flushBufferToFile writerState.outputFileHandle columnChunkState.buffer+ writeIORef+ writerState.currentFileOffsetRef+ (offset + fromIntegral compressedSize)+ writeIORef 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 -> IO 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 -> IO ColumnChunkState+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+ -- ColumnChunk Buffers start at page size and grow to their+ -- actual size over the course of building out the first row+ -- group.+ -- Each column chunk in a row group must have the same number+ -- of rows, but each column chunk is liable to fit the same+ -- number of rows in varying amounts of data depending on the+ -- encoding and the compression characteristics of the data.+ -- So the optimal buffer size of each column chunk is liable+ -- to vary+ -- As a result while one specific column chunk in a row group+ -- is likely to hit the page limit, the others are liable to be+ -- much smaller than the limit.+ buffer_ <- mallocBuffer bufferSize+ uncompressedBufferSize_ <- newIORef 0+ pageState_ <- initPageState bufferSize+ pure+ ColumnChunkState+ { columnName = columnName_+ , nullable = nullable_+ , schema = schema_+ , encoder = encoder_+ , buffer = buffer_+ , uncompressedBufferSize = uncompressedBufferSize_+ , pageState = pageState_+ }++initPageState :: Int -> IO PageState+initPageState bufferSize = do+ pageBuffer_ <- mallocBuffer bufferSize+ definitionLevels_ <- newDefLevels+ currentRowCount_ <- newIORef 0+ pure+ PageState+ { pageBuffer = pageBuffer_+ , definitionLevels = definitionLevels_+ , currentRowCount = currentRowCount_+ }
+ src/DataFrame/IO/Parquet/Writer/DefLevels.hs view
@@ -0,0 +1,61 @@+{-# LANGUAGE OverloadedRecordDot #-}++module DataFrame.IO.Parquet.Writer.DefLevels (+ DefLevels (..),+ newDefLevels,+ pushDef,+ flushDef,+) where++import Control.Monad (when)+import Data.Bits (shiftL, shiftR, (.&.), (.|.))+import Data.IORef (IORef, newIORef, readIORef, writeIORef)+import Data.Word (Word64)+import DataFrame.IO.Utils.RandomAccess (MemoryBuffer, mallocBuffer, writeWord8)++data DefLevels = DefLevels+ { dlBuf :: !MemoryBuffer+ , dlValue :: !(IORef Int)+ , dlCount :: !(IORef Int)+ }++newDefLevels :: IO DefLevels+newDefLevels = DefLevels <$> mallocBuffer 64 <*> newIORef 0 <*> newIORef 0++pushDef :: DefLevels -> Int -> IO ()+pushDef dl value = do+ count <- readIORef dl.dlCount+ if count == 0+ then writeIORef dl.dlValue value >> writeIORef dl.dlCount 1+ else do+ current <- readIORef dl.dlValue+ if current == value+ then writeIORef dl.dlCount (count + 1)+ else do+ writeDefRun dl current count+ writeIORef dl.dlValue value+ writeIORef dl.dlCount 1+{-# INLINE pushDef #-}++flushDef :: DefLevels -> IO ()+flushDef dl = do+ count <- readIORef dl.dlCount+ when (count > 0) $ do+ value <- readIORef dl.dlValue+ writeDefRun dl value count+ writeIORef dl.dlCount 0+{-# INLINE flushDef #-}++writeDefRun :: DefLevels -> Int -> Int -> IO ()+writeDefRun dl value count = do+ writeLeb128 dl.dlBuf (fromIntegral (count `shiftL` 1))+ writeWord8 dl.dlBuf (fromIntegral value)+{-# INLINE writeDefRun #-}++writeLeb128 :: MemoryBuffer -> Word64 -> IO ()+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,394 @@+{-# 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 (stToIO)+import Data.Bits (shiftL, (.|.))+import Data.IORef (newIORef, readIORef, writeIORef)+import Data.Int (Int32, Int64)+import Data.Primitive.ByteArray (+ withMutableByteArrayContents,+ writeByteArray,+ )+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 = Encoder+ { encType :: !ThriftType+ , convertedType :: !(Maybe ConvertedType)+ , logicalType :: !(Maybe LogicalType)+ , encodeValue :: !(MemoryBuffer -> Int -> Int -> IO (Int, Bool))+ , finishValues :: !(MemoryBuffer -> Int -> IO Int)+ }++buildEncoder :: Column -> IO Encoder+buildEncoder col+ | hasElemType @Int32 col =+ pure $+ scalarEncoder @Int32+ (INT32 enum)+ Nothing+ Nothing+ (\buffer pos v -> writeWord32At buffer pos (fromIntegral v) >> pure (pos + 4))+ col+ | hasElemType @Int64 col =+ pure $+ scalarEncoder @Int64+ (INT64 enum)+ Nothing+ Nothing+ (\buffer pos v -> writeWord64At buffer pos (fromIntegral v) >> pure (pos + 8))+ col+ -- Ints in GHC can be 32 bit or 64 bit integers depending on the+ -- underlying computers architecture. So we'll do 64bit integers+ -- to cover all our bases+ | hasElemType @Int col =+ pure $+ scalarEncoder @Int+ (INT64 enum)+ Nothing+ Nothing+ (\buffer pos v -> writeWord64At buffer pos (fromIntegral v) >> pure (pos + 8))+ col+ | hasElemType @Integer col =+ pure $+ scalarEncoder @Integer+ (INT64 enum)+ Nothing+ Nothing+ writeInteger64At+ col+ | hasElemType @Float col =+ pure $+ scalarEncoder @Float+ (FLOAT enum)+ Nothing+ Nothing+ ( \buffer pos v -> writeWord32At buffer pos (castFloatToWord32 v) >> pure (pos + 4)+ )+ col+ | hasElemType @Double col =+ pure $+ scalarEncoder @Double+ (DOUBLE enum)+ Nothing+ Nothing+ ( \buffer pos v -> writeWord64At buffer pos (castDoubleToWord64 v) >> pure (pos + 8)+ )+ col+ | 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)++scalarEncoder ::+ forall a.+ (Columnable a) =>+ ThriftType ->+ Maybe ConvertedType ->+ Maybe LogicalType ->+ (MemoryBuffer -> Int -> a -> IO Int) ->+ Column ->+ Encoder+scalarEncoder tt conv logical writePrim col =+ Encoder tt conv logical (columnWriter @a col writePrim) (\_ pos -> pure pos)+{-# INLINEABLE scalarEncoder #-}+{-# SPECIALIZE scalarEncoder ::+ ThriftType ->+ Maybe ConvertedType ->+ Maybe LogicalType ->+ (MemoryBuffer -> Int -> Int32 -> IO Int) ->+ Column ->+ Encoder+ #-}+{-# SPECIALIZE scalarEncoder ::+ ThriftType ->+ Maybe ConvertedType ->+ Maybe LogicalType ->+ (MemoryBuffer -> Int -> Int64 -> IO Int) ->+ Column ->+ Encoder+ #-}+{-# SPECIALIZE scalarEncoder ::+ ThriftType ->+ Maybe ConvertedType ->+ Maybe LogicalType ->+ (MemoryBuffer -> Int -> Float -> IO Int) ->+ Column ->+ Encoder+ #-}+{-# SPECIALIZE scalarEncoder ::+ ThriftType ->+ Maybe ConvertedType ->+ Maybe LogicalType ->+ (MemoryBuffer -> Int -> Double -> IO Int) ->+ Column ->+ Encoder+ #-}+{-# SPECIALIZE scalarEncoder ::+ ThriftType ->+ Maybe ConvertedType ->+ Maybe LogicalType ->+ (MemoryBuffer -> Int -> Int -> IO Int) ->+ Column ->+ Encoder+ #-}+{-# SPECIALIZE scalarEncoder ::+ ThriftType ->+ Maybe ConvertedType ->+ Maybe LogicalType ->+ (MemoryBuffer -> Int -> Integer -> IO Int) ->+ Column ->+ Encoder+ #-}++columnWriter ::+ forall a.+ (Columnable a) =>+ Column ->+ (MemoryBuffer -> Int -> a -> IO Int) ->+ MemoryBuffer ->+ Int ->+ Int ->+ IO (Int, Bool)+columnWriter 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 (VU.unsafeIndex values)+ 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)+{-# INLINEABLE columnWriter #-}+{-# SPECIALIZE columnWriter ::+ Column ->+ (MemoryBuffer -> Int -> Int32 -> IO Int) ->+ MemoryBuffer ->+ Int ->+ Int ->+ IO (Int, Bool)+ #-}+{-# SPECIALIZE columnWriter ::+ Column ->+ (MemoryBuffer -> Int -> Int64 -> IO Int) ->+ MemoryBuffer ->+ Int ->+ Int ->+ IO (Int, Bool)+ #-}+{-# SPECIALIZE columnWriter ::+ Column ->+ (MemoryBuffer -> Int -> Float -> IO Int) ->+ MemoryBuffer ->+ Int ->+ Int ->+ IO (Int, Bool)+ #-}+{-# SPECIALIZE columnWriter ::+ Column ->+ (MemoryBuffer -> Int -> Double -> IO Int) ->+ MemoryBuffer ->+ Int ->+ Int ->+ IO (Int, Bool)+ #-}+{-# SPECIALIZE columnWriter ::+ Column ->+ (MemoryBuffer -> Int -> Bool -> IO Int) ->+ MemoryBuffer ->+ Int ->+ Int ->+ IO (Int, Bool)+ #-}+{-# SPECIALIZE columnWriter ::+ Column ->+ (MemoryBuffer -> Int -> UTCTime -> IO Int) ->+ MemoryBuffer ->+ Int ->+ Int ->+ IO (Int, Bool)+ #-}+{-# SPECIALIZE columnWriter ::+ Column ->+ (MemoryBuffer -> Int -> Int -> IO Int) ->+ MemoryBuffer ->+ Int ->+ Int ->+ IO (Int, Bool)+ #-}+{-# SPECIALIZE columnWriter ::+ Column ->+ (MemoryBuffer -> Int -> Integer -> IO Int) ->+ MemoryBuffer ->+ Int ->+ Int ->+ IO (Int, Bool)+ #-}++isPresent :: Maybe Bitmap -> Int -> Bool+isPresent Nothing _ = True+isPresent (Just bitmap) row = bitmapTestBit bitmap row+{-# INLINE isPresent #-}++boolEncoder :: Column -> IO Encoder+boolEncoder col = do+ bitsRef <- newIORef (0 :: Word8)+ countRef <- newIORef (0 :: Int)+ let addBit buffer pos value = do+ bits <- readIORef bitsRef+ count <- readIORef countRef+ let bits' = if value then bits .|. ((1 :: Word8) `shiftL` count) else bits+ count' = count + 1+ if count' == 8+ then do+ arr <- readIORef buffer.arrayRef+ writeByteArray arr pos bits'+ writeIORef bitsRef 0+ writeIORef countRef 0+ pure (pos + 1)+ else do+ writeIORef bitsRef bits'+ writeIORef countRef count'+ pure pos+ finish buffer pos = do+ count <- readIORef countRef+ pos' <-+ if count > 0+ then do+ bits <- readIORef bitsRef+ arr <- readIORef buffer.arrayRef+ writeByteArray arr pos bits+ pure (pos + 1)+ else pure pos+ writeIORef bitsRef 0+ writeIORef countRef 0+ pure pos'+ pure+ (Encoder (BOOLEAN enum) Nothing Nothing (columnWriter @Bool col addBit) finish)++textEncoder :: Column -> Encoder+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+ writeIORef buffer.positionRef pos+ _ <- ensureCapacity buffer (pos + 4 + count)+ writeWord32At buffer pos (fromIntegral count)+ arr <- readIORef buffer.arrayRef+ withMutableByteArrayContents arr $ \ptr ->+ stToIO (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+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,175 @@+{-# LANGUAGE OverloadedStrings #-}++module DataFrame.IO.Parquet.Writer.Metadata (+ mkSchemaElem,+ rootSchemaElement,+ mkDataPageHeader,+ mkColumnChunk,+ mkRowGroup,+ writeFooter,+ magic,+) where++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 <- mallocBuffer (BS.length footer + 8)+ writeByteString buffer footer+ 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,83 @@ {-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE OverloadedRecordDot #-} -module DataFrame.IO.Utils.RandomAccess where+module DataFrame.IO.Utils.RandomAccess (+ uncurry3,+ Range (..),+ RandomAccess (..),+ ReaderIO (runReaderIO),+ LocalFile,+ MMappedFile,+ unsafeToByteString,+ WritableBinaryHandle,+ openWritableBinaryFile,+ withWritableBinaryFile,+ atomicallyWriteFile,+ MemoryBuffer (..),+ ensureCapacity,+ mallocBuffer,+ writeByteString,+ appendTextArraySlice,+ writeWord8,+ writeWord32LE,+ writeWord64LE,+ writeInteger64,+ writeWord32At,+ writeWord64At,+ writeInteger64At,+ writeFloatLE,+ writeDoubleLE,+ bufferResidency,+ bufferToByteString,+ flushBufferToBuffer,+ resetPosition,+ flushBufferToFile,+ writeByteStringToFile,+) where +import Control.Exception (bracket, bracketOnError, finally)+import Control.Monad (when) import Control.Monad.IO.Class (MonadIO (..))-import Data.ByteString (ByteString)-import Data.ByteString.Internal (ByteString (PS))+import Control.Monad.Primitive (RealWorld)+import Control.Monad.ST (stToIO)+import Data.Bits (shiftR)+import qualified Data.ByteString as BS+import Data.ByteString.Internal (ByteString (PS), create)+import qualified Data.ByteString.Unsafe as BU+import Data.IORef (IORef, newIORef, readIORef, writeIORef)+import Data.Int (Int64)+import Data.Primitive.ByteArray (+ MutableByteArray,+ copyMutableByteArray,+ getSizeofMutableByteArray,+ newPinnedByteArray,+ withMutableByteArrayContents,+ writeByteArray,+ )+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 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@@ -76,3 +138,274 @@ unsafeToByteString v = PS (castForeignPtr ptr) offset' len where (ptr, offset', len) = VS.unsafeToForeignPtr v++-- Writer Buffer -----------------------------------------------------------------++-- Refer to DataFrame.IO.Parquet.Writer for a justification of what we're doing here+-- There's some overlap here with what's going on in Seeking.hs, so, if this bothers+-- us, eventually someone will have to come back and reconcile the writer buffer+-- approach with the reader oriented patterns in Seeking.hs.+--+-- We're using MutableByteArrays here for convenience and because we don't need+-- the more powerful abstractions vector provides (which uses ByteArrays internally)+--+-- since we want to use hPutBuf, we're going to need a Ptr, which means are ByteArrya+-- must be pinned. Now growing pinned arrays can be problematic, but in the vast majority+-- of cases we shouldn't be growing more than once, if that. See the docs for+-- Data.Primitive.ByteArray.byteArrayContents.++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 = MemoryBuffer+ { arrayRef :: !(IORef (MutableByteArray RealWorld))+ , positionRef :: !(IORef Int)+ }++mallocBuffer :: Int -> IO MemoryBuffer+mallocBuffer capacity+ | capacity < 0 = ioError $ userError "mallocBuffer: negative capacity"+ | otherwise = do+ array <- newPinnedByteArray capacity+ MemoryBuffer <$> newIORef array <*> newIORef 0++-- We're using pinned ByteArrays so we must+-- not use the grow function brovided by Data.Primitive+-- instead we must alloocate a new pinned ByteArray.+-- We might have been worried about heap fragmentation+-- because a single pinned object in a 4KB GHC block can+-- keep the whole plock alive but oyr buffers will tend to+-- be much larger than that.+-- But the memory usage will temporarily spike to 2.5x the size of+-- the buffer, but it should be fine since the current writer is single threaded+-- and grows *should* be rare.+-- If it becomes an issue we should start tracking an array of pointers+-- to buffers intsead of replacing them wholesale so grwoing a buffer+-- is just a matter of adding a new buffer to the array (which we can+-- pre-allocate to three elements to begin with and grow it only on the+-- off chance that a buffer required more than three grows).+ensureCapacity :: MemoryBuffer -> Int -> IO (MutableByteArray RealWorld)+ensureCapacity buffer needed = do+ array <- readIORef buffer.arrayRef+ maxSize <- getSizeofMutableByteArray array+ if needed <= maxSize+ then pure array+ else do+ position <- readIORef buffer.positionRef+ grown <- newPinnedByteArray (needed + (needed `div` 2))+ copyMutableByteArray grown 0 array 0 position+ writeIORef buffer.arrayRef grown+ pure grown+{-# INLINE ensureCapacity #-}++writeWord8 :: MemoryBuffer -> Word8 -> IO ()+writeWord8 buffer b = do+ position <- readIORef buffer.positionRef+ array <- ensureCapacity buffer (position + 1)+ writeByteArray array position b+ writeIORef buffer.positionRef (position + 1)+{-# INLINE writeWord8 #-}++writeByteString :: MemoryBuffer -> ByteString -> IO ()+writeByteString buffer bs =+ BU.unsafeUseAsCStringLen bs $ \(source, len) -> do+ position <- readIORef buffer.positionRef+ array <- ensureCapacity buffer (position + len)+ withMutableByteArrayContents array $ \dst ->+ copyBytes (dst `plusPtr` position) (castPtr source) len+ writeIORef buffer.positionRef (position + len)+{-# INLINE writeByteString #-}++writeWord32LE :: MemoryBuffer -> Word32 -> IO ()+writeWord32LE buffer w = do+ position <- readIORef buffer.positionRef+ writeWord32At buffer position w+ writeIORef buffer.positionRef (position + 4)+{-# INLINE writeWord32LE #-}++writeWord64LE :: MemoryBuffer -> Word64 -> IO ()+writeWord64LE buffer w = do+ position <- readIORef buffer.positionRef+ writeWord64At buffer position w+ writeIORef buffer.positionRef (position + 8)+{-# INLINE writeWord64LE #-}++writeWord32At :: MemoryBuffer -> Int -> Word32 -> IO ()+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 -> Int -> Word64 -> IO ()+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 -> Integer -> IO ()+writeInteger64 buffer value = do+ position <- readIORef buffer.positionRef+ newPosition <- writeInteger64At buffer position value+ writeIORef buffer.positionRef newPosition+{-# INLINE writeInteger64 #-}++writeInteger64At :: MemoryBuffer -> Int -> Integer -> IO 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 =+ ioError (userError "writeParquet: Integer value is outside the INT64 range")+{-# INLINE writeInteger64At #-}++writeFloatLE :: MemoryBuffer -> Float -> IO ()+writeFloatLE buffer = writeWord32LE buffer . castFloatToWord32+{-# INLINE writeFloatLE #-}++writeDoubleLE :: MemoryBuffer -> Double -> IO ()+writeDoubleLE buffer = writeWord64LE buffer . castDoubleToWord64+{-# INLINE writeDoubleLE #-}++flushBufferToBuffer :: MemoryBuffer -> MemoryBuffer -> IO ()+flushBufferToBuffer source destination+ | source.arrayRef == destination.arrayRef = pure ()+ | otherwise = do+ sourceArray <- readIORef source.arrayRef+ sourcePosition <- readIORef source.positionRef+ destinationPosition <- readIORef destination.positionRef+ destinationArray <-+ ensureCapacity destination (destinationPosition + sourcePosition)+ copyMutableByteArray+ destinationArray+ destinationPosition+ sourceArray+ 0+ sourcePosition+ writeIORef destination.positionRef (destinationPosition + sourcePosition)+ writeIORef source.positionRef 0+{-# INLINE flushBufferToBuffer #-}++bufferToByteString :: MemoryBuffer -> IO ByteString+bufferToByteString buffer = do+ array <- readIORef buffer.arrayRef+ position <- readIORef buffer.positionRef+ create position $ \dst ->+ withMutableByteArrayContents array $ \src ->+ copyBytes dst (castPtr src) position++bufferResidency :: MemoryBuffer -> IO Int+bufferResidency buffer = readIORef buffer.positionRef+{-# INLINE bufferResidency #-}++resetPosition :: MemoryBuffer -> IO ()+resetPosition buffer = writeIORef buffer.positionRef 0+{-# INLINE resetPosition #-}++-- I tested write speeds by doing (on Apple Silicon)+-- `dd if=/dev/zero of=test bs={$n}k oflag=direct conv=fdatasync+-- Results:+--+-- ```+-- | block size | data (GiB) | time (s) | GiB/s |+-- |------------|------------|-----------|-------|+-- | 4k | 4.00 | 2.371 | 1.69 |+-- | 8k | 4.00 | 1.486 | 2.69 |+-- | 16k | 4.00 | 1.045 | 3.83 |+-- | 32k | 4.00 | 0.740 | 5.40 |+-- | 64k | 4.00 | 0.675 | 5.92 |+-- | 128k | 4.00 | 0.669 | 5.98 |+-- | 256k | 4.00 | 0.664 | 6.03 |+-- | 512k | 4.00 | 0.670 | 5.97 |+-- | 1024k | 4.00 | 0.664 | 6.02 |+-- | 4096k | 4.00 | 0.668 | 5.99 |+-- ```+-- So when writing to a file to minimize syscall overhead while+-- trying not to create dirty pages in the kernel page cache, we'll+-- be flushing in 256 KiB chunks.+flushBufferToFile :: WritableBinaryHandle -> MemoryBuffer -> IO ()+flushBufferToFile (WritableBinaryHandle h) buffer = do+ array <- readIORef buffer.arrayRef+ position <- readIORef 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+ writeIORef buffer.positionRef 0++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 -> TA.Array -> Int -> Int -> IO ()+appendTextArraySlice buffer source offset count+ | count < 0 = ioError $ userError "appendTextArraySlice: negative length"+ | otherwise = do+ position <- readIORef buffer.positionRef+ array <- ensureCapacity buffer (position + count)+ withMutableByteArrayContents array $ \destination ->+ stToIO (TA.copyToPointer source offset (destination `plusPtr` position) count)+ writeIORef buffer.positionRef (position + count)+{-# INLINE appendTextArraySlice #-}
+ src/DataFrame/Typed/IO/Parquet.hs view
@@ -0,0 +1,111 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | Typed Parquet reading.++The reader validates the file against a type-level schema as it loads, supplied+by type application:++@+type Trips = '[ '(\"id\", Int), '(\"fare\", Double)]++trips <- readParquet \@Trips \"trips.parquet\" -- IO (TypedDataFrame Trips)+@++'readParquet' (and 'readParquetWithOpts'\/'readParquetFiles') throw a+'DataFrameException' on schema mismatch; 'readParquetWithError' returns the+mismatch as an 'Either' instead.+-}+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 @cols@ names, unless the caller selected columns+themselves. Parquet supplies the element types, so the schema contributes the+projection and the freeze checks the types.+-}+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 schema mismatch.+Reads only the columns @cols@ names.++==== __Example__+@+ghci> trips <- readParquet \@Trips \"trips.parquet\"+@+-}+readParquet ::+ forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols)+readParquet = readParquetWithOpts @cols defaultParquetReadOptions++{- | Read a Parquet file, returning a descriptive error on schema mismatch or+a missing column instead of throwing.++==== __Example__+@+ghci> readParquetWithError \@Trips \"trips.parquet\"+Right (TDF ...)+@+-}+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 schema mismatch. The+schema still supplies the column selection; an explicit 'selectedColumns'+takes precedence.++==== __Example__+@+ghci> trips <- readParquetWithOpts \@Trips defaultParquetReadOptions{rowRange = Just (0, 10)} \"trips.parquet\"+@+-}+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.++==== __Example__+@+ghci> trips <- readParquetFiles \@Trips \".\/data\/trips\/*.parquet\"+@+-}+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, fromNamedColumns)++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 = fromNamedColumns (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,283 @@+{-# LANGUAGE OverloadedStrings #-}++-- | 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 Data.Int (Int32, Int64)+import Data.Maybe (fromJust)+import qualified Data.Text as T+import DataFrame.IO.Parquet (readParquet, readParquetFiles)+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,+ fromNamedColumns,+ getColumn,+ )+import System.Directory (listDirectory)++directWrites :: Test+directWrites = TestCase $ do+ buffer <- mallocBuffer 1+ writeWord8 buffer 0xaa+ writeWord32LE buffer 0x78563412+ writeWord64LE buffer 0x0807060504030201+ writeFloatLE buffer 1+ writeDoubleLE buffer 1+ writeByteString buffer (BS.pack [0xfe, 0xff])+ residency <- 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 <- mallocBuffer 0+ destination <- mallocBuffer 0+ writeByteString destination (BS.pack [1, 2])+ writeByteString source (BS.pack [3, 4, 5])+ flushBufferToBuffer source destination+ sourceResidency <- bufferResidency source+ destinationBytes <- bufferToByteString destination+ assertEqual "source cleared" 0 sourceResidency+ assertEqual "destination appended" (BS.pack [1, 2, 3, 4, 5]) destinationBytes+ flushBufferToBuffer destination destination+ selfFlushedBytes <- bufferToByteString destination+ assertEqual "self flush is a no-op" destinationBytes selfFlushedBytes+ resetPosition destination+ destinationResidency <- 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 <- mallocBuffer 1+ writeByteString buffer payload+ withWritableBinaryFile outPath $ \output ->+ flushBufferToFile output buffer+ residency <- 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'++writerRoundTripLargeText :: Test+writerRoundTripLargeText = TestCase $+ withSystemTempDirectory "dfpq-writer" $ \dir -> do+ let payload = T.replicate 4096 "abcdefgh"+ df = fromNamedColumns [("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: @maxRowsPerFile@ splits the frame across a glob pattern,+and reading the shards back reproduces the original frame.+-}+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 '*'")++{- | Parquet has one 64-bit integer type, so @Int@, @Int64@ and @Integer@+columns all land in the file as @INT64@. The writer stamps the original+Haskell type in the footer so the reader can put it back; without that, a+CSV-inferred @Int@ column silently widens to @Int64@ on a round trip.+-}+writerRoundTripNativeIntTypes :: Test+writerRoundTripNativeIntTypes = TestCase $+ withSystemTempDirectory "dfpq-writer" $ \dir -> do+ let df =+ fromNamedColumns+ [ ("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+ }++tests :: Test+tests =+ TestList+ [ 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+ "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+ ]++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)