Use jldopen in collect_results! when scanning JLD2 files. (#293)

* Use `jldopen` in `collect_results!` when scanning JLD2 files.

Employ lazy loading via Requires.jl to allow special load functions for particular file types in `collect_results!`, and falling back to `wload` as a default. This potentially makes scanning collections of large files much faster when using the `black_list` argument of `collect_results!`.

An example use case are simulations that generate amounts of data where the parameters are stored in a Dict, together with the data under the keyword `data`, all saved in a JLD2 file. If we only want to enumerate all the parameters that have been run, we would use `collect_results(datadir("simulations"), black_list = ("data",))`. With `wload`, this still loads the whole file into memory before ignoring the `data` key. With `jldopen`, only the parameters are loaded into memory. In my tests, this leads to around 400x speed up when scanning large collections of files.

For now, only `jldopen` for JLD2 files is implemented, but this could potentially be extended to other data storage formats that allow memory mapped access to storage keys.

* Fix typo.

* Use `Base.require` to try to load JLD2 instead of relying on Requires.jl.

* Add JLD2 as dependency, simplify code.

* Update minor version number and changelog.

Author: ulido

CHANGELOG.md

@@ -1,3 +1,5 @@
+# 2.6.0
+* Use `JLD2`'s jldopen in `collect_results!` to speed up loading of metadata.
 # 2.5.0
 * Add an `update` option of `collect_results!` allowing the updating of an existing results collection if data files were modified or deleted.
 # 2.4.1

Project.toml

@@ -1,11 +1,12 @@
 name = "DrWatson"
 uuid = "634d3b9d-ee7a-5ddf-bec9-22491ea816e1"
 repo = "https://github.com/JuliaDynamics/DrWatson.jl.git"
-version = "2.5.0"
+version = "2.6.0"
 
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LibGit2 = "76f85450-5226-5b5a-8eaa-529ad045b433"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
@@ -15,6 +16,7 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
 [compat]
 FileIO = "1.0.6"
+JLD2 = "0.4.15"
 MacroTools = "0.5"
 Requires = "0.5.2, 0.6, 1"
 UnPack = "1.0.1"

src/DrWatson.jl

@@ -17,6 +17,7 @@ export @pack!, @unpack
 
 # Functionality that saves/loads
 using FileIO
+using JLD2
 export save, load
 export wsave, wload
 

src/result_collection.jl

@@ -148,8 +148,8 @@ function collect_results!(filename, folder;
         # Now update the mtime of the new or modified file
         mtimes[file] = mtime_file
 
-        data = rpath === nothing ? wload(file) : wload(joinpath(rpath, file))
-        df_new = to_data_row(data, file; kwargs...)
+        fpath = rpath === nothing ? file : joinpath(rpath, file)
+        df_new = to_data_row(FileIO.query(fpath); kwargs...)
         #add filename
         df_new[!, :path] .= file
         if replace_entry
@@ -209,6 +209,20 @@ end
 is_valid_file(file, valid_filetypes) =
     any(endswith(file, v) for v in valid_filetypes)
 
+# Use wload per default when nothing else is available
+function to_data_row(file::File; kwargs...)
+    fpath = filename(file)
+    @debug "Opening $(filename(file)) with fallback wload."
+    return to_data_row(wload(fpath), fpath; kwargs...)
+end
+# Specialize for JLD2 files, can do much faster mmapped access
+function to_data_row(file::File{format"JLD2"}; kwargs...)
+    fpath = filename(file)
+    @debug "Opening $(filename(file)) with jldopen."
+    JLD2.jldopen(filename(file), "r") do data
+        return to_data_row(data, fpath; kwargs...)
+    end
+end
 function to_data_row(data, file;
         white_list = collect(keys(data)),
         black_list = keytype(data).((:gitcommit, :gitpatch, :script)),

test/update_results_tests.jl

@@ -195,6 +195,27 @@ subfolders = true, special_list=special_list, black_list = black_list)
     end
 end
 
+###############################################################################
+#                           test jldopen                                      #
+###############################################################################
+
+mktempdir(datadir()) do folder
+    # Create a data file
+    d = Dict("idx" => 1, "value" => rand(100000))
+    fname = joinpath(folder, savename(d, ending, ignores = ("value",)))
+    DrWatson.wsave(fname, d)
+
+    if ending == "jld2"
+        msg_re = r"Opening .* with jldopen."
+    else
+        msg_re = r"Opening .* with fallback wload."
+    end
+    @test_logs (:debug, msg_re) min_level=Base.CoreLogging.Debug match_mode=:any cres = collect_results(folder, black_list = ("value",))
+
+    @test cres.idx[1] == 1 # It's what we've saved above.
+    @test size(cres,1) == 1 # only one file
+    @test size(cres,2) == 2 # idx and path
+end
 
 ###############################################################################
 #                              Quickactivate macro                            #