special path for categoricalarrays/pooledarrays

Project.toml

@@ -1,16 +1,18 @@
 name = "FixedEffects"
 uuid = "c8885935-8500-56a7-9867-7708b20db0eb"
-version = "2.0.2"
+version = "2.0.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
 
 [compat]
 Requires = "1"
 StatsBase = "0.33"
+DataAPI = "1"
 julia = "1.3"
 
 [extras]

src/FixedEffect.jl

@@ -62,17 +62,23 @@ mutable struct GroupedArray{N} <: AbstractArray{UInt32, N}
 	refs::Array{UInt32, N}   # refs must be between 0 and n. 0 means missing
 	n::Int                   # Number of potential values (= maximum(refs))
 end
-
 Base.size(g::GroupedArray) = size(g.refs)
-@propagate_inbounds Base.getindex(g::GroupedArray, i::Int) = getindex(g.refs, i)
-@propagate_inbounds Base.getindex(g::GroupedArray, I) = getindex(g.refs, I)
+@propagate_inbounds Base.getindex(g::GroupedArray, i::Number) = getindex(g.refs, i::Number)
+@propagate_inbounds Base.getindex(g::GroupedArray, I...) = getindex(g.refs, I...)
+Base.firstindex(g::GroupedArray) = firstindex(g.refs)
+Base.lastindex(g::GroupedArray) = lastindex(g.refs)
 
 group(xs::GroupedArray) = xs
 
 function group(xs::AbstractArray)
+	_group(DataAPI.refarray(xs), DataAPI.refpool(xs))
+end
+
+
+function _group(xs, ::Nothing)
 	refs = Array{UInt32}(undef, size(xs))
 	invpool = Dict{eltype(xs), UInt32}()
-	n = 0
+	n = UInt32(0)
 	z = UInt32(0)
 	@inbounds for i in eachindex(xs)
 		x = xs[i]
@@ -92,6 +98,26 @@ function group(xs::AbstractArray)
 	return GroupedArray{ndims(xs)}(refs, n)
 end
 
+function _group(ra, rp)
+	refs = Array{UInt32}(undef, size(ra))
+	hashes = Array{UInt32}(undef, length(rp))
+	firp = firstindex(rp)
+	n = 0
+	for i in eachindex(hashes)
+		if rp[i+firp-1] === missing
+			hashes[i] = UInt32(0)
+		else
+			n += 1
+			hashes[i] = n
+		end
+	end
+	fira = firstindex(ra)
+	@inbounds for i in eachindex(refs)
+		refs[i] = hashes[ra[i+fira-1]-firp+1]
+	end
+	return GroupedArray{ndims(refs)}(refs, n)
+end
+
 function group(args...)
 	g1 = deepcopy(group(args[1]))
 	for j = 2:length(args)

src/FixedEffects.jl

@@ -9,6 +9,7 @@ using Base: @propagate_inbounds
 using LinearAlgebra
 using StatsBase
 using Requires
+using DataAPI
 using Printf
 
 ##############################################################################

test/types.jl

@@ -2,6 +2,7 @@ using Test
 using FixedEffects
 using FixedEffects: GroupedArray, group, factorize!
 using StatsBase
+using PooledArrays, CategoricalArrays
 import Base: ==
 
 ==(x::FixedEffect{R,I}, y::FixedEffect{R,I}) where {R,I} =
@@ -86,4 +87,18 @@ end
     a = rand(N)
     g = group(a)
     @test factorize!(g) == g
+
+
+
+    g = [0, 1, 2, 3, 1, 2, 0]
+    @test group(g) == group(categorical(g))
+    @test group(g) == group(PooledArray(g))
+
+    g = [missing, 1, 2, 3, 1, 2, missing]
+    @test group(g) == group(categorical(g))
+    @test group(g) == group(PooledArray(g))
+
+    g = [missing, "a", "b", "c", "a", "a", "a"]
+    @test group(g) == group(categorical(g))
+    @test group(g) == group(PooledArray(g))
 end