Merge pull request #80 from alan-turing-institute/scitpye-for-ca-views

Make scitpye for categorical array views efficient
JuliaAI · Jun 1, 2021 · f62033d · f62033d
2 parents eb039c4 + 0b0cf79
commit f62033d
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "MLJScientificTypes"
 uuid = "2e2323e0-db8b-457b-ae0d-bdfb3bc63afd"
 authors = ["Anthony D. Blaom <[email protected]>"]
-version = "0.4.7"
+version = "0.4.8"
 
 [deps]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"

diff --git a/src/convention/scitype.jl b/src/convention/scitype.jl
@@ -21,7 +21,10 @@ function ST.scitype(c::Cat, ::MLJ)
     return ifelse(c.pool.ordered, OrderedFactor{nc}, Multiclass{nc})
 end
 
-function ST.scitype(A::CArr{T,N}, ::MLJ) where {T,N}
+const CatArrOrSub{T,N} =
+    Union{CategoricalArray{T,N},SubArray{<:Any,<:Any,<:CategoricalArray{T,N}}}
+
+function ST.scitype(A::CatArrOrSub{T,N}, ::MLJ) where {T,N}
     nlevels = length(levels(A))
     S = ifelse(isordered(A), OrderedFactor{nlevels}, Multiclass{nlevels})
     T >: Missing && (S = Union{S,Missing})

diff --git a/test/basic_tests.jl b/test/basic_tests.jl
@@ -50,6 +50,7 @@ end
 
     # MULTICLASS
     @test scitype(categorical(1:4))         == Vec{Multiclass{4}}
+    @test scitype(view(categorical(1:4), 1:3)) == Vec{Multiclass{4}}
     @test scitype(Any[categorical(1:4)...]) == Vec{Multiclass{4}}
     @test scitype(categorical([1, missing, 3])) ==
         Vec{Union{Multiclass{2},Missing}}