numobs and getobs support for Tables.jl's tables

Project.toml

@@ -1,34 +1,41 @@
 name = "MLUtils"
 uuid = "f1d291b0-491e-4a28-83b9-f70985020b54"
 authors = ["Carlo Lucibello <[email protected]> and contributors"]
-version = "0.2.11"
+version = "0.2.12"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"
 FoldsThreads = "9c68100b-dfe1-47cf-94c8-95104e173443"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ShowCases = "605ecd9f-84a6-4c9e-81e2-4798472b76a3"
+SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999"
 
 [compat]
 ChainRulesCore = "1.0"
+DataAPI = "1.0"
 DelimitedFiles = "1.0"
 FLoops = "0.2"
 FoldsThreads = "0.1"
+SimpleTraits = "0.9"
 ShowCases = "0.1"
 StatsBase = "0.33"
+Tables = "1.10"
 Transducers = "0.4"
 julia = "1.6"
 
 [extras]
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["ChainRulesTestUtils", "SparseArrays", "Test", "Zygote"]
+test = ["ChainRulesTestUtils", "DataFrames", "SparseArrays", "Test", "Zygote"]

src/MLUtils.jl

@@ -8,12 +8,19 @@ using FLoops.Transducers: Executor, ThreadedEx
 using FoldsThreads: TaskPoolEx
 import StatsBase: sample
 using Transducers
+using Tables
+using DataAPI
 using Base: @propagate_inbounds
 using Random: AbstractRNG, shuffle!, GLOBAL_RNG, rand!, randn!
 import ChainRulesCore: rrule
 using ChainRulesCore: @non_differentiable, unthunk, AbstractZero,
                       NoTangent, ZeroTangent, ProjectTo
 
+using SimpleTraits
+
+@traitdef IsTable{X}
+@traitimpl IsTable{X} <- Tables.istable(X)
+
 
 include("observation.jl")
 export numobs,

src/observation.jl

@@ -16,7 +16,9 @@ See also [`getobs`](@ref)
 function numobs end
 
 # Generic Fallbacks
-numobs(data) = length(data)
+@traitfn numobs(data::X) where {X; IsTable{X}} = DataAPI.nrow(data)
+@traitfn numobs(data::X) where {X; !IsTable{X}} = length(data)
+
 
 """
     getobs(data, [idx])
@@ -46,7 +48,11 @@ function getobs end
 
 # Generic Fallbacks
 getobs(data) = data
-getobs(data, idx) = data[idx]
+# getobs(data, idx) = data[idx]
+
+@traitfn getobs(data::X, idx) where {X; IsTable{X}} = Tables.subset(data, idx, viewhint=false)
+@traitfn getobs(data::X, idx) where {X; !IsTable{X}} = data[idx]
+
 
 """
     getobs!(buffer, data, idx)
@@ -161,3 +167,5 @@ function getobs!(buffers, data::Dict, i)
 
     return buffers
 end
+
+

test/observation.jl

@@ -190,4 +190,11 @@ end
         @test getobs!((nothing,xbuf),(Xs,X), 3:4) == (getobs(Xs,3:4),xbuf)
         @test xbuf == getobs(X,3:4)
     end
+
+    @testset "tables" begin
+        df = DataFrame(a=[1,2,3], y=["a","b","c"])
+        @test getobs(df, 1) == df[1,:]
+        @test getobs(df, 2:3) == df[2:3,:]
+        @test numobs(df) == 3
+    end
 end

test/runtests.jl

@@ -1,6 +1,7 @@
 using MLUtils
 using MLUtils.Datasets
 using MLUtils: RingBuffer, eachobsparallel
+using MLUtils: flatten, stack, unstack # also exported by DataFrames.jl
 using SparseArrays
 using Random, Statistics
 using Test
@@ -9,6 +10,7 @@ using FoldsThreads: TaskPoolEx
 using ChainRulesTestUtils: test_rrule
 using Zygote: ZygoteRuleConfig
 using ChainRulesCore: rrule_via_ad
+using DataFrames
 
 showcompact(io, x) = show(IOContext(io, :compact => true), x)