Reduce round trips for prepared bulk writes

benchmark/executemany.jl

@@ -0,0 +1,65 @@
+# Compare grouped executemany with DBInterface's row-by-row method on the same
+# statement and parameters, against a disposable local PostgreSQL server with a
+# trust-authenticated postgres account:
+# julia --project benchmark/executemany.jl PORT [ROWS] [BYTES]
+using Postgres
+const DBI = Postgres.DBInterface
+
+serial(stmt, params) = invoke(DBI.executemany, Tuple{DBI.Statement, Any}, stmt, params)
+batched(stmt, params) = DBI.executemany(stmt, params)
+middle(values) = sort(values)[cld(length(values), 2)]
+cpu_seconds() = Sys.isunix() ? ccall(:clock, Clong, ()) / 1_000_000 : NaN
+
+function sample(f, conn, stmt, params)
+    DBI.execute(conn, "TRUNCATE executemany_benchmark")
+    cpu = cpu_seconds()
+    measured = @timed f(stmt, params)
+    cpu = cpu_seconds() - cpu
+    @assert measured.value === nothing
+    nrows = length(first(params))
+    row = only(DBI.execute(conn,
+        "SELECT count(*) AS n, sum(id) AS total, sum(octet_length(value)) AS bytes FROM executemany_benchmark"))
+    @assert row.n == nrows && row.total == nrows * (nrows + 1) ÷ 2
+    @assert row.bytes == nrows * ncodeunits(params[2][1])
+    return (wall=measured.time, cpu=cpu, bytes=measured.bytes,
+            compilation=hasproperty(measured, :compile_time) ? measured.compile_time : NaN)
+end
+
+function main()
+    port = parse(Int, ARGS[1])
+    nrows = length(ARGS) >= 2 ? parse(Int, ARGS[2]) : 128
+    payload = length(ARGS) >= 3 ? parse(Int, ARGS[3]) : 12
+    conn = DBI.connect(Postgres.Connection, "127.0.0.1", "postgres", nothing;
+        port, dbname="postgres", sslmode="disable", gssencmode="disable")
+    try
+        DBI.execute(conn, "CREATE TEMP TABLE executemany_benchmark(id int PRIMARY KEY, value text)")
+        stmt = DBI.prepare(conn, raw"INSERT INTO executemany_benchmark VALUES ($1, $2)")
+        params = (collect(1:nrows), fill(repeat("v", payload), nrows))
+        for _ in 1:3, f in (serial, batched)
+            sample(f, conn, stmt, params)
+        end
+        results = Dict(:serial => NamedTuple[], :batched => NamedTuple[])
+        for trial in 1:9
+            order = isodd(trial) ? (:serial, :batched) : (:batched, :serial)
+            for variant in order
+                result = sample(variant == :serial ? serial : batched, conn, stmt, params)
+                push!(results[variant], result)
+                println("sample=", trial, " variant=", variant, " ", result)
+            end
+        end
+        println("Julia=", VERSION, " DBInterface=", pkgversion(DBI), " threads=", Threads.nthreads(),
+                " rows=", nrows, " payload_bytes=", payload)
+        for variant in (:serial, :batched)
+            values = results[variant]
+            wall = [x.wall * 1000 for x in values]
+            println(variant, " wall_ms median=", middle(wall), " range=", extrema(wall),
+                " client_cpu_ms=", middle([x.cpu * 1000 for x in values]),
+                " allocated_bytes=", middle([x.bytes for x in values]),
+                " compilation_s=", sum(x.compilation for x in values))
+        end
+        DBI.close!(stmt)
+    finally
+        DBI.close!(conn)
+    end
+end
+main()

docs/src/manual.md

