Rename PartionedSerializer to DestinationSerializer

Make the destination aware serializer more generic
Fix bug w.r.t. pids having the correct shape

Author: amitmurthy

src/core.jl

@@ -445,8 +445,10 @@ function distribute(A::AbstractArray;
     dist = defaultdist(size(A), procs))
     idxs, _ = chunk_idxs([size(A)...], dist)
 
-    pas = PartitionedSerializer(A, procs, idxs)
-    return DArray(I->verify_and_get(pas, I), size(A), procs, dist)
+    s = verified_destination_serializer(reshape(procs, size(idxs)), size(idxs)) do pididx
+        A[idxs[pididx]...]
+    end
+    return DArray(I->localpart(s), size(A), procs, dist)
 end
 
 """
@@ -458,8 +460,10 @@ Distribute a local array `A` like the distributed array `DA`.
 function distribute(A::AbstractArray, DA::DArray)
     size(DA) == size(A) || throw(DimensionMismatch("Distributed array has size $(size(DA)) but array has $(size(A))"))
 
-    pas = PartitionedSerializer(A, procs(DA), DA.indexes)
-    return DArray(I->verify_and_get(pas, I), DA)
+    s = verified_destination_serializer(procs(DA), size(DA.indexes)) do pididx
+        A[DA.indexes[pididx]...]
+    end
+    return DArray(I->localpart(s), DA)
 end
 
 Base.convert{T,N,S<:AbstractArray}(::Type{DArray{T,N,S}}, A::S) = distribute(convert(AbstractArray{T,N}, A))
@@ -589,7 +593,7 @@ indexin_mask(a, b) = [i in b for i in a]
 import Base: tail
 # Given a tuple of indices and a tuple of masks, restrict the indices to the
 # valid regions. This is, effectively, reversing Base.setindex_shape_check.
-# We can't just use indexing into MergedIndices here because getindex is much 
+# We can't just use indexing into MergedIndices here because getindex is much
 # pickier about singleton dimensions than setindex! is.
 restrict_indices(::Tuple{}, ::Tuple{}) = ()
 function restrict_indices(a::Tuple{Any, Vararg{Any}}, b::Tuple{Any, Vararg{Any}})
@@ -639,7 +643,7 @@ end
 Base.size(M::MergedIndices) = M.sz
 Base.@propagate_inbounds Base.getindex{_,N}(M::MergedIndices{_,N}, I::Vararg{Int, N}) =
     CartesianIndex(map(propagate_getindex, M.indices, I))
-# Additionally, we optimize bounds checking when using MergedIndices as an 
+# Additionally, we optimize bounds checking when using MergedIndices as an
 # array index since checking, e.g., A[1:500, 1:500] is *way* faster than
 # checking an array of 500^2 elements of CartesianIndex{2}. This optimization
 # also applies to reshapes of MergedIndices since the outer shape of the

src/mapreduce.jl

@@ -162,16 +162,20 @@ map_localparts(f::Callable, d1::DArray, d2::DArray) = DArray(d1) do I
 end
 
 function map_localparts(f::Callable, DA::DArray, A::Array)
-    pas = PartitionedSerializer(A, procs(DA), DA.indexes)
+    s = verified_destination_serializer(procs(DA), size(DA.indexes)) do pididx
+        A[DA.indexes[pididx]...]
+    end
     DArray(DA) do I
-        f(localpart(DA), verify_and_get(pas, I))
+        f(localpart(DA), localpart(s))
     end
 end
 
 function map_localparts(f::Callable, A::Array, DA::DArray)
-    pas = PartitionedSerializer(A, procs(DA), DA.indexes)
+    s = verified_destination_serializer(procs(DA), size(DA.indexes)) do pididx
+        A[DA.indexes[pididx]...]
+    end
     DArray(DA) do I
-        f(verify_and_get(pas, I), localpart(DA))
+        f(localpart(s), localpart(DA))
     end
 end
 

src/serialize.jl

@@ -40,47 +40,46 @@ function Base.deserialize{T<:DArray}(S::AbstractSerializer, t::Type{T})
 end
 
 # Serialize only those parts of the object as required by the destination worker.
-type PartitionedSerializer
-   indexable_obj                      # An indexable object, Array, SparseMatrix, etc.
-                                      # Complete object on the serializing side.
-                                      # Part object on the deserialized side.
-   pids::Nullable{Array}
-   idxs::Nullable{Array}
-   local_idxs::Nullable{Tuple}
+type DestinationSerializer
+    generate::Nullable{Function}     # Function to generate the part to be serialized
+    pids::Nullable{Array}            # MUST have the same shape as the distribution
 
-   PartitionedSerializer(obj, local_idxs::Tuple) = new(obj, Nullable{Array}(), Nullable{Array}(), local_idxs)
-   function PartitionedSerializer(obj, pids::Array, idxs::Array)
-        pas = new(obj,pids,idxs,Nullable{Tuple}())
+    deser_obj::Nullable{Any}         # Deserialized part
 
-        if myid() in pids
-            pas.local_idxs = idxs[findfirst(pids, myid())]
-        end
-        return pas
-    end
+    DestinationSerializer(f,p,d) = new(f,p,d)
 end
 
-function Base.serialize(S::AbstractSerializer, pas::PartitionedSerializer)
+DestinationSerializer(f::Function, pids::Array) = DestinationSerializer(f, pids, Nullable{Any}())
+
+# contructs a DestinationSerializer after verifying that the shape of pids.
+function verified_destination_serializer(f::Function, pids::Array, verify_size)
+    @assert size(pids) == verify_size
+    return DestinationSerializer(f, pids)
+end
+
+DestinationSerializer(deser_obj::Any) = DestinationSerializer(Nullable{Function}(), Nullable{Array}(), deser_obj)
+
+function Base.serialize(S::AbstractSerializer, s::DestinationSerializer)
     pid = Base.worker_id_from_socket(S.io)
-    I = get(pas.idxs)[findfirst(get(pas.pids), pid)]
-    Serializer.serialize_type(S, typeof(pas))
-    serialize(S, pas.indexable_obj[I...])
-    serialize(S, I)
+    pididx = findfirst(get(s.pids), pid)
+    Serializer.serialize_type(S, typeof(s))
+    serialize(S, get(s.generate)(pididx))
 end
 
-function Base.deserialize{T<:PartitionedSerializer}(S::AbstractSerializer, t::Type{T})
-    obj_part = deserialize(S)
-    I = deserialize(S)
-    return PartitionedSerializer(obj_part, I)
+function Base.deserialize{T<:DestinationSerializer}(S::AbstractSerializer, t::Type{T})
+    lpart = deserialize(S)
+    return DestinationSerializer(lpart)
 end
 
-function verify_and_get(pas::PartitionedSerializer, I)
-    # Handle the special case where myid() is part of pas.pids.
-    # For this case serialize/deserialize is not called as the remotecall is executed locally
-    if myid() in get(pas.pids, [])
-        @assert I == get(pas.idxs)[findfirst(get(pas.pids),myid())]
-        return pas.indexable_obj[I...]
+
+function localpart(s::DestinationSerializer)
+    if !isnull(s.deser_obj)
+        return get(s.deser_obj)
+    elseif  !isnull(s.generate) && (myid() in get(s.pids))
+        # Handle the special case where myid() is part of s.pids.
+        # In this case serialize/deserialize is not called as the remotecall is executed locally
+        return get(s.generate)(findfirst(get(s.pids), myid()))
     else
-        @assert I == get(pas.local_idxs, ())
-        return pas.indexable_obj
+        throw(ErrorException(string("Invalid state in DestinationSerializer.")))
     end
 end