<commit_before>#include <stdlib.h>
#include <stdio.h>
#include <string>
#include <unistd.h>
#include <algorithm>
#include <sstream>
#include <iostream>
#include <gflags/gflags.h>
#include <set>
#include <boost/algorithm/string.hpp>
#include <boost/bind.hpp>
#include <boost/scoped_ptr.hpp>
#include "sort_file.h"
#include "logging.h"
#include "common/filesystem.h"
#include "common/tools_util.h"
#include "thread_pool.h"
#include "mutex.h"

DEFINE_int32(total, 0, "total numbers of map tasks");
DEFINE_int32(reduce_no, 0, "the reduce number of this reduce task");
DEFINE_string(work_dir, "/tmp", "the shuffle work dir");
DEFINE_int32(batch, 300, "merge how many maps output at the same time");
DEFINE_int32(attempt_id, 0, "the attempt_id of this reduce task");
DEFINE_string(dfs_host, "", "host name of dfs master");
DEFINE_string(dfs_port, "", "port of dfs master");
DEFINE_string(dfs_user, "", "user name of dfs master");
DEFINE_string(dfs_password, "", "password of dfs master");
DEFINE_string(pipe, "streaming", "pipe style: streaming/bistreaming");
DEFINE_bool(skip_merge, false, "whether skip merge phase");

using baidu::common::Log;
using baidu::common::FATAL;
using baidu::common::INFO;
using baidu::common::WARNING;
using namespace baidu;
using namespace baidu::shuttle;

std::set<std::string> g_merged;
int32_t g_file_no(0);
FileSystem* g_fs(NULL);
Mutex g_mu;

void FillParam(FileSystem::Param& param) {
    if (!FLAGS_dfs_user.empty()) {
        param["user"] = FLAGS_dfs_user;
    }
    if (!FLAGS_dfs_password.empty()) {
        param["password"] = <PASSWORD>;
    }
    if (!FLAGS_dfs_host.empty()) {
        param["host"] = FLAGS_dfs_host;
    }
    if (!FLAGS_dfs_port.empty()) {
        param["port"] = FLAGS_dfs_port;
    }
}

void CollectFilesToMerge(std::vector<std::string>* maps_to_merge) {
    assert(maps_to_merge);
    std::vector<std::string> candidates;
    for (int i = 0; i < FLAGS_total; i++) {
        std::stringstream ss;
        ss << FLAGS_work_dir << "/map_" << i;
        std::string map_dir = ss.str();
        if (g_merged.find(map_dir) != g_merged.end()) {
            continue;
        }
        candidates.push_back(map_dir);
    }
    std::vector<std::string>::const_iterator it;
    for (it = candidates.begin(); it != candidates.end(); it++) {
        const std::string& map_dir = *it;
        if (g_fs->Exist(map_dir)) {
            LOG(INFO, "maps_to_merge: %s", map_dir.c_str());
            maps_to_merge->push_back(map_dir);
        }
        if (maps_to_merge->size() >= (size_t)FLAGS_batch) {
            break;
        }
    }
}

void AddSortFiles(const std::string map_dir, std::vector<std::string>* file_names) {
    assert(file_names);
    std::vector<FileInfo> sort_files;
    if (g_fs->List(map_dir, &sort_files)) {
        std::vector<FileInfo>::iterator jt;
        for (jt = sort_files.begin(); jt != sort_files.end(); jt++) {
            const std::string& file_name = jt->name;
            if (boost::ends_with(file_name, ".sort")) {
                MutexLock lock(&g_mu);
                file_names->push_back(file_name);
            }
        }
    } else {
        LOG(FATAL, "fail to list %s", map_dir.c_str());
    }
}

