Merge pull request #23 from richiejp/read-perf

Improve read performance & benchmark

Author: MikeInnes

src/BSON.jl

@@ -14,16 +14,16 @@ const Primitive = Union{Nothing,Bool,Int32,Int64,Float64,String,Vector{UInt8},BS
   datetime, null, regex, dbpointer, javascript, symbol, javascript_scoped,
   int32, timestamp, int64, decimal128, minkey=0xFF, maxkey=0x7F)
 
-applychildren!(f, x) = x
+applychildren!(::Function, x) = x
 
-function applychildren!(f, x::BSONDict)
+function applychildren!(f::Function, x::BSONDict)::BSONDict
   for k in keys(x)
     x[k] = f(x[k])
   end
   return x
 end
 
-function applychildren!(f, x::BSONArray)
+function applychildren!(f::Function, x::BSONArray)::BSONArray
   for i = 1:length(x)
     x[i] = f(x[i])
   end

src/extensions.jl

@@ -99,7 +99,16 @@ function newstruct(T, xs...)
     flds = Any[convert(fieldtype(T, i), x) for (i,x) in enumerate(xs)]
     return ccall(:jl_new_structv, Any, (Any,Ptr{Cvoid},UInt32), T, flds, length(flds))
   else
-    newstruct!(initstruct(T), xs...)
+    # Manual inline of newstruct! to work around bug
+    # https://github.com/MikeInnes/BSON.jl/issues/2#issuecomment-452204339
+    x = initstruct(T)
+
+    for (i, f) = enumerate(xs)
+      f = convert(fieldtype(typeof(x),i), f)
+      ccall(:jl_set_nth_field, Nothing, (Any, Csize_t, Any), x, i-1, f)
+    end
+    x
+
   end
 end
 

src/read.jl

@@ -1,26 +1,20 @@
-jtype(tag) =
+jtype(tag::BSONType)::DataType =
   tag == null ? Nothing :
   tag == boolean ? Bool :
   tag == int32 ? Int32 :
   tag == int64 ? Int64 :
   tag == double ? Float64 :
   error("Unsupported tag $tag")
 
-function parse_cstr(io::IO)
-  buf = IOBuffer()
-  while (ch = read(io, UInt8)) != 0x00
-    write(buf, ch)
-  end
-  return String(read(seek(buf, 0)))
-end
+parse_cstr(io::IO) = readuntil(io, '\0')
 
-function parse_tag(io::IO, tag)
+function parse_tag(io::IO, tag::BSONType)
   if tag == null
     nothing
   elseif tag == document
     parse_doc(io)
   elseif tag == array
-    Any[map(x->x[2], parse_pairs(io))...]
+    parse_array(io)
   elseif tag == string
     len = read(io, Int32)-1
     s = String(read(io, len))
@@ -35,18 +29,32 @@ function parse_tag(io::IO, tag)
   end
 end
 
-function parse_pairs(io::IO)
+function parse_array(io::IO)::BSONArray
   len = read(io, Int32)
-  ps = []
+  ps = BSONArray()
+
   while (tag = read(io, BSONType)) ≠ eof
-    k = Symbol(parse_cstr(io))
-    v = parse_tag(io::IO, tag)
-    push!(ps, k => v)
+    # Note that arrays are dicts with the index as the key
+    while read(io, UInt8) != 0x00
+      nothing
+    end
+    push!(ps, parse_tag(io::IO, tag))
   end
-  return ps
+
+  ps
 end
 
-parse_doc(io::IO) = BSONDict(parse_pairs(io))
+function parse_doc(io::IO)::BSONDict
+  len = read(io, Int32)
+  dic = BSONDict()
+
+  while (tag = read(io, BSONType)) ≠ eof
+    k = Symbol(parse_cstr(io))
+    dic[k] = parse_tag(io::IO, tag)
+  end
+
+  dic
+end
 
 backrefs!(x, refs) = applychildren!(x -> backrefs!(x, refs), x)
 

test/benchmark.jl

@@ -0,0 +1,69 @@
+module Benchmark
+
+using Profile
+using BSON
+
+chars = [x for x in '0':'z']
+strings = [String([rand(chars) for _ in 1:20]) for _ in 1:20]
+rstr(n::Int)::String = rand(strings)[1:n]
+
+struct Baz
+  going::String
+  deeper::String
+end
+
+Baz() = Baz(rstr(20), rstr(1))
+
+struct Bar
+  level::Int64
+  bazes::Vector{Baz}
+end
+
+Bar() = Bar(rand(Int64), [Baz() for _ in 1:50])
+
+struct Foo
+  agile::String
+  software::String
+  management::String
+  consultant::String
+  training::String
+  projects::Vector{Bar}
+end
+
+Foo() = Foo(rstr(5), rstr(7), rstr(17), rstr(11), rstr(13), [Bar() for _ in 1:2000])
+
+foos = Dict(:Foo => Foo())
+io = IOBuffer()
+
+@info "Bench Save BSON"
+@timev bson(io, foos)
+seek(io, 0)
+
+@info "Bench Parse BSON"
+dict = @timev BSON.parse(io)
+
+@info "Bench Raise BSON to Julia types"
+@timev BSON.raise_recursive(dict)
+
+@info "Profile Save BSON"
+@profile bson(io, foos)
+Profile.print(;noisefloor=2)
+Profile.clear()
+seek(io, 0)
+
+@info "Profile Parse BSON"
+dict = @profile BSON.parse(io)
+Profile.print(;noisefloor=2)
+Profile.clear()
+
+@info "Profile Raise BSON to Julia types"
+rfoos = @profile BSON.raise_recursive(dict)
+Profile.print(;noisefloor=2)
+Profile.clear()
+
+# Sanity check the results
+rfoos[:Foo]::Foo
+rfoos[:Foo].projects[1]::Bar
+rfoos[:Foo].projects[1].bazes[1]::Baz
+
+end