git.iko.soy/typeable/postgresql-simple
on master help sign in
Add Vector query
commit1a70d75db1722226c6c502401e6535713416c97e authorOleg Grenrus <oleg.grenrus@iki.fi> dateMon 15 Oct 2018, 17:08 (8 years ago) parent5af199a
7 files changed+180−6
compare with master
unifiedsplit
CHANGES.md+5 bench/Select.hs+32 postgresql-simple.cabal+13 src/Database/PostgreSQL/Simple.hs+1 src/Database/PostgreSQL/Simple/Internal/PQResultUtils.hs+40−6 src/Database/PostgreSQL/Simple/Vector.hs+45 src/Database/PostgreSQL/Simple/Vector/Unboxed.hs+44
CHANGES.md+5
@@ -28,6 +28,11 @@
2828 Thanks to Bardur Arantsson.
2929 https://github.com/lpsmith/postgresql-simple/pull/229
3030
31+ * Add `Vector` and `Vector.Unboxed` `query` variants.
32+ These are more memory efficient
33+ (especially, if you anyway will convert to some vector)
34+ https://github.com/phadej/1
35+
3136 * Documentation improvements
3237 https://github.com/lpsmith/postgresql-simple/pull/227
3338 https://github.com/lpsmith/postgresql-simple/pull/236
bench/Select.hsnew file+32
@@ -0,0 +1,32 @@
1+{-# LANGUAGE OverloadedStrings #-}
2+module Main where
3+
4+import Database.PostgreSQL.Simple
5+import qualified Database.PostgreSQL.Simple.Vector as V
6+import qualified Database.PostgreSQL.Simple.Vector.Unboxed as VU
7+
8+import System.Environment (getArgs)
9+import Data.Foldable (Foldable, foldl')
10+import qualified Data.Vector.Unboxed as VU
11+
12+main :: IO ()
13+main = do
14+ args <- getArgs
15+ conn <- connectPostgreSQL ""
16+ case args of
17+ ("vector":_) -> do
18+ result <- V.query_ conn "SELECT * FROM generate_series(1, 10000000);"
19+ print (process result)
20+ ("unboxed":_) -> do
21+ -- dummy column
22+ result <- VU.query_ conn "SELECT (NULL :: VOID), * FROM generate_series(1, 10000000);"
23+ print (process' result)
24+ _ -> do
25+ result <- query_ conn "SELECT * FROM generate_series(1, 10000000);"
26+ print (process result)
27+
28+process :: Foldable f => f (Only Int) -> Int
29+process = foldl' (\x (Only y) -> max x y) 0
30+
31+process' :: VU.Vector ((), Int) -> Int
32+process' = VU.foldl' (\x (_, y) -> max x y) 0
postgresql-simple.cabal+13
@@ -56,6 +56,8 @@ Library
5656 Database.PostgreSQL.Simple.TypeInfo.Static
5757 Database.PostgreSQL.Simple.Types
5858 Database.PostgreSQL.Simple.Errors
59+ Database.PostgreSQL.Simple.Vector
60+ Database.PostgreSQL.Simple.Vector.Unboxed
5961
6062-- Other-modules:
6163 Database.PostgreSQL.Simple.Internal
@@ -160,3 +162,14 @@ test-suite test
160162 if !impl(ghc >= 7.6)
161163 build-depends:
162164 ghc-prim
165+
166+benchmark select
167+ default-language: Haskell2010
168+ type: exitcode-stdio-1.0
169+
170+ hs-source-dirs: bench
171+ main-is: Select.hs
172+
173+ build-depends: base
174+ , postgresql-simple
175+ , vector
src/Database/PostgreSQL/Simple.hs+1
@@ -355,6 +355,7 @@ executeMany conn q qs = do
355355returning :: (ToRow q, FromRow r) => Connection -> Query -> [q] -> IO [r]
356356returning = returningWith fromRow
357357
358+-- | A version of 'returning' taking parser as argument
358359returningWith :: (ToRow q) => RowParser r -> Connection -> Query -> [q] -> IO [r]
359360returningWith _ _ _ [] = return []
360361returningWith parser conn q qs = do
src/Database/PostgreSQL/Simple/Internal/PQResultUtils.hs+40−6
@@ -14,11 +14,14 @@
1414
1515module Database.PostgreSQL.Simple.Internal.PQResultUtils
1616 ( finishQueryWith
17+ , finishQueryWithV
18+ , finishQueryWithVU
1719 , getRowWith
1820 ) where
1921
2022import Control.Exception as E
2123import Data.ByteString (ByteString)
24+import Data.Foldable (for_)
2225import Database.PostgreSQL.Simple.FromField (ResultError(..))
2326import Database.PostgreSQL.Simple.Ok
2427import Database.PostgreSQL.Simple.Types (Query(..))
@@ -26,18 +29,49 @@ import Database.PostgreSQL.Simple.Internal as Base
2629import Database.PostgreSQL.Simple.TypeInfo
2730import qualified Database.PostgreSQL.LibPQ as PQ
2831import qualified Data.ByteString.Char8 as B
32+import qualified Data.Vector as V
33+import qualified Data.Vector.Mutable as MV
34+import qualified Data.Vector.Unboxed as VU
35+import qualified Data.Vector.Unboxed.Mutable as MVU
2936import Control.Monad.Trans.Reader
3037import Control.Monad.Trans.State.Strict
3138
3239finishQueryWith :: RowParser r -> Connection -> Query -> PQ.Result -> IO [r]
33−finishQueryWith parser conn q result = do
40+finishQueryWith parser conn q result = finishQueryWith' q result $ do
41+ nrows <- PQ.ntuples result
42+ ncols <- PQ.nfields result
43+ forM' 0 (nrows-1) $ \row ->
44+ getRowWith parser row ncols conn result
45+
46+finishQueryWithV :: RowParser r -> Connection -> Query -> PQ.Result -> IO (V.Vector r)
47+finishQueryWithV parser conn q result = finishQueryWith' q result $ do
48+ nrows <- PQ.ntuples result
49+ let PQ.Row nrows' = nrows
50+ ncols <- PQ.nfields result
51+ mv <- MV.unsafeNew (fromIntegral nrows')
52+ for_ [ 0 .. nrows-1 ] $ \row -> do
53+ let PQ.Row row' = row
54+ value <- getRowWith parser row ncols conn result
55+ MV.unsafeWrite mv (fromIntegral row') value
56+ V.unsafeFreeze mv
57+
58+finishQueryWithVU :: VU.Unbox r => RowParser r -> Connection -> Query -> PQ.Result -> IO (VU.Vector r)
59+finishQueryWithVU parser conn q result = finishQueryWith' q result $ do
60+ nrows <- PQ.ntuples result
61+ let PQ.Row nrows' = nrows
62+ ncols <- PQ.nfields result
63+ mv <- MVU.unsafeNew (fromIntegral nrows')
64+ for_ [ 0 .. nrows-1 ] $ \row -> do
65+ let PQ.Row row' = row
66+ value <- getRowWith parser row ncols conn result
67+ MVU.unsafeWrite mv (fromIntegral row') value
68+ VU.unsafeFreeze mv
69+
70+finishQueryWith' :: Query -> PQ.Result -> IO a -> IO a
71+finishQueryWith' q result k = do
3472 status <- PQ.resultStatus result
3573 case status of
36− PQ.TuplesOk -> do
37− nrows <- PQ.ntuples result
38− ncols <- PQ.nfields result
39− forM' 0 (nrows-1) $ \row ->
40− getRowWith parser row ncols conn result
74+ PQ.TuplesOk -> k
4175 PQ.EmptyQuery -> queryErr "query: Empty query"
4276 PQ.CommandOk -> queryErr "query resulted in a command response"
4377 PQ.CopyOut -> queryErr "query: COPY TO is not supported"
src/Database/PostgreSQL/Simple/Vector.hsnew file+45
@@ -0,0 +1,45 @@
1+-- | 'query' variants returning 'V.Vector'.
2+module Database.PostgreSQL.Simple.Vector where
3+
4+import Database.PostgreSQL.Simple (Connection, formatQuery, formatMany)
5+import Database.PostgreSQL.Simple.FromRow (FromRow(..))
6+import Database.PostgreSQL.Simple.ToRow (ToRow(..))
7+import Database.PostgreSQL.Simple.Internal (RowParser, exec)
8+import Database.PostgreSQL.Simple.Internal.PQResultUtils
9+import Database.PostgreSQL.Simple.Types ( Query (..) )
10+
11+import qualified Data.Vector as V
12+
13+-- | Perform a @SELECT@ or other SQL query that is expected to return
14+-- results. All results are retrieved and converted before this
15+-- function returns.
16+query :: (ToRow q, FromRow r) => Connection -> Query -> q -> IO (V.Vector r)
17+query = queryWith fromRow
18+
19+-- | A version of 'query' that does not perform query substitution.
20+query_ :: (FromRow r) => Connection -> Query -> IO (V.Vector r)
21+query_ = queryWith_ fromRow
22+
23+-- | A version of 'query' taking parser as argument
24+queryWith :: ToRow q => RowParser r -> Connection -> Query -> q -> IO (V.Vector r)
25+queryWith parser conn template qs = do
26+ result <- exec conn =<< formatQuery conn template qs
27+ finishQueryWithV parser conn template result
28+
29+-- | A version of 'query_' taking parser as argument
30+queryWith_ :: RowParser r -> Connection -> Query -> IO (V.Vector r)
31+queryWith_ parser conn q@(Query que) = do
32+ result <- exec conn que
33+ finishQueryWithV parser conn q result
34+
35+-- | Execute @INSERT ... RETURNING@, @UPDATE ... RETURNING@, or other SQL
36+-- query that accepts multi-row input and is expected to return results.
37+returning :: (ToRow q, FromRow r) => Connection -> Query -> [q] -> IO (V.Vector r)
38+returning = returningWith fromRow
39+
40+-- | A version of 'returning' taking parser as argument
41+returningWith :: (ToRow q) => RowParser r -> Connection -> Query -> [q] -> IO (V.Vector r)
42+returningWith _ _ _ [] = return V.empty
43+returningWith parser conn q qs = do
44+ result <- exec conn =<< formatMany conn q qs
45+ finishQueryWithV parser conn q result
src/Database/PostgreSQL/Simple/Vector/Unboxed.hsnew file+44
@@ -0,0 +1,44 @@
1+module Database.PostgreSQL.Simple.Vector.Unboxed where
2+
3+import Database.PostgreSQL.Simple (Connection, formatQuery, formatMany)
4+import Database.PostgreSQL.Simple.FromRow (FromRow(..))
5+import Database.PostgreSQL.Simple.ToRow (ToRow(..))
6+import Database.PostgreSQL.Simple.Internal (RowParser, exec)
7+import Database.PostgreSQL.Simple.Internal.PQResultUtils
8+import Database.PostgreSQL.Simple.Types ( Query (..) )
9+
10+import qualified Data.Vector.Unboxed as VU
11+
12+-- | Perform a @SELECT@ or other SQL query that is expected to return
13+-- results. All results are retrieved and converted before this
14+-- function returns.
15+query :: (ToRow q, FromRow r, VU.Unbox r) => Connection -> Query -> q -> IO (VU.Vector r)
16+query = queryWith fromRow
17+
18+-- | A version of 'query' that does not perform query substitution.
19+query_ :: (FromRow r, VU.Unbox r) => Connection -> Query -> IO (VU.Vector r)
20+query_ = queryWith_ fromRow
21+
22+-- | A version of 'query' taking parser as argument
23+queryWith :: (ToRow q, VU.Unbox r) => RowParser r -> Connection -> Query -> q -> IO (VU.Vector r)
24+queryWith parser conn template qs = do
25+ result <- exec conn =<< formatQuery conn template qs
26+ finishQueryWithVU parser conn template result
27+
28+-- | A version of 'query_' taking parser as argument
29+queryWith_ :: VU.Unbox r => RowParser r -> Connection -> Query -> IO (VU.Vector r)
30+queryWith_ parser conn q@(Query que) = do
31+ result <- exec conn que
32+ finishQueryWithVU parser conn q result
33+
34+-- | Execute @INSERT ... RETURNING@, @UPDATE ... RETURNING@, or other SQL
35+-- query that accepts multi-row input and is expected to return results.
36+returning :: (ToRow q, FromRow r, VU.Unbox r) => Connection -> Query -> [q] -> IO (VU.Vector r)
37+returning = returningWith fromRow
38+
39+-- | A version of 'returning' taking parser as argument
40+returningWith :: (ToRow q, VU.Unbox r) => RowParser r -> Connection -> Query -> [q] -> IO (VU.Vector r)
41+returningWith _ _ _ [] = return VU.empty
42+returningWith parser conn q qs = do
43+ result <- exec conn =<< formatMany conn q qs
44+ finishQueryWithVU parser conn q result