void MergeMapOutput(const std::vector<std::string>& maps_to_merge) {
    std::vector<std::string> * file_names = new std::vector<std::string>();
    boost::scoped_ptr<std::vector<std::string> > file_names_guard(file_names);
    std::vector<std::string>::const_iterator it;
    std::vector<std::string> real_merged_maps;
    ThreadPool pool;
    int map_ct = 0;
    for (it = maps_to_merge.begin(); it != maps_to_merge.end(); it++) {<commit_msg>improve hdfs shuffle performance<commit_after>#include <stdlib.h>
#include <stdio.h>
#include <string>
#include <unistd.h>
#include <algorithm>
#include <sstream>
#include <iostream>
#include <gflags/gflags.h>
#include <set>
#include <boost/algorithm/string.hpp>
#include <boost/bind.hpp>
#include <boost/scoped_ptr.hpp>
#include "sort_file.h"
#include "logging.h"
#include "common/filesystem.h"
#include "common/tools_util.h"
#include "thread_pool.h"
#include "mutex.h"

DEFINE_int32(total, 0, "total numbers of map tasks");
DEFINE_int32(reduce_no, 0, "the reduce number of this reduce task");
DEFINE_string(work_dir, "/tmp", "the shuffle work dir");
DEFINE_int32(batch, 300, "merge how many maps output at the same time");
DEFINE_int32(attempt_id, 0, "the attempt_id of this reduce task");
DEFINE_string(dfs_host, "", "host name of dfs master");
DEFINE_string(dfs_port, "", "port of dfs master");
DEFINE_string(dfs_user, "", "user name of dfs master");
DEFINE_string(dfs_password, "", "password of dfs master");
DEFINE_string(pipe, "streaming", "pipe style: streaming/bistreaming");
DEFINE_bool(skip_merge, false, "whether skip merge phase");

using baidu::common::Log;
using baidu::common::FATAL;
using baidu::common::INFO;
using baidu::common::WARNING;
using namespace baidu;
using namespace baidu::shuttle;

std::set<std::string> g_merged;
int32_t g_file_no(0);
FileSystem* g_fs(NULL);
Mutex g_mu;

const static int32_t sMaxSortFiles = 100;

void FillParam(FileSystem::Param& param) {
    if (!FLAGS_dfs_user.empty()) {
        param["user"] = FLAGS_dfs_user;
    }
    if (!FLAGS_dfs_password.empty()) {
        param["password"] = <PASSWORD>;
    }
    if (!FLAGS_dfs_host.empty()) {
        param["host"] = FLAGS_dfs_host;
    }
    if (!FLAGS_dfs_port.empty()) {
        param["port"] = FLAGS_dfs_port;
    }
}

void CollectFilesToMerge(std::vector<std::string>* maps_to_merge) {
    assert(maps_to_merge);
    std::vector<std::string> candidates;
    for (int i = 0; i < FLAGS_total; i++) {
        std::stringstream ss;
        ss << FLAGS_work_dir << "/map_" << i;
        std::string map_dir = ss.str();
        if (g_merged.find(map_dir) != g_merged.end()) {
            continue;
        }
        candidates.push_back(map_dir);
    }
    std::vector<std::string>::const_iterator it;
    for (it = candidates.begin(); it != candidates.end(); it++) {
        const std::string& map_dir = *it;
        if (g_fs->Exist(map_dir)) {
            LOG(INFO, "maps_to_merge: %s", map_dir.c_str());
            maps_to_merge->push_back(map_dir);
        }
        if (maps_to_merge->size() >= (size_t)FLAGS_batch) {
            break;
        }
    }
}

void AddSortFiles(const std::string map_dir, std::vector<std::string>* file_names) {
    assert(file_names);
    std::vector<FileInfo> sort_files;
    for (int i = 0; i < sMaxSortFiles; i++) {
        std::stringstream ss;
        ss << map_dir << "/" << i << ".sort";
        const std::string f_name = ss.str();
        if (g_fs->Exist(f_name)) {
            MutexLock lock(&g_mu);
            file_names->push_back(f_name);
        } else {
            break;
        }
    }
}

void MergeMapOutput(const std::vector<std::string>& maps_to_merge) {
    std::vector<std::string> * file_names = new std::vector<std::string>();
    boost::scoped_ptr<std::vector<std::string> > file_names_guard(file_names);
    std::vector<std::string>::const_iterator it;
    std::vector<std::string> real_merged_maps;
    ThreadPool pool;
    int map_ct = 0;
    for (it = maps_to_merge.begin(); it != maps_to_merge.end(); it++) {
{'fraction_non_alphanumeric': 0.1592895795754128, 'fraction_numerical': 0.006105175523796309, 'mean_word_length': 3.824631860776439, 'pattern_counts': {'":': 0, '<': 79, '<?xml version=': 0, '>': 69, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7116501', 'n_tokens_mistral': 2613, 'n_tokens_neox': 2525, 'n_words': 530}