Initial port of FastAI dataset containers

Author: darsnack

Project.toml

@@ -4,16 +4,22 @@ version = "0.5.15"
 
 [deps]
 BinDeps = "9e28174c-4ba2-5203-b857-d8d62c4213ee"
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 ColorTypes = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
+FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+FilePathsBase = "48062228-2e41-5def-b9a4-89aafe57970f"
 FixedPointNumbers = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
 GZip = "92fee26a-97fe-5a0c-ad85-20a5f3185b63"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 Pickle = "fbb45041-c46e-462f-888f-7c521cafbc2c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [compat]
 BinDeps = "1"

src/MLDatasets.jl

@@ -7,6 +7,15 @@ using DelimitedFiles: readdlm
 using FixedPointNumbers, ColorTypes
 using Pickle
 using SparseArrays
+using DataFrames, CSV, Tables
+using FilePathsBase
+using FilePathsBase: AbstractPath
+
+import MLUtils
+using MLUtils: getobs, numobs, AbstractDataContainer
+
+export FileDataset, TableDataset
+export getobs, numobs
 
 # Julia 1.0 compatibility
 if !isdefined(Base, :isnothing)
@@ -36,6 +45,8 @@ end
 
 include("download.jl")
 
+include("containers/filedataset.jl")
+include("containers/tabledataset.jl")
 
 # Misc.
 include("BostonHousing/BostonHousing.jl")

src/containers/filedataset.jl

@@ -0,0 +1,38 @@
+# FileDataset
+
+"""
+    rglob(filepattern, dir = pwd(), depth = 4)
+
+Recursive glob up to `depth` layers deep within `dir`.
+"""
+function rglob(filepattern = "*", dir = pwd(), depth = 4)
+    patterns = [repeat("*/", i) * filepattern for i in 0:(depth - 1)]
+
+    return vcat([glob(pattern, dir) for pattern in patterns]...)
+end
+
+"""
+    loadfile(file)
+
+Load a file from disk into the appropriate format.
+"""
+function loadfile(file::String)
+    if isimagefile(file)
+        # faster image loading
+        return FileIO.load(file, view = true)
+    elseif endswith(file, ".csv")
+        return DataFrame(CSV.File(file))
+    else
+        return FileIO.load(file)
+    end
+end
+loadfile(file::AbstractPath) = loadfile(string(file))
+
+struct FileDataset{T} <: AbstractDataContainer
+    paths::T
+end
+
+FileDataset(dir, pattern = "*", depth = 4) = rglob(pattern, string(dir), depth)
+
+MLUtils.getobs(dataset::FileDataset, i) = loadfile(dataset.paths[i])
+MLUtils.numobs(dataset::FileDataset) = length(dataset.paths)

src/containers/tabledataset.jl

@@ -0,0 +1,47 @@
+struct TableDataset{T} <: AbstractDataContainer
+    table::T
+
+    # TableDatasets must implement the Tables.jl interface
+    function TableDataset{T}(table::T) where {T}
+        Tables.istable(table) ||
+            throw(ArgumentError("TableDatasets must implement the Tabels.jl interface"))
+
+        new{T}(table)
+    end
+end
+
+TableDataset(table::T) where {T} = TableDataset{T}(table)
+TableDataset(path::Union{AbstractPath, AbstractString}) =
+    TableDataset(DataFrame(CSV.File(path)))
+
+# slow accesses based on Tables.jl
+function MLUtils.getobs(dataset::TableDataset, i)
+    if Tables.rowaccess(dataset.table)
+        row, _ = Iterators.peel(Iterators.drop(Tables.rows(dataset.table), i - 1))
+        return row
+    elseif Tables.columnaccess(dataset.table)
+        colnames = Tables.columnnames(dataset.table)
+        rowvals = [Tables.getcolumn(dataset.table, j)[i] for j in 1:length(colnames)]
+        return (; zip(colnames, rowvals)...)
+    else
+        error("The Tables.jl implementation used should have either rowaccess or columnaccess.")
+    end
+end
+function MLUtils.numobs(dataset::TableDataset)
+    if Tables.columnaccess(dataset.table)
+        return length(Tables.getcolumn(dataset.table, 1))
+    elseif Tables.rowaccess(dataset.table)
+        # length might not be defined, but has to be for this to work.
+        return length(Tables.rows(dataset.table))
+    else
+        error("The Tables.jl implementation used should have either rowaccess or columnaccess.")
+    end
+end
+
+# fast access for DataFrame
+MLUtils.getobs(dataset::TableDataset{<:DataFrame}, i) = dataset.table[i, :]
+MLUtils.numobs(dataset::TableDataset{<:DataFrame}) = nrow(dataset.table)
+
+# fast access for CSV.File
+MLUtils.getobs(dataset::TableDataset{<:CSV.File}, i) = dataset.table[i]
+MLUtils.numobs(dataset::TableDataset{<:CSV.File}) = length(dataset.table)

