Add support for table types.

JuliaML · May 27, 2022 · 90d3104 · 90d3104
1 parent c1dc3c2
commit 90d3104
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/Project.toml b/Project.toml
@@ -12,6 +12,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ShowCases = "605ecd9f-84a6-4c9e-81e2-4798472b76a3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [compat]
 ChainRulesCore = "1.0"

diff --git a/src/MLUtils.jl b/src/MLUtils.jl
@@ -13,6 +13,7 @@ import ChainRulesCore: rrule
 using ChainRulesCore: @non_differentiable, unthunk, AbstractZero,
                       NoTangent, ZeroTangent, ProjectTo
 
+using Tables: istable, rows
 
 include("observation.jl")
 export numobs,

diff --git a/src/observation.jl b/src/observation.jl
@@ -16,7 +16,7 @@ See also [`getobs`](@ref)
 function numobs end
 
 # Generic Fallbacks
-numobs(data) = length(data)
+numobs(data) = istable(data) ? length(rows(data)) : length(data)
 
 """
     getobs(data, [idx])
@@ -40,13 +40,13 @@ Every author behind some custom data container can make this
 decision themselves.
 The output should be consistent when `idx` is a scalar vs vector.
 
-See also [`getobs!`](@ref) and [`numobs`](@ref) 
+See also [`getobs!`](@ref) and [`numobs`](@ref)
 """
 function getobs end
 
 # Generic Fallbacks
 getobs(data) = data
-getobs(data, idx) = data[idx]
+getobs(data, idx) = istable(data) ? collect(rows(data))[idx] : data[idx]
 
 """
     getobs!(buffer, data, idx)
@@ -82,7 +82,7 @@ Base.lastindex(x::AbstractDataContainer) = numobs(x)
 # --------------------------------------------------------------------
 # Arrays
 # We are very opinionated with arrays: the observation dimension
-# is th last dimension. For different behavior wrap the array in 
+# is th last dimension. For different behavior wrap the array in
 # a custom type, e.g. with Tables.table.