| 10 | using namespace std; |
| 11 | |
| 12 | int main( int argc, char** argv) |
| 13 | { |
| 14 | if(argc < 4) { |
| 15 | cout << "Usage .... exec input_file output_file_basename num_partitions" << endl; |
| 16 | exit(-1); |
| 17 | } |
| 18 | |
| 19 | std::string input_file = argv[1]; |
| 20 | std::string output_file = argv[2]; |
| 21 | int num_partition = atoi(argv[3]); |
| 22 | |
| 23 | std::ifstream infile(input_file); |
| 24 | std::vector<std::pair<std::string, int> > lines; |
| 25 | std::string line; |
| 26 | size_t pos; |
| 27 | int label; |
| 28 | while (std::getline(infile, line)) { |
| 29 | pos = line.find_last_of(' '); |
| 30 | label = atoi(line.substr(pos + 1).c_str()); |
| 31 | lines.push_back(std::make_pair(line.substr(0, pos), label)); |
| 32 | } |
| 33 | |
| 34 | std::cout << "Input file " << input_file << std::endl; |
| 35 | std::cout << "Shuffling data" << std::endl; |
| 36 | unsigned seed = std::chrono::system_clock::now().time_since_epoch().count(); |
| 37 | std::shuffle(lines.begin(), lines.end(),std::default_random_engine(seed)); |
| 38 | |
| 39 | int num = lines.size(); |
| 40 | int batch = int(num/num_partition); |
| 41 | |
| 42 | std::cout << "A total of " << num << " rows"; |
| 43 | std::cout << " Num of partition " << num_partition; |
| 44 | std::cout << " Data per partition " << batch << std::endl; |
| 45 | |
| 46 | std::cout << "Original file shuffled and save as: " << output_file << std::endl; |
| 47 | ofstream base_ofs(output_file.c_str()); |
| 48 | if (!base_ofs) { std::cout << "\n In write_file can't open file : " << output_file; exit(1); } |
| 49 | for(auto& l0 : lines) base_ofs << l0.first << " " << l0.second << std::endl; |
| 50 | |
| 51 | std::cout << " Start partitioning " << std::endl; |
| 52 | for(int p = 0; p < num_partition; p++) { |
| 53 | int start = p * batch; |
| 54 | int end = start + batch; |
| 55 | std::string partition_file = output_file + ".p" + to_string(p); |
| 56 | std::cout << "Partitioned file name " << partition_file << std::endl; |
| 57 | ofstream ofs(partition_file.c_str()); |
| 58 | if (!ofs) { std::cout << "\n In write_file: can't open file : " << partition_file; exit(1); } |
| 59 | for(int l = start; l < end; l++) |
| 60 | ofs << lines[l].first << " " << lines[l].second << endl; |
| 61 | } |
| 62 | |
| 63 | std::cout << "DONE!" << std::endl; |
| 64 | return 0; |
| 65 | |
| 66 | } |