Add HDF5 dataset

Author: darsnack

Project.toml

@@ -14,6 +14,7 @@ FilePathsBase = "48062228-2e41-5def-b9a4-89aafe57970f"
 FixedPointNumbers = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
 GZip = "92fee26a-97fe-5a0c-ad85-20a5f3185b63"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
@@ -29,8 +30,8 @@ ColorTypes = "0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.10, 0.11"
 DataDeps = "0.3, 0.4, 0.5, 0.6, 0.7"
 DataFrames = "1.3"
 FixedPointNumbers = "0.3, 0.4, 0.5, 0.6, 0.7, 0.8"
-Glob = "1.3"
 GZip = "0.5"
+Glob = "1.3"
 ImageCore = "0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8"
 JSON3 = "1"
 MAT = "0.7, 0.8, 0.9, 0.10"

src/MLDatasets.jl

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

src/containers/filedataset.jl

@@ -39,8 +39,8 @@ Wrap a set of file `paths` as a dataset (traversed in the same order as `paths`)
 Alternatively, specify a `dir` and collect all paths that match a glob `pattern`
 (recursively globbing by `depth`). The glob order determines the traversal order.
 """
-struct FileDataset{T} <: AbstractDataContainer
-    paths::T
+struct FileDataset{T<:Union{AbstractPath, AbstractString}} <: AbstractDataContainer
+    paths::Vector{T}
 end
 
 FileDataset(dir, pattern = "*", depth = 4) = FileDataset(rglob(pattern, string(dir), depth))

src/containers/hdf5dataset.jl

@@ -0,0 +1,50 @@
+function _check_hdf5_shapes(shapes)
+    nobs = map(last, filter(!isempty, shapes))
+
+    return all(==(first(nobs)), nobs[2:end])
+end
+
+"""
+    HDF5Dataset(file::Union{AbstractString, AbstractPath}, paths)
+    HDF5Dataset(fid::HDF5.File, paths::Vector{HDF5.Dataset})
+    HDF5Dataset(fid::HDF5.File, paths::Vector{<:AbstractString})
+    HDF5Dataset(fid::HDF5.File, paths::Vector{HDF5.Dataset}, shapes)
+
+Wrap several HDF5 datasets (`paths`) as a single dataset container.
+Each dataset `p` in `paths` should be accessible as `fid[p]`.
+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
+    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."))
+
+        new(fid, paths, shapes)
+    end
+end
+
+HDF5Dataset(fid::HDF5.File, paths::Vector{HDF5.Dataset}) =
+    HDF5Dataset(fid, paths, map(size, paths))
+HDF5Dataset(fid::HDF5.File, paths::Vector{<:AbstractString}) =
+    HDF5Dataset(fid, map(p -> fid[p], paths))
+HDF5Dataset(file::Union{AbstractString, AbstractPath}, paths) =
+    HDF5Dataset(h5open(file, "r"), paths)
+
+MLUtils.getobs(dataset::HDF5Dataset, i) = Tuple(map(dataset.paths, dataset.shapes) do path, shape
+    if isempty(shape)
+        return read(path)
+    else
+        I = map(s -> 1:s, shape[1:(end - 1)])
+        return path[I..., i]
+    end
+end)
+MLUtils.numobs(dataset::HDF5Dataset) = last(first(filter(!isempty, dataset.shapes)))
+
+Base.close(dataset::HDF5Dataset) = close(dataset.fid)

test/containers/hdf5dataset.jl

@@ -0,0 +1,28 @@
+function setup_hdf5dataset_test()
+    datasets = [
+        ("d1", rand(2, 10)),
+        ("g1/d1", rand(10)),
+        # these are broken at HDF5 level
+        # ("g1/d2", string.('a':'j')),
+        # ("g2/g1/d1", "test")
+    ]
+
+    fid = h5open("test.h5", "w")
+    for (path, data) in datasets
+        fid[path] = data
+    end
+    close(fid)
+
+    return first.(datasets), last.(datasets)
+end
+
+@testset "HDF5Dataset" begin
+    paths, datas = setup_hdf5dataset_test()
+    dataset = HDF5Dataset("test.h5", paths)
+    for i in 1:10
+        data = Tuple(map(x -> getobs(x, i), datas))
+        @test getobs(dataset, i) == data
+    end
+    @test numobs(dataset) == 10
+    rm("test.h5")
+end

test/runtests.jl

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