getindex(::Column) performance (#252)

* Initial refactor

* Remove allocations

* Remove commented code

* Revert code movement

* Bump patch

* Remove indirection

* Revert wrapper

* Sort out formatting

* Remove formatting issues

* Initial performance regression tests

* Uncommnt tests

* Make allocation tests more robust

* Naming and count_allocs and binary_format

* Add string test

* Uncomment code

* Remove test code

* Revert test data extraction

* Formatting

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Merge reversion

* Update test/runtests.jl

Co-authored-by: Eric Davies <[email protected]>

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Eric Davies <[email protected]>

Author: 3 people

Project.toml

@@ -1,7 +1,7 @@
 name = "LibPQ"
 uuid = "194296ae-ab2e-5f79-8cd4-7183a0a5a0d1"
 license = "MIT"
-version = "1.13.0"
+version = "1.14.0"
 
 [deps]
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"

src/parsing.jl

@@ -683,12 +683,13 @@ end
 
 struct FallbackConversion <: AbstractDict{Tuple{Oid,Type},Base.Callable} end
 
+struct ParseType{T} <: Function end
+
+(::ParseType{typ})(pqv::PQValue) where {typ} = parse(typ, pqv)
+
 function Base.getindex(cmap::FallbackConversion, oid_typ::Tuple{Integer,Type})
     _, typ = oid_typ
-
-    return function parse_type(pqv::PQValue)
-        return parse(typ, pqv)
-    end
+    return ParseType{typ}()
 end
 
 Base.haskey(cmap::FallbackConversion, oid_typ::Tuple{Integer,Type}) = true

src/tables.jl

@@ -51,14 +51,12 @@ end
 
 # Columns
 
-struct Column{T} <: AbstractVector{T}
-    result::Result
+struct Column{T,oid,typ,Tresult<:Result,Tfunc<:Base.Callable} <: AbstractVector{T}
+    result::Tresult
     col::Int
     col_name::Symbol
-    oid::Oid
     not_null::Bool
-    typ::Type
-    func::Base.Callable
+    func::Tfunc
 end
 
 struct Columns <: AbstractVector{Column}
@@ -86,7 +84,9 @@ function Tables.schema(cs::Columns)
     return Tables.Schema(map(Symbol, column_names(jl_result)), types)
 end
 
-function Column(jl_result::Result, col::Integer, name=Symbol(column_name(jl_result, col)))
+function Column(
+    jl_result::Tresult, col::Integer, name=Symbol(column_name(jl_result, col))
+) where {Tresult<:Result}
     @boundscheck if !checkindex(Bool, Base.OneTo(num_columns(jl_result)), col)
         throw(BoundsError(Columns(jl_result), col))
     end
@@ -95,8 +95,10 @@ function Column(jl_result::Result, col::Integer, name=Symbol(column_name(jl_resu
     typ = column_types(jl_result)[col]
     func = jl_result.column_funcs[col]
     not_null = jl_result.not_null[col]
-    element_type = not_null ? typ : Union{typ, Missing}
-    return Column{element_type}(jl_result, col, name, oid, not_null, typ, func)
+    element_type = not_null ? typ : Union{typ,Missing}
+    return Column{element_type,oid,typ,Tresult,typeof(func)}(
+        jl_result, col, name, not_null, func
+    )
 end
 
 function Column(jl_result::Result, name::Symbol, col=column_number(jl_result, name))
@@ -107,13 +109,13 @@ result(c::Column) = getfield(c, :result)
 column_number(c::Column) = getfield(c, :col)
 column_name(c::Column) = getfield(c, :col_name)
 
-function Base.getindex(c::Column{T}, row::Integer)::T where T
+function Base.getindex(c::Column{T,oid,typ}, row::Integer)::T where {T,oid,typ}
     jl_result = result(c)
     col = column_number(c)
     if isnull(jl_result, row, col)
         return missing
     else
-        return c.func(PQValue{c.oid}(jl_result, row, col))::c.typ
+        return c.func(PQValue{oid}(jl_result, row, col))::typ
     end
 end
 

test/runtests.jl

@@ -32,6 +32,12 @@ macro test_nolog_on_windows(ex...)
     end
 end
 
+# Copied from `@time`.
+function count_allocs(f, args...)
+    stats = @timed f(args...)
+    return Base.gc_alloc_count(stats.gcstats)
+end
+
 @testset "LibPQ" begin
 
 @testset "ConninfoDisplay" begin
@@ -1612,6 +1618,48 @@ end
             close(conn)
             @test !isopen(conn)
         end
+
+        @testset "getindex(::Column) performance" begin
+            @testset "$in_val, bin=$bin_fmt" for (
+                in_val, out_val, bin_fmt, num_allocs
+            ) in [
+                ("5::float8", 5.0, true, 0),
+                ("5::float4", 5.0f0, true, 0),
+                ("3::int8", Int64(3), true, 0),
+                ("3::int4", Int32(3), true, 0),
+                ("3::int2", Int16(3), true, 0),
+                ("'hello'::varchar", "hello", true, 1),
+                ("5::float8", 5.0, false, 2),
+                ("5::float4", 5.0f0, false, 2),
+                ("3::int8", Int64(3), false, 2),
+                ("3::int4", Int32(3), false, 2),
+                ("3::int2", Int16(3), false, 2),
+                ("'hello'::varchar", "hello", false, 3),
+            ]
+
+                # Establish connection and construct temporary table.
+                conn = LibPQ.Connection("dbname=postgres user=$DATABASE_USER")
+
+                # Get the column.
+                result = execute(
+                    conn, "SELECT $in_val AS my_column;"; binary_format=bin_fmt,
+                )
+                col = columntable(result).my_column
+
+                # Ensure that element is of expected type and value.
+                @test typeof(col[1]) == typeof(out_val)
+                @test col[1] == out_val
+
+                # Ensure that getting an element from the column produces num_allocs allocs.
+                foo(col) = [col[1] for _ in 1:100]
+                count_allocs(foo, col)
+                max_expected_allocs = num_allocs * 100 + 5
+                @test count_allocs(foo, col) < max_expected_allocs
+
+                close(result)
+                close(conn)
+            end
+        end
     end
 
     @testset "Statements" begin