Add Vector query

CHANGES.md

@@ -28,6 +28,11 @@
     Thanks to Bardur Arantsson.
     https://github.com/lpsmith/postgresql-simple/pull/229
 
+  * Add `Vector` and `Vector.Unboxed` `query` variants.
+    These are more memory efficient
+    (especially, if you anyway will convert to some vector)
+    https://github.com/phadej/1
+
   * Documentation improvements
     https://github.com/lpsmith/postgresql-simple/pull/227
     https://github.com/lpsmith/postgresql-simple/pull/236

bench/Select.hs

@@ -0,0 +1,32 @@
+{-# LANGUAGE OverloadedStrings #-}
+module Main where
+
+import           Database.PostgreSQL.Simple
+import qualified Database.PostgreSQL.Simple.Vector as V
+import qualified Database.PostgreSQL.Simple.Vector.Unboxed as VU
+
+import           System.Environment (getArgs)
+import           Data.Foldable (Foldable, foldl')
+import qualified Data.Vector.Unboxed as VU
+
+main :: IO ()
+main = do
+    args <- getArgs
+    conn <- connectPostgreSQL ""
+    case args of
+        ("vector":_) -> do
+            result <- V.query_ conn "SELECT * FROM generate_series(1, 10000000);"
+            print (process result)
+        ("unboxed":_) -> do
+            -- dummy column
+            result <- VU.query_ conn "SELECT (NULL :: VOID), * FROM generate_series(1, 10000000);"
+            print (process' result)
+        _ -> do
+            result <- query_ conn "SELECT * FROM generate_series(1, 10000000);"
+            print (process result)
+
+process :: Foldable f => f (Only Int) ->  Int
+process = foldl' (\x (Only y) -> max x y) 0
+
+process' :: VU.Vector ((), Int) ->  Int
+process' = VU.foldl' (\x (_, y) -> max x y) 0

postgresql-simple.cabal

@@ -56,6 +56,8 @@ Library
      Database.PostgreSQL.Simple.TypeInfo.Static
      Database.PostgreSQL.Simple.Types
      Database.PostgreSQL.Simple.Errors
+     Database.PostgreSQL.Simple.Vector
+     Database.PostgreSQL.Simple.Vector.Unboxed
 
 -- Other-modules:
      Database.PostgreSQL.Simple.Internal
@@ -160,3 +162,14 @@ test-suite test
   if !impl(ghc >= 7.6)
     build-depends:
       ghc-prim
+
+benchmark select
+  default-language: Haskell2010
+  type:             exitcode-stdio-1.0
+
+  hs-source-dirs:  bench
+  main-is: Select.hs
+
+  build-depends: base
+               , postgresql-simple
+               , vector

src/Database/PostgreSQL/Simple.hs

@@ -355,6 +355,7 @@ executeMany conn q qs = do
 returning :: (ToRow q, FromRow r) => Connection -> Query -> [q] -> IO [r]
 returning = returningWith fromRow
 
+-- | A version of 'returning' taking parser as argument
 returningWith :: (ToRow q) => RowParser r -> Connection -> Query -> [q] -> IO [r]
 returningWith _ _ _ [] = return []
 returningWith parser conn q qs = do

src/Database/PostgreSQL/Simple/Internal/PQResultUtils.hs

@@ -14,30 +14,64 @@
 
 module Database.PostgreSQL.Simple.Internal.PQResultUtils
     ( finishQueryWith
+    , finishQueryWithV
+    , finishQueryWithVU
     , getRowWith
     ) where
 
 import           Control.Exception as E
 import           Data.ByteString (ByteString)
+import           Data.Foldable (for_)
 import           Database.PostgreSQL.Simple.FromField (ResultError(..))
 import           Database.PostgreSQL.Simple.Ok
 import           Database.PostgreSQL.Simple.Types (Query(..))
 import           Database.PostgreSQL.Simple.Internal as Base
 import           Database.PostgreSQL.Simple.TypeInfo
 import qualified Database.PostgreSQL.LibPQ as PQ
 import qualified Data.ByteString.Char8 as B
+import qualified Data.Vector as V
+import qualified Data.Vector.Mutable as MV
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as MVU
 import           Control.Monad.Trans.Reader
 import           Control.Monad.Trans.State.Strict
 
 finishQueryWith :: RowParser r -> Connection -> Query -> PQ.Result -> IO [r]
-finishQueryWith parser conn q result = do
+finishQueryWith parser conn q result = finishQueryWith' q result $ do
+    nrows <- PQ.ntuples result
+    ncols <- PQ.nfields result
+    forM' 0 (nrows-1) $ \row ->
+        getRowWith parser row ncols conn result
+
+finishQueryWithV :: RowParser r -> Connection -> Query -> PQ.Result -> IO (V.Vector r)
+finishQueryWithV parser conn q result = finishQueryWith' q result $ do
+    nrows <- PQ.ntuples result
+    let PQ.Row nrows' = nrows
+    ncols <- PQ.nfields result
+    mv <- MV.unsafeNew (fromIntegral nrows')
+    for_ [ 0 .. nrows-1 ] $ \row -> do
+        let PQ.Row row' = row
+        value <- getRowWith parser row ncols conn result
+        MV.unsafeWrite mv (fromIntegral row') value
+    V.unsafeFreeze mv
+
+finishQueryWithVU :: VU.Unbox r => RowParser r -> Connection -> Query -> PQ.Result -> IO (VU.Vector r)
+finishQueryWithVU parser conn q result = finishQueryWith' q result $ do
+    nrows <- PQ.ntuples result
+    let PQ.Row nrows' = nrows
+    ncols <- PQ.nfields result
+    mv <- MVU.unsafeNew (fromIntegral nrows')
+    for_ [ 0 .. nrows-1 ] $ \row -> do
+        let PQ.Row row' = row
+        value <- getRowWith parser row ncols conn result
+        MVU.unsafeWrite mv (fromIntegral row') value
+    VU.unsafeFreeze mv
+
+finishQueryWith' :: Query -> PQ.Result -> IO a -> IO a
+finishQueryWith' q result k = do
   status <- PQ.resultStatus result
   case status of
-    PQ.TuplesOk -> do
-        nrows <- PQ.ntuples result
-        ncols <- PQ.nfields result
-        forM' 0 (nrows-1) $ \row ->
-            getRowWith parser row ncols conn result
+    PQ.TuplesOk -> k
     PQ.EmptyQuery    -> queryErr "query: Empty query"
     PQ.CommandOk     -> queryErr "query resulted in a command response"
     PQ.CopyOut       -> queryErr "query: COPY TO is not supported"

src/Database/PostgreSQL/Simple/Vector.hs

@@ -0,0 +1,45 @@
+-- | 'query' variants returning 'V.Vector'.
+module Database.PostgreSQL.Simple.Vector where
+
+import           Database.PostgreSQL.Simple (Connection, formatQuery, formatMany)
+import           Database.PostgreSQL.Simple.FromRow (FromRow(..))
+import           Database.PostgreSQL.Simple.ToRow (ToRow(..))
+import           Database.PostgreSQL.Simple.Internal (RowParser, exec)
+import           Database.PostgreSQL.Simple.Internal.PQResultUtils
+import           Database.PostgreSQL.Simple.Types ( Query (..) )
+
+import qualified Data.Vector as V
+
+-- | Perform a @SELECT@ or other SQL query that is expected to return
+-- results. All results are retrieved and converted before this
+-- function returns.
+query :: (ToRow q, FromRow r) => Connection -> Query -> q -> IO (V.Vector r)
+query = queryWith fromRow
+
+-- | A version of 'query' that does not perform query substitution.
+query_ :: (FromRow r) => Connection -> Query -> IO (V.Vector r)
+query_ = queryWith_ fromRow
+
+-- | A version of 'query' taking parser as argument
+queryWith :: ToRow q => RowParser r -> Connection -> Query -> q -> IO (V.Vector r)
+queryWith parser conn template qs = do
+  result <- exec conn =<< formatQuery conn template qs
+  finishQueryWithV parser conn template result
+
+-- | A version of 'query_' taking parser as argument
+queryWith_ :: RowParser r -> Connection -> Query -> IO (V.Vector r)
+queryWith_ parser conn q@(Query que) = do
+  result <- exec conn que
+  finishQueryWithV parser conn q result
+
+-- | Execute @INSERT ... RETURNING@, @UPDATE ... RETURNING@, or other SQL
+-- query that accepts multi-row input and is expected to return results.
+returning :: (ToRow q, FromRow r) => Connection -> Query -> [q] -> IO (V.Vector r)
+returning = returningWith fromRow
+
+-- | A version of 'returning' taking parser as argument
+returningWith :: (ToRow q) => RowParser r -> Connection -> Query -> [q] -> IO (V.Vector r)
+returningWith _ _ _ [] = return V.empty
+returningWith parser conn q qs = do
+  result <- exec conn =<< formatMany conn q qs
+  finishQueryWithV parser conn q result

src/Database/PostgreSQL/Simple/Vector/Unboxed.hs

@@ -0,0 +1,44 @@
+module Database.PostgreSQL.Simple.Vector.Unboxed where
+
+import           Database.PostgreSQL.Simple (Connection, formatQuery, formatMany)
+import           Database.PostgreSQL.Simple.FromRow (FromRow(..))
+import           Database.PostgreSQL.Simple.ToRow (ToRow(..))
+import           Database.PostgreSQL.Simple.Internal (RowParser, exec)
+import           Database.PostgreSQL.Simple.Internal.PQResultUtils
+import           Database.PostgreSQL.Simple.Types ( Query (..) )
+
+import qualified Data.Vector.Unboxed as VU
+
+-- | Perform a @SELECT@ or other SQL query that is expected to return
+-- results. All results are retrieved and converted before this
+-- function returns.
+query :: (ToRow q, FromRow r, VU.Unbox r) => Connection -> Query -> q -> IO (VU.Vector r)
+query = queryWith fromRow
+
+-- | A version of 'query' that does not perform query substitution.
+query_ :: (FromRow r, VU.Unbox r) => Connection -> Query -> IO (VU.Vector r)
+query_ = queryWith_ fromRow
+
+-- | A version of 'query' taking parser as argument
+queryWith :: (ToRow q, VU.Unbox r) => RowParser r -> Connection -> Query -> q -> IO (VU.Vector r)
+queryWith parser conn template qs = do
+  result <- exec conn =<< formatQuery conn template qs
+  finishQueryWithVU parser conn template result
+
+-- | A version of 'query_' taking parser as argument
+queryWith_ :: VU.Unbox r => RowParser r -> Connection -> Query -> IO (VU.Vector r)
+queryWith_ parser conn q@(Query que) = do
+  result <- exec conn que
+  finishQueryWithVU parser conn q result
+
+-- | Execute @INSERT ... RETURNING@, @UPDATE ... RETURNING@, or other SQL
+-- query that accepts multi-row input and is expected to return results.
+returning :: (ToRow q, FromRow r, VU.Unbox r) => Connection -> Query -> [q] -> IO (VU.Vector r)
+returning = returningWith fromRow
+
+-- | A version of 'returning' taking parser as argument
+returningWith :: (ToRow q, VU.Unbox r) => RowParser r -> Connection -> Query -> [q] -> IO (VU.Vector r)
+returningWith _ _ _ [] = return VU.empty
+returningWith parser conn q qs = do
+  result <- exec conn =<< formatMany conn q qs
+  finishQueryWithVU parser conn q result