dataset_loader.cpp 46.1 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
  const int* tmp_feature_map = reinterpret_cast<const int*>(mem_ptr);
  dataset->used_feature_map_.clear();
Guolin Ke's avatar
Guolin Ke committed
331
  for (int i = 0; i < dataset->num_total_features_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
332
333
    dataset->used_feature_map_.push_back(tmp_feature_map[i]);
  }
Guolin Ke's avatar
Guolin Ke committed
334
335
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
  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();
370
  for (int i = 0; i < dataset->num_groups_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
371
372
373
374
375
376
377
378
379
380
381
382
    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_);

383
  if(!config_.monotone_constraints.empty()) {
384
    CHECK(static_cast<size_t>(dataset->num_total_features_) == config_.monotone_constraints.size());
385
386
387
388
389
390
391
392
393
394
395
396
397
398
    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
399
400
401
402
403
404
405
  }
  mem_ptr += sizeof(int8_t) * (dataset->num_features_);

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

406
  if(!config_.feature_contri.empty()) {
407
    CHECK(static_cast<size_t>(dataset->num_total_features_) == config_.feature_contri.size());
408
409
410
411
412
413
414
415
416
417
418
419
420
421
    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
422
423
424
425
426
427
428
  }
  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
429
  // get feature names
Guolin Ke's avatar
Guolin Ke committed
430
  dataset->feature_names_.clear();
Guolin Ke's avatar
Guolin Ke committed
431
432
433
434
435
436
437
438
439
440
  // 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
441
    dataset->feature_names_.emplace_back(str_buf.str());
Guolin Ke's avatar
Guolin Ke committed
442
443
444
  }

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

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

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

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

  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
465
  dataset->metadata_.LoadFromMemory(buffer.data());
Guolin Ke's avatar
Guolin Ke committed
466

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

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

    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
524
    dataset->feature_groups_.emplace_back(std::unique_ptr<FeatureGroup>(
525
526
527
528
      new FeatureGroup(buffer.data(),
                       *num_global_data,
                       *used_data_indices)
      ));
Guolin Ke's avatar
Guolin Ke committed
529
  }
Guolin Ke's avatar
Guolin Ke committed
530
  dataset->feature_groups_.shrink_to_fit();
Guolin Ke's avatar
Guolin Ke committed
531
  dataset->is_finish_load_ = true;
Guolin Ke's avatar
Guolin Ke committed
532
  return dataset.release();
Guolin Ke's avatar
Guolin Ke committed
533
534
}

535
536
537
538
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);
539
540
  // fill feature_names_ if not header
  if (feature_names_.empty()) {
541
    for (int i = 0; i < num_col; ++i) {
542
543
544
545
546
      std::stringstream str_buf;
      str_buf << "Column_" << i;
      feature_names_.push_back(str_buf.str());
    }
  }
Guolin Ke's avatar
Guolin Ke committed
547
548

  const data_size_t filter_cnt = static_cast<data_size_t>(
Guolin Ke's avatar
Guolin Ke committed
549
    static_cast<double>(config_.min_data_in_leaf * total_sample_size) / num_data);
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
  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
566
                              config_.max_bin, config_.min_data_in_bin, filter_cnt, bin_type, config_.use_missing, config_.zero_as_missing);
567
568
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
      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
603
                              config_.max_bin, config_.min_data_in_bin, filter_cnt, bin_type, config_.use_missing, config_.zero_as_missing);
604
605
606
607
608
609
610
611
      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
612
    }
613
614
615
616
    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
617
    comm_size_t buffer_size = type_size * total_num_feature;
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
    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
634
635
    std::vector<comm_size_t> size_start(num_machines);
    std::vector<comm_size_t> size_len(num_machines);
636
637
    // convert to binary size
    for (int i = 0; i < num_machines; ++i) {
Guolin Ke's avatar
Guolin Ke committed
638
639
      size_start[i] = start[i] * static_cast<comm_size_t>(type_size);
      size_len[i] = len[i] * static_cast<comm_size_t>(type_size);
640
641
    }
    // gather global feature bin mappers
Guolin Ke's avatar
Guolin Ke committed
642
    Network::Allgather(input_buffer.data(), size_start.data(), size_len.data(), output_buffer.data(), buffer_size);
643
644
645
646
647
648
649
650
    // 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);
651
    }
Guolin Ke's avatar
Guolin Ke committed
652
  }
Guolin Ke's avatar
Guolin Ke committed
653
  auto dataset = std::unique_ptr<Dataset>(new Dataset(num_data));