@@ -177,6 +177,16 @@ DBInterface.executemany(stmt, ([1, 2, 3], ["a", "b", "c"]))
 DBInterface.close!(stmt)
 ```
 
+`executemany` runs all rows in one transaction, or in a savepoint when a
+transaction is already open, so an error rolls back every row.
+
+For `INSERT`, `UPDATE`, `DELETE`, and `MERGE` statements that return no rows,
+with parameters given as `Vector` columns of fixed-width integers and floats,
+`Bool`, strings, dates, times, `UUID`s, timestamps, decimals, `missing`, or
+`Vector`s of these, rows are sent in groups of up to 64 without waiting for the
+server after each row. Other statements and inputs, custom styles, debug mode,
+and GSS-encrypted connections run one row at a time.
+
 ## Transactions
 
 Postgres.jl supports both its connection-passing helper and the DBInterface transaction API.

src/api/API.jl

@@ -1184,6 +1184,79 @@ function exec(style::S, socket::BufferedConn, stmtname::String,
                    Ref{Union{Nothing, Int}}(nothing), Ref{UInt8}(UInt8('I')))
 end
 
+# A group contains only no-result Bind/Execute pairs followed by one Sync.
+# Reading while writing is necessary even for a bounded group: a NOTICE may
+# fill the server's output while a large parameter fills its input.
+function exec_batch(style::AbstractPostgresStyle, socket::BufferedConn,
+                    messages::Vector{UInt8}, count::Int,
+                    server_parameters::Dict{String, String}, debug::Bool)
+    # Each side records its error before aborting the connection, so the first
+    # entry is the cause; a later one only reports that abort.
+    failures = Channel{Any}(2)
+    writer = errormonitor(Threads.@spawn try
+        write(socket, messages)
+        flush(socket)
+    catch err
+        put!(failures, err)
+        try
+            abort(socket)
+        catch
+        end
+    end)
+    server_error = nothing
+    completed = 0
+    bound = false
+    status = UInt8('I')
+    try
+        while true
+            mt, len = readheader(socket, debug)
+            if mt == UInt8('2')
+                (len == 0 && !bound && completed < count && server_error === nothing) ||
+                    throw(Error("unexpected BindComplete in executemany response"))
+                bound = true
+            elseif mt == UInt8('C')
+                (bound && server_error === nothing) ||
+                    throw(Error("unexpected CommandComplete in executemany response"))
+                commandComplete(len, socket)
+                completed += 1
+                bound = false
+            elseif mt == UInt8('E')
+                server_error = something(server_error, errorResponse(len, socket, debug))
+            elseif mt == UInt8('N')
+                notice_callback(style, noticeResponse(len, socket))
+            elseif mt == UInt8('A')
+                notification_callback(style, notificationResponse(len, socket))
+            elseif mt == UInt8('S')
+                parameterStatus!(server_parameters, len, socket)
+            elseif mt == UInt8('Z')
+                status = read_ready_status(socket, len)
+                (server_error !== nothing || (!bound && completed == count)) ||
+                    throw(Error("incomplete executemany response from server"))
+                status == (server_error === nothing ? UInt8('T') : UInt8('E')) ||
+                    throw(Error("unexpected transaction status in executemany response"))
+                break
+            else
+                throw(Error("unexpected message type '$(Char(mt))' in executemany response"))
+            end
+        end
+        wait(writer)
+        isready(failures) && throw(take!(failures))
+        return status, server_error
+    catch err
+        # The caller still owns the connection lock. Stop and join the writer
+        # before that lock can be released or the connection can be reused.
+        put!(failures, err)
+        try
+            abort(socket)
+        catch
+        end
+        wait(writer)
+        server_error === nothing || throw(server_error)
+        cause = take!(failures)
+        cause === err ? rethrow() : throw(cause)
+    end
+end
+
 # Execute an unnamed statement as one extended-query segment. There is one
 # Sync, after Parse/Describe/Bind/Execute. A transaction-mode pooler therefore
 # cannot return the backend between dependent protocol messages and replace the

src/execute.jl

@@ -560,9 +560,108 @@ function DBInterface.execute(conn::Connection, sql::AbstractString, params=nothi
     end
 end
 
-# DBInterface's generic connection overload prepares a statement without
-# closing it. That leaks named server statements when this driver's cache is
-# disabled. Keep the handle lifetime explicit for both bulk fallbacks.
+# Only native values in ordinary owned columns may be converted ahead of the
+# server. Custom indexing/conversion can have observable per-row side effects.
+function batch_parameter_type(::Type{T}) where {T}
+    T <: Union{Missing, Nothing, Bool, Int8, Int16, Int32, Int64, Int128,
+               UInt8, UInt16, UInt32, UInt64, UInt128, Float16, Float32, Float64,
+               String, SubString{String}, Dates.Date, Dates.DateTime, Dates.Time,
+               UUID, API.PGTimestamp, Durations.Timestamp{Dates.Nanosecond},
+               API.PGDecimal} && return true
+    T <: Vector && return batch_parameter_type(eltype(T))
+    T isa Union && return all(batch_parameter_type, Base.uniontypes(T))
+    return false
+end
+
+function batch_columns(stmt::Statement, params)
+    params isa Union{Tuple, NamedTuple, Vector} || return nothing
+    columns = params isa NamedTuple ? values(params) : params
+    isempty(columns) && return nothing
+    length(columns) == stmt.nparams || return nothing
+    all(x -> x isa Vector && batch_parameter_type(eltype(x)), columns) || return nothing
+    nrows = length(first(columns))
+    nrows > 1 || return nothing
+    all(x -> length(x) == nrows, columns) || return nothing
+    return columns
+end
+
+function flush_batch!(stmt::Statement, buf::IOBuffer, count::Int)
+    conn = stmt.conn
+    API._write_message_to_buffer(buf, conn.debug, ('S',))
+    status, server_error = API.exec_batch(conn.style, conn.socket, take!(buf), count,
+                                         conn.server_parameters, conn.debug)
+    conn.server_in_transaction = API.in_transaction_status(status)
+    server_error === nothing || throw(server_error)
+    return nothing
+end
+
+function execute_batch!(stmt::Statement, columns, row::Int, nrows::Int)
+    @lock stmt.conn.lock begin
+        checkstmt(stmt)
+        buf = IOBuffer()
+        count = 0
+        complete_bytes = 0
+        try
+            while row <= nrows && count < 64
+                try
+                    bind_params!(stmt.params, DBInterface.LazyIndex(columns, row), stmt.sql)
+                    API._write_message_to_buffer(buf, stmt.conn.debug,
+                        ('B', "", stmt.name, Int16(0), Int16(stmt.nparams), API.Params(stmt.params), Int16(0)))
+                    API._write_message_to_buffer(buf, stmt.conn.debug, ('E', "", Int32(0)))
+                finally
+                    fill!(stmt.params, missing)
+                end
+                complete_bytes = position(buf)
+                count += 1
+                row += 1
+                # End the group at 64 KiB; the row that crosses it still goes whole.
+                complete_bytes >= 65536 && break
+            end
+        catch
+            # Send the rows already encoded before reporting a conversion
+            # error: as in serial execution, their server error wins.
+            truncate(buf, complete_bytes)
+            seekend(buf)
+            count > 0 && flush_batch!(stmt, buf, count)
+            rethrow()
+        end
+        flush_batch!(stmt, buf, count)
+        return row
+    end
+end
+
+function DBInterface.executemany(stmt::Statement, params)
+    conn = stmt.conn
+    columns = batch_columns(stmt, params)
+    # Custom styles can observe per-row callbacks and query logs; keep them serial.
+    if columns === nothing || !(conn.style isa API.PostgresStyle) || conn.debug || stmt.nfields != 0
+        return invoke(DBInterface.executemany, Tuple{DBInterface.Statement, Any}, stmt, params)
+    end
+    nrows = length(first(columns))
+    DBInterface.transaction(conn) do
+        result = DBInterface.execute(stmt, DBInterface.LazyIndex(columns, 1))
+        tag = command_tag(result)
+        # Use prepared metadata and the actual command tag, not a SQL lexer.
+        # GSS wrap/unwrap share one security context and stay on the serial path.
+        pipeline = tag !== nothing &&
+            first(split(tag; limit=2)) in ("INSERT", "UPDATE", "DELETE", "MERGE") &&
+            !(@lock conn.lock conn.socket.transport isa API.GSSConn)
+        if pipeline
+            row = 2
+            while row <= nrows
+                row = execute_batch!(stmt, columns, row, nrows)
+            end
+        else
+            for row in 2:nrows
+                DBInterface.close!(DBInterface.execute(stmt, DBInterface.LazyIndex(columns, row)))
+            end
+        end
+    end
+    return nothing
+end
+
+# Keep the statement lifetime explicit on every supported DBInterface version,
+# including versions whose generic connection overload does not close it.
 function DBInterface.executemany(conn::Connection, sql::AbstractString, params)
     stmt = DBInterface.prepare(conn, sql)
     try