Add CachedDataset

Author: darsnack

src/MLDatasets.jl

@@ -16,8 +16,6 @@ using JLD2
 
 import MLUtils
 using MLUtils: getobs, numobs, AbstractDataContainer
-
-export FileDataset, TableDataset, HDF5Dataset, JLD2Dataset
 export getobs, numobs
 
 # Julia 1.0 compatibility
@@ -49,9 +47,15 @@ end
 include("download.jl")
 
 include("containers/filedataset.jl")
+export FileDataset
 include("containers/tabledataset.jl")
+export TableDataset
 include("containers/hdf5dataset.jl")
+export HDF5Dataset
 include("containers/jld2dataset.jl")
+export JLD2Dataset
+include("containers/cacheddataset.jl")
+export CachedDataset
 
 # Misc.
 include("BostonHousing/BostonHousing.jl")

src/containers/cacheddataset.jl

@@ -0,0 +1,39 @@
+"""
+    make_cache(source, cacheidx)
+
+Return a in-memory copy of `source` at observation indices `cacheidx`.
+Defaults to `getobs(source, cacheidx)`.
+"""
+make_cache(source, cacheidx) = getobs(source, cacheidx)
+
+"""
+    CachedDataset(source, cachesize = numbobs(source))
+    CachedDataset(source, cacheidx, cache)
+
+Wrap a `source` data container and cache `cachesize` samples in memory.
+This can be useful for improving read speeds when `source` is a lazy data container,
+but your system memory is large enough to store a sizeable chunk of it.
+
+By default the observation indices `1:cachesize` are cached.
+You can manually pass in a `cache` and set of `cacheidx` as well.
+
+See also [`make_cache`](@ref) for customizing the default cache creation for `source`.
+"""
+struct CachedDataset{T, S}
+    source::T
+    cacheidx::Vector{Int}
+    cache::S
+end
+
+function CachedDataset(source, cachesize::Int = numobs(source))
+    cacheidx = 1:cachesize
+
+    CachedDataset(source, collect(cacheidx), make_cache(source, cacheidx))
+end
+
+function MLUtils.getobs(dataset::CachedDataset, i::Integer)
+    _i = findfirst(==(i), dataset.cacheidx)
+
+    return isnothing(_i) ? getobs(dataset.source, i) : getobs(dataset.cache, _i)
+end
+MLUtils.numobs(dataset::CachedDataset) = numobs(dataset.source)

src/containers/filedataset.jl

@@ -52,5 +52,6 @@ FileDataset(loadfn,
 FileDataset(dir::Union{AbstractPath, AbstractString}, pattern::AbstractString = "*", depth = 4) =
     FileDataset(loadfile, dir, pattern, depth)
 
-MLUtils.getobs(dataset::FileDataset, i) = loadfile(dataset.paths[i])
+MLUtils.getobs(dataset::FileDataset, i::Integer) = loadfile(dataset.paths[i])
+MLUtils.getobs(dataset::FileDataset, is::AbstractVector) = map(Base.Fix1(getobs, dataset), is)
 MLUtils.numobs(dataset::FileDataset) = length(dataset.paths)

test/containers/cacheddataset.jl

@@ -0,0 +1,25 @@
+@testset "CachedDataset" begin
+    @testset "CachedDataset(::FileDataset)" begin
+        files = setup_filedataset_test()
+        fdataset = FileDataset("root", "*.csv")
+        cdataset = CachedDataset(fdataset)
+
+        @test numobs(cdataset) == numobs(fdataset)
+        @test cdataset.cache == getobs(fdataset, 1:numobs(fdataset))
+        @test all(getobs(cdataset, i) == getobs(fdataset, i) for i in 1:numobs(fdataset))
+
+        cleanup_filedataset_test()
+    end
+
+    @testset "CachedDataset(::HDF5Dataset)" begin
+        paths, datas = setup_hdf5dataset_test()
+        hdataset = HDF5Dataset("test.h5", ["d1"])
+        cdataset = CachedDataset(hdataset, 5)
+
+        @test numobs(cdataset) == numobs(hdataset)
+        @test cdataset.cache == getobs(hdataset, 1:5)
+        @test all(getobs(cdataset, i) == getobs(hdataset, i) for i in 1:10)
+
+        cleanup_hdf5dataset_test()
+    end
+end

test/containers/filedataset.jl

@@ -23,6 +23,7 @@ function setup_filedataset_test()
 
     return files
 end
+cleanup_filedataset_test() = rm("root"; recursive = true)
 
 @testset "FileDataset" begin
     files = setup_filedataset_test()
@@ -32,5 +33,5 @@ end
         true_obs = MLDatasets.loadfile(file)
         @test getobs(dataset, i) == true_obs
     end
-    rm("root"; recursive = true)
+    cleanup_filedataset_test()
 end

test/containers/hdf5dataset.jl

@@ -14,6 +14,7 @@ function setup_hdf5dataset_test()
 
     return first.(datasets), last.(datasets)
 end
+cleanup_hdf5dataset_test() = rm("test.h5")
 
 @testset "HDF5Dataset" begin
     paths, datas = setup_hdf5dataset_test()
@@ -24,5 +25,5 @@ end
     end
     @test numobs(dataset) == 10
     close(dataset)
-    rm("test.h5")
+    cleanup_hdf5dataset_test()
 end

test/containers/jld2dataset.jl

@@ -14,6 +14,7 @@ function setup_jld2dataset_test()
 
     return first.(datasets), Tuple(last.(datasets))
 end
+cleanup_jld2dataset_test() = rm("test.jld2")
 
 @testset "JLD2Dataset" begin
     paths, datas = setup_jld2dataset_test()
@@ -23,5 +24,5 @@ end
     end
     @test numobs(dataset) == 10
     close(dataset)
-    rm("test.jld2")
+    cleanup_jld2dataset_test()
 end

test/runtests.jl

@@ -2,7 +2,6 @@ using Test
 using MLDatasets
 using ImageCore
 using DataDeps
-using MLUtils: getobs, numobs
 using DataFrames, CSV, Tables
 using HDF5
 using JLD2
@@ -34,6 +33,7 @@ container_tests = [
     "containers/tabledataset.jl",
     "containers/hdf5dataset.jl",
     "containers/jld2dataset.jl",
+    "containers/cacheddataset.jl",
 ]
 
 @testset "Datasets" for t in dataset_tests