From 5afa2156520bd6d2881f854316b43613035ab841 Mon Sep 17 00:00:00 2001 From: Raghav Sharma Date: Tue, 1 Sep 2026 23:32:57 +0530 Subject: [PATCH 1/2] Parquet Writer Optimization --- .../src/DataFrame/IO/Parquet/Writer.hs | 84 ++++--- .../DataFrame/IO/Parquet/Writer/Encoder.hs | 208 ++++++++++++------ .../src/DataFrame/IO/Utils/RandomAccess.hs | 46 +++- dataframe-parquet/tests/Main.hs | 15 ++ 4 files changed, 251 insertions(+), 102 deletions(-) diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs index 18a41627..2e8d4e9f 100644 --- a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE OverloadedRecordDot #-} {-# LANGUAGE OverloadedStrings #-} @@ -16,6 +17,7 @@ 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) @@ -41,11 +43,12 @@ import DataFrame.IO.Parquet.Writer.Options ( defaultParquetWriteOptions, ) import DataFrame.IO.Utils.RandomAccess ( - MemoryBuffer, + MemoryBuffer (..), WritableBinaryHandle, atomicallyWriteFile, bufferResidency, bufferToByteString, + ensureCapacity, flushBufferToBuffer, flushBufferToFile, mallocBuffer, @@ -170,16 +173,24 @@ writeShard options path_ df startRow endRow = do 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 - VB.forM_ columnChunks_ (writeRow options scratchBuffer_ rowNum) - modifyIORef' rowNumberRef_ (+ 1) - when ((rowNum - startRow + 1) `mod` interval == 0) $ do - size <- bufferedSize columnChunks_ - when (size >= options.rowGroupSize) (flushRowGroup options writerState) - loop (rowNum + 1) + 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_ @@ -204,30 +215,53 @@ nativeTypeKeyValues names df = , Just col <- [getColumn name df] ] -writeRow :: - ParquetWriteOptions -> MemoryBuffer -> Int -> ColumnChunkState -> IO () -writeRow options scratch rowNum columnChunkState = do - let page = columnChunkState.pageState - notNull <- columnChunkState.encoder.writeValue page.pageBuffer rowNum - when columnChunkState.nullable $ - pushDef page.definitionLevels (if notNull then 1 else 0) - modifyIORef' page.currentRowCount (+ 1) - pageRowCount <- readIORef page.currentRowCount - let subInterval = max 1 options.subBatchRows - when (pageRowCount `mod` subInterval == 0) $ do - flushDef page.definitionLevels - pageBufferResidency <- bufferResidency page.pageBuffer - defLevelsResidency <- bufferResidency page.definitionLevels.dlBuf - when - (pageBufferResidency + defLevelsResidency >= options.pageSize) - (flushPage options scratch columnChunkState) +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 - columnChunkState.encoder.finishValues page.pageBuffer + 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 diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs index 0a6615ed..ba26cd7e 100644 --- a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs @@ -10,10 +10,14 @@ module DataFrame.IO.Parquet.Writer.Encoder ( buildEncoder, ) where -import Control.Monad (when) +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)) @@ -25,14 +29,11 @@ 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, - writeInteger64, - writeWord32LE, - writeWord64LE, - writeWord8, + MemoryBuffer (..), + ensureCapacity, + writeInteger64At, + writeWord32At, + writeWord64At, ) import DataFrame.Internal.Column ( Column (..), @@ -49,6 +50,8 @@ import DataFrame.Internal.Data.PackedText ( offAt, selAt, ) +import Foreign (plusPtr) +import GHC.Float (castDoubleToWord64, castFloatToWord32) import Pinch (enum, putField) import Type.Reflection (typeRep) @@ -56,8 +59,8 @@ data Encoder = Encoder { encType :: !ThriftType , convertedType :: !(Maybe ConvertedType) , logicalType :: !(Maybe LogicalType) - , writeValue :: !(MemoryBuffer -> Int -> IO Bool) - , finishValues :: !(MemoryBuffer -> IO ()) + , encodeValue :: !(MemoryBuffer -> Int -> Int -> IO (Int, Bool)) + , finishValues :: !(MemoryBuffer -> Int -> IO Int) } buildEncoder :: Column -> IO Encoder @@ -68,7 +71,7 @@ buildEncoder col (INT32 enum) Nothing Nothing - (\buffer -> writeWord32LE buffer . fromIntegral) + (\buffer pos v -> writeWord32At buffer pos (fromIntegral v) >> pure (pos + 4)) col | hasElemType @Int64 col = pure $ @@ -76,7 +79,7 @@ buildEncoder col (INT64 enum) Nothing Nothing - (\buffer -> writeWord64LE buffer . fromIntegral) + (\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 @@ -87,7 +90,7 @@ buildEncoder col (INT64 enum) Nothing Nothing - (\buffer -> writeWord64LE buffer . fromIntegral) + (\buffer pos v -> writeWord64At buffer pos (fromIntegral v) >> pure (pos + 8)) col | hasElemType @Integer col = pure $ @@ -95,12 +98,24 @@ buildEncoder col (INT64 enum) Nothing Nothing - writeInteger64 + writeInteger64At col | hasElemType @Float col = - pure $ scalarEncoder @Float (FLOAT enum) Nothing Nothing writeFloatLE 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 writeDoubleLE 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) @@ -113,17 +128,17 @@ scalarEncoder :: ThriftType -> Maybe ConvertedType -> Maybe LogicalType -> - (MemoryBuffer -> a -> IO ()) -> + (MemoryBuffer -> Int -> a -> IO Int) -> Column -> Encoder -scalarEncoder tt conv logical writeValue col = - Encoder tt conv logical (columnWriter @a col writeValue) (const (pure ())) +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 -> Int32 -> IO ()) -> + (MemoryBuffer -> Int -> Int32 -> IO Int) -> Column -> Encoder #-} @@ -131,7 +146,7 @@ scalarEncoder tt conv logical writeValue col = ThriftType -> Maybe ConvertedType -> Maybe LogicalType -> - (MemoryBuffer -> Int64 -> IO ()) -> + (MemoryBuffer -> Int -> Int64 -> IO Int) -> Column -> Encoder #-} @@ -139,7 +154,7 @@ scalarEncoder tt conv logical writeValue col = ThriftType -> Maybe ConvertedType -> Maybe LogicalType -> - (MemoryBuffer -> Float -> IO ()) -> + (MemoryBuffer -> Int -> Float -> IO Int) -> Column -> Encoder #-} @@ -147,7 +162,7 @@ scalarEncoder tt conv logical writeValue col = ThriftType -> Maybe ConvertedType -> Maybe LogicalType -> - (MemoryBuffer -> Double -> IO ()) -> + (MemoryBuffer -> Int -> Double -> IO Int) -> Column -> Encoder #-} @@ -155,7 +170,7 @@ scalarEncoder tt conv logical writeValue col = ThriftType -> Maybe ConvertedType -> Maybe LogicalType -> - (MemoryBuffer -> Int -> IO ()) -> + (MemoryBuffer -> Int -> Int -> IO Int) -> Column -> Encoder #-} @@ -163,7 +178,7 @@ scalarEncoder tt conv logical writeValue col = ThriftType -> Maybe ConvertedType -> Maybe LogicalType -> - (MemoryBuffer -> Integer -> IO ()) -> + (MemoryBuffer -> Int -> Integer -> IO Int) -> Column -> Encoder #-} @@ -172,11 +187,12 @@ columnWriter :: forall a. (Columnable a) => Column -> - (MemoryBuffer -> a -> IO ()) -> + (MemoryBuffer -> Int -> a -> IO Int) -> MemoryBuffer -> Int -> - IO Bool -columnWriter col writeValue = case col of + 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) @@ -187,36 +203,78 @@ columnWriter col writeValue = case col of Nothing -> mismatch _ -> mismatch where - writeFrom bitmap at buffer row - | isPresent bitmap row = writeValue buffer (at row) >> pure True - | otherwise = pure False + 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 -> Int32 -> IO ()) -> MemoryBuffer -> Int -> IO Bool + Column -> + (MemoryBuffer -> Int -> Int32 -> IO Int) -> + MemoryBuffer -> + Int -> + Int -> + IO (Int, Bool) #-} {-# SPECIALIZE columnWriter :: - Column -> (MemoryBuffer -> Int64 -> IO ()) -> MemoryBuffer -> Int -> IO Bool + Column -> + (MemoryBuffer -> Int -> Int64 -> IO Int) -> + MemoryBuffer -> + Int -> + Int -> + IO (Int, Bool) #-} {-# SPECIALIZE columnWriter :: - Column -> (MemoryBuffer -> Float -> IO ()) -> MemoryBuffer -> Int -> IO Bool + Column -> + (MemoryBuffer -> Int -> Float -> IO Int) -> + MemoryBuffer -> + Int -> + Int -> + IO (Int, Bool) #-} {-# SPECIALIZE columnWriter :: - Column -> (MemoryBuffer -> Double -> IO ()) -> MemoryBuffer -> Int -> IO Bool + Column -> + (MemoryBuffer -> Int -> Double -> IO Int) -> + MemoryBuffer -> + Int -> + Int -> + IO (Int, Bool) #-} {-# SPECIALIZE columnWriter :: - Column -> (MemoryBuffer -> Bool -> IO ()) -> MemoryBuffer -> Int -> IO Bool + Column -> + (MemoryBuffer -> Int -> Bool -> IO Int) -> + MemoryBuffer -> + Int -> + Int -> + IO (Int, Bool) #-} {-# SPECIALIZE columnWriter :: - Column -> (MemoryBuffer -> UTCTime -> IO ()) -> MemoryBuffer -> Int -> IO Bool + Column -> + (MemoryBuffer -> Int -> UTCTime -> IO Int) -> + MemoryBuffer -> + Int -> + Int -> + IO (Int, Bool) #-} {-# SPECIALIZE columnWriter :: - Column -> (MemoryBuffer -> Int -> IO ()) -> MemoryBuffer -> Int -> IO Bool + Column -> + (MemoryBuffer -> Int -> Int -> IO Int) -> + MemoryBuffer -> + Int -> + Int -> + IO (Int, Bool) #-} {-# SPECIALIZE columnWriter :: - Column -> (MemoryBuffer -> Integer -> IO ()) -> MemoryBuffer -> Int -> IO Bool + Column -> + (MemoryBuffer -> Int -> Integer -> IO Int) -> + MemoryBuffer -> + Int -> + Int -> + IO (Int, Bool) #-} isPresent :: Maybe Bitmap -> Int -> Bool @@ -228,19 +286,35 @@ boolEncoder :: Column -> IO Encoder boolEncoder col = do bitsRef <- newIORef (0 :: Word8) countRef <- newIORef (0 :: Int) - let addBit buffer value = do + 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 writeWord8 buffer bits' >> writeIORef bitsRef 0 >> writeIORef countRef 0 - else writeIORef bitsRef bits' >> writeIORef countRef count' - finish buffer = do + 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 - when (count > 0) (readIORef bitsRef >>= writeWord8 buffer) + 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) @@ -251,7 +325,7 @@ textEncoder col = (Just (UTF8 enum)) (Just (LT_STRING (putField StringType))) writePresent - (const (pure ())) + (\_ pos -> pure pos) where writePresent = case col of BoxedColumn bitmap (values :: VB.Vector a) -> @@ -260,32 +334,32 @@ textEncoder col = 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 + 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) - writeTextSlice buffer packed.ptBytes start (end - start) - pure True - | otherwise = pure False + 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) -writeText :: MemoryBuffer -> T.Text -> IO () -writeText buffer (Text bytes offset count) = writeTextSlice buffer bytes offset count -{-# INLINE writeText #-} - -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 @@ -293,9 +367,11 @@ timestampEncoder col = (Just (TIMESTAMP_MICROS enum)) (Just timestampLogical) (columnWriter @UTCTime col writeMicros) - (const (pure ())) + (\_ pos -> pure pos) where - writeMicros buffer t = writeWord64LE buffer (fromIntegral (utcToMicros t)) + writeMicros buffer pos t = do + writeWord64At buffer pos (fromIntegral (utcToMicros t)) + pure (pos + 8) timestampLogical :: LogicalType timestampLogical = diff --git a/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs b/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs index 3bc600cf..7acc7e8a 100644 --- a/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs +++ b/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs @@ -14,6 +14,7 @@ module DataFrame.IO.Utils.RandomAccess ( withWritableBinaryFile, atomicallyWriteFile, MemoryBuffer (..), + ensureCapacity, mallocBuffer, writeByteString, appendTextArraySlice, @@ -21,6 +22,9 @@ module DataFrame.IO.Utils.RandomAccess ( writeWord32LE, writeWord64LE, writeInteger64, + writeWord32At, + writeWord64At, + writeInteger64At, writeFloatLE, writeDoubleLE, bufferResidency, @@ -252,17 +256,28 @@ writeByteString buffer bs = 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) + 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) @@ -272,17 +287,26 @@ writeWord64LE buffer w = do 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) -{-# INLINE writeWord64LE #-} +{-# INLINE writeWord64At #-} writeInteger64 :: MemoryBuffer -> Integer -> IO () -writeInteger64 buffer value +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 = writeWord64LE buffer (fromIntegral value) + | otherwise = do + writeWord64At buffer position (fromIntegral value) + pure (position + 8) where - outOfRange = ioError (userError "writeParquet: Integer value is outside the INT64 range") -{-# INLINE writeInteger64 #-} + outOfRange = + ioError (userError "writeParquet: Integer value is outside the INT64 range") +{-# INLINE writeInteger64At #-} writeFloatLE :: MemoryBuffer -> Float -> IO () writeFloatLE buffer = writeWord32LE buffer . castFloatToWord32 diff --git a/dataframe-parquet/tests/Main.hs b/dataframe-parquet/tests/Main.hs index 4e2f2ae2..c1c7ea7c 100644 --- a/dataframe-parquet/tests/Main.hs +++ b/dataframe-parquet/tests/Main.hs @@ -128,6 +128,20 @@ writerRoundTripTiny label path = TestCase $ 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. -} @@ -258,6 +272,7 @@ tests = "alltypes_plain multi-page" "tests/data/alltypes_plain.parquet" ) + , TestLabel "writer roundtrip: large text" writerRoundTripLargeText ] main :: IO () From 23a7b27dea6de6a3149c9c26007a5f9c472225da Mon Sep 17 00:00:00 2001 From: Raghav Sharma Date: Tue, 1 Sep 2026 23:47:59 +0530 Subject: [PATCH 2/2] Fourmolu --- .../src/DataFrame/IO/Parquet/Writer/Encoder.hs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs index ba26cd7e..0df7b76d 100644 --- a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs @@ -106,7 +106,8 @@ buildEncoder col (FLOAT enum) Nothing Nothing - (\buffer pos v -> writeWord32At buffer pos (castFloatToWord32 v) >> pure (pos + 4)) + ( \buffer pos v -> writeWord32At buffer pos (castFloatToWord32 v) >> pure (pos + 4) + ) col | hasElemType @Double col = pure $ @@ -114,7 +115,8 @@ buildEncoder col (DOUBLE enum) Nothing Nothing - (\buffer pos v -> writeWord64At buffer pos (castDoubleToWord64 v) >> pure (pos + 8)) + ( \buffer pos v -> writeWord64At buffer pos (castDoubleToWord64 v) >> pure (pos + 8) + ) col | hasElemType @Bool col = boolEncoder col | hasElemType @T.Text col = pure (textEncoder col)