dataset_loader.cpp 46.2 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
#include <LightGBM/dataset_loader.h>

3
#include <LightGBM/utils/openmp_wrapper.h>
Guolin Ke's avatar
Guolin Ke committed
4
#include <LightGBM/utils/log.h>
Guolin Ke's avatar
Guolin Ke committed
5
6
#include <LightGBM/utils/array_args.h>

Guolin Ke's avatar
Guolin Ke committed
7
8
9
10
11
#include <LightGBM/network.h>


namespace LightGBM {

Guolin Ke's avatar
Guolin Ke committed
12
13
DatasetLoader::DatasetLoader(const Config& io_config, const PredictFunction& predict_fun, int num_class, const char* filename)
  :config_(io_config), random_(config_.data_random_seed), predict_fun_(predict_fun), num_class_(num_class) {
Guolin Ke's avatar
Guolin Ke committed
14
15
16
17
  label_idx_ = 0;
  weight_idx_ = NO_SPECIFIC;
  group_idx_ = NO_SPECIFIC;
  SetHeader(filename);
Guolin Ke's avatar
Guolin Ke committed
18
19
20
21
22
23
}

DatasetLoader::~DatasetLoader() {

}

Guolin Ke's avatar
Guolin Ke committed
24
void DatasetLoader::SetHeader(const char* filename) {
Guolin Ke's avatar
Guolin Ke committed
25
  std::unordered_map<std::string, int> name2idx;
Guolin Ke's avatar
Guolin Ke committed
26
27
  std::string name_prefix("name:");
  if (filename != nullptr) {
Guolin Ke's avatar
Guolin Ke committed
28
    TextReader<data_size_t> text_reader(filename, config_.header);
Guolin Ke's avatar
Guolin Ke committed
29

Guolin Ke's avatar
Guolin Ke committed
30
    // get column names
Guolin Ke's avatar
Guolin Ke committed
31
    if (config_.header) {
Guolin Ke's avatar
Guolin Ke committed
32
      std::string first_line = text_reader.first_line();
33
      feature_names_ = Common::Split(first_line.c_str(), "\t,");
Guolin Ke's avatar
Guolin Ke committed
34
35
    }

Guolin Ke's avatar
Guolin Ke committed
36
    // load label idx first
Guolin Ke's avatar
Guolin Ke committed
37
38
39
    if (config_.label_column.size() > 0) {
      if (Common::StartsWith(config_.label_column, name_prefix)) {
        std::string name = config_.label_column.substr(name_prefix.size());
Guolin Ke's avatar
Guolin Ke committed
40
41
42
43
44
45
46
47
48
49
        label_idx_ = -1;
        for (int i = 0; i < static_cast<int>(feature_names_.size()); ++i) {
          if (name == feature_names_[i]) {
            label_idx_ = i;
            break;
          }
        }
        if (label_idx_ >= 0) {
          Log::Info("Using column %s as label", name.c_str());
        } else {
50
51
          Log::Fatal("Could not find label column %s in data file \n"
                     "or data file doesn't contain header", name.c_str());
Guolin Ke's avatar
Guolin Ke committed
52
        }
Guolin Ke's avatar
Guolin Ke committed
53
      } else {
Guolin Ke's avatar
Guolin Ke committed
54
        if (!Common::AtoiAndCheck(config_.label_column.c_str(), &label_idx_)) {
55
56
57
          Log::Fatal("label_column is not a number,\n"
                     "if you want to use a column name,\n"
                     "please add the prefix \"name:\" to the column name");
Guolin Ke's avatar
Guolin Ke committed
58
59
        }
        Log::Info("Using column number %d as label", label_idx_);
Guolin Ke's avatar
Guolin Ke committed
60
61
      }
    }
Guolin Ke's avatar
Guolin Ke committed
62

Guolin Ke's avatar
Guolin Ke committed
63
    if (!feature_names_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
64
65
66
67
      // erase label column name
      feature_names_.erase(feature_names_.begin() + label_idx_);
      for (size_t i = 0; i < feature_names_.size(); ++i) {
        name2idx[feature_names_[i]] = static_cast<int>(i);
Guolin Ke's avatar
Guolin Ke committed
68
      }
Guolin Ke's avatar
Guolin Ke committed
69
70
71
    }

    // load ignore columns
Guolin Ke's avatar
Guolin Ke committed
72
73
74
    if (config_.ignore_column.size() > 0) {
      if (Common::StartsWith(config_.ignore_column, name_prefix)) {
        std::string names = config_.ignore_column.substr(name_prefix.size());
Guolin Ke's avatar
Guolin Ke committed
75
76
77
78
79
80
81
82
83
        for (auto name : Common::Split(names.c_str(), ',')) {
          if (name2idx.count(name) > 0) {
            int tmp = name2idx[name];
            ignore_features_.emplace(tmp);
          } else {
            Log::Fatal("Could not find ignore column %s in data file", name.c_str());
          }
        }
      } else {
Guolin Ke's avatar
Guolin Ke committed
84
        for (auto token : Common::Split(config_.ignore_column.c_str(), ',')) {
Guolin Ke's avatar
Guolin Ke committed
85
86
          int tmp = 0;
          if (!Common::AtoiAndCheck(token.c_str(), &tmp)) {
87
88
89
            Log::Fatal("ignore_column is not a number,\n"
                       "if you want to use a column name,\n"
                       "please add the prefix \"name:\" to the column name");
Guolin Ke's avatar
Guolin Ke committed
90
91
          }
          ignore_features_.emplace(tmp);
Guolin Ke's avatar
Guolin Ke committed
92
93
94
        }
      }
    }
Guolin Ke's avatar
Guolin Ke committed
95
    // load weight idx
Guolin Ke's avatar
Guolin Ke committed
96
97
98
    if (config_.weight_column.size() > 0) {
      if (Common::StartsWith(config_.weight_column, name_prefix)) {
        std::string name = config_.weight_column.substr(name_prefix.size());
Guolin Ke's avatar
Guolin Ke committed
99
100
101
102
103
104
        if (name2idx.count(name) > 0) {
          weight_idx_ = name2idx[name];
          Log::Info("Using column %s as weight", name.c_str());
        } else {
          Log::Fatal("Could not find weight column %s in data file", name.c_str());
        }
Guolin Ke's avatar
Guolin Ke committed
105
      } else {
Guolin Ke's avatar
Guolin Ke committed
106
        if (!Common::AtoiAndCheck(config_.weight_column.c_str(), &weight_idx_)) {
107
108
109
          Log::Fatal("weight_column is not a number,\n"
                     "if you want to use a column name,\n"
                     "please add the prefix \"name:\" to the column name");
Guolin Ke's avatar
Guolin Ke committed
110
111
        }
        Log::Info("Using column number %d as weight", weight_idx_);
Guolin Ke's avatar
Guolin Ke committed
112
      }
Guolin Ke's avatar
Guolin Ke committed
113
      ignore_features_.emplace(weight_idx_);
Guolin Ke's avatar
Guolin Ke committed
114
    }
Guolin Ke's avatar
Guolin Ke committed
115
    // load group idx
Guolin Ke's avatar
Guolin Ke committed
116
117
118
    if (config_.group_column.size() > 0) {
      if (Common::StartsWith(config_.group_column, name_prefix)) {
        std::string name = config_.group_column.substr(name_prefix.size());
Guolin Ke's avatar
Guolin Ke committed
119
120
121
122
123
124
125
        if (name2idx.count(name) > 0) {
          group_idx_ = name2idx[name];
          Log::Info("Using column %s as group/query id", name.c_str());
        } else {
          Log::Fatal("Could not find group/query column %s in data file", name.c_str());
        }
      } else {
Guolin Ke's avatar
Guolin Ke committed
126
        if (!Common::AtoiAndCheck(config_.group_column.c_str(), &group_idx_)) {
127
128
129
          Log::Fatal("group_column is not a number,\n"
                     "if you want to use a column name,\n"
                     "please add the prefix \"name:\" to the column name");
Guolin Ke's avatar
Guolin Ke committed
130
131
132
133
        }
        Log::Info("Using column number %d as group/query id", group_idx_);
      }
      ignore_features_.emplace(group_idx_);
Guolin Ke's avatar
Guolin Ke committed
134
135
    }
  }
Guolin Ke's avatar
Guolin Ke committed
136
137
138
  if (config_.categorical_feature.size() > 0) {
    if (Common::StartsWith(config_.categorical_feature, name_prefix)) {
      std::string names = config_.categorical_feature.substr(name_prefix.size());
139
140
141
142
143
      for (auto name : Common::Split(names.c_str(), ',')) {
        if (name2idx.count(name) > 0) {
          int tmp = name2idx[name];
          categorical_features_.emplace(tmp);
        } else {
Guolin Ke's avatar
Guolin Ke committed
144
          Log::Fatal("Could not find categorical_feature %s in data file", name.c_str());
145
146
147
        }
      }
    } else {
Guolin Ke's avatar
Guolin Ke committed
148
      for (auto token : Common::Split(config_.categorical_feature.c_str(), ',')) {
149
150
        int tmp = 0;
        if (!Common::AtoiAndCheck(token.c_str(), &tmp)) {
Guolin Ke's avatar
Guolin Ke committed
151
          Log::Fatal("categorical_feature is not a number,\n"
152
153
                     "if you want to use a column name,\n"
                     "please add the prefix \"name:\" to the column name");
154
155
156
157
158
        }
        categorical_features_.emplace(tmp);
      }
    }
  }
Guolin Ke's avatar
Guolin Ke committed
159
160
}

161
Dataset* DatasetLoader::LoadFromFile(const char* filename, const char* initscore_file, int rank, int num_machines) {
Guolin Ke's avatar
Guolin Ke committed
162
  // don't support query id in data file when training in parallel
Guolin Ke's avatar
Guolin Ke committed
163
  if (num_machines > 1 && !config_.pre_partition) {
Guolin Ke's avatar
Guolin Ke committed
164
    if (group_idx_ > 0) {
165
166
      Log::Fatal("Using a query id without pre-partitioning the data file is not supported for parallel training.\n"
                 "Please use an additional query file or pre-partition the data");
Guolin Ke's avatar
Guolin Ke committed
167
168
    }
  }
Guolin Ke's avatar
Guolin Ke committed
169
  auto dataset = std::unique_ptr<Dataset>(new Dataset());
Guolin Ke's avatar
Guolin Ke committed
170
171
  data_size_t num_global_data = 0;
  std::vector<data_size_t> used_data_indices;
172
173
  auto bin_filename = CheckCanLoadFromBin(filename);
  if (bin_filename.size() == 0) {
Guolin Ke's avatar
Guolin Ke committed
174
    auto parser = std::unique_ptr<Parser>(Parser::CreateParser(filename, config_.header, 0, label_idx_));
Guolin Ke's avatar
Guolin Ke committed
175
176
177
178
    if (parser == nullptr) {
      Log::Fatal("Could not recognize data format of %s", filename);
    }
    dataset->data_filename_ = filename;
Guolin Ke's avatar
Guolin Ke committed
179
    dataset->label_idx_ = label_idx_;
180
    dataset->metadata_.Init(filename, initscore_file);
Guolin Ke's avatar
Guolin Ke committed
181
    if (!config_.two_round) {
Guolin Ke's avatar
Guolin Ke committed
182
      // read data to memory
183
      auto text_data = LoadTextDataToMemory(filename, dataset->metadata_, rank, num_machines, &num_global_data, &used_data_indices);
Guolin Ke's avatar
Guolin Ke committed
184
185
186
187
      dataset->num_data_ = static_cast<data_size_t>(text_data.size());
      // sample data
      auto sample_data = SampleTextDataFromMemory(text_data);
      // construct feature bin mappers
Guolin Ke's avatar
Guolin Ke committed
188
      ConstructBinMappersFromTextData(rank, num_machines, sample_data, parser.get(), dataset.get());
Guolin Ke's avatar
Guolin Ke committed
189
      // initialize label
190
      dataset->metadata_.Init(dataset->num_data_, weight_idx_, group_idx_);
Guolin Ke's avatar
Guolin Ke committed
191
      // extract features
Guolin Ke's avatar
Guolin Ke committed
192
      ExtractFeaturesFromMemory(text_data, parser.get(), dataset.get());
Guolin Ke's avatar
Guolin Ke committed
193
194
195
196
197
198
199
200
201
202
      text_data.clear();
    } else {
      // sample data from file
      auto sample_data = SampleTextDataFromFile(filename, dataset->metadata_, rank, num_machines, &num_global_data, &used_data_indices);
      if (used_data_indices.size() > 0) {
        dataset->num_data_ = static_cast<data_size_t>(used_data_indices.size());
      } else {
        dataset->num_data_ = num_global_data;
      }
      // construct feature bin mappers
Guolin Ke's avatar
Guolin Ke committed
203
      ConstructBinMappersFromTextData(rank, num_machines, sample_data, parser.get(), dataset.get());
Guolin Ke's avatar
Guolin Ke committed
204
      // initialize label
205
      dataset->metadata_.Init(dataset->num_data_, weight_idx_, group_idx_);
Guolin Ke's avatar
Guolin Ke committed
206
207

      // extract features
Guolin Ke's avatar
Guolin Ke committed
208
      ExtractFeaturesFromFile(filename, parser.get(), used_data_indices, dataset.get());
Guolin Ke's avatar
Guolin Ke committed
209
210
211
    }
  } else {
    // load data from binary file
212
    dataset.reset(LoadFromBinFile(filename, bin_filename.c_str(), rank, num_machines, &num_global_data, &used_data_indices));
Guolin Ke's avatar
Guolin Ke committed
213
214
215
216
  }
  // check meta data
  dataset->metadata_.CheckOrPartition(num_global_data, used_data_indices);
  // need to check training data
Guolin Ke's avatar
Guolin Ke committed
217
218
  CheckDataset(dataset.get());
  return dataset.release();
Guolin Ke's avatar
Guolin Ke committed
219
220
221
222
}



223
Dataset* DatasetLoader::LoadFromFileAlignWithOtherDataset(const char* filename, const char* initscore_file, const Dataset* train_data) {
Guolin Ke's avatar
Guolin Ke committed
224
225
  data_size_t num_global_data = 0;
  std::vector<data_size_t> used_data_indices;
Guolin Ke's avatar
Guolin Ke committed
226
  auto dataset = std::unique_ptr<Dataset>(new Dataset());
227
228
  auto bin_filename = CheckCanLoadFromBin(filename);
  if (bin_filename.size() == 0) {
Guolin Ke's avatar
Guolin Ke committed
229
    auto parser = std::unique_ptr<Parser>(Parser::CreateParser(filename, config_.header, 0, label_idx_));
Guolin Ke's avatar
Guolin Ke committed
230
231
232
233
    if (parser == nullptr) {
      Log::Fatal("Could not recognize data format of %s", filename);
    }
    dataset->data_filename_ = filename;
Guolin Ke's avatar
Guolin Ke committed
234
    dataset->label_idx_ = label_idx_;
235
    dataset->metadata_.Init(filename, initscore_file);
Guolin Ke's avatar
Guolin Ke committed
236
    if (!config_.two_round) {
Guolin Ke's avatar
Guolin Ke committed
237
238
239
240
      // read data in memory
      auto text_data = LoadTextDataToMemory(filename, dataset->metadata_, 0, 1, &num_global_data, &used_data_indices);
      dataset->num_data_ = static_cast<data_size_t>(text_data.size());
      // initialize label
241
      dataset->metadata_.Init(dataset->num_data_, weight_idx_, group_idx_);
Guolin Ke's avatar
Guolin Ke committed
242
      dataset->CreateValid(train_data);
Guolin Ke's avatar
Guolin Ke committed
243
      // extract features
Guolin Ke's avatar
Guolin Ke committed
244
      ExtractFeaturesFromMemory(text_data, parser.get(), dataset.get());
Guolin Ke's avatar
Guolin Ke committed
245
246
      text_data.clear();
    } else {
Guolin Ke's avatar
Guolin Ke committed
247
      TextReader<data_size_t> text_reader(filename, config_.header);
Guolin Ke's avatar
Guolin Ke committed
248
249
250
251
      // Get number of lines of data file
      dataset->num_data_ = static_cast<data_size_t>(text_reader.CountLine());
      num_global_data = dataset->num_data_;
      // initialize label
252
      dataset->metadata_.Init(dataset->num_data_, weight_idx_, group_idx_);
Guolin Ke's avatar
Guolin Ke committed
253
      dataset->CreateValid(train_data);
Guolin Ke's avatar
Guolin Ke committed
254
      // extract features
Guolin Ke's avatar
Guolin Ke committed
255
      ExtractFeaturesFromFile(filename, parser.get(), used_data_indices, dataset.get());
Guolin Ke's avatar
Guolin Ke committed
256
257
258
    }
  } else {
    // load data from binary file
259
    dataset.reset(LoadFromBinFile(filename, bin_filename.c_str(), 0, 1, &num_global_data, &used_data_indices));
Guolin Ke's avatar
Guolin Ke committed
260
261
262
263
  }
  // not need to check validation data
  // check meta data
  dataset->metadata_.CheckOrPartition(num_global_data, used_data_indices);
Guolin Ke's avatar
Guolin Ke committed
264
  return dataset.release();
Guolin Ke's avatar
Guolin Ke committed
265
266
}

267
Dataset* DatasetLoader::LoadFromBinFile(const char* data_filename, const char* bin_filename, int rank, int num_machines, int* num_global_data, std::vector<data_size_t>* used_data_indices) {
Guolin Ke's avatar
Guolin Ke committed
268
  auto dataset = std::unique_ptr<Dataset>(new Dataset());
269
  auto reader = VirtualFileReader::Make(bin_filename);
Guolin Ke's avatar
Guolin Ke committed
270
  dataset->data_filename_ = data_filename;
271
  if (!reader->Init()) {
Guolin Ke's avatar
Guolin Ke committed
272
273
274
275
276
    Log::Fatal("Could not read binary data from %s", bin_filename);
  }

  // buffer to read binary file
  size_t buffer_size = 16 * 1024 * 1024;
Guolin Ke's avatar
Guolin Ke committed
277
  auto buffer = std::vector<char>(buffer_size);
278

279
280
  // check token
  size_t size_of_token = std::strlen(Dataset::binary_file_token);
281
282
  size_t read_cnt = reader->Read(buffer.data(), sizeof(char) * size_of_token);
  if (read_cnt != sizeof(char) * size_of_token) {
283
284
285
    Log::Fatal("Binary file error: token has the wrong size");
  }
  if (std::string(buffer.data()) != std::string(Dataset::binary_file_token)) {
286
    Log::Fatal("Input file is not LightGBM binary file");
287
  }
Guolin Ke's avatar
Guolin Ke committed
288
289

  // read size of header
290
  read_cnt = reader->Read(buffer.data(), sizeof(size_t));
Guolin Ke's avatar
Guolin Ke committed
291

292
  if (read_cnt != sizeof(size_t)) {
Guolin Ke's avatar
Guolin Ke committed
293
294
295
    Log::Fatal("Binary file error: header has the wrong size");
  }

Guolin Ke's avatar
Guolin Ke committed
296
  size_t size_of_head = *(reinterpret_cast<size_t*>(buffer.data()));
Guolin Ke's avatar
Guolin Ke committed
297
298
299
300

  // re-allocmate space if not enough
  if (size_of_head > buffer_size) {
    buffer_size = size_of_head;
Guolin Ke's avatar
Guolin Ke committed
301
    buffer.resize(buffer_size);
Guolin Ke's avatar
Guolin Ke committed
302
303
  }
  // read header
304
  read_cnt = reader->Read(buffer.data(), size_of_head);
Guolin Ke's avatar
Guolin Ke committed
305
306
307
308
309

  if (read_cnt != size_of_head) {
    Log::Fatal("Binary file error: header is incorrect");
  }
  // get header
Guolin Ke's avatar
Guolin Ke committed
310
  const char* mem_ptr = buffer.data();
Guolin Ke's avatar
Guolin Ke committed
311
312
313
314
315
316
  dataset->num_data_ = *(reinterpret_cast<const data_size_t*>(mem_ptr));
  mem_ptr += sizeof(dataset->num_data_);
  dataset->num_features_ = *(reinterpret_cast<const int*>(mem_ptr));
  mem_ptr += sizeof(dataset->num_features_);
  dataset->num_total_features_ = *(reinterpret_cast<const int*>(mem_ptr));
  mem_ptr += sizeof(dataset->num_total_features_);
Guolin Ke's avatar
Guolin Ke committed
317
318
  dataset->label_idx_ = *(reinterpret_cast<const int*>(mem_ptr));
  mem_ptr += sizeof(dataset->label_idx_);
319
320
321
322
323
324
325
326
327
328
  dataset->max_bin_ = *(reinterpret_cast<const int*>(mem_ptr));
  mem_ptr += sizeof(dataset->max_bin_);
  dataset->bin_construct_sample_cnt_ = *(reinterpret_cast<const int*>(mem_ptr));
  mem_ptr += sizeof(dataset->bin_construct_sample_cnt_);
  dataset->min_data_in_bin_ = *(reinterpret_cast<const int*>(mem_ptr));
  mem_ptr += sizeof(dataset->min_data_in_bin_);
  dataset->use_missing_ = *(reinterpret_cast<const bool*>(mem_ptr));
  mem_ptr += sizeof(dataset->use_missing_);
  dataset->zero_as_missing_ = *(reinterpret_cast<const bool*>(mem_ptr));
  mem_ptr += sizeof(dataset->zero_as_missing_);
Guolin Ke's avatar
Guolin Ke committed
329
330
  dataset->sparse_threshold_ = *(reinterpret_cast<const double*>(mem_ptr));
  mem_ptr += sizeof(dataset->sparse_threshold_);
Guolin Ke's avatar
Guolin Ke committed
331
332
  const int* tmp_feature_map = reinterpret_cast<const int*>(mem_ptr);
  dataset->used_feature_map_.clear();
Guolin Ke's avatar
Guolin Ke committed
333
  for (int i = 0; i < dataset->num_total_features_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
334
335
    dataset->used_feature_map_.push_back(tmp_feature_map[i]);
  }
Guolin Ke's avatar
Guolin Ke committed
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
  mem_ptr += sizeof(int) * dataset->num_total_features_;
  // num_groups
  dataset->num_groups_ = *(reinterpret_cast<const int*>(mem_ptr));
  mem_ptr += sizeof(dataset->num_groups_);
  // real_feature_idx_
  const int* tmp_ptr_real_feature_idx_ = reinterpret_cast<const int*>(mem_ptr);
  dataset->real_feature_idx_.clear();
  for (int i = 0; i < dataset->num_features_; ++i) {
    dataset->real_feature_idx_.push_back(tmp_ptr_real_feature_idx_[i]);
  }
  mem_ptr += sizeof(int) * dataset->num_features_;
  // feature2group
  const int* tmp_ptr_feature2group = reinterpret_cast<const int*>(mem_ptr);
  dataset->feature2group_.clear();
  for (int i = 0; i < dataset->num_features_; ++i) {
    dataset->feature2group_.push_back(tmp_ptr_feature2group[i]);
  }
  mem_ptr += sizeof(int) * dataset->num_features_;
  // feature2subfeature
  const int* tmp_ptr_feature2subfeature = reinterpret_cast<const int*>(mem_ptr);
  dataset->feature2subfeature_.clear();
  for (int i = 0; i < dataset->num_features_; ++i) {
    dataset->feature2subfeature_.push_back(tmp_ptr_feature2subfeature[i]);
  }
  mem_ptr += sizeof(int) * dataset->num_features_;
  // group_bin_boundaries
  const uint64_t* tmp_ptr_group_bin_boundaries = reinterpret_cast<const uint64_t*>(mem_ptr);
  dataset->group_bin_boundaries_.clear();
  for (int i = 0; i < dataset->num_groups_ + 1; ++i) {
    dataset->group_bin_boundaries_.push_back(tmp_ptr_group_bin_boundaries[i]);
  }
  mem_ptr += sizeof(uint64_t) * (dataset->num_groups_ + 1);

  // group_feature_start_
  const int* tmp_ptr_group_feature_start = reinterpret_cast<const int*>(mem_ptr);
  dataset->group_feature_start_.clear();
372
  for (int i = 0; i < dataset->num_groups_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
373
374
375
376
377
378
379
380
381
382
383
384
    dataset->group_feature_start_.push_back(tmp_ptr_group_feature_start[i]);
  }
  mem_ptr += sizeof(int) * (dataset->num_groups_);

  // group_feature_cnt_
  const int* tmp_ptr_group_feature_cnt = reinterpret_cast<const int*>(mem_ptr);
  dataset->group_feature_cnt_.clear();
  for (int i = 0; i < dataset->num_groups_; ++i) {
    dataset->group_feature_cnt_.push_back(tmp_ptr_group_feature_cnt[i]);
  }
  mem_ptr += sizeof(int) * (dataset->num_groups_);

385
  if(!config_.monotone_constraints.empty()) {
386
    CHECK(static_cast<size_t>(dataset->num_total_features_) == config_.monotone_constraints.size());
387
388
389
390
391
392
393
394
395
396
397
398
399
400
    dataset->monotone_types_.resize(dataset->num_features_);
    for(int i = 0; i < dataset->num_total_features_; ++i){
      int inner_fidx = dataset->InnerFeatureIndex(i);
      if(inner_fidx >= 0) {
        dataset->monotone_types_[inner_fidx] = config_.monotone_constraints[i];
      }
    }
  }
  else {
    const int8_t* tmp_ptr_monotone_type = reinterpret_cast<const int8_t*>(mem_ptr);
    dataset->monotone_types_.clear();
    for (int i = 0; i < dataset->num_features_; ++i) {
      dataset->monotone_types_.push_back(tmp_ptr_monotone_type[i]);
    }
Guolin Ke's avatar
Guolin Ke committed
401
402
403
404
405
406
407
  }
  mem_ptr += sizeof(int8_t) * (dataset->num_features_);

  if (ArrayArgs<int8_t>::CheckAllZero(dataset->monotone_types_)) {
    dataset->monotone_types_.clear();
  }

408
  if(!config_.feature_contri.empty()) {
409
    CHECK(static_cast<size_t>(dataset->num_total_features_) == config_.feature_contri.size());
410
411
412
413
414
415
416
417
418
419
420
421
422
423
    dataset->feature_penalty_.resize(dataset->num_features_);
    for(int i = 0; i < dataset->num_total_features_; ++i){
      int inner_fidx = dataset->InnerFeatureIndex(i);
      if(inner_fidx >= 0) {
        dataset->feature_penalty_[inner_fidx] = config_.feature_contri[i];
      }
    }
  }
  else {
    const double* tmp_ptr_feature_penalty = reinterpret_cast<const double*>(mem_ptr);
    dataset->feature_penalty_.clear();
    for (int i = 0; i < dataset->num_features_; ++i) {
      dataset->feature_penalty_.push_back(tmp_ptr_feature_penalty[i]);
    }
Guolin Ke's avatar
Guolin Ke committed
424
425
426
427
428
429
430
  }
  mem_ptr += sizeof(double) * (dataset->num_features_);

  if (ArrayArgs<double>::CheckAll(dataset->feature_penalty_, 1)) {
    dataset->feature_penalty_.clear();
  }

Guolin Ke's avatar
Guolin Ke committed
431
  // get feature names
Guolin Ke's avatar
Guolin Ke committed
432
  dataset->feature_names_.clear();
Guolin Ke's avatar
Guolin Ke committed
433
434
435
436
437
438
439
440
441
442
  // write feature names
  for (int i = 0; i < dataset->num_total_features_; ++i) {
    int str_len = *(reinterpret_cast<const int*>(mem_ptr));
    mem_ptr += sizeof(int);
    std::stringstream str_buf;
    for (int j = 0; j < str_len; ++j) {
      char tmp_char = *(reinterpret_cast<const char*>(mem_ptr));
      mem_ptr += sizeof(char);
      str_buf << tmp_char;
    }
Guolin Ke's avatar
Guolin Ke committed
443
    dataset->feature_names_.emplace_back(str_buf.str());
Guolin Ke's avatar
Guolin Ke committed
444
445
446
  }

  // read size of meta data
447
  read_cnt = reader->Read(buffer.data(), sizeof(size_t));
Guolin Ke's avatar
Guolin Ke committed
448

449
  if (read_cnt != sizeof(size_t)) {
Guolin Ke's avatar
Guolin Ke committed
450
451
452
    Log::Fatal("Binary file error: meta data has the wrong size");
  }

Guolin Ke's avatar
Guolin Ke committed
453
  size_t size_of_metadata = *(reinterpret_cast<size_t*>(buffer.data()));
Guolin Ke's avatar
Guolin Ke committed
454
455
456
457

  // re-allocate space if not enough
  if (size_of_metadata > buffer_size) {
    buffer_size = size_of_metadata;
Guolin Ke's avatar
Guolin Ke committed
458
    buffer.resize(buffer_size);
Guolin Ke's avatar
Guolin Ke committed
459
460
  }
  //  read meta data
461
  read_cnt = reader->Read(buffer.data(), size_of_metadata);
Guolin Ke's avatar
Guolin Ke committed
462
463
464
465
466

  if (read_cnt != size_of_metadata) {
    Log::Fatal("Binary file error: meta data is incorrect");
  }
  // load meta data
Guolin Ke's avatar
Guolin Ke committed
467
  dataset->metadata_.LoadFromMemory(buffer.data());
Guolin Ke's avatar
Guolin Ke committed
468

469
470
  *num_global_data = dataset->num_data_;
  used_data_indices->clear();
Guolin Ke's avatar
Guolin Ke committed
471
  // sample local used data if need to partition
Guolin Ke's avatar
Guolin Ke committed
472
  if (num_machines > 1 && !config_.pre_partition) {
Guolin Ke's avatar
Guolin Ke committed
473
474
475
476
    const data_size_t* query_boundaries = dataset->metadata_.query_boundaries();
    if (query_boundaries == nullptr) {
      // if not contain query file, minimal sample unit is one record
      for (data_size_t i = 0; i < dataset->num_data_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
477
        if (random_.NextShort(0, num_machines) == rank) {
478
          used_data_indices->push_back(i);
Guolin Ke's avatar
Guolin Ke committed
479
480
481
482
483
484
485
486
487
        }
      }
    } else {
      // if contain query file, minimal sample unit is one query
      data_size_t num_queries = dataset->metadata_.num_queries();
      data_size_t qid = -1;
      bool is_query_used = false;
      for (data_size_t i = 0; i < dataset->num_data_; ++i) {
        if (qid >= num_queries) {
488
489
          Log::Fatal("Current query exceeds the range of the query file,\n"
                     "please ensure the query file is correct");
Guolin Ke's avatar
Guolin Ke committed
490
491
492
493
        }
        if (i >= query_boundaries[qid + 1]) {
          // if is new query
          is_query_used = false;
Guolin Ke's avatar
Guolin Ke committed
494
          if (random_.NextShort(0, num_machines) == rank) {
Guolin Ke's avatar
Guolin Ke committed
495
496
497
498
499
            is_query_used = true;
          }
          ++qid;
        }
        if (is_query_used) {
500
          used_data_indices->push_back(i);
Guolin Ke's avatar
Guolin Ke committed
501
502
503
        }
      }
    }
504
    dataset->num_data_ = static_cast<data_size_t>((*used_data_indices).size());
Guolin Ke's avatar
Guolin Ke committed
505
  }
506
  dataset->metadata_.PartitionLabel(*used_data_indices);
Guolin Ke's avatar
Guolin Ke committed
507
  // read feature data
Guolin Ke's avatar
Guolin Ke committed
508
  for (int i = 0; i < dataset->num_groups_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
509
    // read feature size
510
511
    read_cnt = reader->Read(buffer.data(), sizeof(size_t));
    if (read_cnt != sizeof(size_t)) {
Guolin Ke's avatar
Guolin Ke committed
512
513
      Log::Fatal("Binary file error: feature %d has the wrong size", i);
    }
Guolin Ke's avatar
Guolin Ke committed
514
    size_t size_of_feature = *(reinterpret_cast<size_t*>(buffer.data()));
Guolin Ke's avatar
Guolin Ke committed
515
516
517
    // re-allocate space if not enough
    if (size_of_feature > buffer_size) {
      buffer_size = size_of_feature;
Guolin Ke's avatar
Guolin Ke committed
518
      buffer.resize(buffer_size);
Guolin Ke's avatar
Guolin Ke committed
519
520
    }

521
    read_cnt = reader->Read(buffer.data(), size_of_feature);
Guolin Ke's avatar
Guolin Ke committed
522
523
524
525

    if (read_cnt != size_of_feature) {
      Log::Fatal("Binary file error: feature %d is incorrect, read count: %d", i, read_cnt);
    }
Guolin Ke's avatar
Guolin Ke committed
526
    dataset->feature_groups_.emplace_back(std::unique_ptr<FeatureGroup>(
527
528
529
530
      new FeatureGroup(buffer.data(),
                       *num_global_data,
                       *used_data_indices)
      ));
Guolin Ke's avatar
Guolin Ke committed
531
  }
Guolin Ke's avatar
Guolin Ke committed
532
  dataset->feature_groups_.shrink_to_fit();
Guolin Ke's avatar
Guolin Ke committed
533
  dataset->is_finish_load_ = true;
Guolin Ke's avatar
Guolin Ke committed
534
  return dataset.release();
Guolin Ke's avatar
Guolin Ke committed
535
536
}

537
538
539
540
Dataset* DatasetLoader::CostructFromSampleData(double** sample_values,
                                               int** sample_indices, int num_col, const int* num_per_col,
                                               size_t total_sample_size, data_size_t num_data) {
  std::vector<std::unique_ptr<BinMapper>> bin_mappers(num_col);
541
542
  // fill feature_names_ if not header
  if (feature_names_.empty()) {
543
    for (int i = 0; i < num_col; ++i) {
544
545
546
547
548
      std::stringstream str_buf;
      str_buf << "Column_" << i;
      feature_names_.push_back(str_buf.str());
    }
  }
Guolin Ke's avatar
Guolin Ke committed
549
550

  const data_size_t filter_cnt = static_cast<data_size_t>(
Guolin Ke's avatar
Guolin Ke committed
551
    static_cast<double>(config_.min_data_in_leaf * total_sample_size) / num_data);
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
  if (Network::num_machines() == 1) {
    // if only one machine, find bin locally
    OMP_INIT_EX();
    #pragma omp parallel for schedule(guided)
    for (int i = 0; i < num_col; ++i) {
      OMP_LOOP_EX_BEGIN();
      if (ignore_features_.count(i) > 0) {
        bin_mappers[i] = nullptr;
        continue;
      }
      BinType bin_type = BinType::NumericalBin;
      if (categorical_features_.count(i)) {
        bin_type = BinType::CategoricalBin;
      }
      bin_mappers[i].reset(new BinMapper());
      bin_mappers[i]->FindBin(sample_values[i], num_per_col[i], total_sample_size,
Guolin Ke's avatar
Guolin Ke committed
568
                              config_.max_bin, config_.min_data_in_bin, filter_cnt, bin_type, config_.use_missing, config_.zero_as_missing);
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
      OMP_LOOP_EX_END();
    }
    OMP_THROW_EX();
  } else {
    // if have multi-machines, need to find bin distributed
    // different machines will find bin for different features
    int num_machines = Network::num_machines();
    int rank = Network::rank();
    int total_num_feature = num_col;
    total_num_feature = Network::GlobalSyncUpByMin(total_num_feature);
    // start and len will store the process feature indices for different machines
    // machine i will find bins for features in [ start[i], start[i] + len[i] )
    std::vector<int> start(num_machines);
    std::vector<int> len(num_machines);
    int step = (total_num_feature + num_machines - 1) / num_machines;
    if (step < 1) { step = 1; }

    start[0] = 0;
    for (int i = 0; i < num_machines - 1; ++i) {
      len[i] = std::min(step, total_num_feature - start[i]);
      start[i + 1] = start[i] + len[i];
    }
    len[num_machines - 1] = total_num_feature - start[num_machines - 1];
    OMP_INIT_EX();
    #pragma omp parallel for schedule(guided)
    for (int i = 0; i < len[rank]; ++i) {
      OMP_LOOP_EX_BEGIN();
      if (ignore_features_.count(start[rank] + i) > 0) {
        continue;
      }
      BinType bin_type = BinType::NumericalBin;
      if (categorical_features_.count(start[rank] + i)) {
        bin_type = BinType::CategoricalBin;
      }
      bin_mappers[i].reset(new BinMapper());
      bin_mappers[i]->FindBin(sample_values[start[rank] + i], num_per_col[start[rank] + i], total_sample_size,
Guolin Ke's avatar
Guolin Ke committed
605
                              config_.max_bin, config_.min_data_in_bin, filter_cnt, bin_type, config_.use_missing, config_.zero_as_missing);
606
607
608
609
610
611
612
613
      OMP_LOOP_EX_END();
    }
    OMP_THROW_EX();
    int max_bin = 0;
    for (int i = 0; i < len[rank]; ++i) {
      if (bin_mappers[i] != nullptr) {
        max_bin = std::max(max_bin, bin_mappers[i]->num_bin());
      }
Guolin Ke's avatar
Guolin Ke committed
614
    }
615
616
617
618
    max_bin = Network::GlobalSyncUpByMax(max_bin);
    // get size of bin mapper with max_bin size
    int type_size = BinMapper::SizeForSpecificBin(max_bin);
    // since sizes of different feature may not be same, we expand all bin mapper to type_size
Guolin Ke's avatar
Guolin Ke committed
619
    comm_size_t buffer_size = type_size * total_num_feature;
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
    auto input_buffer = std::vector<char>(buffer_size);
    auto output_buffer = std::vector<char>(buffer_size);

    // find local feature bins and copy to buffer
    #pragma omp parallel for schedule(guided)
    for (int i = 0; i < len[rank]; ++i) {
      OMP_LOOP_EX_BEGIN();
      if (ignore_features_.count(start[rank] + i) > 0) {
        continue;
      }
      bin_mappers[i]->CopyTo(input_buffer.data() + i * type_size);
      // free
      bin_mappers[i].reset(nullptr);
      OMP_LOOP_EX_END();
    }
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
636
637
    std::vector<comm_size_t> size_start(num_machines);
    std::vector<comm_size_t> size_len(num_machines);
638
639
    // convert to binary size
    for (int i = 0; i < num_machines; ++i) {
Guolin Ke's avatar
Guolin Ke committed
640
641
      size_start[i] = start[i] * static_cast<comm_size_t>(type_size);
      size_len[i] = len[i] * static_cast<comm_size_t>(type_size);
642
643
    }
    // gather global feature bin mappers
Guolin Ke's avatar
Guolin Ke committed
644
    Network::Allgather(input_buffer.data(), size_start.data(), size_len.data(), output_buffer.data(), buffer_size);
645
646
647
648
649
650
651
652
    // restore features bins from buffer
    for (int i = 0; i < total_num_feature; ++i) {
      if (ignore_features_.count(i) > 0) {
        bin_mappers[i] = nullptr;
        continue;
      }
      bin_mappers[i].reset(new BinMapper());
      bin_mappers[i]->CopyFrom(output_buffer.data() + i * type_size);
653
    }
Guolin Ke's avatar
Guolin Ke committed
654
  }
Guolin Ke's avatar
Guolin Ke committed
655
  auto dataset = std::unique_ptr<Dataset>(new Dataset(num_data));
Guolin Ke's avatar
Guolin Ke committed
656
  dataset->Construct(bin_mappers, sample_indices, num_per_col, total_sample_size, config_);
657
  dataset->set_feature_names(feature_names_);
Guolin Ke's avatar
Guolin Ke committed
658
  return dataset.release();
Guolin Ke's avatar
Guolin Ke committed
659
}
Guolin Ke's avatar
Guolin Ke committed
660
661
662
663
664
665


// ---- private functions ----

void DatasetLoader::CheckDataset(const Dataset* dataset) {
  if (dataset->num_data_ <= 0) {
Guolin Ke's avatar
Guolin Ke committed
666
    Log::Fatal("Data file %s is empty", dataset->data_filename_.c_str());
Guolin Ke's avatar
Guolin Ke committed
667
  }
668
669
  if (dataset->feature_names_.size() != static_cast<size_t>(dataset->num_total_features_)) {
    Log::Fatal("Size of feature name error, should be %d, got %d", dataset->num_total_features_,
670
               static_cast<int>(dataset->feature_names_.size()));
671
  }
Guolin Ke's avatar
Guolin Ke committed
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
  bool is_feature_order_by_group = true;
  int last_group = -1;
  int last_sub_feature = -1;
  // if features are ordered, not need to use hist_buf
  for (int i = 0; i < dataset->num_features_; ++i) {
    int group = dataset->feature2group_[i];
    int sub_feature = dataset->feature2subfeature_[i];
    if (group < last_group) {
      is_feature_order_by_group = false;
    } else if (group == last_group) {
      if (sub_feature <= last_sub_feature) {
        is_feature_order_by_group = false;
        break;
      }
    }
    last_group = group;
    last_sub_feature = sub_feature;
  }
  if (!is_feature_order_by_group) {
691
    Log::Fatal("Features in dataset should be ordered by group");
Guolin Ke's avatar
Guolin Ke committed
692
  }
Guolin Ke's avatar
Guolin Ke committed
693
694
695
}

std::vector<std::string> DatasetLoader::LoadTextDataToMemory(const char* filename, const Metadata& metadata,
696
697
                                                             int rank, int num_machines, int* num_global_data,
                                                             std::vector<data_size_t>* used_data_indices) {
Guolin Ke's avatar
Guolin Ke committed
698
  TextReader<data_size_t> text_reader(filename, config_.header);
Guolin Ke's avatar
Guolin Ke committed
699
  used_data_indices->clear();
Guolin Ke's avatar
Guolin Ke committed
700
  if (num_machines == 1 || config_.pre_partition) {
Guolin Ke's avatar
Guolin Ke committed
701
702
703
704
705
706
707
708
709
    // read all lines
    *num_global_data = text_reader.ReadAllLines();
  } else {  // need partition data
            // get query data
    const data_size_t* query_boundaries = metadata.query_boundaries();

    if (query_boundaries == nullptr) {
      // if not contain query data, minimal sample unit is one record
      *num_global_data = text_reader.ReadAndFilterLines([this, rank, num_machines](data_size_t) {
Guolin Ke's avatar
Guolin Ke committed
710
        if (random_.NextShort(0, num_machines) == rank) {
Guolin Ke's avatar
Guolin Ke committed
711
712
713
714
715
716
717
718
719
720
721
722
723
724
          return true;
        } else {
          return false;
        }
      }, used_data_indices);
    } else {
      // if contain query data, minimal sample unit is one query
      data_size_t num_queries = metadata.num_queries();
      data_size_t qid = -1;
      bool is_query_used = false;
      *num_global_data = text_reader.ReadAndFilterLines(
        [this, rank, num_machines, &qid, &query_boundaries, &is_query_used, num_queries]
      (data_size_t line_idx) {
        if (qid >= num_queries) {
725
726
          Log::Fatal("Current query exceeds the range of the query file,\n"
                     "please ensure the query file is correct");
Guolin Ke's avatar
Guolin Ke committed
727
728
729
730
        }
        if (line_idx >= query_boundaries[qid + 1]) {
          // if is new query
          is_query_used = false;
Guolin Ke's avatar
Guolin Ke committed
731
          if (random_.NextShort(0, num_machines) == rank) {
Guolin Ke's avatar
Guolin Ke committed
732
733
734
735
736
737
738
739
740
741
742
743
            is_query_used = true;
          }
          ++qid;
        }
        return is_query_used;
      }, used_data_indices);
    }
  }
  return std::move(text_reader.Lines());
}

std::vector<std::string> DatasetLoader::SampleTextDataFromMemory(const std::vector<std::string>& data) {
Guolin Ke's avatar
Guolin Ke committed
744
  int sample_cnt = config_.bin_construct_sample_cnt;
745
746
  if (static_cast<size_t>(sample_cnt) > data.size()) {
    sample_cnt = static_cast<int>(data.size());
747
  }
748
  auto sample_indices = random_.Sample(static_cast<int>(data.size()), sample_cnt);
Guolin Ke's avatar
Guolin Ke committed
749
  std::vector<std::string> out(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
750
751
  for (size_t i = 0; i < sample_indices.size(); ++i) {
    const size_t idx = sample_indices[i];
Guolin Ke's avatar
Guolin Ke committed
752
    out[i] = data[idx];
Guolin Ke's avatar
Guolin Ke committed
753
754
755
756
757
  }
  return out;
}

std::vector<std::string> DatasetLoader::SampleTextDataFromFile(const char* filename, const Metadata& metadata, int rank, int num_machines, int* num_global_data, std::vector<data_size_t>* used_data_indices) {
Guolin Ke's avatar
Guolin Ke committed
758
759
  const data_size_t sample_cnt = static_cast<data_size_t>(config_.bin_construct_sample_cnt);
  TextReader<data_size_t> text_reader(filename, config_.header);
Guolin Ke's avatar
Guolin Ke committed
760
  std::vector<std::string> out_data;
Guolin Ke's avatar
Guolin Ke committed
761
  if (num_machines == 1 || config_.pre_partition) {
Guolin Ke's avatar
Guolin Ke committed
762
763
764
765
766
767
768
769
    *num_global_data = static_cast<data_size_t>(text_reader.SampleFromFile(random_, sample_cnt, &out_data));
  } else {  // need partition data
            // get query data
    const data_size_t* query_boundaries = metadata.query_boundaries();
    if (query_boundaries == nullptr) {
      // if not contain query file, minimal sample unit is one record
      *num_global_data = text_reader.SampleAndFilterFromFile([this, rank, num_machines]
      (data_size_t) {
Guolin Ke's avatar
Guolin Ke committed
770
        if (random_.NextShort(0, num_machines) == rank) {
Guolin Ke's avatar
Guolin Ke committed
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
          return true;
        } else {
          return false;
        }
      }, used_data_indices, random_, sample_cnt, &out_data);
    } else {
      // if contain query file, minimal sample unit is one query
      data_size_t num_queries = metadata.num_queries();
      data_size_t qid = -1;
      bool is_query_used = false;
      *num_global_data = text_reader.SampleAndFilterFromFile(
        [this, rank, num_machines, &qid, &query_boundaries, &is_query_used, num_queries]
      (data_size_t line_idx) {
        if (qid >= num_queries) {
          Log::Fatal("Query id exceeds the range of the query file, \
                      please ensure the query file is correct");
        }
        if (line_idx >= query_boundaries[qid + 1]) {
          // if is new query
          is_query_used = false;
Guolin Ke's avatar
Guolin Ke committed
791
          if (random_.NextShort(0, num_machines) == rank) {
Guolin Ke's avatar
Guolin Ke committed
792
793
794
795
796
797
798
799
800
801
802
803
            is_query_used = true;
          }
          ++qid;
        }
        return is_query_used;
      }, used_data_indices, random_, sample_cnt, &out_data);
    }
  }
  return out_data;
}

void DatasetLoader::ConstructBinMappersFromTextData(int rank, int num_machines, const std::vector<std::string>& sample_data, const Parser* parser, Dataset* dataset) {
Guolin Ke's avatar
Guolin Ke committed
804

Guolin Ke's avatar
Guolin Ke committed
805
  std::vector<std::vector<double>> sample_values;
Guolin Ke's avatar
Guolin Ke committed
806
  std::vector<std::vector<int>> sample_indices;
Guolin Ke's avatar
Guolin Ke committed
807
808
  std::vector<std::pair<int, double>> oneline_features;
  double label;
Guolin Ke's avatar
Guolin Ke committed
809
  for (int i = 0; i < static_cast<int>(sample_data.size()); ++i) {
Guolin Ke's avatar
Guolin Ke committed
810
811
812
813
    oneline_features.clear();
    // parse features
    parser->ParseOneLine(sample_data[i].c_str(), &oneline_features, &label);
    for (std::pair<int, double>& inner_data : oneline_features) {
814
      if (static_cast<size_t>(inner_data.first) >= sample_values.size()) {
Guolin Ke's avatar
Guolin Ke committed
815
816
        sample_values.resize(inner_data.first + 1);
        sample_indices.resize(inner_data.first + 1);
817
      }
Guolin Ke's avatar
Guolin Ke committed
818
      if (std::fabs(inner_data.second) > kZeroThreshold || std::isnan(inner_data.second)) {
Guolin Ke's avatar
Guolin Ke committed
819
820
        sample_values[inner_data.first].emplace_back(inner_data.second);
        sample_indices[inner_data.first].emplace_back(i);
Guolin Ke's avatar
Guolin Ke committed
821
822
823
824
      }
    }
  }

Guolin Ke's avatar
Guolin Ke committed
825
  dataset->feature_groups_.clear();
Guolin Ke's avatar
Guolin Ke committed
826

827
828
  if (feature_names_.empty()) {
    // -1 means doesn't use this feature
Guolin Ke's avatar
Guolin Ke committed
829
830
    dataset->num_total_features_ = std::max(static_cast<int>(sample_values.size()), parser->TotalColumns() - 1);
    dataset->used_feature_map_ = std::vector<int>(dataset->num_total_features_, -1);
831
832
833
834
  } else {
    dataset->used_feature_map_ = std::vector<int>(feature_names_.size(), -1);
    dataset->num_total_features_ = static_cast<int>(feature_names_.size());
  }
Guolin Ke's avatar
Guolin Ke committed
835
836
837
838
839
840
841

  // check the range of label_idx, weight_idx and group_idx
  CHECK(label_idx_ >= 0 && label_idx_ <= dataset->num_total_features_);
  CHECK(weight_idx_ < 0 || weight_idx_ < dataset->num_total_features_);
  CHECK(group_idx_ < 0 || group_idx_ < dataset->num_total_features_);

  // fill feature_names_ if not header
Guolin Ke's avatar
Guolin Ke committed
842
  if (feature_names_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
843
844
845
846
847
848
    for (int i = 0; i < dataset->num_total_features_; ++i) {
      std::stringstream str_buf;
      str_buf << "Column_" << i;
      feature_names_.push_back(str_buf.str());
    }
  }
849
  dataset->set_feature_names(feature_names_);
Guolin Ke's avatar
Guolin Ke committed
850
  std::vector<std::unique_ptr<BinMapper>> bin_mappers(dataset->num_total_features_);
Guolin Ke's avatar
Guolin Ke committed
851
  const data_size_t filter_cnt = static_cast<data_size_t>(
Guolin Ke's avatar
Guolin Ke committed
852
    static_cast<double>(config_.min_data_in_leaf* sample_data.size()) / dataset->num_data_);
Guolin Ke's avatar
Guolin Ke committed
853

Guolin Ke's avatar
Guolin Ke committed
854
855
856
  // start find bins
  if (num_machines == 1) {
    // if only one machine, find bin locally
857
    OMP_INIT_EX();
858
    #pragma omp parallel for schedule(guided)
Guolin Ke's avatar
Guolin Ke committed
859
    for (int i = 0; i < static_cast<int>(sample_values.size()); ++i) {
860
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
861
      if (ignore_features_.count(i) > 0) {
Guolin Ke's avatar
Guolin Ke committed
862
        bin_mappers[i] = nullptr;
Guolin Ke's avatar
Guolin Ke committed
863
864
        continue;
      }
865
866
867
868
      BinType bin_type = BinType::NumericalBin;
      if (categorical_features_.count(i)) {
        bin_type = BinType::CategoricalBin;
      }
Guolin Ke's avatar
Guolin Ke committed
869
      bin_mappers[i].reset(new BinMapper());
870
      bin_mappers[i]->FindBin(sample_values[i].data(), static_cast<int>(sample_values[i].size()),
Guolin Ke's avatar
Guolin Ke committed
871
                              sample_data.size(), config_.max_bin, config_.min_data_in_bin, filter_cnt, bin_type, config_.use_missing, config_.zero_as_missing);
872
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
873
    }
874
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
875
  } else {
876
    // if have multi-machines, need to find bin distributed
Guolin Ke's avatar
Guolin Ke committed
877
878
    // different machines will find bin for different features

Guolin Ke's avatar
Guolin Ke committed
879
880
881
    int num_total_features = dataset->num_total_features_;
    num_total_features = Network::GlobalSyncUpByMin(num_total_features);
    dataset->num_total_features_ = num_total_features;
Guolin Ke's avatar
Guolin Ke committed
882
    // start and len will store the process feature indices for different machines
883
    // machine i will find bins for features in [ start[i], start[i] + len[i] )
Guolin Ke's avatar
Guolin Ke committed
884
885
    std::vector<int> start(num_machines);
    std::vector<int> len(num_machines);
Guolin Ke's avatar
Guolin Ke committed
886
    int step = (num_total_features + num_machines - 1) / num_machines;
Guolin Ke's avatar
Guolin Ke committed
887
888
889
890
    if (step < 1) { step = 1; }

    start[0] = 0;
    for (int i = 0; i < num_machines - 1; ++i) {
Guolin Ke's avatar
Guolin Ke committed
891
      len[i] = std::min(step, num_total_features - start[i]);
Guolin Ke's avatar
Guolin Ke committed
892
893
      start[i + 1] = start[i] + len[i];
    }
Guolin Ke's avatar
Guolin Ke committed
894
    len[num_machines - 1] = num_total_features - start[num_machines - 1];
895
    OMP_INIT_EX();
896
    #pragma omp parallel for schedule(guided)
897
    for (int i = 0; i < len[rank]; ++i) {
898
      OMP_LOOP_EX_BEGIN();
899
900
901
902
903
904
905
906
      if (ignore_features_.count(start[rank] + i) > 0) {
        continue;
      }
      BinType bin_type = BinType::NumericalBin;
      if (categorical_features_.count(start[rank] + i)) {
        bin_type = BinType::CategoricalBin;
      }
      bin_mappers[i].reset(new BinMapper());
Guolin Ke's avatar
Guolin Ke committed
907
      bin_mappers[i]->FindBin(sample_values[start[rank] + i].data(), static_cast<int>(sample_values[start[rank] + i].size()),
Guolin Ke's avatar
Guolin Ke committed
908
                              sample_data.size(), config_.max_bin, config_.min_data_in_bin, filter_cnt, bin_type, config_.use_missing, config_.zero_as_missing);
909
      OMP_LOOP_EX_END();
910
    }
911
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
912
913
914
915
916
917
918
    int max_bin = 0;
    for (int i = 0; i < len[rank]; ++i) {
      if (bin_mappers[i] != nullptr) {
        max_bin = std::max(max_bin, bin_mappers[i]->num_bin());
      }
    }
    max_bin = Network::GlobalSyncUpByMax(max_bin);
919
    // get size of bin mapper with max_bin size
920
    int type_size = BinMapper::SizeForSpecificBin(max_bin);
Guolin Ke's avatar
Guolin Ke committed
921
    // since sizes of different feature may not be same, we expand all bin mapper to type_size
Guolin Ke's avatar
Guolin Ke committed
922
    comm_size_t buffer_size = type_size * num_total_features;
Guolin Ke's avatar
Guolin Ke committed
923
924
    auto input_buffer = std::vector<char>(buffer_size);
    auto output_buffer = std::vector<char>(buffer_size);
Guolin Ke's avatar
Guolin Ke committed
925
926

    // find local feature bins and copy to buffer
927
    #pragma omp parallel for schedule(guided)
Guolin Ke's avatar
Guolin Ke committed
928
    for (int i = 0; i < len[rank]; ++i) {
929
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
930
931
932
      if (ignore_features_.count(start[rank] + i) > 0) {
        continue;
      }
933
934
935
      bin_mappers[i]->CopyTo(input_buffer.data() + i * type_size);
      // free
      bin_mappers[i].reset(nullptr);
936
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
937
    }
938
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
939
940
    std::vector<comm_size_t> size_start(num_machines);
    std::vector<comm_size_t> size_len(num_machines);
Guolin Ke's avatar
Guolin Ke committed
941
942
    // convert to binary size
    for (int i = 0; i < num_machines; ++i) {
Guolin Ke's avatar
Guolin Ke committed
943
944
      size_start[i] = start[i] * static_cast<comm_size_t>(type_size);
      size_len[i] = len[i] * static_cast<comm_size_t>(type_size);
Guolin Ke's avatar
Guolin Ke committed
945
946
    }
    // gather global feature bin mappers
Guolin Ke's avatar
Guolin Ke committed
947
    Network::Allgather(input_buffer.data(), size_start.data(), size_len.data(), output_buffer.data(), buffer_size);
Guolin Ke's avatar
Guolin Ke committed
948
    // restore features bins from buffer
Guolin Ke's avatar
Guolin Ke committed
949
    for (int i = 0; i < num_total_features; ++i) {
Guolin Ke's avatar
Guolin Ke committed
950
      if (ignore_features_.count(i) > 0) {
Guolin Ke's avatar
Guolin Ke committed
951
        bin_mappers[i] = nullptr;
Guolin Ke's avatar
Guolin Ke committed
952
953
        continue;
      }
Guolin Ke's avatar
Guolin Ke committed
954
955
      bin_mappers[i].reset(new BinMapper());
      bin_mappers[i]->CopyFrom(output_buffer.data() + i * type_size);
Guolin Ke's avatar
Guolin Ke committed
956
957
    }
  }
Guolin Ke's avatar
Guolin Ke committed
958
  sample_values.clear();
Guolin Ke's avatar
Guolin Ke committed
959
  dataset->Construct(bin_mappers, Common::Vector2Ptr<int>(sample_indices).data(),
Guolin Ke's avatar
Guolin Ke committed
960
                     Common::VectorSize<int>(sample_indices).data(), sample_data.size(), config_);
Guolin Ke's avatar
Guolin Ke committed
961
962
963
964
965
966
967
}

/*! \brief Extract local features from memory */
void DatasetLoader::ExtractFeaturesFromMemory(std::vector<std::string>& text_data, const Parser* parser, Dataset* dataset) {
  std::vector<std::pair<int, double>> oneline_features;
  double tmp_label = 0.0f;
  if (predict_fun_ == nullptr) {
968
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
969
    // if doesn't need to prediction with initial model
970
    #pragma omp parallel for schedule(static) private(oneline_features) firstprivate(tmp_label)
Guolin Ke's avatar
Guolin Ke committed
971
    for (data_size_t i = 0; i < dataset->num_data_; ++i) {
972
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
973
974
975
976
977
      const int tid = omp_get_thread_num();
      oneline_features.clear();
      // parser
      parser->ParseOneLine(text_data[i].c_str(), &oneline_features, &tmp_label);
      // set label
978
      dataset->metadata_.SetLabelAt(i, static_cast<label_t>(tmp_label));
Guolin Ke's avatar
Guolin Ke committed
979
980
981
982
983
984
      // free processed line:
      text_data[i].clear();
      // shrink_to_fit will be very slow in linux, and seems not free memory, disable for now
      // text_reader_->Lines()[i].shrink_to_fit();
      // push data
      for (auto& inner_data : oneline_features) {
985
        if (inner_data.first >= dataset->num_total_features_) { continue; }
Guolin Ke's avatar
Guolin Ke committed
986
987
988
        int feature_idx = dataset->used_feature_map_[inner_data.first];
        if (feature_idx >= 0) {
          // if is used feature
Guolin Ke's avatar
Guolin Ke committed
989
990
991
          int group = dataset->feature2group_[feature_idx];
          int sub_feature = dataset->feature2subfeature_[feature_idx];
          dataset->feature_groups_[group]->PushData(tid, sub_feature, i, inner_data.second);
Guolin Ke's avatar
Guolin Ke committed
992
993
        } else {
          if (inner_data.first == weight_idx_) {
994
            dataset->metadata_.SetWeightAt(i, static_cast<label_t>(inner_data.second));
Guolin Ke's avatar
Guolin Ke committed
995
996
997
998
999
          } else if (inner_data.first == group_idx_) {
            dataset->metadata_.SetQueryAt(i, static_cast<data_size_t>(inner_data.second));
          }
        }
      }
1000
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1001
    }
1002
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1003
  } else {
1004
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1005
    // if need to prediction with initial model
1006
    std::vector<double> init_score(dataset->num_data_ * num_class_);
1007
    #pragma omp parallel for schedule(static) private(oneline_features) firstprivate(tmp_label)
Guolin Ke's avatar
Guolin Ke committed
1008
    for (data_size_t i = 0; i < dataset->num_data_; ++i) {
1009
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1010
1011
1012
1013
1014
      const int tid = omp_get_thread_num();
      oneline_features.clear();
      // parser
      parser->ParseOneLine(text_data[i].c_str(), &oneline_features, &tmp_label);
      // set initial score
Guolin Ke's avatar
Guolin Ke committed
1015
1016
      std::vector<double> oneline_init_score(num_class_);
      predict_fun_(oneline_features, oneline_init_score.data());
1017
      for (int k = 0; k < num_class_; ++k) {
1018
        init_score[k * dataset->num_data_ + i] = static_cast<double>(oneline_init_score[k]);
Guolin Ke's avatar
Guolin Ke committed
1019
1020
      }
      // set label
1021
      dataset->metadata_.SetLabelAt(i, static_cast<label_t>(tmp_label));
Guolin Ke's avatar
Guolin Ke committed
1022
1023
1024
1025
1026
1027
      // free processed line:
      text_data[i].clear();
      // shrink_to_fit will be very slow in linux, and seems not free memory, disable for now
      // text_reader_->Lines()[i].shrink_to_fit();
      // push data
      for (auto& inner_data : oneline_features) {
1028
        if (inner_data.first >= dataset->num_total_features_) { continue; }
Guolin Ke's avatar
Guolin Ke committed
1029
1030
1031
        int feature_idx = dataset->used_feature_map_[inner_data.first];
        if (feature_idx >= 0) {
          // if is used feature
Guolin Ke's avatar
Guolin Ke committed
1032
1033
          int group = dataset->feature2group_[feature_idx];
          int sub_feature = dataset->feature2subfeature_[feature_idx];
1034
          dataset->feature_groups_[group]->PushData(tid, sub_feature, i, inner_data.second);
Guolin Ke's avatar
Guolin Ke committed
1035
1036
        } else {
          if (inner_data.first == weight_idx_) {
1037
            dataset->metadata_.SetWeightAt(i, static_cast<label_t>(inner_data.second));
Guolin Ke's avatar
Guolin Ke committed
1038
1039
1040
1041
1042
          } else if (inner_data.first == group_idx_) {
            dataset->metadata_.SetQueryAt(i, static_cast<data_size_t>(inner_data.second));
          }
        }
      }
1043
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1044
    }
1045
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1046
    // metadata_ will manage space of init_score
1047
    dataset->metadata_.SetInitScore(init_score.data(), dataset->num_data_ * num_class_);
Guolin Ke's avatar
Guolin Ke committed
1048
  }
Guolin Ke's avatar
Guolin Ke committed
1049
  dataset->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1050
1051
1052
1053
1054
1055
  // text data can be free after loaded feature values
  text_data.clear();
}

/*! \brief Extract local features from file */
void DatasetLoader::ExtractFeaturesFromFile(const char* filename, const Parser* parser, const std::vector<data_size_t>& used_data_indices, Dataset* dataset) {
1056
  std::vector<double> init_score;
Guolin Ke's avatar
Guolin Ke committed
1057
  if (predict_fun_ != nullptr) {
1058
    init_score = std::vector<double>(dataset->num_data_ * num_class_);
Guolin Ke's avatar
Guolin Ke committed
1059
1060
1061
1062
1063
1064
  }
  std::function<void(data_size_t, const std::vector<std::string>&)> process_fun =
    [this, &init_score, &parser, &dataset]
  (data_size_t start_idx, const std::vector<std::string>& lines) {
    std::vector<std::pair<int, double>> oneline_features;
    double tmp_label = 0.0f;
1065
    OMP_INIT_EX();
1066
    #pragma omp parallel for schedule(static) private(oneline_features) firstprivate(tmp_label)
Guolin Ke's avatar
Guolin Ke committed
1067
    for (data_size_t i = 0; i < static_cast<data_size_t>(lines.size()); ++i) {
1068
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1069
1070
1071
1072
1073
      const int tid = omp_get_thread_num();
      oneline_features.clear();
      // parser
      parser->ParseOneLine(lines[i].c_str(), &oneline_features, &tmp_label);
      // set initial score
Guolin Ke's avatar
Guolin Ke committed
1074
      if (!init_score.empty()) {
Guolin Ke's avatar
Guolin Ke committed
1075
1076
        std::vector<double> oneline_init_score(num_class_);
        predict_fun_(oneline_features, oneline_init_score.data());
1077
        for (int k = 0; k < num_class_; ++k) {
1078
          init_score[k * dataset->num_data_ + start_idx + i] = static_cast<double>(oneline_init_score[k]);
Guolin Ke's avatar
Guolin Ke committed
1079
1080
1081
        }
      }
      // set label
1082
      dataset->metadata_.SetLabelAt(start_idx + i, static_cast<label_t>(tmp_label));
Guolin Ke's avatar
Guolin Ke committed
1083
1084
      // push data
      for (auto& inner_data : oneline_features) {
1085
        if (inner_data.first >= dataset->num_total_features_) { continue; }
Guolin Ke's avatar
Guolin Ke committed
1086
1087
1088
        int feature_idx = dataset->used_feature_map_[inner_data.first];
        if (feature_idx >= 0) {
          // if is used feature
Guolin Ke's avatar
Guolin Ke committed
1089
1090
1091
          int group = dataset->feature2group_[feature_idx];
          int sub_feature = dataset->feature2subfeature_[feature_idx];
          dataset->feature_groups_[group]->PushData(tid, sub_feature, start_idx + i, inner_data.second);
Guolin Ke's avatar
Guolin Ke committed
1092
1093
        } else {
          if (inner_data.first == weight_idx_) {
1094
            dataset->metadata_.SetWeightAt(start_idx + i, static_cast<label_t>(inner_data.second));
Guolin Ke's avatar
Guolin Ke committed
1095
1096
1097
1098
1099
          } else if (inner_data.first == group_idx_) {
            dataset->metadata_.SetQueryAt(start_idx + i, static_cast<data_size_t>(inner_data.second));
          }
        }
      }
1100
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1101
    }
1102
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1103
  };
Guolin Ke's avatar
Guolin Ke committed
1104
  TextReader<data_size_t> text_reader(filename, config_.header);
Guolin Ke's avatar
Guolin Ke committed
1105
  if (!used_data_indices.empty()) {
Guolin Ke's avatar
Guolin Ke committed
1106
1107
1108
1109
1110
1111
1112
1113
    // only need part of data
    text_reader.ReadPartAndProcessParallel(used_data_indices, process_fun);
  } else {
    // need full data
    text_reader.ReadAllAndProcessParallel(process_fun);
  }

  // metadata_ will manage space of init_score
Guolin Ke's avatar
Guolin Ke committed
1114
  if (!init_score.empty()) {
1115
    dataset->metadata_.SetInitScore(init_score.data(), dataset->num_data_ * num_class_);
Guolin Ke's avatar
Guolin Ke committed
1116
  }
Guolin Ke's avatar
Guolin Ke committed
1117
  dataset->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1118
1119
1120
}

/*! \brief Check can load from binary file */
1121
std::string DatasetLoader::CheckCanLoadFromBin(const char* filename) {
Guolin Ke's avatar
Guolin Ke committed
1122
1123
1124
  std::string bin_filename(filename);
  bin_filename.append(".bin");

1125
  auto reader = VirtualFileReader::Make(bin_filename.c_str());
Guolin Ke's avatar
Guolin Ke committed
1126

1127
  if (!reader->Init()) {
1128
    bin_filename = std::string(filename);
1129
1130
    reader = VirtualFileReader::Make(bin_filename.c_str());
    if (!reader->Init()) {
1131
      Log::Fatal("Cannot open data file %s", bin_filename.c_str());
1132
    }
1133
  }
1134
1135
1136
1137
1138

  size_t buffer_size = 256;
  auto buffer = std::vector<char>(buffer_size);
  // read size of token
  size_t size_of_token = std::strlen(Dataset::binary_file_token);
1139
  size_t read_cnt = reader->Read(buffer.data(), size_of_token);
1140
1141
  if (read_cnt == size_of_token
      && std::string(buffer.data()) == std::string(Dataset::binary_file_token)) {
1142
    return bin_filename;
Guolin Ke's avatar
Guolin Ke committed
1143
  } else {
1144
    return std::string();
Guolin Ke's avatar
Guolin Ke committed
1145
  }
1146

Guolin Ke's avatar
Guolin Ke committed
1147
1148
}

Guolin Ke's avatar
Guolin Ke committed
1149
}