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 d86ff474..201e2f78 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,13 @@ library DataFrame.IO.Parquet.Thrift DataFrame.IO.Parquet.Time DataFrame.IO.Parquet.Utils + DataFrame.IO.Parquet.Writer + DataFrame.IO.Parquet.Writer.ColumnChunkWriter + DataFrame.IO.Parquet.Writer.DefLevels + DataFrame.IO.Parquet.Writer.Encoder + DataFrame.IO.Parquet.Writer.Metadata + DataFrame.IO.Parquet.Writer.Options + DataFrame.IO.Parquet.Writer.PageWriter DataFrame.IO.Utils.RandomAccess DataFrame.Typed.IO.Parquet build-depends: base >= 4 && < 5, @@ -53,6 +66,7 @@ library dataframe-core >= 2.4 && < 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, @@ -65,3 +79,54 @@ 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, + primitive >= 0.7 && < 0.11, + 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..b4a46709 --- /dev/null +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs @@ -0,0 +1,203 @@ +{-# 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.Vector as VB +import DataFrame.IO.Parquet.Thrift +import DataFrame.IO.Parquet.Writer.ColumnChunkWriter ( + ColumnChunkState (..), + bufferedSize, + finalizePage, + initColumnState, + runColumnChunkWriter, + writeRowAndMaybeFinalize, + ) +import DataFrame.IO.Parquet.Writer.Encoder (Encoder (..)) +import DataFrame.IO.Parquet.Writer.Metadata (magic, rootSchemaElement) +import DataFrame.IO.Parquet.Writer.Options ( + ParquetWriteOptions (..), + WriterStrategy (..), + defaultParquetWriteOptions, + ) +import DataFrame.IO.Utils.RandomAccess ( + WritableBinaryHandle, + bufferResidency, + flushBufferToFile, + mallocBuffer, + onBuffer, + putByteString, + putWord32LE, + withWritableBinaryFile, + writeByteStringToFile, + ) +import DataFrame.Internal.DataFrame ( + DataFrame, + columnNames, + dataframeDimensions, + getColumn, + ) +import Pinch (enum, putField) +import qualified Pinch + +writeParquet :: FilePath -> DataFrame -> IO () +writeParquet = writeParquetWithOptions defaultParquetWriteOptions + +writeParquetWithOptions :: ParquetWriteOptions -> FilePath -> DataFrame -> IO () +writeParquetWithOptions opts path df = do + when (opts.strategy == TwoPass) $ + error "writeParquet: TwoPass strategy is not yet implemented" + let (nRows, _) = dataframeDimensions df + names = columnNames df + cols <- + VB.fromList + <$> mapM (\n -> initColumnState opts n (fromJust (getColumn n df))) names + withWritableBinaryFile path $ \out -> do + writeByteStringToFile out magic + fileOff <- newIORef 4 + rowGroupsRef <- newIORef [] -- Row Group Metadata + rgRowsRef <- newIORef 0 + let st = WriterState out cols fileOff rowGroupsRef rgRowsRef + interval = max 1 opts.batchRows + loop row + | row >= nRows = pure () + | otherwise = do + -- go row by row and append to a pgee + -- When a page is full (frome the page size writer option) flush it to its ColumnChunk + -- When all the columnChunks combined match or exceed the row group size option + -- flush all the columnchunks to file one by one + VB.forM_ cols (writeRowAndMaybeFinalize opts row) + modifyIORef' rgRowsRef (+ 1) + when ((row + 1) `mod` interval == 0) $ do + size <- bufferedSize cols + when (size >= opts.rowGroupSize) (finalizeRowGroup opts st) + loop (row + 1) + loop 0 + finalizeRowGroup opts st + writeFooter st nRows + +data WriterState = WriterState + { wsOut :: !WritableBinaryHandle + , wsCols :: !(VB.Vector ColumnChunkState) + , wsFileOffset :: !(IORef Int64) + , wsRowGroups :: !(IORef [RowGroup]) + , wsRgRows :: !(IORef Int) + } + +finalizeRowGroup :: ParquetWriteOptions -> WriterState -> IO () +finalizeRowGroup opts st = do + rgRows <- readIORef st.wsRgRows + when (rgRows > 0) $ do + -- flush the page buffer of each columnChunk into their respective ColumnChunk Buffer + -- before flushing the entire row group + VB.mapM_ (runColumnChunkWriter (finalizePage opts.compressionCodec)) st.wsCols + (chunksRev, totalCompressed, totalUncompressed) <- + VB.foldM' + ( \(acc, totalCompressedSize, totalUncompressedSize) cs -> do + offset <- readIORef st.wsFileOffset + compressedSize <- bufferResidency (ckBuffer cs) + uncompressedSize <- readIORef (ckUncompressed cs) + flushBufferToFile st.wsOut (ckBuffer cs) + writeIORef st.wsFileOffset (offset + fromIntegral compressedSize) + writeIORef (ckUncompressed cs) 0 + let chunk = mkColumnChunk opts offset compressedSize uncompressedSize rgRows cs + pure + ( chunk : acc + , totalCompressedSize + fromIntegral compressedSize + , totalUncompressedSize + fromIntegral uncompressedSize + ) + ) + ([], 0 :: Int64, 0 :: Int64) + st.wsCols + modifyIORef' + st.wsRowGroups + (mkRowGroup (reverse chunksRev) totalCompressed totalUncompressed rgRows :) + writeIORef st.wsRgRows 0 + +mkColumnChunk :: + ParquetWriteOptions -> + Int64 -> + Int -> + Int64 -> + Int -> + ColumnChunkState -> + ColumnChunk +mkColumnChunk opts offset compressedSize uncompressedSize rgRows cs = + 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 (ckEncoder cs).encType + , cmd_encodings = putField [PLAIN enum, RLE enum] + , cmd_path_in_schema = putField [ckName cs] + , cmd_codec = putField opts.compressionCodec + , 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 :: WriterState -> Int -> IO () +writeFooter st nRows = do + rowGroups <- reverse <$> readIORef st.wsRowGroups + let schemaElements = + rootSchemaElement (VB.length st.wsCols) : VB.toList (VB.map ckSchema st.wsCols) + metadata = + FileMetadata + { version = putField 1 + , schema = putField schemaElements + , num_rows = putField (fromIntegral nRows) + , row_groups = putField rowGroups + , 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) + onBuffer buffer $ do + putByteString footer + putWord32LE (fromIntegral (BS.length footer)) + putByteString magic + flushBufferToFile st.wsOut buffer diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/ColumnChunkWriter.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/ColumnChunkWriter.hs new file mode 100644 index 00000000..b0bc354e --- /dev/null +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/ColumnChunkWriter.hs @@ -0,0 +1,175 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} + +module DataFrame.IO.Parquet.Writer.ColumnChunkWriter ( + ColumnChunkWriter (..), + ColumnChunkState (..), + page, + askColumnChunk, + initColumnState, + writeRowAndMaybeFinalize, + finalizePage, + bufferedSize, +) where + +import Control.Monad (when) +import Control.Monad.IO.Class (MonadIO (..)) +import qualified Data.ByteString as BS +import Data.IORef (IORef, modifyIORef', newIORef, readIORef) +import Data.Int (Int64) +import qualified Data.Text as T +import qualified Data.Vector as VB +import DataFrame.IO.Parquet.Thrift +import DataFrame.IO.Parquet.Writer.DefLevels (DefLevels (..), pushDef) +import DataFrame.IO.Parquet.Writer.Encoder (Encoder (..), buildEncoder) +import DataFrame.IO.Parquet.Writer.Metadata (mkDataPageHeader, mkSchemaElem) +import DataFrame.IO.Parquet.Writer.Options (ParquetWriteOptions (..)) +import DataFrame.IO.Parquet.Writer.PageWriter ( + PageState (..), + PageWriter (..), + assemblePageBody, + newPageState, + pageRows, + resetPage, + ) +import DataFrame.IO.Utils.RandomAccess ( + HasBuffer (..), + MemoryBuffer (..), + ReaderIO (runReaderIO), + bufferResidency, + bufferToByteString, + copyBuffer, + mallocBuffer, + putByteString, + ) +import DataFrame.Internal.Column (Column, hasMissing) +import qualified Pinch +import qualified Snappy + +data ColumnChunkState = ColumnChunkState + { ckName :: !T.Text + , ckNullable :: !Bool + , ckSchema :: !SchemaElement + , ckEncoder :: !Encoder + , ckBuffer :: !MemoryBuffer + , ckUncompressed :: !(IORef Int64) + , ckPage :: !PageState + } + +newtype ColumnChunkWriter a = ColumnChunkWriter {runColumnChunkWriter :: ColumnChunkState -> IO a} + +instance Functor ColumnChunkWriter where + fmap f (ColumnChunkWriter g) = ColumnChunkWriter (fmap f . g) + +instance Applicative ColumnChunkWriter where + pure x = ColumnChunkWriter (const (pure x)) + ColumnChunkWriter f <*> ColumnChunkWriter g = ColumnChunkWriter (\s -> f s <*> g s) + +instance Monad ColumnChunkWriter where + ColumnChunkWriter m >>= k = ColumnChunkWriter (\s -> m s >>= \a -> runColumnChunkWriter (k a) s) + +instance MonadIO ColumnChunkWriter where + liftIO io = ColumnChunkWriter (const io) + +instance HasBuffer ColumnChunkWriter where + type Buffer ColumnChunkWriter = MemoryBuffer + askBuffer = ColumnChunkWriter (pure . ckBuffer) + residency = ColumnChunkWriter (bufferResidency . ckBuffer) + writeBytes bytes = ColumnChunkWriter (runReaderIO (writeBytes bytes) . ckBuffer) + flushTo sink = ColumnChunkWriter (runReaderIO (flushTo sink) . ckBuffer) + +page :: PageWriter a -> ColumnChunkWriter a +page (PageWriter f) = ColumnChunkWriter (f . ckPage) + +askColumnChunk :: ColumnChunkWriter ColumnChunkState +askColumnChunk = ColumnChunkWriter pure + +initColumnState :: + ParquetWriteOptions -> T.Text -> Column -> IO ColumnChunkState +initColumnState opts name col = do + encoder <- buildEncoder col + let nullable = hasMissing col + schemaElem = + mkSchemaElem + name + encoder.encType + nullable + encoder.encConverted + encoder.encLogical + cap = max 1 opts.pageSize + chunk <- mallocBuffer cap + uncompressed <- newIORef 0 + pageState <- newPageState cap nullable + pure + ColumnChunkState + { ckName = name + , ckNullable = nullable + , ckSchema = schemaElem + , ckEncoder = encoder + , ckBuffer = chunk + , ckUncompressed = uncompressed + , ckPage = pageState + } + +writeRowIO :: Int -> ColumnChunkState -> IO () +writeRowIO row st = do + let pageState = ckPage st + present <- encWriteValue (ckEncoder st) pageState.psValues row + when (ckNullable st) (pushDef pageState.psDefs (if present then 1 else 0)) + modifyIORef' pageState.psRows (+ 1) +{-# INLINE writeRowIO #-} + +writeRowAndMaybeFinalize :: + ParquetWriteOptions -> Int -> ColumnChunkState -> IO () +writeRowAndMaybeFinalize opts row st = do + writeRowIO row st + valuesSize <- bufferResidency st.ckPage.psValues + levelsSize <- bufferResidency st.ckPage.psDefs.dlBuf + let size = valuesSize + levelsSize + when + (size >= opts.pageSize) + (runColumnChunkWriter (finalizePage opts.compressionCodec) st) +{-# INLINE writeRowAndMaybeFinalize #-} + +finalizePage :: CompressionCodec -> ColumnChunkWriter () +finalizePage codec = do + st <- askColumnChunk + rows <- page pageRows + when (rows > 0) $ do + liftIO (encFinishValues (ckEncoder st) st.ckPage.psValues) + body <- page (assemblePageBody (ckNullable st)) + writeDataPage codec rows body + page resetPage + +writeDataPage :: CompressionCodec -> Int -> MemoryBuffer -> ColumnChunkWriter () +writeDataPage codec rows body = do + uncompressedSize <- liftIO (bufferResidency body) + (compressedSize, emit) <- case codec of + UNCOMPRESSED _ -> pure (uncompressedSize, copyBuffer body) + SNAPPY _ -> do + compressed <- liftIO (Snappy.compress <$> bufferToByteString body) + pure (BS.length compressed, putByteString compressed) + other -> error ("writeParquet: unsupported codec " <> show other) + let headerBytes = + Pinch.encode + Pinch.compactProtocol + (mkDataPageHeader rows uncompressedSize compressedSize) + putByteString headerBytes + emit + bumpUncompressed (fromIntegral (BS.length headerBytes + uncompressedSize)) + +bumpUncompressed :: Int64 -> ColumnChunkWriter () +bumpUncompressed n = ColumnChunkWriter (\st -> modifyIORef' (ckUncompressed st) (+ n)) + +bufferedSize :: VB.Vector ColumnChunkState -> IO Int +bufferedSize = + VB.foldM' + ( \total st -> do + chunk <- bufferResidency (ckBuffer st) + values <- bufferResidency (psValues (ckPage st)) + defs <- bufferResidency (ckPage st).psDefs.dlBuf + pure (total + chunk + values + defs) + ) + 0 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..db406bce --- /dev/null +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs @@ -0,0 +1,270 @@ +{-# 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 ( + Bitmap, + Column (..), + Columnable, + bitmapTestBit, + columnTypeString, + hasElemType, + ) +import DataFrame.Internal.PackedText ( + PackedTextData (..), + offAt, + selAt, + ) +import Pinch (enum, putField) +import Type.Reflection (typeRep) + +data Encoder = Encoder + { encType :: !ThriftType + , encConverted :: !(Maybe ConvertedType) + , encLogical :: !(Maybe LogicalType) + , encWriteValue :: !(MemoryBuffer -> Int -> IO Bool) + , encFinishValues :: !(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..586632d6 --- /dev/null +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Metadata.hs @@ -0,0 +1,75 @@ +{-# LANGUAGE OverloadedStrings #-} + +module DataFrame.IO.Parquet.Writer.Metadata ( + mkSchemaElem, + rootSchemaElement, + mkDataPageHeader, + magic, +) where + +import qualified Data.ByteString as BS +import qualified Data.Text as T +import DataFrame.IO.Parquet.Thrift +import Pinch (enum, putField) + +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 + } + +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/Parquet/Writer/PageWriter.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/PageWriter.hs new file mode 100644 index 00000000..1ced3c34 --- /dev/null +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/PageWriter.hs @@ -0,0 +1,115 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE TypeFamilies #-} + +module DataFrame.IO.Parquet.Writer.PageWriter ( + PageWriter (..), + PageState (..), + newPageState, + askPage, + recordDef, + bumpRows, + pageRows, + assemblePageBody, + resetPage, +) where + +import Control.Monad.IO.Class (MonadIO (..)) +import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef) +import DataFrame.IO.Parquet.Writer.DefLevels ( + DefLevels (..), + flushDef, + newDefLevels, + pushDef, + ) +import DataFrame.IO.Utils.RandomAccess ( + HasBuffer (..), + MemoryBuffer, + ReaderIO (runReaderIO), + Sink (..), + bufferResidency, + copyBuffer, + mallocBuffer, + onBuffer, + putWord32LE, + resetPosition, + ) + +data PageState = PageState + { psValues :: !MemoryBuffer + , psScratch :: !MemoryBuffer + , psDefs :: !DefLevels + , psRows :: !(IORef Int) + , psNullable :: !Bool + } + +newPageState :: Int -> Bool -> IO PageState +newPageState cap nullable = do + values <- mallocBuffer cap + scratch <- mallocBuffer cap + defs <- newDefLevels + rows <- newIORef 0 + pure + PageState + { psValues = values + , psScratch = scratch + , psDefs = defs + , psRows = rows + , psNullable = nullable + } + +newtype PageWriter a = PageWriter {runPageWriter :: PageState -> IO a} + +instance Functor PageWriter where + fmap f (PageWriter g) = PageWriter (fmap f . g) + +instance Applicative PageWriter where + pure x = PageWriter (const (pure x)) + PageWriter f <*> PageWriter g = PageWriter (\s -> f s <*> g s) + +instance Monad PageWriter where + PageWriter m >>= k = PageWriter (\s -> m s >>= \a -> runPageWriter (k a) s) + +instance MonadIO PageWriter where + liftIO io = PageWriter (const io) + +instance HasBuffer PageWriter where + type Buffer PageWriter = MemoryBuffer + askBuffer = PageWriter (pure . psValues) + residency = PageWriter (bufferResidency . psValues) + writeBytes bytes = PageWriter (runReaderIO (writeBytes bytes) . psValues) + flushTo sink = PageWriter (runReaderIO (flushTo sink) . psValues) + +askPage :: PageWriter PageState +askPage = PageWriter pure + +recordDef :: Bool -> PageWriter () +recordDef present = PageWriter (\ps -> pushDef (psDefs ps) (if present then 1 else 0)) + +bumpRows :: PageWriter () +bumpRows = PageWriter (\ps -> modifyIORef' (psRows ps) (+ 1)) + +pageRows :: PageWriter Int +pageRows = PageWriter (readIORef . psRows) + +assemblePageBody :: Bool -> PageWriter MemoryBuffer +assemblePageBody nullable = do + ps <- askPage + if not nullable + then pure (psValues ps) + else do + liftIO (flushDef (psDefs ps)) + liftIO (resetPosition (psScratch ps)) + defSize <- liftIO (bufferResidency ps.psDefs.dlBuf) + onBuffer (psScratch ps) $ do + putWord32LE (fromIntegral defSize) + copyBuffer ps.psDefs.dlBuf + copyBuffer (psValues ps) + pure (psScratch ps) + +resetPage :: PageWriter () +resetPage = PageWriter $ \ps -> do + resetPosition (psValues ps) + resetPosition ps.psDefs.dlBuf + resetPosition (psScratch ps) + writeIORef (psRows ps) 0 diff --git a/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs b/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs index c6b84655..90f9fc83 100644 --- a/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs +++ b/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs @@ -1,21 +1,95 @@ +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE TypeFamilies #-} -module DataFrame.IO.Utils.RandomAccess where +module DataFrame.IO.Utils.RandomAccess ( + uncurry3, + Range (..), + RandomAccess (..), + ReaderIO (runReaderIO), + LocalFile, + MMappedFile, + unsafeToByteString, + WritableBinaryHandle, + openWritableBinaryFile, + withWritableBinaryFile, + HasBuffer (..), + Sink (..), + MemoryBuffer (..), + mallocBuffer, + BufferHandle, + withFileBuffer, + appendByteString, + appendGeneratedBytes, + appendTextArraySlice, + appendByteStringHandle, + writeWord8, + writeWord32LE, + writeWord64LE, + writeFloatLE, + writeDoubleLE, + bufferResidency, + bufferToByteString, + copyBufferInto, + resetPosition, + flushBufferToFile, + writeByteStringToFile, + MemoryWriter, + onBuffer, + putWord8, + putWord32LE, + putWord64LE, + putFloatLE, + putDoubleLE, + putByteString, + putGenerated, + copyBuffer, +) where +import Control.Exception (bracket) +import Control.Monad (foldM) 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 qualified Data.Foldable as Foldable +import Data.IORef (IORef, newIORef, readIORef, writeIORef) +import Data.Primitive.ByteArray ( + MutableByteArray, + copyMutableByteArray, + getSizeofMutableByteArray, + mutableByteArrayContents, + 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 (ReadWriteMode, WriteMode), SeekMode (AbsoluteSeek), + hClose, + hGetBuf, + hPutBuf, + hSeek, + hSetBinaryMode, + hSetBuffering, + openBinaryFile, ) uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d @@ -76,3 +150,370 @@ 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) + +class (Monad m) => HasBuffer m where + type Buffer m + askBuffer :: m (Buffer m) + residency :: m Int -- number of bytes currently in the buffer + writeBytes :: (Foldable f) => f Word8 -> m () + flushTo :: Sink -> m () + +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 + +data Sink = MemorySink MemoryBuffer | FileSink WritableBinaryHandle + +instance HasBuffer (ReaderIO MemoryBuffer) where + type Buffer (ReaderIO MemoryBuffer) = MemoryBuffer + + askBuffer = ReaderIO pure + + residency = ReaderIO $ \buffer -> readIORef buffer.positionRef + + writeBytes bytes = ReaderIO $ \buffer -> do + position <- readIORef buffer.positionRef + array <- ensureCapacity buffer (position + Foldable.length bytes) + newPosition <- + foldM (\i byte -> writeByteArray array i byte >> pure (i + 1)) position bytes + writeIORef buffer.positionRef newPosition + + flushTo (MemorySink destination) = ReaderIO $ \source -> + if destination.arrayRef == source.arrayRef + then pure () + else do + sourceArray <- readIORef source.arrayRef + sourcePosition <- readIORef source.positionRef + destinationPosition <- readIORef destination.positionRef + let newDestinationPosition = destinationPosition + sourcePosition + destinationArray <- ensureCapacity destination newDestinationPosition + copyMutableByteArray + destinationArray + destinationPosition + sourceArray + 0 -- offset + sourcePosition -- number of bytes + writeIORef destination.positionRef newDestinationPosition + writeIORef source.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. + flushTo (FileSink (WritableBinaryHandle h)) = ReaderIO $ \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 + +-- We're using pinned ByteArrays so we must +-- not use the grow fuynction brovided by primitive +-- instead we must alloocatie a new pinned byteArray. +-- We might have been worried about heap fragmentation +-- becasue 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 useage 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 (we also allocate a little extra space to +-- begin with). +-- 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). The extra +-- ceremony of handling writes and flushes can be encapsulated well in +-- HasBuffer instances. +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 + +appendByteString :: MemoryBuffer -> ByteString -> IO () +appendByteString 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) + +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) + +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 + +copyBufferInto :: MemoryBuffer -> MemoryBuffer -> IO () +copyBufferInto destination source = 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) + +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 + +appendGeneratedBytes :: MemoryBuffer -> Int -> (Int -> Word8) -> IO () +appendGeneratedBytes buffer count at + | count < 0 = ioError $ userError "appendGeneratedBytes: negative length" + | otherwise = do + position <- readIORef buffer.positionRef + array <- ensureCapacity buffer (position + count) + let go i + | i >= count = pure () + | otherwise = writeByteArray array (position + i) (at i) >> go (i + 1) + go 0 + writeIORef buffer.positionRef (position + count) + +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 #-} + +flushBufferToFile :: WritableBinaryHandle -> MemoryBuffer -> IO () +flushBufferToFile handle = runReaderIO (flushTo (FileSink handle)) + +writeByteStringToFile :: WritableBinaryHandle -> ByteString -> IO () +writeByteStringToFile handle bs = do + buffer <- mallocBuffer (max 1 (BS.length bs)) + appendByteString buffer bs + flushBufferToFile handle buffer + +type MemoryWriter m = (HasBuffer m, Buffer m ~ MemoryBuffer, MonadIO m) + +onBuffer :: (MonadIO m) => MemoryBuffer -> ReaderIO MemoryBuffer a -> m a +onBuffer buffer action = liftIO (runReaderIO action buffer) + +putWord8 :: (MemoryWriter m) => Word8 -> m () +putWord8 value = askBuffer >>= \buffer -> liftIO (writeWord8 buffer value) + +putWord32LE :: (MemoryWriter m) => Word32 -> m () +putWord32LE value = askBuffer >>= \buffer -> liftIO (writeWord32LE buffer value) + +putWord64LE :: (MemoryWriter m) => Word64 -> m () +putWord64LE value = askBuffer >>= \buffer -> liftIO (writeWord64LE buffer value) + +putFloatLE :: (MemoryWriter m) => Float -> m () +putFloatLE value = askBuffer >>= \buffer -> liftIO (writeFloatLE buffer value) + +putDoubleLE :: (MemoryWriter m) => Double -> m () +putDoubleLE value = askBuffer >>= \buffer -> liftIO (writeDoubleLE buffer value) + +putByteString :: (MemoryWriter m) => ByteString -> m () +putByteString bytes = askBuffer >>= \buffer -> liftIO (appendByteString buffer bytes) + +putGenerated :: (MemoryWriter m) => Int -> (Int -> Word8) -> m () +putGenerated count at = askBuffer >>= \buffer -> liftIO (appendGeneratedBytes buffer count at) + +copyBuffer :: (MemoryWriter m) => MemoryBuffer -> m () +copyBuffer source = askBuffer >>= \destination -> liftIO (copyBufferInto destination source) + +data BufferHandle = BufferHandle + { bufferPath :: !FilePath + , bufferHandle :: !WritableBinaryHandle + , residencyRef :: !(IORef Int) + , flushedRef :: !(IORef Int) + } + +withFileBuffer :: FilePath -> (BufferHandle -> IO a) -> IO a +withFileBuffer filepath = + bracket open (hClose . unHandle . bufferHandle) + where + open = do + h <- openBinaryFile filepath ReadWriteMode + hSetBinaryMode h True + hSetBuffering h NoBuffering + res <- newIORef 0 + flushed <- newIORef 0 + pure (BufferHandle filepath (WritableBinaryHandle h) res flushed) + +instance HasBuffer (ReaderIO BufferHandle) where + type Buffer (ReaderIO BufferHandle) = BufferHandle + + askBuffer = ReaderIO pure + + residency = ReaderIO $ \bh -> readIORef bh.residencyRef + + writeBytes bytes = ReaderIO $ \bh -> do + let WritableBinaryHandle h = bh.bufferHandle + scratch <- newPinnedByteArray 1 + n <- + foldM + ( \count byte -> do + writeByteArray scratch 0 byte + hPutBuf h (mutableByteArrayContents scratch) 1 + pure (count + 1) + ) + 0 + bytes + position <- readIORef bh.residencyRef + writeIORef bh.residencyRef (position + n) + + -- Flushing from a file to a sink happens using + -- a reusable 256 KiB buffer closed over this function + -- Usually for this instance we shoulc be diong only + -- file to file but file to buffer is also bossible + -- if you should want to do that, for whatever reason + -- (don't let me tell you how to live your life) + flushTo sink = ReaderIO $ \bh -> do + count <- readIORef bh.residencyRef + offset <- readIORef bh.flushedRef + let WritableBinaryHandle h = bh.bufferHandle + chunkSize = 262144 + chunk <- newPinnedByteArray (min chunkSize count) + let ptr = mutableByteArrayContents chunk + pushChunk n = case sink of + FileSink (WritableBinaryHandle out) -> hPutBuf out ptr n + MemorySink dest -> do + destinationPosition <- readIORef dest.positionRef + let newDestinationPosition = destinationPosition + n + destinationArray <- ensureCapacity dest newDestinationPosition + copyMutableByteArray destinationArray destinationPosition chunk 0 n + writeIORef dest.positionRef newDestinationPosition + go remaining + | remaining <= 0 = pure () + | otherwise = do + actual <- hGetBuf h ptr (min chunkSize remaining) + if actual <= 0 + then pure () + else pushChunk actual >> go (remaining - actual) + hSeek h AbsoluteSeek (fromIntegral offset) + go count + writeIORef bh.residencyRef 0 + writeIORef bh.flushedRef (offset + count) + +appendByteStringHandle :: BufferHandle -> ByteString -> IO () +appendByteStringHandle bh bs = + BU.unsafeUseAsCStringLen bs $ \(source, len) -> do + let WritableBinaryHandle h = bh.bufferHandle + hPutBuf h source len + position <- readIORef bh.residencyRef + writeIORef bh.residencyRef (position + len) diff --git a/dataframe-parquet/stress/DataFrame10GB.hs b/dataframe-parquet/stress/DataFrame10GB.hs new file mode 100644 index 00000000..824c35e5 --- /dev/null +++ b/dataframe-parquet/stress/DataFrame10GB.hs @@ -0,0 +1,136 @@ +{-# 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 (Bitmap, Column (..)) +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..efec900a --- /dev/null +++ b/dataframe-parquet/tests/Main.hs @@ -0,0 +1,327 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Tests for the writer-buffer logic in "DataFrame.IO.Utils.RandomAccess". +module Main where + +import Control.Monad (forM) +import qualified Data.ByteString as BS +import Data.IORef (readIORef) +import Data.Primitive.ByteArray (readByteArray) +import Data.Word (Word8) +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 + +withTempFileBuffer :: FilePath -> String -> (BufferHandle -> IO a) -> IO a +withTempFileBuffer dir name = withFileBuffer (dir name) + +memoryBufferBytes :: MemoryBuffer -> IO [Word8] +memoryBufferBytes buf = do + array <- readIORef (arrayRef buf) + n <- readIORef (positionRef buf) + forM [0 .. n - 1] (readByteArray array) + +residencyTracksWrites :: Test +residencyTracksWrites = TestCase $ + withSystemTempDirectory "dfpq-buffer" $ \dir -> + withTempFileBuffer dir "residency.bin" $ \bh -> do + runReaderIO (writeBytes [1, 2, 3, 4, 5 :: Word8]) bh + r1 <- runReaderIO residency bh + assertEqual "residency after first write" 5 r1 + runReaderIO (writeBytes [6, 7, 8 :: Word8]) bh + r2 <- runReaderIO residency bh + assertEqual "residency accumulates across writes" 8 r2 + +-- Flushing a file buffer into a FileSink writes exactly the buffered bytes to +-- the output handle and empties the buffer. +flushToFileSink :: Test +flushToFileSink = TestCase $ + withSystemTempDirectory "dfpq-buffer" $ \dir -> do + let outPath = dir "out.bin" + payload = [10, 20, 30, 40, 50, 60 :: Word8] + withTempFileBuffer dir "buf.bin" $ \bh -> do + runReaderIO (writeBytes payload) bh + withWritableBinaryFile outPath $ \out -> + runReaderIO (flushTo (FileSink out)) bh + r <- runReaderIO residency bh + assertEqual "residency reset after flush" 0 r + contents <- BS.readFile outPath + assertEqual "flushed file content" (BS.pack payload) contents + +-- Flushing a file buffer into a MemorySink copies the buffered bytes into the +-- destination memory buffer and empties the source. +flushToMemorySink :: Test +flushToMemorySink = TestCase $ + withSystemTempDirectory "dfpq-buffer" $ \dir -> do + let payload = [7, 6, 5, 4, 3, 2, 1 :: Word8] + destination <- mallocBuffer 0 + withTempFileBuffer dir "buf.bin" $ \bh -> do + runReaderIO (writeBytes payload) bh + runReaderIO (flushTo (MemorySink destination)) bh + r <- runReaderIO residency bh + assertEqual "residency reset after flush" 0 r + memBytes <- memoryBufferBytes destination + assertEqual "flushed memory content" payload memBytes + +-- Successive flushes into the same FileSink append rather than overwrite. +flushAppendsToFileSink :: Test +flushAppendsToFileSink = TestCase $ + withSystemTempDirectory "dfpq-buffer" $ \dir -> do + let outPath = dir "out.bin" + chunkA = [1, 2, 3 :: Word8] + chunkB = [4, 5, 6, 7 :: Word8] + withWritableBinaryFile outPath $ \out -> do + withTempFileBuffer dir "a.bin" $ \bh -> do + runReaderIO (writeBytes chunkA) bh + runReaderIO (flushTo (FileSink out)) bh + withTempFileBuffer dir "b.bin" $ \bh -> do + runReaderIO (writeBytes chunkB) bh + runReaderIO (flushTo (FileSink out)) bh + contents <- BS.readFile outPath + assertEqual "appended flushes" (BS.pack (chunkA ++ chunkB)) contents + +-- After a flush the buffer is emptied, so re-using it only re-flushes the +-- bytes written since the previous flush (old bytes are not resurrected). +flushEmptiesForReuse :: Test +flushEmptiesForReuse = TestCase $ + withSystemTempDirectory "dfpq-buffer" $ \dir -> do + let firstPath = dir "first.bin" + secondPath = dir "second.bin" + chunkA = [11, 22, 33 :: Word8] + chunkB = [44, 55 :: Word8] + withTempFileBuffer dir "buf.bin" $ \bh -> do + runReaderIO (writeBytes chunkA) bh + withWritableBinaryFile firstPath $ \out -> + runReaderIO (flushTo (FileSink out)) bh + runReaderIO (writeBytes chunkB) bh + r <- runReaderIO residency bh + assertEqual "residency reflects only post-flush bytes" 2 r + withWritableBinaryFile secondPath $ \out -> + runReaderIO (flushTo (FileSink out)) bh + first <- BS.readFile firstPath + second <- BS.readFile secondPath + assertEqual "first flush" (BS.pack chunkA) first + assertEqual "second flush excludes first chunk" (BS.pack chunkB) second + +-- A payload larger than the 256 KiB flush chunk round-trips intact, exercising +-- the chunked flush loop. +flushLargePayload :: Test +flushLargePayload = TestCase $ + withSystemTempDirectory "dfpq-buffer" $ \dir -> do + let outPath = dir "big.bin" + payload = take 300000 (cycle [0 .. 255]) :: [Word8] + withTempFileBuffer dir "big-buf.bin" $ \bh -> do + runReaderIO (writeBytes payload) bh + r <- runReaderIO residency bh + assertEqual "residency for large payload" 300000 r + withWritableBinaryFile outPath $ \out -> + runReaderIO (flushTo (FileSink out)) bh + contents <- BS.readFile outPath + assertEqual "large payload round-trips" (BS.pack payload) contents + +-- Memory buffer instance ------------------------------------------------------ +-- The same observable contract as the file buffer, plus the memory-specific +-- self-flush no-op. + +-- residency reports the running byte count and accumulates across writes. +memResidencyTracksWrites :: Test +memResidencyTracksWrites = TestCase $ do + buf <- mallocBuffer 0 + runReaderIO (writeBytes [1, 2, 3, 4, 5 :: Word8]) buf + r1 <- runReaderIO residency buf + assertEqual "residency after first write" 5 r1 + runReaderIO (writeBytes [6, 7, 8 :: Word8]) buf + r2 <- runReaderIO residency buf + assertEqual "residency accumulates across writes" 8 r2 + +-- Flushing a memory buffer into a FileSink writes the buffered bytes to the +-- output handle and empties the buffer. +memFlushToFileSink :: Test +memFlushToFileSink = TestCase $ + withSystemTempDirectory "dfpq-buffer" $ \dir -> do + let outPath = dir "out.bin" + payload = [10, 20, 30, 40, 50, 60 :: Word8] + buf <- mallocBuffer 0 + runReaderIO (writeBytes payload) buf + withWritableBinaryFile outPath $ \out -> + runReaderIO (flushTo (FileSink out)) buf + r <- runReaderIO residency buf + assertEqual "residency reset after flush" 0 r + contents <- BS.readFile outPath + assertEqual "flushed file content" (BS.pack payload) contents + +-- Flushing a memory buffer into another MemorySink copies the buffered bytes +-- into the destination and empties the source. +memFlushToMemorySink :: Test +memFlushToMemorySink = TestCase $ do + let payload = [7, 6, 5, 4, 3, 2, 1 :: Word8] + source <- mallocBuffer 0 + destination <- mallocBuffer 0 + runReaderIO (writeBytes payload) source + runReaderIO (flushTo (MemorySink destination)) source + rSrc <- runReaderIO residency source + assertEqual "source residency reset after flush" 0 rSrc + destBytes <- memoryBufferBytes destination + assertEqual "flushed memory content" payload destBytes + +-- Successive flushes into the same FileSink append rather than overwrite. +memFlushAppendsToFileSink :: Test +memFlushAppendsToFileSink = TestCase $ + withSystemTempDirectory "dfpq-buffer" $ \dir -> do + let outPath = dir "out.bin" + chunkA = [1, 2, 3 :: Word8] + chunkB = [4, 5, 6, 7 :: Word8] + withWritableBinaryFile outPath $ \out -> do + bufA <- mallocBuffer 0 + runReaderIO (writeBytes chunkA) bufA + runReaderIO (flushTo (FileSink out)) bufA + bufB <- mallocBuffer 0 + runReaderIO (writeBytes chunkB) bufB + runReaderIO (flushTo (FileSink out)) bufB + contents <- BS.readFile outPath + assertEqual "appended flushes" (BS.pack (chunkA ++ chunkB)) contents + +-- After a flush the buffer is emptied, so re-using it only re-flushes the +-- bytes written since the previous flush. +memFlushEmptiesForReuse :: Test +memFlushEmptiesForReuse = TestCase $ + withSystemTempDirectory "dfpq-buffer" $ \dir -> do + let firstPath = dir "first.bin" + secondPath = dir "second.bin" + chunkA = [11, 22, 33 :: Word8] + chunkB = [44, 55 :: Word8] + buf <- mallocBuffer 0 + runReaderIO (writeBytes chunkA) buf + withWritableBinaryFile firstPath $ \out -> + runReaderIO (flushTo (FileSink out)) buf + runReaderIO (writeBytes chunkB) buf + r <- runReaderIO residency buf + assertEqual "residency reflects only post-flush bytes" 2 r + withWritableBinaryFile secondPath $ \out -> + runReaderIO (flushTo (FileSink out)) buf + first <- BS.readFile firstPath + second <- BS.readFile secondPath + assertEqual "first flush" (BS.pack chunkA) first + assertEqual "second flush excludes first chunk" (BS.pack chunkB) second + +-- A memory buffer flushing to a MemorySink backed by itself is a no-op: the +-- bytes and residency are left untouched (identity is compared via arrayRef). +memSelfFlushIsNoop :: Test +memSelfFlushIsNoop = TestCase $ do + let payload = [3, 1, 4, 1, 5, 9, 2, 6 :: Word8] + buf <- mallocBuffer 0 + runReaderIO (writeBytes payload) buf + runReaderIO (flushTo (MemorySink buf)) buf + r <- runReaderIO residency buf + assertEqual "self-flush leaves residency unchanged" (Prelude.length payload) r + bytes <- memoryBufferBytes buf + assertEqual "self-flush leaves content unchanged" payload bytes + +-- A payload larger than the 256 KiB flush chunk round-trips intact. +memFlushLargePayload :: Test +memFlushLargePayload = TestCase $ + withSystemTempDirectory "dfpq-buffer" $ \dir -> do + let outPath = dir "big.bin" + payload = take 300000 (cycle [0 .. 255]) :: [Word8] + buf <- mallocBuffer 0 + runReaderIO (writeBytes payload) buf + r <- runReaderIO residency buf + assertEqual "residency for large payload" 300000 r + withWritableBinaryFile outPath $ \out -> + runReaderIO (flushTo (FileSink out)) buf + contents <- BS.readFile outPath + assertEqual "large payload round-trips" (BS.pack 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 "file buffer: residency tracks writes" residencyTracksWrites + , TestLabel "file buffer: flush to file sink" flushToFileSink + , TestLabel "file buffer: flush to memory sink" flushToMemorySink + , TestLabel "file buffer: flush appends to file sink" flushAppendsToFileSink + , TestLabel "file buffer: flush empties for reuse" flushEmptiesForReuse + , TestLabel "file buffer: flush large payload" flushLargePayload + , TestLabel "memory buffer: residency tracks writes" memResidencyTracksWrites + , TestLabel "memory buffer: flush to file sink" memFlushToFileSink + , TestLabel "memory buffer: flush to memory sink" memFlushToMemorySink + , TestLabel "memory buffer: flush appends to file sink" memFlushAppendsToFileSink + , TestLabel "memory buffer: flush empties for reuse" memFlushEmptiesForReuse + , TestLabel "memory buffer: self-flush is a no-op" memSelfFlushIsNoop + , TestLabel "memory buffer: flush large payload" memFlushLargePayload + , 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 b7310c6f..21395729 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,