metadata.cpp 17 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
4
5
6
7
8
9
#include <LightGBM/dataset.h>

#include <LightGBM/utils/common.h>

#include <vector>
#include <string>

namespace LightGBM {

Guolin Ke's avatar
Guolin Ke committed
10
Metadata::Metadata() {
Guolin Ke's avatar
Guolin Ke committed
11
12
}

Guolin Ke's avatar
Guolin Ke committed
13
void Metadata::Init(const char * data_filename, const int num_class) {
Guolin Ke's avatar
Guolin Ke committed
14
  data_filename_ = data_filename;
15
  num_class_ = num_class;
Guolin Ke's avatar
Guolin Ke committed
16
17
18
19
20
21
22
23
24
25
26
27
28
  // for lambdarank, it needs query data for partition data in parallel learning
  LoadQueryBoundaries();
  LoadWeights();
  LoadQueryWeights();
  LoadInitialScore();
}



Metadata::~Metadata() {
}


29
void Metadata::Init(data_size_t num_data, int num_class, int weight_idx, int query_idx) {
Guolin Ke's avatar
Guolin Ke committed
30
  num_data_ = num_data;
31
  num_class_ = num_class;
Guolin Ke's avatar
Guolin Ke committed
32
  label_ = std::vector<float>(num_data_);
Guolin Ke's avatar
Guolin Ke committed
33
  if (weight_idx >= 0) {
Guolin Ke's avatar
Guolin Ke committed
34
    if (weights_.size() > 0) {
35
      Log::Info("Using weights in data file, ignoring the additional weights file");
Guolin Ke's avatar
Guolin Ke committed
36
      weights_.clear();
Guolin Ke's avatar
Guolin Ke committed
37
    }
Guolin Ke's avatar
Guolin Ke committed
38
    weights_ = std::vector<float>(num_data_);
Guolin Ke's avatar
Guolin Ke committed
39
    num_weights_ = num_data_;
Guolin Ke's avatar
Guolin Ke committed
40
    std::fill(weights_.begin(), weights_.end(), 0.0f);
Guolin Ke's avatar
Guolin Ke committed
41
42
  }
  if (query_idx >= 0) {
Guolin Ke's avatar
Guolin Ke committed
43
    if (query_boundaries_.size() >  0) {
44
      Log::Info("Using query id in data file, ignoring the additional query file");
Guolin Ke's avatar
Guolin Ke committed
45
      query_boundaries_.clear();
Guolin Ke's avatar
Guolin Ke committed
46
    }
Guolin Ke's avatar
Guolin Ke committed
47
48
49
    if (query_weights_.size() > 0) { query_weights_.clear(); }
    queries_ = std::vector<data_size_t>(num_data_);
    std::fill(queries_.begin(), queries_.end(), 0);
Guolin Ke's avatar
Guolin Ke committed
50
  }
Guolin Ke's avatar
Guolin Ke committed
51
52
}

wxchan's avatar
wxchan committed
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
void Metadata::Init(const Metadata& fullset, const data_size_t* used_indices, data_size_t num_used_indices) {
  num_data_ = num_used_indices;
  num_class_ = fullset.num_class_;

  label_ = std::vector<float>(num_used_indices);
  for (data_size_t i = 0; i < num_used_indices; i++) {
    label_[i] = fullset.label_[used_indices[i]];
  }

  if (fullset.weights_.size() > 0) {
    weights_ = std::vector<float>(num_used_indices);
    num_weights_ = num_used_indices;
    for (data_size_t i = 0; i < num_used_indices; i++) {
      weights_[i] = fullset.weights_[used_indices[i]];
    }
  } else {
    num_weights_ = 0;
  }

  if (fullset.init_score_.size() > 0) {
    init_score_ = std::vector<float>(num_used_indices);
    num_init_score_ = num_used_indices;
    for (data_size_t i = 0; i < num_used_indices; i++) {
      init_score_[i] = fullset.init_score_[used_indices[i]];
    }
  } else {
    num_init_score_ = 0;
  }

  if (fullset.query_boundaries_.size() > 0) {
    std::vector<data_size_t> used_query;
    data_size_t data_idx = 0;
    for (data_size_t qid = 0; qid < num_queries_ && data_idx < num_used_indices; ++qid) {
      data_size_t start = fullset.query_boundaries_[qid];
      data_size_t end = fullset.query_boundaries_[qid + 1];
      data_size_t len = end - start;
      if (used_indices[data_idx] > start) {
        continue;
      } else if (used_indices[data_idx] == start) {
        if (num_used_indices >= data_idx + len && used_indices[data_idx + len - 1] == end - 1) {
          used_query.push_back(qid);
          data_idx += len;
        } else {
          Log::Fatal("Data partition error, data didn't match queries");
        }
      } else {
        Log::Fatal("Data partition error, data didn't match queries");
      }
    }
    query_boundaries_ = std::vector<data_size_t>(used_query.size() + 1);
    num_queries_ = static_cast<data_size_t>(used_query.size());
    query_boundaries_[0] = 0;
    for (data_size_t i = 0; i < num_queries_; ++i) {
      data_size_t qid = used_query[i];
      data_size_t len = fullset.query_boundaries_[qid + 1] - fullset.query_boundaries_[qid];
      query_boundaries_[i + 1] = query_boundaries_[i] + len;
    }
  } else {
    num_queries_ = 0;
  }

}

Guolin Ke's avatar
Guolin Ke committed
116
117
118
119
void Metadata::PartitionLabel(const std::vector<data_size_t>& used_indices) {
  if (used_indices.size() <= 0) {
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
120
  auto old_label = label_;
Guolin Ke's avatar
Guolin Ke committed
121
  num_data_ = static_cast<data_size_t>(used_indices.size());
Guolin Ke's avatar
Guolin Ke committed
122
  label_ = std::vector<float>(num_data_);
Guolin Ke's avatar
Guolin Ke committed
123
124
125
  for (data_size_t i = 0; i < num_data_; ++i) {
    label_[i] = old_label[used_indices[i]];
  }
Guolin Ke's avatar
Guolin Ke committed
126
  old_label.clear();
Guolin Ke's avatar
Guolin Ke committed
127
128
129
130
}

void Metadata::CheckOrPartition(data_size_t num_all_data, const std::vector<data_size_t>& used_data_indices) {
  if (used_data_indices.size() == 0) {
Guolin Ke's avatar
Guolin Ke committed
131
    if (queries_.size() > 0) {
Guolin Ke's avatar
Guolin Ke committed
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
      // need convert query_id to boundaries
      std::vector<data_size_t> tmp_buffer;
      data_size_t last_qid = -1;
      data_size_t cur_cnt = 0;
      for (data_size_t i = 0; i < num_data_; ++i) {
        if (last_qid != queries_[i]) {
          if (cur_cnt > 0) {
            tmp_buffer.push_back(cur_cnt);
          }
          cur_cnt = 0;
          last_qid = queries_[i];
        }
        ++cur_cnt;
      }
      tmp_buffer.push_back(cur_cnt);
Guolin Ke's avatar
Guolin Ke committed
147
      query_boundaries_ = std::vector<data_size_t>(tmp_buffer.size() + 1);
Guolin Ke's avatar
Guolin Ke committed
148
149
150
151
152
153
      num_queries_ = static_cast<data_size_t>(tmp_buffer.size());
      query_boundaries_[0] = 0;
      for (size_t i = 0; i < tmp_buffer.size(); ++i) {
        query_boundaries_[i + 1] = query_boundaries_[i] + tmp_buffer[i];
      }
      LoadQueryWeights();
Guolin Ke's avatar
Guolin Ke committed
154
      queries_.clear();
Guolin Ke's avatar
Guolin Ke committed
155
    }
Guolin Ke's avatar
Guolin Ke committed
156
    // check weights
Guolin Ke's avatar
Guolin Ke committed
157
158
    if (weights_.size() > 0 && num_weights_ != num_data_) {
      weights_.clear();
Guolin Ke's avatar
Guolin Ke committed
159
      num_weights_ = 0;
Guolin Ke's avatar
Guolin Ke committed
160
      Log::Fatal("Weights size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
161
162
163
    }

    // check query boundries
Guolin Ke's avatar
Guolin Ke committed
164
165
    if (query_boundaries_.size() > 0 && query_boundaries_[num_queries_] != num_data_) {
      query_boundaries_.clear();
Guolin Ke's avatar
Guolin Ke committed
166
      num_queries_ = 0;
Guolin Ke's avatar
Guolin Ke committed
167
      Log::Fatal("Query size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
168
169
170
    }

    // contain initial score file
Guolin Ke's avatar
Guolin Ke committed
171
172
    if (init_score_.size() > 0 && num_init_score_ != num_data_) {
      init_score_.clear();
Guolin Ke's avatar
Guolin Ke committed
173
      num_init_score_ = 0;
Guolin Ke's avatar
Guolin Ke committed
174
      Log::Fatal("Initial score size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
175
176
177
178
    }
  } else {
    data_size_t num_used_data = static_cast<data_size_t>(used_data_indices.size());
    // check weights
Guolin Ke's avatar
Guolin Ke committed
179
180
    if (weights_.size() > 0 && num_weights_ != num_all_data) {
      weights_.clear();
Guolin Ke's avatar
Guolin Ke committed
181
      num_weights_ = 0;
Guolin Ke's avatar
Guolin Ke committed
182
      Log::Fatal("Weights size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
183
184
    }
    // check query boundries
Guolin Ke's avatar
Guolin Ke committed
185
186
    if (query_boundaries_.size() > 0 && query_boundaries_[num_queries_] != num_all_data) {
      query_boundaries_.clear();
Guolin Ke's avatar
Guolin Ke committed
187
      num_queries_ = 0;
Guolin Ke's avatar
Guolin Ke committed
188
      Log::Fatal("Query size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
189
190
191
    }

    // contain initial score file
Guolin Ke's avatar
Guolin Ke committed
192
193
    if (init_score_.size() > 0 && num_init_score_ != num_all_data) {
      init_score_.clear();
Guolin Ke's avatar
Guolin Ke committed
194
      num_init_score_ = 0;
Guolin Ke's avatar
Guolin Ke committed
195
      Log::Fatal("Initial score size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
196
197
198
    }

    // get local weights
Guolin Ke's avatar
Guolin Ke committed
199
200
    if (weights_.size() > 0) {
      auto old_weights = weights_;
Guolin Ke's avatar
Guolin Ke committed
201
      num_weights_ = num_data_;
Guolin Ke's avatar
Guolin Ke committed
202
      weights_ = std::vector<float>(num_data_);
Guolin Ke's avatar
Guolin Ke committed
203
204
205
      for (size_t i = 0; i < used_data_indices.size(); ++i) {
        weights_[i] = old_weights[used_data_indices[i]];
      }
Guolin Ke's avatar
Guolin Ke committed
206
      old_weights.clear();
Guolin Ke's avatar
Guolin Ke committed
207
208
209
    }

    // get local query boundaries
Guolin Ke's avatar
Guolin Ke committed
210
    if (query_boundaries_.size() > 0) {
Guolin Ke's avatar
Guolin Ke committed
211
212
213
214
215
216
217
218
219
220
221
222
223
      std::vector<data_size_t> used_query;
      data_size_t data_idx = 0;
      for (data_size_t qid = 0; qid < num_queries_ && data_idx < num_used_data; ++qid) {
        data_size_t start = query_boundaries_[qid];
        data_size_t end = query_boundaries_[qid + 1];
        data_size_t len = end - start;
        if (used_data_indices[data_idx] > start) {
          continue;
        } else if (used_data_indices[data_idx] == start) {
          if (num_used_data >= data_idx + len && used_data_indices[data_idx + len - 1] == end - 1) {
            used_query.push_back(qid);
            data_idx += len;
          } else {
Guolin Ke's avatar
Guolin Ke committed
224
            Log::Fatal("Data partition error, data didn't match queries");
Guolin Ke's avatar
Guolin Ke committed
225
226
          }
        } else {
Guolin Ke's avatar
Guolin Ke committed
227
          Log::Fatal("Data partition error, data didn't match queries");
Guolin Ke's avatar
Guolin Ke committed
228
229
        }
      }
Guolin Ke's avatar
Guolin Ke committed
230
231
      auto old_query_boundaries = query_boundaries_;
      query_boundaries_ = std::vector<data_size_t>(used_query.size() + 1);
Guolin Ke's avatar
Guolin Ke committed
232
233
234
235
236
237
238
      num_queries_ = static_cast<data_size_t>(used_query.size());
      query_boundaries_[0] = 0;
      for (data_size_t i = 0; i < num_queries_; ++i) {
        data_size_t qid = used_query[i];
        data_size_t len = old_query_boundaries[qid + 1] - old_query_boundaries[qid];
        query_boundaries_[i + 1] = query_boundaries_[i] + len;
      }
Guolin Ke's avatar
Guolin Ke committed
239
      old_query_boundaries.clear();
Guolin Ke's avatar
Guolin Ke committed
240
241
242
    }

    // get local initial scores
Guolin Ke's avatar
Guolin Ke committed
243
244
    if (init_score_.size() > 0) {
      auto old_scores = init_score_;
Guolin Ke's avatar
Guolin Ke committed
245
      num_init_score_ = num_data_;
Guolin Ke's avatar
Guolin Ke committed
246
      init_score_ = std::vector<float>(num_init_score_ * num_class_);
247
248
249
      for (int k = 0; k < num_class_; ++k){
        for (size_t i = 0; i < used_data_indices.size(); ++i) {
          init_score_[k * num_data_ + i] = old_scores[k * num_all_data + used_data_indices[i]];
250
        }
Guolin Ke's avatar
Guolin Ke committed
251
      }
Guolin Ke's avatar
Guolin Ke committed
252
      old_scores.clear();
Guolin Ke's avatar
Guolin Ke committed
253
254
255
256
257
258
259
260
    }

    // re-load query weight
    LoadQueryWeights();
  }
}


Guolin Ke's avatar
Guolin Ke committed
261
void Metadata::SetInitScore(const float* init_score, data_size_t len) {
wxchan's avatar
wxchan committed
262
263
264
265
266
267
268
  std::lock_guard<std::mutex> lock(mutex_);
  // save to nullptr
  if (init_score == nullptr || len == 0) {
    init_score_.clear();
    num_init_score_ = 0;
    return;
  }
269
  if (len != num_data_ * num_class_) {
270
    Log::Fatal("Initial score size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
271
  }
Guolin Ke's avatar
Guolin Ke committed
272
  if (init_score_.size() > 0) { init_score_.clear(); }
Guolin Ke's avatar
Guolin Ke committed
273
  num_init_score_ = num_data_;
Guolin Ke's avatar
Guolin Ke committed
274
  init_score_ = std::vector<float>(len);
275
  for (data_size_t i = 0; i < len; ++i) {
Guolin Ke's avatar
Guolin Ke committed
276
277
    init_score_[i] = init_score[i];
  }
Guolin Ke's avatar
Guolin Ke committed
278
279
}

Guolin Ke's avatar
Guolin Ke committed
280
void Metadata::SetLabel(const float* label, data_size_t len) {
wxchan's avatar
wxchan committed
281
282
283
284
  std::lock_guard<std::mutex> lock(mutex_);
  if (label == nullptr) {
    Log::Fatal("label cannot be nullptr");
  }
Guolin Ke's avatar
Guolin Ke committed
285
286
287
  if (num_data_ != len) {
    Log::Fatal("len of label is not same with #data");
  }
Guolin Ke's avatar
Guolin Ke committed
288
289
  if (label_.size() > 0) { label_.clear(); }
  label_ = std::vector<float>(num_data_);
Guolin Ke's avatar
Guolin Ke committed
290
291
292
293
294
295
  for (data_size_t i = 0; i < num_data_; ++i) {
    label_[i] = label[i];
  }
}

void Metadata::SetWeights(const float* weights, data_size_t len) {
wxchan's avatar
wxchan committed
296
297
298
299
300
301
302
  std::lock_guard<std::mutex> lock(mutex_);
  // save to nullptr
  if (weights == nullptr || len == 0) {
    weights_.clear();
    num_weights_ = 0;
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
303
304
305
  if (num_data_ != len) {
    Log::Fatal("len of weights is not same with #data");
  }
Guolin Ke's avatar
Guolin Ke committed
306
  if (weights_.size() > 0) { weights_.clear(); }
Guolin Ke's avatar
Guolin Ke committed
307
  num_weights_ = num_data_;
Guolin Ke's avatar
Guolin Ke committed
308
  weights_ = std::vector<float>(num_weights_);
Guolin Ke's avatar
Guolin Ke committed
309
310
311
312
313
314
315
  for (data_size_t i = 0; i < num_weights_; ++i) {
    weights_[i] = weights[i];
  }
  LoadQueryWeights();
}

void Metadata::SetQueryBoundaries(const data_size_t* query_boundaries, data_size_t len) {
wxchan's avatar
wxchan committed
316
317
318
319
320
321
322
  std::lock_guard<std::mutex> lock(mutex_);
  // save to nullptr
  if (query_boundaries == nullptr || len == 0) {
    query_boundaries_.clear();
    num_queries_ = 0;
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
323
324
325
326
327
328
329
  data_size_t sum = 0;
  for (data_size_t i = 0; i < len; ++i) {
    sum += query_boundaries[i];
  }
  if (num_data_ != sum) {
    Log::Fatal("sum of query counts is not same with #data");
  }
Guolin Ke's avatar
Guolin Ke committed
330
  if (query_boundaries_.size() > 0) { query_boundaries_.clear(); }
Guolin Ke's avatar
Guolin Ke committed
331
  num_queries_ = len;
Guolin Ke's avatar
Guolin Ke committed
332
  query_boundaries_ = std::vector<data_size_t>(num_queries_);
Guolin Ke's avatar
Guolin Ke committed
333
334
335
336
337
338
  for (data_size_t i = 0; i < num_queries_; ++i) {
    query_boundaries_[i] = query_boundaries[i];
  }
  LoadQueryWeights();
}

wxchan's avatar
wxchan committed
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
372
373
374
375
376
377
378
379
void Metadata::SetQueryId(const data_size_t* query_id, data_size_t len) {
  std::lock_guard<std::mutex> lock(mutex_);
  // save to nullptr
  if (query_id == nullptr || len == 0) {
    query_boundaries_.clear();
    queries_.clear();
    num_queries_ = 0;
    return;
  }
  if (num_data_ != len) {
    Log::Fatal("len of query id is not same with #data");
  }
  if (queries_.size() > 0) { queries_.clear(); }
  queries_ = std::vector<data_size_t>(num_data_);
  for (data_size_t i = 0; i < num_weights_; ++i) {
    queries_[i] = query_id[i];
  }
  // need convert query_id to boundaries
  std::vector<data_size_t> tmp_buffer;
  data_size_t last_qid = -1;
  data_size_t cur_cnt = 0;
  for (data_size_t i = 0; i < num_data_; ++i) {
    if (last_qid != queries_[i]) {
      if (cur_cnt > 0) {
        tmp_buffer.push_back(cur_cnt);
      }
      cur_cnt = 0;
      last_qid = queries_[i];
    }
    ++cur_cnt;
  }
  tmp_buffer.push_back(cur_cnt);
  query_boundaries_ = std::vector<data_size_t>(tmp_buffer.size() + 1);
  num_queries_ = static_cast<data_size_t>(tmp_buffer.size());
  query_boundaries_[0] = 0;
  for (size_t i = 0; i < tmp_buffer.size(); ++i) {
    query_boundaries_[i + 1] = query_boundaries_[i] + tmp_buffer[i];
  }
  queries_.clear();
  LoadQueryWeights();
}
Guolin Ke's avatar
Guolin Ke committed
380

Guolin Ke's avatar
Guolin Ke committed
381
382
383
384
385
void Metadata::LoadWeights() {
  num_weights_ = 0;
  std::string weight_filename(data_filename_);
  // default weight file name
  weight_filename.append(".weight");
Guolin Ke's avatar
Guolin Ke committed
386
  TextReader<size_t> reader(weight_filename.c_str(), false);
Guolin Ke's avatar
Guolin Ke committed
387
388
389
390
  reader.ReadAllLines();
  if (reader.Lines().size() <= 0) {
    return;
  }
391
  Log::Info("Loading weights...");
Guolin Ke's avatar
Guolin Ke committed
392
  num_weights_ = static_cast<data_size_t>(reader.Lines().size());
Guolin Ke's avatar
Guolin Ke committed
393
  weights_ = std::vector<float>(num_weights_);
Guolin Ke's avatar
Guolin Ke committed
394
  for (data_size_t i = 0; i < num_weights_; ++i) {
395
    double tmp_weight = 0.0f;
Guolin Ke's avatar
Guolin Ke committed
396
    Common::Atof(reader.Lines()[i].c_str(), &tmp_weight);
397
    weights_[i] = static_cast<float>(tmp_weight);
Guolin Ke's avatar
Guolin Ke committed
398
399
400
401
402
  }
}

void Metadata::LoadInitialScore() {
  num_init_score_ = 0;
Guolin Ke's avatar
Guolin Ke committed
403
404
405
406
  std::string init_score_filename(data_filename_);
  // default weight file name
  init_score_filename.append(".init");
  TextReader<size_t> reader(init_score_filename.c_str(), false);
Guolin Ke's avatar
Guolin Ke committed
407
  reader.ReadAllLines();
Guolin Ke's avatar
Guolin Ke committed
408
409
410
  if (reader.Lines().size() <= 0) {
    return;
  }
411
  Log::Info("Loading initial scores...");
Guolin Ke's avatar
Guolin Ke committed
412
  num_init_score_ = static_cast<data_size_t>(reader.Lines().size());
413

Guolin Ke's avatar
Guolin Ke committed
414
  init_score_ = std::vector<float>(num_init_score_ * num_class_);
415
  double tmp = 0.0f;
416

417
418
419
420
421
422
423
424
  if (num_class_ == 1){
      for (data_size_t i = 0; i < num_init_score_; ++i) {
        Common::Atof(reader.Lines()[i].c_str(), &tmp);
        init_score_[i] = static_cast<float>(tmp);
      }
  } else {
      std::vector<std::string> oneline_init_score;
      for (data_size_t i = 0; i < num_init_score_; ++i) {
425
        oneline_init_score = Common::Split(reader.Lines()[i].c_str(), '\t');
426
427
428
429
430
431
432
433
        if (static_cast<int>(oneline_init_score.size()) != num_class_){
            Log::Fatal("Invalid initial score file. Redundant or insufficient columns.");
        }
        for (int k = 0; k < num_class_; ++k) {
          Common::Atof(oneline_init_score[k].c_str(), &tmp);
          init_score_[k * num_init_score_ + i] = static_cast<float>(tmp);
        }
      }
Guolin Ke's avatar
Guolin Ke committed
434
435
436
437
438
439
440
441
  }
}

void Metadata::LoadQueryBoundaries() {
  num_queries_ = 0;
  std::string query_filename(data_filename_);
  // default query file name
  query_filename.append(".query");
Guolin Ke's avatar
Guolin Ke committed
442
  TextReader<size_t> reader(query_filename.c_str(), false);
Guolin Ke's avatar
Guolin Ke committed
443
444
445
446
  reader.ReadAllLines();
  if (reader.Lines().size() <= 0) {
    return;
  }
447
  Log::Info("Loading query boundaries...");
Guolin Ke's avatar
Guolin Ke committed
448
  query_boundaries_ = std::vector<data_size_t>(reader.Lines().size() + 1);
Guolin Ke's avatar
Guolin Ke committed
449
450
451
452
453
454
455
456
457
458
  num_queries_ = static_cast<data_size_t>(reader.Lines().size());
  query_boundaries_[0] = 0;
  for (size_t i = 0; i < reader.Lines().size(); ++i) {
    int tmp_cnt;
    Common::Atoi(reader.Lines()[i].c_str(), &tmp_cnt);
    query_boundaries_[i + 1] = query_boundaries_[i] + static_cast<data_size_t>(tmp_cnt);
  }
}

void Metadata::LoadQueryWeights() {
Guolin Ke's avatar
Guolin Ke committed
459
  if (weights_.size() == 0 || query_boundaries_.size() == 0) {
Guolin Ke's avatar
Guolin Ke committed
460
461
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
462
  query_weights_.clear();
463
  Log::Info("Loading query weights...");
Guolin Ke's avatar
Guolin Ke committed
464
  query_weights_ = std::vector<float>(num_queries_);
Guolin Ke's avatar
Guolin Ke committed
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
  for (data_size_t i = 0; i < num_queries_; ++i) {
    query_weights_[i] = 0.0f;
    for (data_size_t j = query_boundaries_[i]; j < query_boundaries_[i + 1]; ++j) {
      query_weights_[i] += weights_[j];
    }
    query_weights_[i] /= (query_boundaries_[i + 1] - query_boundaries_[i]);
  }
}

void Metadata::LoadFromMemory(const void* memory) {
  const char* mem_ptr = reinterpret_cast<const char*>(memory);

  num_data_ = *(reinterpret_cast<const data_size_t*>(mem_ptr));
  mem_ptr += sizeof(num_data_);
  num_weights_ = *(reinterpret_cast<const data_size_t*>(mem_ptr));
  mem_ptr += sizeof(num_weights_);
  num_queries_ = *(reinterpret_cast<const data_size_t*>(mem_ptr));
  mem_ptr += sizeof(num_queries_);

Guolin Ke's avatar
Guolin Ke committed
484
485
486
  if (label_.size() > 0) { label_.clear(); }
  label_ = std::vector<float>(num_data_);
  std::memcpy(label_.data(), mem_ptr, sizeof(float)*num_data_);
Guolin Ke's avatar
Guolin Ke committed
487
  mem_ptr += sizeof(float)*num_data_;
Guolin Ke's avatar
Guolin Ke committed
488
489

  if (num_weights_ > 0) {
Guolin Ke's avatar
Guolin Ke committed
490
491
492
    if (weights_.size() > 0) { weights_.clear(); }
    weights_ = std::vector<float>(num_weights_);
    std::memcpy(weights_.data(), mem_ptr, sizeof(float)*num_weights_);
Guolin Ke's avatar
Guolin Ke committed
493
494
495
    mem_ptr += sizeof(float)*num_weights_;
  }
  if (num_queries_ > 0) {
Guolin Ke's avatar
Guolin Ke committed
496
497
498
    if (query_boundaries_.size() > 0) { query_boundaries_.clear(); }
    query_boundaries_ = std::vector<data_size_t>(num_queries_ + 1);
    std::memcpy(query_boundaries_.data(), mem_ptr, sizeof(data_size_t)*(num_queries_ + 1));
Guolin Ke's avatar
Guolin Ke committed
499
500
    mem_ptr += sizeof(data_size_t)*(num_queries_ + 1);
  }
Guolin Ke's avatar
Guolin Ke committed
501
  LoadQueryWeights();
Guolin Ke's avatar
Guolin Ke committed
502
503
504
505
506
507
}

void Metadata::SaveBinaryToFile(FILE* file) const {
  fwrite(&num_data_, sizeof(num_data_), 1, file);
  fwrite(&num_weights_, sizeof(num_weights_), 1, file);
  fwrite(&num_queries_, sizeof(num_queries_), 1, file);
Guolin Ke's avatar
Guolin Ke committed
508
509
510
  fwrite(label_.data(), sizeof(float), num_data_, file);
  if (weights_.size() > 0) {
    fwrite(weights_.data(), sizeof(float), num_weights_, file);
Guolin Ke's avatar
Guolin Ke committed
511
  }
Guolin Ke's avatar
Guolin Ke committed
512
513
  if (query_boundaries_.size() > 0) {
    fwrite(query_boundaries_.data(), sizeof(data_size_t), num_queries_ + 1, file);
Guolin Ke's avatar
Guolin Ke committed
514
515
516
517
518
519
520
521
  }

}

size_t Metadata::SizesInByte() const  {
  size_t size = sizeof(num_data_) + sizeof(num_weights_)
    + sizeof(num_queries_);
  size += sizeof(float) * num_data_;
Guolin Ke's avatar
Guolin Ke committed
522
  if (weights_.size() > 0) {
Guolin Ke's avatar
Guolin Ke committed
523
524
    size += sizeof(float) * num_weights_;
  }
Guolin Ke's avatar
Guolin Ke committed
525
  if (query_boundaries_.size() > 0) {
Guolin Ke's avatar
Guolin Ke committed
526
527
528
529
530
531
532
    size += sizeof(data_size_t) * (num_queries_ + 1);
  }
  return size;
}


}  // namespace LightGBM