Add JLD2Dataset

Author: darsnack

Project.toml

@@ -15,6 +15,7 @@ FixedPointNumbers = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
 GZip = "92fee26a-97fe-5a0c-ad85-20a5f3185b63"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"

src/MLDatasets.jl

@@ -12,11 +12,12 @@ using FilePathsBase
 using FilePathsBase: AbstractPath
 using Glob
 using HDF5
+using JLD2
 
 import MLUtils
 using MLUtils: getobs, numobs, AbstractDataContainer
 
-export FileDataset, TableDataset, HDF5Dataset
+export FileDataset, TableDataset, HDF5Dataset, JLD2Dataset
 export getobs, numobs
 
 # Julia 1.0 compatibility
@@ -50,6 +51,7 @@ include("download.jl")
 include("containers/filedataset.jl")
 include("containers/tabledataset.jl")
 include("containers/hdf5dataset.jl")
+include("containers/jld2dataset.jl")
 
 # Misc.
 include("BostonHousing/BostonHousing.jl")

src/containers/hdf5dataset.jl

@@ -19,14 +19,14 @@ See [`close(::HDF5Dataset)`](@ref) for closing the underlying HDF5 file pointer.
 For array datasets, the last dimension is assumed to be the observation dimension.
 For scalar datasets, the stored value is returned by `getobs` for any index.
 """
-struct HDF5Dataset
+struct HDF5Dataset <: AbstractDataContainer
     fid::HDF5.File
     paths::Vector{HDF5.Dataset}
     shapes::Vector{Tuple}
 
     function HDF5Dataset(fid::HDF5.File, paths::Vector{HDF5.Dataset}, shapes::Vector)
         _check_hdf5_shapes(shapes) ||
-            throw(ArgumentError("Cannot create HDF5Dataset for datasets with mismatch number of observations."))
+            throw(ArgumentError("Cannot create HDF5Dataset for datasets with mismatched number of observations."))
 
         new(fid, paths, shapes)
     end

src/containers/jld2dataset.jl

@@ -0,0 +1,37 @@
+_check_jld2_nobs(nobs) = all(==(first(nobs)), nobs[2:end])
+
+"""
+    JLD2Dataset(file::Union{AbstractString, AbstractPath}, paths)
+    JLD2Dataset(fid::JLD2.JLDFile, paths)
+
+Wrap several JLD2 datasets (`paths`) as a single dataset container.
+Each dataset `p` in `paths` should be accessible as `fid[p]`.
+Calling `getobs` on a `JLD2Dataset` is equivalent to mapping `getobs` on
+each dataset in `paths`.
+See [`close(::JLD2Dataset)`](@ref) for closing the underlying JLD2 file pointer.
+"""
+struct JLD2Dataset{T<:JLD2.JLDFile} <: AbstractDataContainer
+    fid::T
+    paths::Vector{String}
+
+    function JLD2Dataset(fid::JLD2.JLDFile, paths::Vector{String})
+        nobs = map(p -> numobs(fid[p]), paths)
+        _check_jld2_nobs(nobs) ||
+            throw(ArgumentError("Cannot create JLD2Dataset for datasets with mismatched number of observations."))
+
+        new{typeof(fid)}(fid, paths)
+    end
+end
+
+JLD2Dataset(file::Union{AbstractString, AbstractPath}, paths) =
+    JLD2Dataset(jldopen(file, "r"), paths)
+
+MLUtils.getobs(dataset::JLD2Dataset, i) = Tuple(map(p -> getobs(dataset.fid[p], i), dataset.paths))
+MLUtils.numobs(dataset::JLD2Dataset) = numobs(dataset.fid[dataset.paths[1]])
+
+"""
+    close(dataset::JLD2Dataset)
+
+Close the underlying JLD2 file pointer for `dataset`.
+"""
+Base.close(dataset::JLD2Dataset) = close(dataset.fid)

test/containers/jld2dataset.jl

@@ -0,0 +1,27 @@
+function setup_jld2dataset_test()
+    datasets = [
+        ("d1", rand(2, 10)),
+        ("g1/d1", rand(10)),
+        ("g1/d2", string.('a':'j')),
+        ("g2/g1/d1", rand(Bool, 3, 3, 10))
+    ]
+
+    fid = jldopen("test.jld2", "w")
+    for (path, data) in datasets
+        fid[path] = data
+    end
+    close(fid)
+
+    return first.(datasets), Tuple(last.(datasets))
+end
+
+@testset "JLD2Dataset" begin
+    paths, datas = setup_jld2dataset_test()
+    dataset = JLD2Dataset("test.jld2", paths)
+    for i in 1:10
+        @test getobs(dataset, i) == getobs(datas, i)
+    end
+    @test numobs(dataset) == 10
+    close(dataset)
+    rm("test.jld2")
+end

test/runtests.jl

@@ -5,6 +5,7 @@ using DataDeps
 using MLUtils: getobs, numobs
 using DataFrames, CSV, Tables
 using HDF5
+using JLD2
 
 ENV["DATADEPS_ALWAYS_ACCEPT"] = true
 
@@ -32,6 +33,7 @@ container_tests = [
     "containers/filedataset.jl",
     "containers/tabledataset.jl",
     "containers/hdf5dataset.jl",
+    "containers/jld2dataset.jl",
 ]
 
 @testset "Datasets" for t in dataset_tests