int MatrixWorkerTable<T>::Partition(const std::vector<Blob>& kv,
std::unordered_map<int, std::vector<Blob>>* out) {
CHECK(kv.size() == 1 || kv.size() == 2 || kv.size() == 3);
CHECK_NOTNULL(out);
size_t keys_size = kv[0].size<integer_t>();
integer_t *keys = reinterpret_cast<integer_t*>(kv[0].data());
if (keys_size == 1 && keys[0] == -1){
for (auto i = 0; i < num_server_; ++i){
int rank = MV_ServerIdToRank(i);
(*out)[rank].push_back(kv[0]);
}
if (kv.size() >= 2){ //process add values
for (integer_t i = 0; i < num_server_; ++i){
int rank = MV_ServerIdToRank(i);
Blob blob(kv[1].data() + server_offsets_[i] * row_size_,
(server_offsets_[i + 1] - server_offsets_[i]) * row_size_);
(*out)[rank].push_back(blob);
if (kv.size() == 3) {// update option blob
(*out)[rank].push_back(kv[2]);
}
}
}
else {
CHECK(get_reply_count_ == 0);
get_reply_count_ = static_cast<int>(out->size());
}
return static_cast<int>(out->size());
}
//count row number in each server
std::vector<int> dest;
std::vector<integer_t> count;
count.resize(num_server_, 0);
integer_t num_row_each = num_row_ / num_server_; // num_server_;
for (auto i = 0; i < keys_size; ++i){
int dst = keys[i] / num_row_each;
dst = (dst >= num_server_ ? num_server_ - 1 : dst);
dest.push_back(dst);
++count[dst];
}
for (auto i = 0; i < num_server_; i++) { // allocate memory for blobs
int rank = MV_ServerIdToRank(i);
if (count[i] != 0) {
std::vector<Blob>& vec = (*out)[rank];
vec.push_back(Blob(count[i] * sizeof(integer_t)));
if (kv.size() >= 2) vec.push_back(Blob(count[i] * row_size_));
}
}
count.clear();
count.resize(num_server_, 0);
integer_t offset = 0;
for (auto i = 0; i < keys_size; ++i) {
int dst = dest[i];
int rank = MV_ServerIdToRank(dst);
(*out)[rank][0].As<integer_t>(count[dst]) = keys[i];
if (kv.size() >= 2){ // copy add values
memcpy(&((*out)[rank][1].As<T>(count[dst] * num_col_)),
kv[1].data() + offset, row_size_);
offset += row_size_;
}
++count[dst];
}
for (int i = 0; i < num_server_; ++i){
int rank = MV_ServerIdToRank(i);
if (count[i] != 0) {
if (kv.size() == 3) {// update option blob
(*out)[rank].push_back(kv[2]);
}
}
}
if (kv.size() == 1){
CHECK(get_reply_count_ == 0);
get_reply_count_ = static_cast<int>(out->size());
}
return static_cast<int>(out->size());
}
int SparseMatrixWorkerTable<T>::Partition(const std::vector<Blob>& kv,
MsgType partition_type,
std::unordered_map<int, std::vector<Blob>>* out) {
int res;
CHECK(kv.size() == 1 || kv.size() == 2 || kv.size() == 3);
CHECK_NOTNULL(out);
if (kv.size() == 2) { // processing Get()
size_t keys_size = kv[0].size<integer_t>();
integer_t* keys = reinterpret_cast<integer_t*>(kv[0].data());
if (keys[0] == -1) {
for (auto i = 0; i < this->num_server_; ++i) {
int rank = MV_ServerIdToRank(i);
(*out)[rank].push_back(kv[0]);
}
for (auto i = 0; i < this->num_server_; ++i){
int rank = MV_ServerIdToRank(i);
if (kv.size() == 2) {// general option blob
(*out)[rank].push_back(kv[1]);
}
}
CHECK(this->get_reply_count_ == 0);
this->get_reply_count_ = static_cast<int>(out->size());
res = static_cast<int>(out->size());
} else {
// count row number in each server
//std::unordered_map<int, integer_t> count;
std::vector<integer_t> count;
std::vector<int> dest;
count.resize(this->num_server_, 0);
integer_t num_row_each = this->num_row_ / this->num_server_; // num_server_;
for (auto i = 0; i < keys_size; ++i) {
int dst = keys[i] / num_row_each;
dst = (dst >= this->num_server_ ? this->num_server_ - 1 : dst);
dest.push_back(dst);
++count[dst];
}
for (auto i = 0; i < this->num_server_; i++) { // allocate memory for blobs
int rank = MV_ServerIdToRank(i);
if (count[i] != 0) {
std::vector<Blob>& vec = (*out)[rank];
vec.push_back(Blob(count[i] * sizeof(integer_t)));
}
}
count.clear();
count.resize(this->num_server_, 0);
for (auto i = 0; i < keys_size; ++i) {
int dst = dest[i];
int rank = MV_ServerIdToRank(dst);
(*out)[rank][0].As<integer_t>(count[dst]) = keys[i];
++count[dst];
}
for (auto i = 0; i < this->num_server_; ++i){
int rank = MV_ServerIdToRank(i);
if (count[i] != 0) {
if (kv.size() == 2) {// add option blob
(*out)[rank].push_back(kv[1]);
}
}
}
CHECK(this->get_reply_count_ == 0);
this->get_reply_count_ = static_cast<int>(out->size());
res = static_cast<int>(out->size());
}
} else { // processing Add()
// call base class's Partition
res = MatrixWorkerTable<T>::Partition(kv, partition_type, out);
}
// only have effect when adding elements
SparseFilter<T, int32_t> filter(0, true);
for (auto& pair : *out) {
std::vector<Blob> compressed_blobs;
filter.FilterIn(pair.second, &compressed_blobs);
pair.second.swap(compressed_blobs);
}
return res;
}