Add unified matrix table add interface.

huxuan · huxuan · commit 7c36fdd860ea · 2016-05-23T18:15:49.000+09:00
diff --git a/include/multiverso/table/matrix_table.h b/include/multiverso/table/matrix_table.h
@@ -29,16 +29,16 @@ class MatrixWorkerTable : public WorkerTable {
   void Get(const std::vector<integer_t>& row_ids,
            const std::vector<T*>& data_vec, size_t size);
 
-  // Add whole table
-  void Add(T* data, size_t size, const AddOption* option = nullptr);
-
   void Add(integer_t row_id, T* data, size_t size, 
            const AddOption* option = nullptr);
 
   void Add(const std::vector<integer_t>& row_ids,
            const std::vector<T*>& data_vec, size_t size, 
            const AddOption* option = nullptr);
 
+  void Add(T* data, size_t size, integer_t* row_ids = nullptr,
+           integer_t row_ids_size = 0, const AddOption* option = nullptr);
+
   int Partition(const std::vector<Blob>& kv,
     std::unordered_map<int, std::vector<Blob>>* out) override;
 
diff --git a/src/table/matrix_table.cpp b/src/table/matrix_table.cpp
@@ -89,13 +89,6 @@ void MatrixWorkerTable<T>::Get(const std::vector<integer_t>& row_ids,
   Log::Debug("[Get] worker = %d, #rows_set = %d\n", MV_Rank(), row_ids.size());
 }
 
-template <typename T>
-void MatrixWorkerTable<T>::Add(T* data, size_t size, const AddOption* option) {
-  CHECK(size == num_col_ * num_row_);
-  integer_t whole_table = -1;
-  Add(whole_table, data, size, option);
-}
-
 template <typename T>
 void MatrixWorkerTable<T>::Add(integer_t row_id, T* data, size_t size, 
                                const AddOption* option) {
@@ -122,6 +115,30 @@ void MatrixWorkerTable<T>::Add(const std::vector<integer_t>& row_ids,
   Log::Debug("[Add] worker = %d, #rows_set = %d\n", MV_Rank(), row_ids.size());
 }
 
+template <typename T>
+void MatrixWorkerTable<T>::Add(T* data, size_t size, integer_t* row_ids,
+                               integer_t row_ids_size,
+                               const AddOption* option) {
+  if (row_ids_size == 0) {
+    CHECK(size == num_col_ * num_row_);
+    integer_t row_id = -1;
+    Blob ids_blob(&row_id, sizeof(integer_t));
+    Blob data_blob(data, size * sizeof(T));
+    WorkerTable::Add(ids_blob, data_blob, option);
+    Log::Debug("[Add] worker = %d, #row = %d\n", MV_Rank(), row_id);
+  } else {
+    CHECK(size == num_col_);
+    Blob ids_blob(row_ids, sizeof(integer_t) * row_ids_size);
+    Blob data_blob(row_ids_size * row_size_);
+    //copy each row
+    for (auto i = 0; i < row_ids_size; ++i){
+      memcpy(data_blob.data() + i * row_size_, &data[i * num_col_], row_size_);
+    }
+    WorkerTable::Add(ids_blob, data_blob, option);
+    Log::Debug("[Add] worker = %d, #rows_set = %d\n", MV_Rank(), row_ids_size);
+  }
+}
+
 template <typename T>
 int MatrixWorkerTable<T>::Partition(const std::vector<Blob>& kv,
   std::unordered_map<int, std::vector<Blob>>* out) {
