diff --git a/benchmark/Writer10GB.hs b/benchmark/Writer10GB.hs new file mode 100644 --- /dev/null +++ b/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.cabal b/dataframe-parquet.cabal --- a/dataframe-parquet.cabal +++ b/dataframe-parquet.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: dataframe-parquet -version: 1.5.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 @@ -19,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: @@ -28,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 @@ -44,6 +50,11 @@ 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 @@ 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, @@ -64,3 +76,56 @@ 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 diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs --- a/src/DataFrame/IO/Parquet.hs +++ b/src/DataFrame/IO/Parquet.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MonoLocalBinds #-} @@ -7,8 +8,62 @@ {-# 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 (..)) @@ -48,7 +103,7 @@ int96Decoder, ) import DataFrame.IO.Parquet.Seeking ( - FileBufferedOrSeekable, + FileBufferedOrSeekable (..), ForceNonSeekable, withFileBufferedOrSeekable, ) @@ -56,6 +111,7 @@ ColumnChunk (..), DecimalType (..), FileMetadata (..), + KeyValue (..), LogicalType (..), RowGroup (..), ThriftType (..), @@ -74,6 +130,14 @@ generateColumnDescriptions, getColumnNames, ) +import DataFrame.IO.Parquet.Writer ( + ParquetWriteOptions (..), + WriterStrategy (..), + defaultParquetWriteOptions, + nativeTypeKeyPrefix, + writeParquet, + writeParquetWithOptions, + ) import DataFrame.IO.Utils.RandomAccess ( RandomAccess (..), ReaderIO (runReaderIO), @@ -89,6 +153,12 @@ 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 (()) @@ -207,7 +277,7 @@ matches <- glob pat - files <- filterM (fmap not . doesDirectoryExist) matches + files <- L.sort <$> filterM (fmap not . doesDirectoryExist) matches case files of [] -> @@ -299,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) @@ -613,6 +688,55 @@ Left _ -> col else col applyLogicalType _ col = col + +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. diff --git a/src/DataFrame/IO/Parquet/Writer.hs b/src/DataFrame/IO/Parquet/Writer.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Writer.hs @@ -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_ + } diff --git a/src/DataFrame/IO/Parquet/Writer/DefLevels.hs b/src/DataFrame/IO/Parquet/Writer/DefLevels.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Writer/DefLevels.hs @@ -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 #-} diff --git a/src/DataFrame/IO/Parquet/Writer/Encoder.hs b/src/DataFrame/IO/Parquet/Writer/Encoder.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Writer/Encoder.hs @@ -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 #-} diff --git a/src/DataFrame/IO/Parquet/Writer/Metadata.hs b/src/DataFrame/IO/Parquet/Writer/Metadata.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Writer/Metadata.hs @@ -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" diff --git a/src/DataFrame/IO/Parquet/Writer/Options.hs b/src/DataFrame/IO/Parquet/Writer/Options.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Writer/Options.hs @@ -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 + } diff --git a/src/DataFrame/IO/Utils/RandomAccess.hs b/src/DataFrame/IO/Utils/RandomAccess.hs --- a/src/DataFrame/IO/Utils/RandomAccess.hs +++ b/src/DataFrame/IO/Utils/RandomAccess.hs @@ -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 #-} diff --git a/stress/DataFrame10GB.hs b/stress/DataFrame10GB.hs new file mode 100644 --- /dev/null +++ b/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.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 diff --git a/stress/StressMain.hs b/stress/StressMain.hs new file mode 100644 --- /dev/null +++ b/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/tests/Main.hs b/tests/Main.hs new file mode 100644 --- /dev/null +++ b/tests/Main.hs @@ -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 diff --git a/tests/data/alltypes_dictionary.parquet b/tests/data/alltypes_dictionary.parquet new file mode 100644 Binary files /dev/null and b/tests/data/alltypes_dictionary.parquet differ diff --git a/tests/data/alltypes_plain.parquet b/tests/data/alltypes_plain.parquet new file mode 100644 Binary files /dev/null and b/tests/data/alltypes_plain.parquet differ diff --git a/tests/data/alltypes_plain.snappy.parquet b/tests/data/alltypes_plain.snappy.parquet new file mode 100644 Binary files /dev/null and b/tests/data/alltypes_plain.snappy.parquet differ diff --git a/tests/data/alltypes_tiny_pages.parquet b/tests/data/alltypes_tiny_pages.parquet new file mode 100644 Binary files /dev/null and b/tests/data/alltypes_tiny_pages.parquet differ diff --git a/tests/data/int32_decimal.parquet b/tests/data/int32_decimal.parquet new file mode 100644 Binary files /dev/null and b/tests/data/int32_decimal.parquet differ diff --git a/tests/data/int64_decimal.parquet b/tests/data/int64_decimal.parquet new file mode 100644 Binary files /dev/null and b/tests/data/int64_decimal.parquet differ diff --git a/tests/data/mtcars.parquet b/tests/data/mtcars.parquet new file mode 100644 Binary files /dev/null and b/tests/data/mtcars.parquet differ diff --git a/tests/data/transactions.parquet b/tests/data/transactions.parquet new file mode 100644 Binary files /dev/null and b/tests/data/transactions.parquet differ