Benchmark for decoding DataRows

bench/Bench.hs

@@ -24,9 +24,9 @@ import qualified Database.PostgreSQL.LibPQ as LibPQ
 import Database.PostgreSQL.Protocol.Types
 import Database.PostgreSQL.Protocol.Encoders
 import Database.PostgreSQL.Protocol.Decoders
+import Database.PostgreSQL.Protocol.DataRows
 import Database.PostgreSQL.Protocol.Store.Decode
 import Database.PostgreSQL.Protocol.Codecs.Decoders
-import Database.PostgreSQL.Protocol.ExtractDataRows
 import Database.PostgreSQL.Driver.Connection
 import Database.PostgreSQL.Driver
 import Criterion.Main
@@ -48,23 +48,30 @@ import Criterion.Main
 --  (SELECT repeat('a', 100)::bytea FROM generate_series(1, 300));
 
 -- main = benchMultiPw
--- main = defaultMain
---     [ bgroup "Requests"
---         [ 
---             -- env createConnection (\c -> bench "100 of 1k" . nfIO $ requestAction c)
---             bench "parser" $ nf parse bs
---         ]
---     ]
-main = benchMultiPw
+main = defaultMain
+    -- [ bgroup "Requests"
+    --     [ 
+    --         -- env createConnection (\c -> bench "100 of 1k" . nfIO $ requestAction c)
+    --         bench "parser" $ nf parse bs
+    --     ]
+    -- ]
+    [ bgroup "Decoder"
+        [ bench "datarow" $ nf benchDataRowDecoder bs
+        ]
+    ]
+-- main = benchMultiPw
+
+benchDataRowDecoder bs = decodeManyRows decodeDataRow $ 
+    DataRows (DataChunk 350 bs) Empty
+  where
+    decodeDataRow = do
+        (Header _ len) <- decodeHeader
+        getByteString len
 
 {-# NOINLINE bs #-}
 bs :: B.ByteString
 bs = unsafePerformIO $ B.readFile "1.txt"
 
-parse bs | B.null bs = ()
-         | otherwise = let (rest, v) = runDecode getCustomRow bs
-                       in v `seq` parse rest
-
 benchLoop :: IO ()
 benchLoop = do
     ref <- newIORef 0 :: IO (IORef Word)

src/Database/PostgreSQL/Protocol/DataRows.hs

@@ -13,6 +13,7 @@ import qualified Data.ByteString as B
 import qualified Data.ByteString.Unsafe as B
 import qualified Data.Vector as V
 import qualified Data.Vector.Mutable as MV
+import qualified Data.List as L
 import Data.Foldable
 import System.IO.Unsafe
 
@@ -141,19 +142,18 @@ decodeOneRow dec (DataRows (DataChunk _ bs) _) = snd $ runDecode dec bs
 
 decodeManyRows :: Decode a -> DataRows -> V.Vector a
 decodeManyRows dec dr = unsafePerformIO $ do
-    vec <- MV.unsafeNew count
-    go vec 0 dr
+    vec <- MV.unsafeNew . fromIntegral $ countDataRows dr
+    let go startInd Empty = pure ()
+        go startInd (DataRows (DataChunk len bs) nextDr) = do
+            let endInd = startInd + fromIntegral len
+            runDecodeIO 
+                (traverse_ (writeDec vec) [startInd .. (endInd  -1)]) 
+                bs
+            go endInd nextDr
+    go 0 dr
     V.unsafeFreeze vec
   where
-    go vec startInd Empty = pure ()
-    go vec startInd (DataRows (DataChunk len bs) nextDr) = do
-        let endInd = startInd + fromIntegral len
-        runDecodeIO 
-            (traverse_ (writeDec vec) [startInd .. (endInd  -1)]) 
-            bs
-        go vec endInd nextDr
-
-    count = fromIntegral $ countDataRows dr
+    {-# INLINE writeDec #-}
     writeDec vec pos = dec >>= embedIO . MV.unsafeWrite vec pos
 
 ---

src/Database/PostgreSQL/Protocol/Types.hs

@@ -52,7 +52,7 @@ data DataChunk = DataChunk
     deriving (Show, Eq)
 
 -- | Helper types that contains only raw DataRows messages.
-data DataRows = Empty | DataRows {-# UNPACK #-} DataChunk DataRows
+data DataRows = Empty | DataRows {-# UNPACK #-} !DataChunk DataRows
     deriving (Show, Eq)
 
 -- | Ad-hoc type only for data rows.