test/containers/tabledataset.jl

@@ -0,0 +1,53 @@
+@testset "TableDataset" begin
+    @testset "TableDataset from rowaccess table" begin
+        Tables.columnaccess(::Type{<:Tables.MatrixTable}) = false
+        Tables.rowaccess(::Type{<:Tables.MatrixTable}) = true
+
+        testtable = Tables.table([1 4.0 "7"; 2 5.0 "8"; 3 6.0 "9"])
+        td = TableDataset(testtable)
+
+        @test all(getobs(td, 1) .== [1, 4.0, "7"])
+        @test numobs(td) == 3
+    end
+
+    @testset "TableDataset from columnaccess table" begin
+        Tables.columnaccess(::Type{<:Tables.MatrixTable}) = true
+        Tables.rowaccess(::Type{<:Tables.MatrixTable}) = false
+
+        testtable = Tables.table([1 4.0 "7"; 2 5.0 "8"; 3 6.0 "9"])
+        td = TableDataset(testtable)
+
+        @test [data for data in getobs(td, 2)] == [2, 5.0, "8"]
+        @test numobs(td) == 3
+
+        @test getobs(td, 1) isa NamedTuple
+    end
+
+    @testset "TableDataset from DataFrames" begin
+        testtable = DataFrame(
+            col1 = [1, 2, 3, 4, 5],
+            col2 = ["a", "b", "c", "d", "e"],
+            col3 = [10, 20, 30, 40, 50],
+            col4 = ["A", "B", "C", "D", "E"],
+            col5 = [100.0, 200.0, 300.0, 400.0, 500.0],
+            split = ["train", "train", "train", "valid", "valid"],
+        )
+        td = TableDataset(testtable)
+        @test td isa TableDataset{<:DataFrame}
+
+        @test [data for data in getobs(td, 1)] == [1, "a", 10, "A", 100.0, "train"]
+        @test numobs(td) == 5
+    end
+
+    @testset "TableDataset from CSV" begin
+        open("test.csv", "w") do io
+            write(io, "col1,col2,col3,col4,col5, split\n1,a,10,A,100.,train")
+        end
+        testtable = CSV.File("test.csv")
+        td = TableDataset(testtable)
+        @test td isa TableDataset{<:CSV.File}
+        @test [data for data in getobs(td, 1)] == [1, "a", 10, "A", 100.0, "train"]
+        @test numobs(td) == 1
+        rm("test.csv")
+    end
+end

test/runtests.jl

@@ -2,11 +2,12 @@ using Test
 using MLDatasets
 using ImageCore
 using DataDeps
-
+using MLUtils: getobs, numobs
+using DataFrames, CSV, Tables
 
 ENV["DATADEPS_ALWAYS_ACCEPT"] = true
 
-tests = [
+dataset_tests = [
     # misc
     "tst_iris.jl",
     "tst_boston_housing.jl",
@@ -26,12 +27,20 @@ tests = [
     "tst_tudataset.jl",
 ]
 
-for t in tests
+container_tests = [
+    "containers/tabledataset.jl",
+]
+
+@testset "Datasets" for t in dataset_tests
     @testset "$t" begin
         include(t)
     end
 end
 
+@testset "Containers" for t in container_tests
+    include(t)
+end
+
 #temporary to not stress CI
 if !parse(Bool, get(ENV, "CI", "false"))
     @testset "other tests" begin