Guolin Ke's avatar
Guolin Ke committed
654
  dataset->Construct(bin_mappers, sample_indices, num_per_col, total_sample_size, config_);
655
  dataset->set_feature_names(feature_names_);
Guolin Ke's avatar
Guolin Ke committed
656
  return dataset.release();
Guolin Ke's avatar
Guolin Ke committed
657
}
Guolin Ke's avatar
Guolin Ke committed
658
659
660
661
662
663


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

void DatasetLoader::CheckDataset(const Dataset* dataset) {
  if (dataset->num_data_ <= 0) {
Guolin Ke's avatar
Guolin Ke committed
664
    Log::Fatal("Data file %s is empty", dataset->data_filename_.c_str());
Guolin Ke's avatar
Guolin Ke committed
665
  }
666
667
  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_,
668
               static_cast<int>(dataset->feature_names_.size()));
669
  }
Guolin Ke's avatar
Guolin Ke committed
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
  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) {
689
    Log::Fatal("Features in dataset should be ordered by group");
Guolin Ke's avatar
Guolin Ke committed
690
  }
Guolin Ke's avatar
Guolin Ke committed
691
692
693
}

std::vector<std::string> DatasetLoader::LoadTextDataToMemory(const char* filename, const Metadata& metadata,
694
695
                                                             int rank, int num_machines, int* num_global_data,
                                                             std::vector<data_size_t>* used_data_indices) {
Guolin Ke's avatar
Guolin Ke committed
696
  TextReader<data_size_t> text_reader(filename, config_.header);
Guolin Ke's avatar
Guolin Ke committed
697
  used_data_indices->clear();
Guolin Ke's avatar
Guolin Ke committed
698
  if (num_machines == 1 || config_.pre_partition) {
Guolin Ke's avatar
Guolin Ke committed
699
700
701
702
703
704
705
706
707
    // 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
708
        if (random_.NextShort(0, num_machines) == rank) {
Guolin Ke's avatar
Guolin Ke committed
709
710
711
712
713
714
715
716
717
718
719
720
721
722
          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) {
723
724
          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
725
726
727
728
        }
        if (line_idx >= query_boundaries[qid + 1]) {
          // if is new query
          is_query_used = false;
Guolin Ke's avatar
Guolin Ke committed
729
          if (random_.NextShort(0, num_machines) == rank) {
Guolin Ke's avatar
Guolin Ke committed
730
731
732
733
734
735
736
737
738
739
740
741
            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
742
  int sample_cnt = config_.bin_construct_sample_cnt;
743
744
  if (static_cast<size_t>(sample_cnt) > data.size()) {
    sample_cnt = static_cast<int>(data.size());
745
  }
746
  auto sample_indices = random_.Sample(static_cast<int>(data.size()), sample_cnt);
Guolin Ke's avatar
Guolin Ke committed
747
  std::vector<std::string> out(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
748
749
  for (size_t i = 0; i < sample_indices.size(); ++i) {
    const size_t idx = sample_indices[i];
Guolin Ke's avatar
Guolin Ke committed
750
    out[i] = data[idx];
Guolin Ke's avatar
Guolin Ke committed
751
752
753
754
755
  }
  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
756
757
  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
758
  std::vector<std::string> out_data;
Guolin Ke's avatar
Guolin Ke committed
759
  if (num_machines == 1 || config_.pre_partition) {
Guolin Ke's avatar
Guolin Ke committed
760
761
762
763
764
765
766
767
    *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
768
        if (random_.NextShort(0, num_machines) == rank) {
Guolin Ke's avatar
Guolin Ke committed
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
          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
789
          if (random_.NextShort(0, num_machines) == rank) {
Guolin Ke's avatar
Guolin Ke committed
790
791
792
793
794
795
796
797
798
799
800
801
            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
802

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

Guolin Ke's avatar
Guolin Ke committed
823
  dataset->feature_groups_.clear();
Guolin Ke's avatar
Guolin Ke committed
824

825
826
  if (feature_names_.empty()) {
    // -1 means doesn't use this feature
Guolin Ke's avatar
Guolin Ke committed
827
828
    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);
829
830
831
832
  } 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
833
834
835
836
837
838
839

  // 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
840
  if (feature_names_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
841
842
843
844
845
846
    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());
    }
  }
847
  dataset->set_feature_names(feature_names_);
Guolin Ke's avatar
Guolin Ke committed
848
  std::vector<std::unique_ptr<BinMapper>> bin_mappers(dataset->num_total_features_);
Guolin Ke's avatar
Guolin Ke committed
849
  const data_size_t filter_cnt = static_cast<data_size_t>(
Guolin Ke's avatar
Guolin Ke committed
850
    static_cast<double>(config_.min_data_in_leaf* sample_data.size()) / dataset->num_data_);
Guolin Ke's avatar
Guolin Ke committed
851

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

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

    start[0] = 0;
    for (int i = 0; i < num_machines - 1; ++i) {
Guolin Ke's avatar
Guolin Ke committed
889
      len[i] = std::min(step, num_total_features - start[i]);
Guolin Ke's avatar
Guolin Ke committed
890
891
      start[i + 1] = start[i] + len[i];
    }
Guolin Ke's avatar
Guolin Ke committed
892
    len[num_machines - 1] = num_total_features - start[num_machines - 1];
893
    OMP_INIT_EX();
894
    #pragma omp parallel for schedule(guided)
895
    for (int i = 0; i < len[rank]; ++i) {
896
      OMP_LOOP_EX_BEGIN();
897
898
899
900
901
902
903
904
      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
905
      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
906
                              sample_data.size(), config_.max_bin, config_.min_data_in_bin, filter_cnt, bin_type, config_.use_missing, config_.zero_as_missing);
907
      OMP_LOOP_EX_END();
908
    }
909
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
910
911
912
913
914
915
916
    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);
917
    // get size of bin mapper with max_bin size
918
    int type_size = BinMapper::SizeForSpecificBin(max_bin);
Guolin Ke's avatar
Guolin Ke committed
919
    // since sizes of different feature may not be same, we expand all bin mapper to type_size
Guolin Ke's avatar
Guolin Ke committed
920
    comm_size_t buffer_size = type_size * num_total_features;
Guolin Ke's avatar
Guolin Ke committed
921
922
    auto input_buffer = std::vector<char>(buffer_size);
    auto output_buffer = std::vector<char>(buffer_size);
Guolin Ke's avatar
Guolin Ke committed
923
924

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

/*! \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) {
966
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
967
    // if doesn't need to prediction with initial model
968
    #pragma omp parallel for schedule(static) private(oneline_features) firstprivate(tmp_label)
Guolin Ke's avatar
Guolin Ke committed
969
    for (data_size_t i = 0; i < dataset->num_data_; ++i) {
970
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
971
972
973
974
975
      const int tid = omp_get_thread_num();
      oneline_features.clear();
      // parser
      parser->ParseOneLine(text_data[i].c_str(), &oneline_features, &tmp_label);
      // set label
976
      dataset->metadata_.SetLabelAt(i, static_cast<label_t>(tmp_label));
Guolin Ke's avatar
Guolin Ke committed
977
978
979
980
981
982
      // 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) {
983
        if (inner_data.first >= dataset->num_total_features_) { continue; }
Guolin Ke's avatar
Guolin Ke committed
984
985
986
        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
987
988
989
          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
990
991
        } else {
          if (inner_data.first == weight_idx_) {
992
            dataset->metadata_.SetWeightAt(i, static_cast<label_t>(inner_data.second));
Guolin Ke's avatar
Guolin Ke committed
993
994
995
996
997
          } else if (inner_data.first == group_idx_) {
            dataset->metadata_.SetQueryAt(i, static_cast<data_size_t>(inner_data.second));
          }
        }
      }
998
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
999
    }
1000
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1001
  } else {
1002
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1003
    // if need to prediction with initial model
1004
    std::vector<double> init_score(dataset->num_data_ * num_class_);
1005
    #pragma omp parallel for schedule(static) private(oneline_features) firstprivate(tmp_label)
Guolin Ke's avatar
Guolin Ke committed
1006
    for (data_size_t i = 0; i < dataset->num_data_; ++i) {
1007
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1008
1009
1010
1011
1012
      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
1013
1014
      std::vector<double> oneline_init_score(num_class_);
      predict_fun_(oneline_features, oneline_init_score.data());
1015
      for (int k = 0; k < num_class_; ++k) {
1016
        init_score[k * dataset->num_data_ + i] = static_cast<double>(oneline_init_score[k]);
Guolin Ke's avatar
Guolin Ke committed
1017
1018
      }
      // set label
1019
      dataset->metadata_.SetLabelAt(i, static_cast<label_t>(tmp_label));
Guolin Ke's avatar
Guolin Ke committed
1020
1021
1022
1023
1024
1025
      // 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) {
1026
        if (inner_data.first >= dataset->num_total_features_) { continue; }
Guolin Ke's avatar
Guolin Ke committed
1027
1028
1029
        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
1030
1031
          int group = dataset->feature2group_[feature_idx];
          int sub_feature = dataset->feature2subfeature_[feature_idx];
1032
          dataset->feature_groups_[group]->PushData(tid, sub_feature, i, inner_data.second);
Guolin Ke's avatar
Guolin Ke committed
1033
1034
        } else {
          if (inner_data.first == weight_idx_) {
1035
            dataset->metadata_.SetWeightAt(i, static_cast<label_t>(inner_data.second));
Guolin Ke's avatar
Guolin Ke committed
1036
1037
1038
1039
1040
          } else if (inner_data.first == group_idx_) {
            dataset->metadata_.SetQueryAt(i, static_cast<data_size_t>(inner_data.second));
          }
        }
      }
1041
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1042
    }
1043
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1044
    // metadata_ will manage space of init_score
1045
    dataset->metadata_.SetInitScore(init_score.data(), dataset->num_data_ * num_class_);
Guolin Ke's avatar
Guolin Ke committed
1046
  }
Guolin Ke's avatar
Guolin Ke committed
1047
  dataset->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1048
1049
1050
1051
1052
1053
  // 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) {
1054
  std::vector<double> init_score;
Guolin Ke's avatar
Guolin Ke committed
1055
  if (predict_fun_ != nullptr) {
1056
    init_score = std::vector<double>(dataset->num_data_ * num_class_);
Guolin Ke's avatar
Guolin Ke committed
1057
1058
1059
1060
1061
1062
  }
  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;
1063
    OMP_INIT_EX();
1064
    #pragma omp parallel for schedule(static) private(oneline_features) firstprivate(tmp_label)
Guolin Ke's avatar
Guolin Ke committed
1065
    for (data_size_t i = 0; i < static_cast<data_size_t>(lines.size()); ++i) {
1066
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1067
1068
1069
1070
1071
      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
1072
      if (!init_score.empty()) {
Guolin Ke's avatar
Guolin Ke committed
1073
1074
        std::vector<double> oneline_init_score(num_class_);
        predict_fun_(oneline_features, oneline_init_score.data());
1075
        for (int k = 0; k < num_class_; ++k) {
1076
          init_score[k * dataset->num_data_ + start_idx + i] = static_cast<double>(oneline_init_score[k]);
Guolin Ke's avatar
Guolin Ke committed
1077
1078
1079
        }
      }
      // set label
1080
      dataset->metadata_.SetLabelAt(start_idx + i, static_cast<label_t>(tmp_label));
Guolin Ke's avatar
Guolin Ke committed
1081
1082
      // push data
      for (auto& inner_data : oneline_features) {
1083
        if (inner_data.first >= dataset->num_total_features_) { continue; }
Guolin Ke's avatar
Guolin Ke committed
1084
1085
1086
        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
1087
1088
1089
          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
1090
1091
        } else {
          if (inner_data.first == weight_idx_) {
1092
            dataset->metadata_.SetWeightAt(start_idx + i, static_cast<label_t>(inner_data.second));
Guolin Ke's avatar
Guolin Ke committed
1093
1094
1095
1096
1097
          } else if (inner_data.first == group_idx_) {
            dataset->metadata_.SetQueryAt(start_idx + i, static_cast<data_size_t>(inner_data.second));
          }
        }
      }
1098
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1099
    }
1100
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1101
  };
Guolin Ke's avatar
Guolin Ke committed
1102
  TextReader<data_size_t> text_reader(filename, config_.header);
Guolin Ke's avatar
Guolin Ke committed
1103
  if (!used_data_indices.empty()) {
Guolin Ke's avatar
Guolin Ke committed
1104
1105
1106
1107
1108
1109
1110
1111
    // 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
1112
  if (!init_score.empty()) {
1113
    dataset->metadata_.SetInitScore(init_score.data(), dataset->num_data_ * num_class_);
Guolin Ke's avatar
Guolin Ke committed
1114
  }
Guolin Ke's avatar
Guolin Ke committed
1115
  dataset->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1116
1117
1118
}

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

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

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

  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);
1137
  size_t read_cnt = reader->Read(buffer.data(), size_of_token);
1138
1139
  if (read_cnt == size_of_token
      && std::string(buffer.data()) == std::string(Dataset::binary_file_token)) {
1140
    return bin_filename;
Guolin Ke's avatar
Guolin Ke committed
1141
  } else {
1142
    return std::string();
Guolin Ke's avatar
Guolin Ke committed
1143
  }
1144

Guolin Ke's avatar
Guolin Ke committed
1145
1146
}

Guolin Ke's avatar
Guolin Ke committed
1147
}