dataframe-parquet 1.5.0.1 → 1.5.1.0

raw patch · 20 files changed

+2116/−10 lines, 20 filesdep +HUnitdep +criteriondep +dataframe-parquetdep ~directorynew-component:exe:dataframe-parquet-10gb-stressbinary-addedPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependencies added: HUnit, criterion, dataframe-parquet, deepseq, primitive, temporary

Dependency ranges changed: directory

API changes (from Hackage documentation)

Files

• benchmark/Writer10GB.hs +49/−0
• dataframe-parquet.cabal +66/−1
• src/DataFrame/IO/Parquet.hs +128/−4
• src/DataFrame/IO/Parquet/Writer.hs +413/−0
• src/DataFrame/IO/Parquet/Writer/DefLevels.hs +61/−0
• src/DataFrame/IO/Parquet/Writer/Encoder.hs +394/−0
• src/DataFrame/IO/Parquet/Writer/Metadata.hs +175/−0
• src/DataFrame/IO/Parquet/Writer/Options.hs +34/−0
• src/DataFrame/IO/Utils/RandomAccess.hs +338/−5
• stress/DataFrame10GB.hs +137/−0
• stress/StressMain.hs +38/−0
• tests/Main.hs +283/−0
• tests/data/alltypes_dictionary.parquet binary
• tests/data/alltypes_plain.parquet binary
• tests/data/alltypes_plain.snappy.parquet binary
• tests/data/alltypes_tiny_pages.parquet binary
• tests/data/int32_decimal.parquet binary
• tests/data/int64_decimal.parquet binary
• tests/data/mtcars.parquet binary
• tests/data/transactions.parquet binary

@@ -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)+                    )+        ]

@@ -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

@@ -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.

@@ -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_+            }

@@ -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 #-}

@@ -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 #-}

@@ -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"

@@ -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+        }

@@ -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 #-}

@@ -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

@@ -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

@@ -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

binary file changed (absent → 1698 bytes)

binary file changed (absent → 1851 bytes)

binary file changed (absent → 1736 bytes)

binary file changed (absent → 454233 bytes)

binary file changed (absent → 478 bytes)

binary file changed (absent → 591 bytes)

binary file changed (absent → 4564 bytes)

binary file changed (absent → 1746 bytes)