Merge pull request JuliaLang#161 from alyst/pooled_array_permute

Speed up permute!(PooledDataArray)
martinholters · Aug 6, 2015 · b26774c · b26774c
2 parents f783cd1 + cc5e749
commit b26774c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 12 deletions.
diff --git a/src/pooleddataarray.jl b/src/pooleddataarray.jl
@@ -443,6 +443,16 @@ reorder(x::PooledDataArray, y::AbstractVector...) = reorder(mean, x, y...)
 
 Base.reverse(x::PooledDataArray) = PooledDataArray(RefArray(reverse(x.refs)), x.pool)
 
+function Base.permute!!{T<:Integer}(x::PooledDataArray, p::AbstractVector{T})
+    Base.permute!!(x.refs, p)
+    x
+end
+
+function Base.ipermute!!{T<:Integer}(x::PooledDataArray, p::AbstractVector{T})
+    Base.ipermute!!(x.refs, p)
+    x
+end
+
 ##############################################################################
 ##
 ## similar()
@@ -467,7 +477,7 @@ Base.find(pdv::PooledDataVector{Bool}) = find(convert(Vector{Bool}, pdv, false))
 ##
 ##############################################################################
 
-function getpoolidx{T,R<:Union(UInt8, UInt16, Int8, Int16)}(pda::PooledDataArray{T,R}, val::Any)
+function getpoolidx{T,R}(pda::PooledDataArray{T,R}, val::Any)
     val::T = convert(T,val)
     pool_idx = findfirst(pda.pool, val)
     if pool_idx <= 0
@@ -484,17 +494,6 @@ function getpoolidx{T,R<:Union(UInt8, UInt16, Int8, Int16)}(pda::PooledDataArray
     return pool_idx
 end
 
-function getpoolidx{T,R}(pda::PooledDataArray{T,R}, val::Any)
-    val::T = convert(T,val)
-    pool_idx = findfirst(pda.pool, val)
-    if pool_idx <= 0
-        push!(pda.pool, val)
-        pool_idx = length(pda.pool)
-    end
-    return pool_idx
-end
-
-getpoolidx{T,R<:Union(UInt8, UInt16, Int8, Int16)}(pda::PooledDataArray{T,R}, val::NAtype) = zero(R)
 getpoolidx{T,R}(pda::PooledDataArray{T,R}, val::NAtype) = zero(R)
 
 ##############################################################################

diff --git a/test/pooleddataarray.jl b/test/pooleddataarray.jl
@@ -102,4 +102,9 @@ module TestPDA
     da = convert(DataArray{Float32,2}, @pdata(ones(5, 5)))
     @assert isequal(da, @data(ones(Float32, 5, 5)))
     @assert typeof(da) == DataArray{Float32,2}
+
+    # permute
+    pda = @pdata([NA, "A", "B", "C", "A", "B"])
+    @test isequal(Base.permute!!(copy(pda), [2, 5, 3, 6, 4, 1]), @pdata(["A", "A", "B", "B", "C", NA]))
+    @test isequal(Base.ipermute!!(copy(pda), [6, 1, 3, 5, 2, 4]), @pdata(["A", "A", "B", "B", "C", NA]))
 end