diff --git a/dataframe-parquet/benchmark/Writer10GB.hs b/dataframe-parquet/benchmark/Writer10GB.hs new file mode 100644 index 00000000..09be1cba --- /dev/null +++ b/dataframe-parquet/benchmark/Writer10GB.hs @@ -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) + ) + ] diff --git a/dataframe-parquet/dataframe-parquet.cabal b/dataframe-parquet/dataframe-parquet.cabal index 99d5fbbd..1d79e38c 100644 --- a/dataframe-parquet/dataframe-parquet.cabal +++ b/dataframe-parquet/dataframe-parquet.cabal @@ -19,6 +19,7 @@ maintainer: mschavinda@gmail.com 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: @@ -28,6 +29,11 @@ common warnings -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 @@ -44,6 +50,11 @@ library 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, @@ -52,6 +63,7 @@ library dataframe-core >= 2.4.1 && < 2.5, dataframe-operations >= 2.4 && < 2.5, 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, @@ -64,3 +76,53 @@ library 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-parquet, + filepath >= 1.4 && < 2, + temporary >= 1.3 && < 1.5, + 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.4 && < 2.5, + 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.4 && < 2.5, + 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 diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs new file mode 100644 index 00000000..de86e695 --- /dev/null +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs @@ -0,0 +1,319 @@ +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE OverloadedStrings #-} + +module DataFrame.IO.Parquet.Writer ( + writeParquet, + writeParquetWithOptions, + ParquetWriteOptions (..), + WriterStrategy (..), + defaultParquetWriteOptions, +) where + +import Control.Monad (when) +import qualified Data.ByteString as BS +import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef) +import Data.Int (Int64) +import Data.Maybe (fromJust) +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, + bufferResidency, + bufferToByteString, + flushBufferToBuffer, + flushBufferToFile, + mallocBuffer, + resetPosition, + withWritableBinaryFile, + writeByteString, + writeByteStringToFile, + writeWord32LE, + ) +import DataFrame.Internal.Column (Column, hasMissing) +import DataFrame.Internal.DataFrame ( + DataFrame, + columnNames, + dataframeDimensions, + getColumn, + ) +import qualified Pinch +import qualified Snappy + +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 (maxRows, _) = dataframeDimensions df + names = columnNames df + columnChunks_ <- + VB.fromList + <$> mapM + ( \columnName_ -> + initColumnChunkState + options + columnName_ + (fromJust (getColumn columnName_ df)) + ) + names + scratchBuffer_ <- mallocBuffer (max 1 options.pageSize) + 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 + loop :: Int -> IO () + loop rowNum + | rowNum >= maxRows = pure () + | otherwise = do + VB.forM_ columnChunks_ (writeRow options scratchBuffer_ rowNum) + modifyIORef' rowNumberRef_ (+ 1) + when ((rowNum + 1) `mod` interval == 0) $ do + size <- bufferedSize columnChunks_ + when (size >= options.rowGroupSize) (flushRowGroup options writerState) + loop (rowNum + 1) + loop 0 + flushRowGroup options writerState + rowGroupMetadata <- reverse <$> readIORef rowGroupMetadataRef_ + let schemaElements = + rootSchemaElement (VB.length columnChunks_) + : VB.toList (VB.map schema columnChunks_) + writeFooter output schemaElements maxRows rowGroupMetadata + +writeRow :: + ParquetWriteOptions -> MemoryBuffer -> Int -> ColumnChunkState -> IO () +writeRow options scratch rowNum columnChunkState = do + let page = columnChunkState.pageState + notNull <- columnChunkState.encoder.writeValue page.pageBuffer rowNum + when columnChunkState.nullable $ + pushDef page.definitionLevels (if notNull then 1 else 0) + modifyIORef' page.currentRowCount (+ 1) + pageRowCount <- readIORef page.currentRowCount + let subInterval = max 1 options.subBatchRows + when (pageRowCount `mod` subInterval == 0) $ do + flushDef page.definitionLevels + pageBufferResidency <- bufferResidency page.pageBuffer + defLevelsResidency <- bufferResidency page.definitionLevels.dlBuf + when + (pageBufferResidency + defLevelsResidency >= options.pageSize) + (flushPage options scratch columnChunkState) + +flushPage :: ParquetWriteOptions -> MemoryBuffer -> ColumnChunkState -> IO () +flushPage options scratch columnChunkState = do + let page = columnChunkState.pageState + numPageRows <- readIORef page.currentRowCount + when (numPageRows > 0) $ do + columnChunkState.encoder.finishValues page.pageBuffer + 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_ + } diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/DefLevels.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/DefLevels.hs new file mode 100644 index 00000000..c8b6ffca --- /dev/null +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/DefLevels.hs @@ -0,0 +1,57 @@ +{-# 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 + +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 + +writeDefRun :: DefLevels -> Int -> Int -> IO () +writeDefRun dl value count = do + writeLeb128 dl.dlBuf (fromIntegral (count `shiftL` 1)) + writeWord8 dl.dlBuf (fromIntegral value) + +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) diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs new file mode 100644 index 00000000..bd75757e --- /dev/null +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs @@ -0,0 +1,272 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.IO.Parquet.Writer.Encoder ( + Encoder (..), + buildEncoder, +) where + +import Control.Monad (when) +import Data.Bits (shiftL, (.|.)) +import Data.IORef (newIORef, readIORef, writeIORef) +import Data.Int (Int32, Int64) +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, + appendTextArraySlice, + writeDoubleLE, + writeFloatLE, + writeWord32LE, + writeWord64LE, + writeWord8, + ) +import DataFrame.Internal.Column ( + Column (..), + Columnable, + columnTypeString, + hasElemType, + ) +import DataFrame.Internal.Column.Bitmap ( + Bitmap, + bitmapTestBit, + ) +import DataFrame.Internal.PackedText ( + PackedTextData (..), + offAt, + selAt, + ) +import Pinch (enum, putField) +import Type.Reflection (typeRep) + +data Encoder = Encoder + { encType :: !ThriftType + , convertedType :: !(Maybe ConvertedType) + , logicalType :: !(Maybe LogicalType) + , writeValue :: !(MemoryBuffer -> Int -> IO Bool) + , finishValues :: !(MemoryBuffer -> IO ()) + } + +buildEncoder :: Column -> IO Encoder +buildEncoder col + | hasElemType @Int32 col = + pure $ + scalarEncoder @Int32 + (INT32 enum) + Nothing + Nothing + (\buffer -> writeWord32LE buffer . fromIntegral) + col + | hasElemType @Int64 col = + pure $ + scalarEncoder @Int64 + (INT64 enum) + Nothing + Nothing + (\buffer -> writeWord64LE buffer . fromIntegral) + col + | hasElemType @Float col = + pure $ scalarEncoder @Float (FLOAT enum) Nothing Nothing writeFloatLE col + | hasElemType @Double col = + pure $ scalarEncoder @Double (DOUBLE enum) Nothing Nothing writeDoubleLE 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, VU.Unbox a) => + ThriftType -> + Maybe ConvertedType -> + Maybe LogicalType -> + (MemoryBuffer -> a -> IO ()) -> + Column -> + Encoder +scalarEncoder tt conv logical writeValue col = + Encoder tt conv logical (columnWriter @a col writeValue) (const (pure ())) +{-# INLINEABLE scalarEncoder #-} +{-# SPECIALIZE scalarEncoder :: + ThriftType -> + Maybe ConvertedType -> + Maybe LogicalType -> + (MemoryBuffer -> Int32 -> IO ()) -> + Column -> + Encoder + #-} +{-# SPECIALIZE scalarEncoder :: + ThriftType -> + Maybe ConvertedType -> + Maybe LogicalType -> + (MemoryBuffer -> Int64 -> IO ()) -> + Column -> + Encoder + #-} +{-# SPECIALIZE scalarEncoder :: + ThriftType -> + Maybe ConvertedType -> + Maybe LogicalType -> + (MemoryBuffer -> Float -> IO ()) -> + Column -> + Encoder + #-} +{-# SPECIALIZE scalarEncoder :: + ThriftType -> + Maybe ConvertedType -> + Maybe LogicalType -> + (MemoryBuffer -> Double -> IO ()) -> + Column -> + Encoder + #-} + +columnWriter :: + forall a. + (Columnable a) => + Column -> + (MemoryBuffer -> a -> IO ()) -> + MemoryBuffer -> + Int -> + IO Bool +columnWriter col writeValue = 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 row + | isPresent bitmap row = writeValue buffer (at row) >> pure True + | otherwise = pure False + mismatch = + error + ("writeParquet: incompatible column representation for " <> columnTypeString col) +{-# INLINEABLE columnWriter #-} +{-# SPECIALIZE columnWriter :: + Column -> (MemoryBuffer -> Int32 -> IO ()) -> MemoryBuffer -> Int -> IO Bool + #-} +{-# SPECIALIZE columnWriter :: + Column -> (MemoryBuffer -> Int64 -> IO ()) -> MemoryBuffer -> Int -> IO Bool + #-} +{-# SPECIALIZE columnWriter :: + Column -> (MemoryBuffer -> Float -> IO ()) -> MemoryBuffer -> Int -> IO Bool + #-} +{-# SPECIALIZE columnWriter :: + Column -> (MemoryBuffer -> Double -> IO ()) -> MemoryBuffer -> Int -> IO Bool + #-} +{-# SPECIALIZE columnWriter :: + Column -> (MemoryBuffer -> Bool -> IO ()) -> MemoryBuffer -> Int -> IO Bool + #-} +{-# SPECIALIZE columnWriter :: + Column -> (MemoryBuffer -> UTCTime -> IO ()) -> MemoryBuffer -> Int -> IO 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 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 writeWord8 buffer bits' >> writeIORef bitsRef 0 >> writeIORef countRef 0 + else writeIORef bitsRef bits' >> writeIORef countRef count' + finish buffer = do + count <- readIORef countRef + when (count > 0) (readIORef bitsRef >>= writeWord8 buffer) + writeIORef bitsRef 0 + writeIORef countRef 0 + 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 + (const (pure ())) + 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 row + | isPresent bitmap row = + writeText buffer (VB.unsafeIndex values row) >> pure True + | otherwise = pure False + writePacked bitmap packed buffer row + | isPresent bitmap row = do + let baseRow = maybe row (`selAt` row) packed.ptSel + start = offAt packed.ptOffsets baseRow + end = offAt packed.ptOffsets (baseRow + 1) + writeTextSlice buffer packed.ptBytes start (end - start) + pure True + | otherwise = pure False + mismatch = + error + ("writeParquet: incompatible text representation for " <> columnTypeString col) + +writeText :: MemoryBuffer -> T.Text -> IO () +writeText buffer (Text bytes offset count) = writeTextSlice buffer bytes offset count + +writeTextSlice :: MemoryBuffer -> TA.Array -> Int -> Int -> IO () +writeTextSlice buffer bytes offset count = do + writeWord32LE buffer (fromIntegral count) + appendTextArraySlice buffer bytes offset count +{-# INLINE writeTextSlice #-} + +timestampEncoder :: Column -> Encoder +timestampEncoder col = + Encoder + (INT64 enum) + (Just (TIMESTAMP_MICROS enum)) + (Just timestampLogical) + (columnWriter @UTCTime col writeMicros) + (const (pure ())) + where + writeMicros buffer t = writeWord64LE buffer (fromIntegral (utcToMicros t)) + +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 + ) diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Metadata.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Metadata.hs new file mode 100644 index 00000000..3d707852 --- /dev/null +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Metadata.hs @@ -0,0 +1,162 @@ +{-# 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] -> IO () +writeFooter output schemaElements numRows rowGroupMetadata = do + let metadata = + FileMetadata + { version = putField 1 + , schema = putField schemaElements + , num_rows = putField (fromIntegral numRows) + , row_groups = putField rowGroupMetadata + , key_value_metadata = putField Nothing + , 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" diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Options.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Options.hs new file mode 100644 index 00000000..40d730e6 --- /dev/null +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Options.hs @@ -0,0 +1,32 @@ +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 + } + deriving (Eq, Show) + +defaultParquetWriteOptions :: ParquetWriteOptions +defaultParquetWriteOptions = + ParquetWriteOptions + { pageSize = 1048576 + , rowGroupSize = 134217728 + , batchRows = 8192 + , subBatchRows = 2048 + , compressionCodec = SNAPPY enum + , strategy = InMemory + } diff --git a/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs b/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs index c6b84655..eb5f4fba 100644 --- a/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs +++ b/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs @@ -1,21 +1,72 @@ {-# 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, + MemoryBuffer (..), + mallocBuffer, + writeByteString, + appendTextArraySlice, + writeWord8, + writeWord32LE, + writeWord64LE, + writeFloatLE, + writeDoubleLE, + bufferResidency, + bufferToByteString, + flushBufferToBuffer, + resetPosition, + flushBufferToFile, + writeByteStringToFile, +) where +import Control.Exception (bracket) 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.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.IO ( + BufferMode (NoBuffering), + Handle, + IOMode (WriteMode), SeekMode (AbsoluteSeek), + hClose, + hPutBuf, + hSetBinaryMode, + hSetBuffering, + openBinaryFile, ) uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d @@ -76,3 +127,210 @@ unsafeToByteString :: VS.Vector Word8 -> ByteString 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 + +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 + +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) + +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) + +writeWord32LE :: MemoryBuffer -> Word32 -> IO () +writeWord32LE buffer w = do + position <- readIORef buffer.positionRef + 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) + writeIORef buffer.positionRef (position + 4) + +writeWord64LE :: MemoryBuffer -> Word64 -> IO () +writeWord64LE buffer w = do + position <- readIORef buffer.positionRef + 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) + writeIORef buffer.positionRef (position + 8) + +writeFloatLE :: MemoryBuffer -> Float -> IO () +writeFloatLE buffer = writeWord32LE buffer . castFloatToWord32 + +writeDoubleLE :: MemoryBuffer -> Double -> IO () +writeDoubleLE buffer = writeWord64LE buffer . castDoubleToWord64 + +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 + +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 + +resetPosition :: MemoryBuffer -> IO () +resetPosition buffer = writeIORef buffer.positionRef 0 + +-- 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 #-} diff --git a/dataframe-parquet/stress/DataFrame10GB.hs b/dataframe-parquet/stress/DataFrame10GB.hs new file mode 100644 index 00000000..25a2db8a --- /dev/null +++ b/dataframe-parquet/stress/DataFrame10GB.hs @@ -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.DataFrame (DataFrame, fromNamedColumns) +import DataFrame.Internal.PackedText (mkPackedContiguous32) + +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 diff --git a/dataframe-parquet/stress/StressMain.hs b/dataframe-parquet/stress/StressMain.hs new file mode 100644 index 00000000..976ce3d7 --- /dev/null +++ b/dataframe-parquet/stress/StressMain.hs @@ -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 diff --git a/dataframe-parquet/tests/Main.hs b/dataframe-parquet/tests/Main.hs new file mode 100644 index 00000000..ef6a4662 --- /dev/null +++ b/dataframe-parquet/tests/Main.hs @@ -0,0 +1,169 @@ +-- | Tests for the writer-buffer logic in "DataFrame.IO.Utils.RandomAccess". +module Main where + +import qualified Data.ByteString as BS +import qualified System.Exit as Exit +import System.FilePath (()) +import System.IO.Temp (withSystemTempDirectory) +import Test.HUnit + +import DataFrame.IO.Parquet (readParquet) +import DataFrame.IO.Parquet.Writer ( + ParquetWriteOptions (..), + defaultParquetWriteOptions, + writeParquet, + writeParquetWithOptions, + ) +import DataFrame.IO.Utils.RandomAccess + +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' + +tinyWriteOpts :: ParquetWriteOptions +tinyWriteOpts = + defaultParquetWriteOptions + { pageSize = 64 + , rowGroupSize = 512 + , batchRows = 4 + } + +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: alltypes_plain multi-page" + ( writerRoundTripTiny + "alltypes_plain multi-page" + "tests/data/alltypes_plain.parquet" + ) + ] + +main :: IO () +main = do + result <- runTestTT tests + if failures result > 0 || errors result > 0 + then Exit.exitFailure + else Exit.exitSuccess diff --git a/dataframe-parquet/tests/data/alltypes_dictionary.parquet b/dataframe-parquet/tests/data/alltypes_dictionary.parquet new file mode 100644 index 00000000..e6da6ab7 Binary files /dev/null and b/dataframe-parquet/tests/data/alltypes_dictionary.parquet differ diff --git a/dataframe-parquet/tests/data/alltypes_plain.parquet b/dataframe-parquet/tests/data/alltypes_plain.parquet new file mode 100644 index 00000000..a63f5dca Binary files /dev/null and b/dataframe-parquet/tests/data/alltypes_plain.parquet differ diff --git a/dataframe-parquet/tests/data/alltypes_plain.snappy.parquet b/dataframe-parquet/tests/data/alltypes_plain.snappy.parquet new file mode 100644 index 00000000..9809d676 Binary files /dev/null and b/dataframe-parquet/tests/data/alltypes_plain.snappy.parquet differ diff --git a/dataframe-parquet/tests/data/alltypes_tiny_pages.parquet b/dataframe-parquet/tests/data/alltypes_tiny_pages.parquet new file mode 100644 index 00000000..90019d16 Binary files /dev/null and b/dataframe-parquet/tests/data/alltypes_tiny_pages.parquet differ diff --git a/dataframe-parquet/tests/data/int32_decimal.parquet b/dataframe-parquet/tests/data/int32_decimal.parquet new file mode 100644 index 00000000..5bf2d4ea Binary files /dev/null and b/dataframe-parquet/tests/data/int32_decimal.parquet differ diff --git a/dataframe-parquet/tests/data/int64_decimal.parquet b/dataframe-parquet/tests/data/int64_decimal.parquet new file mode 100644 index 00000000..5043bcac Binary files /dev/null and b/dataframe-parquet/tests/data/int64_decimal.parquet differ diff --git a/dataframe-parquet/tests/data/mtcars.parquet b/dataframe-parquet/tests/data/mtcars.parquet new file mode 100644 index 00000000..cbf0c163 Binary files /dev/null and b/dataframe-parquet/tests/data/mtcars.parquet differ diff --git a/dataframe-parquet/tests/data/transactions.parquet b/dataframe-parquet/tests/data/transactions.parquet new file mode 100644 index 00000000..ba194bdc Binary files /dev/null and b/dataframe-parquet/tests/data/transactions.parquet differ diff --git a/dataframe.cabal b/dataframe.cabal index f59dca1e..3714076a 100644 --- a/dataframe.cabal +++ b/dataframe.cabal @@ -153,6 +153,7 @@ library DataFrame.IO.Parquet.Page, DataFrame.IO.Parquet.Schema, DataFrame.IO.Parquet.Utils, + DataFrame.IO.Parquet.Writer, DataFrame.IO.Parquet.Seeking, DataFrame.IO.Parquet.Time, DataFrame.IO.Utils.RandomAccess, diff --git a/examples/examples.cabal b/examples/examples.cabal index f4531c9b..05de1429 100644 --- a/examples/examples.cabal +++ b/examples/examples.cabal @@ -72,6 +72,7 @@ executable examples DataFrame.IO.Parquet.Thrift, DataFrame.IO.Parquet.Time, DataFrame.IO.Parquet.Utils, + DataFrame.IO.Parquet.Writer, DataFrame.IO.Utils.RandomAccess, DataFrame.Lazy.IO.CSV, DataFrame.Lazy.IO.Binary,