metadata.cpp 28.8 KB
Newer Older
1
2
3
4
/*!
 * Copyright (c) 2016 Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See LICENSE file in the project root for license information.
 */
5
6
#include <LightGBM/dataset.h>
#include <LightGBM/utils/common.h>
Guolin Ke's avatar
Guolin Ke committed
7

8
#include <set>
Guolin Ke's avatar
Guolin Ke committed
9
#include <string>
10
#include <vector>
Guolin Ke's avatar
Guolin Ke committed
11
12
13

namespace LightGBM {

Guolin Ke's avatar
Guolin Ke committed
14
Metadata::Metadata() {
Guolin Ke's avatar
Guolin Ke committed
15
16
17
18
  num_weights_ = 0;
  num_init_score_ = 0;
  num_data_ = 0;
  num_queries_ = 0;
19
  num_positions_ = 0;
20
  weight_load_from_file_ = false;
21
  position_load_from_file_ = false;
22
23
  query_load_from_file_ = false;
  init_score_load_from_file_ = false;
24
  #ifdef USE_CUDA
25
  cuda_metadata_ = nullptr;
26
  #endif  // USE_CUDA
Guolin Ke's avatar
Guolin Ke committed
27
28
}

29
void Metadata::Init(const char* data_filename) {
Guolin Ke's avatar
Guolin Ke committed
30
  data_filename_ = data_filename;
31
  // for lambdarank, it needs query data for partition data in distributed learning
Guolin Ke's avatar
Guolin Ke committed
32
33
  LoadQueryBoundaries();
  LoadWeights();
34
  LoadPositions();
35
  CalculateQueryWeights();
36
  LoadInitialScore(data_filename_);
Guolin Ke's avatar
Guolin Ke committed
37
38
39
40
41
}

Metadata::~Metadata() {
}

42
void Metadata::Init(data_size_t num_data, int weight_idx, int query_idx) {
Guolin Ke's avatar
Guolin Ke committed
43
  num_data_ = num_data;
44
  label_ = std::vector<label_t>(num_data_);
Guolin Ke's avatar
Guolin Ke committed
45
  if (weight_idx >= 0) {
Guolin Ke's avatar
Guolin Ke committed
46
    if (!weights_.empty()) {
47
      Log::Info("Using weights in data file, ignoring the additional weights file");
Guolin Ke's avatar
Guolin Ke committed
48
      weights_.clear();
Guolin Ke's avatar
Guolin Ke committed
49
    }
50
    weights_ = std::vector<label_t>(num_data_, 0.0f);
Guolin Ke's avatar
Guolin Ke committed
51
    num_weights_ = num_data_;
52
    weight_load_from_file_ = false;
Guolin Ke's avatar
Guolin Ke committed
53
54
  }
  if (query_idx >= 0) {
Guolin Ke's avatar
Guolin Ke committed
55
    if (!query_boundaries_.empty()) {
56
      Log::Info("Using query id in data file, ignoring the additional query file");
Guolin Ke's avatar
Guolin Ke committed
57
      query_boundaries_.clear();
Guolin Ke's avatar
Guolin Ke committed
58
    }
Guolin Ke's avatar
Guolin Ke committed
59
    if (!query_weights_.empty()) { query_weights_.clear(); }
60
    queries_ = std::vector<data_size_t>(num_data_, 0);
61
    query_load_from_file_ = false;
Guolin Ke's avatar
Guolin Ke committed
62
  }
Guolin Ke's avatar
Guolin Ke committed
63
64
}

65
66
67
68
void Metadata::InitByReference(data_size_t num_data, const Metadata* reference) {
  int has_weights = reference->num_weights_ > 0;
  int has_init_scores = reference->num_init_score_ > 0;
  int has_queries = reference->num_queries_ > 0;
69
  int nclasses = reference->num_init_score_classes();
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
  Init(num_data, has_weights, has_init_scores, has_queries, nclasses);
}

void Metadata::Init(data_size_t num_data, int32_t has_weights, int32_t has_init_scores, int32_t has_queries, int32_t nclasses) {
  num_data_ = num_data;
  label_ = std::vector<label_t>(num_data_);
  if (has_weights) {
    if (!weights_.empty()) {
      Log::Fatal("Calling Init() on Metadata weights that have already been initialized");
    }
    weights_.resize(num_data_, 0.0f);
    num_weights_ = num_data_;
    weight_load_from_file_ = false;
  }
  if (has_init_scores) {
    if (!init_score_.empty()) {
      Log::Fatal("Calling Init() on Metadata initial scores that have already been initialized");
    }
    num_init_score_ = static_cast<int64_t>(num_data) * nclasses;
    init_score_.resize(num_init_score_, 0);
  }
  if (has_queries) {
    if (!query_weights_.empty()) {
      Log::Fatal("Calling Init() on Metadata queries that have already been initialized");
    }
    queries_.resize(num_data_, 0);
    query_load_from_file_ = false;
  }
}

Guolin Ke's avatar
Guolin Ke committed
100
101
102
void Metadata::Init(const Metadata& fullset, const data_size_t* used_indices, data_size_t num_used_indices) {
  num_data_ = num_used_indices;

103
  label_ = std::vector<label_t>(num_used_indices);
104
#pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static, 512) if (num_used_indices >= 1024)
105
  for (data_size_t i = 0; i < num_used_indices; ++i) {
Guolin Ke's avatar
Guolin Ke committed
106
107
108
    label_[i] = fullset.label_[used_indices[i]];
  }

Guolin Ke's avatar
Guolin Ke committed
109
  if (!fullset.weights_.empty()) {
110
    weights_ = std::vector<label_t>(num_used_indices);
Guolin Ke's avatar
Guolin Ke committed
111
    num_weights_ = num_used_indices;
112
#pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static, 512) if (num_used_indices >= 1024)
113
    for (data_size_t i = 0; i < num_used_indices; ++i) {
Guolin Ke's avatar
Guolin Ke committed
114
115
116
117
118
119
      weights_[i] = fullset.weights_[used_indices[i]];
    }
  } else {
    num_weights_ = 0;
  }

Guolin Ke's avatar
Guolin Ke committed
120
  if (!fullset.init_score_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
121
    int num_class = static_cast<int>(fullset.num_init_score_ / fullset.num_data_);
122
    init_score_ = std::vector<double>(static_cast<size_t>(num_used_indices) * num_class);
Guolin Ke's avatar
Guolin Ke committed
123
    num_init_score_ = static_cast<int64_t>(num_used_indices) * num_class;
124
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
125
    for (int k = 0; k < num_class; ++k) {
126
127
128
129
      const size_t offset_dest = static_cast<size_t>(k) * num_data_;
      const size_t offset_src = static_cast<size_t>(k) * fullset.num_data_;
      for (data_size_t i = 0; i < num_used_indices; ++i) {
        init_score_[offset_dest + i] = fullset.init_score_[offset_src + used_indices[i]];
130
      }
Guolin Ke's avatar
Guolin Ke committed
131
132
133
134
135
    }
  } else {
    num_init_score_ = 0;
  }

Guolin Ke's avatar
Guolin Ke committed
136
  if (!fullset.query_boundaries_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
    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
169
void Metadata::PartitionLabel(const std::vector<data_size_t>& used_indices) {
Guolin Ke's avatar
Guolin Ke committed
170
  if (used_indices.empty()) {
Guolin Ke's avatar
Guolin Ke committed
171
172
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
173
  auto old_label = label_;
Guolin Ke's avatar
Guolin Ke committed
174
  num_data_ = static_cast<data_size_t>(used_indices.size());
175
  label_ = std::vector<label_t>(num_data_);
176
#pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static, 512) if (num_data_ >= 1024)
Guolin Ke's avatar
Guolin Ke committed
177
178
179
  for (data_size_t i = 0; i < num_data_; ++i) {
    label_[i] = old_label[used_indices[i]];
  }
Guolin Ke's avatar
Guolin Ke committed
180
  old_label.clear();
Guolin Ke's avatar
Guolin Ke committed
181
182
}

183
184
185
186
187
188
189
190
191
192
void Metadata::CalculateQueryBoundaries() {
  if (!queries_.empty()) {
    // 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);
Guolin Ke's avatar
Guolin Ke committed
193
        }
194
195
        cur_cnt = 0;
        last_qid = queries_[i];
Guolin Ke's avatar
Guolin Ke committed
196
      }
197
      ++cur_cnt;
Guolin Ke's avatar
Guolin Ke committed
198
    }
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
    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];
    }
    CalculateQueryWeights();
    queries_.clear();
  }
}

void Metadata::CheckOrPartition(data_size_t num_all_data, const std::vector<data_size_t>& used_data_indices) {
  if (used_data_indices.empty()) {
     CalculateQueryBoundaries();
Guolin Ke's avatar
Guolin Ke committed
214
    // check weights
Guolin Ke's avatar
Guolin Ke committed
215
    if (!weights_.empty() && num_weights_ != num_data_) {
Guolin Ke's avatar
Guolin Ke committed
216
      weights_.clear();
Guolin Ke's avatar
Guolin Ke committed
217
      num_weights_ = 0;
Guolin Ke's avatar
Guolin Ke committed
218
      Log::Fatal("Weights size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
219
220
    }

221
222
223
224
225
226
227
    // check positions
    if (!positions_.empty() && num_positions_ != num_data_) {
      Log::Fatal("Positions size (%i) doesn't match data size (%i)", num_positions_, num_data_);
      positions_.clear();
      num_positions_ = 0;
    }

Guolin Ke's avatar
Guolin Ke committed
228
    // check query boundries
Guolin Ke's avatar
Guolin Ke committed
229
    if (!query_boundaries_.empty() && query_boundaries_[num_queries_] != num_data_) {
Guolin Ke's avatar
Guolin Ke committed
230
      query_boundaries_.clear();
Guolin Ke's avatar
Guolin Ke committed
231
      num_queries_ = 0;
Guolin Ke's avatar
Guolin Ke committed
232
      Log::Fatal("Query size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
233
234
235
    }

    // contain initial score file
236
    if (!init_score_.empty() && (num_init_score_ % num_data_) != 0) {
Guolin Ke's avatar
Guolin Ke committed
237
      init_score_.clear();
Guolin Ke's avatar
Guolin Ke committed
238
      num_init_score_ = 0;
Guolin Ke's avatar
Guolin Ke committed
239
      Log::Fatal("Initial score size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
240
241
    }
  } else {
242
    if (!queries_.empty()) {
243
      Log::Fatal("Cannot used query_id for distributed training");
244
    }
Guolin Ke's avatar
Guolin Ke committed
245
246
    data_size_t num_used_data = static_cast<data_size_t>(used_data_indices.size());
    // check weights
247
248
249
250
251
252
253
254
255
256
    if (weight_load_from_file_) {
      if (weights_.size() > 0 && num_weights_ != num_all_data) {
        weights_.clear();
        num_weights_ = 0;
        Log::Fatal("Weights size doesn't match data size");
      }
      // get local weights
      if (!weights_.empty()) {
        auto old_weights = weights_;
        num_weights_ = num_data_;
257
        weights_ = std::vector<label_t>(num_data_);
258
#pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static, 512)
259
260
261
262
        for (int i = 0; i < static_cast<int>(used_data_indices.size()); ++i) {
          weights_[i] = old_weights[used_data_indices[i]];
        }
        old_weights.clear();
Guolin Ke's avatar
Guolin Ke committed
263
264
      }
    }
265
266
267
268
269
270
271
272
273
274
275
276
    // check positions
    if (position_load_from_file_) {
      if (positions_.size() > 0 && num_positions_ != num_all_data) {
        positions_.clear();
        num_positions_ = 0;
        Log::Fatal("Positions size (%i) doesn't match data size (%i)", num_positions_, num_data_);
      }
      // get local positions
      if (!positions_.empty()) {
        auto old_positions = positions_;
        num_positions_ = num_data_;
        positions_ = std::vector<data_size_t>(num_data_);
277
        #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static, 512)
278
279
280
281
282
283
        for (int i = 0; i < static_cast<int>(used_data_indices.size()); ++i) {
          positions_[i] = old_positions[used_data_indices[i]];
        }
        old_positions.clear();
      }
    }
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
    if (query_load_from_file_) {
      // check query boundries
      if (!query_boundaries_.empty() && query_boundaries_[num_queries_] != num_all_data) {
        query_boundaries_.clear();
        num_queries_ = 0;
        Log::Fatal("Query size doesn't match data size");
      }
      // get local query boundaries
      if (!query_boundaries_.empty()) {
        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 {
              Log::Fatal("Data partition error, data didn't match queries");
            }
Guolin Ke's avatar
Guolin Ke committed
308
          } else {
Guolin Ke's avatar
Guolin Ke committed
309
            Log::Fatal("Data partition error, data didn't match queries");
Guolin Ke's avatar
Guolin Ke committed
310
311
          }
        }
312
313
314
315
316
317
318
319
320
321
        auto old_query_boundaries = query_boundaries_;
        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 = old_query_boundaries[qid + 1] - old_query_boundaries[qid];
          query_boundaries_[i + 1] = query_boundaries_[i] + len;
        }
        old_query_boundaries.clear();
Guolin Ke's avatar
Guolin Ke committed
322
323
      }
    }
324
325
326
327
328
329
330
    if (init_score_load_from_file_) {
      // contain initial score file
      if (!init_score_.empty() && (num_init_score_ % num_all_data) != 0) {
        init_score_.clear();
        num_init_score_ = 0;
        Log::Fatal("Initial score size doesn't match data size");
      }
Guolin Ke's avatar
Guolin Ke committed
331

332
333
334
335
336
337
      // get local initial scores
      if (!init_score_.empty()) {
        auto old_scores = init_score_;
        int num_class = static_cast<int>(num_init_score_ / num_all_data);
        num_init_score_ = static_cast<int64_t>(num_data_) * num_class;
        init_score_ = std::vector<double>(num_init_score_);
338
#pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
339
        for (int k = 0; k < num_class; ++k) {
340
341
          const size_t offset_dest = static_cast<size_t>(k) * num_data_;
          const size_t offset_src = static_cast<size_t>(k) * num_all_data;
342
          for (size_t i = 0; i < used_data_indices.size(); ++i) {
343
            init_score_[offset_dest + i] = old_scores[offset_src + used_data_indices[i]];
344
          }
345
        }
346
        old_scores.clear();
Guolin Ke's avatar
Guolin Ke committed
347
348
      }
    }
349
350
    // re-calculate query weight
    CalculateQueryWeights();
Guolin Ke's avatar
Guolin Ke committed
351
  }
352
353
354
355
  if (num_queries_ > 0) {
    Log::Debug("Number of queries in %s: %i. Average number of rows per query: %f.",
      data_filename_.c_str(), static_cast<int>(num_queries_), static_cast<double>(num_data_) / num_queries_);
  }
Guolin Ke's avatar
Guolin Ke committed
356
357
}

358
359
360
361
362
363
364
365
366
367
368
void Metadata::SetInitScore(const double* init_score, data_size_t len) {
  std::lock_guard<std::mutex> lock(mutex_);
  // save to nullptr
  if (init_score == nullptr || len == 0) {
    init_score_.clear();
    num_init_score_ = 0;
    return;
  }
  if ((len % num_data_) != 0) {
    Log::Fatal("Initial score size doesn't match data size");
  }
369
  if (init_score_.empty()) { init_score_.resize(len); }
370
  num_init_score_ = len;
371

372
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static, 512) if (num_init_score_ >= 1024)
Guolin Ke's avatar
Guolin Ke committed
373
  for (int64_t i = 0; i < num_init_score_; ++i) {
374
    init_score_[i] = Common::AvoidInf(init_score[i]);
375
  }
376
  init_score_load_from_file_ = false;
377
  #ifdef USE_CUDA
378
379
380
  if (cuda_metadata_ != nullptr) {
    cuda_metadata_->SetInitScore(init_score_.data(), len);
  }
381
  #endif  // USE_CUDA
382
383
}

384
385
386
387
388
389
390
391
392
393
void Metadata::InsertInitScores(const double* init_scores, data_size_t start_index, data_size_t len, data_size_t source_size) {
  if (num_init_score_ <= 0) {
    Log::Fatal("Inserting initial score data into dataset with no initial scores");
  }
  if (start_index + len > num_data_) {
    // Note that len here is row count, not num_init_score, so we compare against num_data
    Log::Fatal("Inserted initial score data is too large for dataset");
  }
  if (init_score_.empty()) { init_score_.resize(num_init_score_); }

394
  int nclasses = num_init_score_classes();
395
396
397
398
399
400
401
402
403
404
405

  for (int32_t col = 0; col < nclasses; ++col) {
    int32_t dest_offset = num_data_ * col + start_index;
    // We need to use source_size here, because len might not equal size (due to a partially loaded dataset)
    int32_t source_offset = source_size * col;
    memcpy(init_score_.data() + dest_offset, init_scores + source_offset, sizeof(double) * len);
  }
  init_score_load_from_file_ = false;
  // CUDA is handled after all insertions are complete
}

406
407
template <typename It>
void Metadata::SetLabelsFromIterator(It first, It last) {
408
  std::lock_guard<std::mutex> lock(mutex_);
409
410
  if (num_data_ != last - first) {
    Log::Fatal("Length of labels differs from the length of #data");
411
  }
412
413
  if (label_.empty()) {
    label_.resize(num_data_);
Guolin Ke's avatar
Guolin Ke committed
414
  }
415

416
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static, 512) if (num_data_ >= 1024)
Guolin Ke's avatar
Guolin Ke committed
417
  for (data_size_t i = 0; i < num_data_; ++i) {
418
    label_[i] = Common::AvoidInf(first[i]);
Guolin Ke's avatar
Guolin Ke committed
419
  }
420

421
  #ifdef USE_CUDA
422
  if (cuda_metadata_ != nullptr) {
423
    cuda_metadata_->SetLabel(label_.data(), label_.size());
424
  }
425
  #endif  // USE_CUDA
Guolin Ke's avatar
Guolin Ke committed
426
427
}

428
429
430
431
432
433
434
435
436
437
438
void Metadata::SetLabel(const label_t* label, data_size_t len) {
  if (label == nullptr) {
    Log::Fatal("label cannot be nullptr");
  }
  SetLabelsFromIterator(label, label + len);
}

void Metadata::SetLabel(const ArrowChunkedArray& array) {
  SetLabelsFromIterator(array.begin<label_t>(), array.end<label_t>());
}

439
440
441
442
443
444
445
446
447
448
449
450
451
452
void Metadata::InsertLabels(const label_t* labels, data_size_t start_index, data_size_t len) {
  if (labels == nullptr) {
    Log::Fatal("label cannot be nullptr");
  }
  if (start_index + len > num_data_) {
    Log::Fatal("Inserted label data is too large for dataset");
  }
  if (label_.empty()) { label_.resize(num_data_); }

  memcpy(label_.data() + start_index, labels, sizeof(label_t) * len);

  // CUDA is handled after all insertions are complete
}

453
void Metadata::SetWeights(const label_t* weights, data_size_t len) {
454
  std::lock_guard<std::mutex> lock(mutex_);
455
456
457
458
459
460
  // save to nullptr
  if (weights == nullptr || len == 0) {
    weights_.clear();
    num_weights_ = 0;
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
461
  if (num_data_ != len) {
462
    Log::Fatal("Length of weights is not same with #data");
Guolin Ke's avatar
Guolin Ke committed
463
  }
464
  if (weights_.empty()) { weights_.resize(num_data_); }
Guolin Ke's avatar
Guolin Ke committed
465
  num_weights_ = num_data_;
466

467
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static, 512) if (num_weights_ >= 1024)
Guolin Ke's avatar
Guolin Ke committed
468
  for (data_size_t i = 0; i < num_weights_; ++i) {
469
    weights_[i] = Common::AvoidInf(weights[i]);
Guolin Ke's avatar
Guolin Ke committed
470
  }
471
  CalculateQueryWeights();
472
  weight_load_from_file_ = false;
473
  #ifdef USE_CUDA
474
475
476
  if (cuda_metadata_ != nullptr) {
    cuda_metadata_->SetWeights(weights_.data(), len);
  }
477
  #endif  // USE_CUDA
Guolin Ke's avatar
Guolin Ke committed
478
479
}

480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
void Metadata::InsertWeights(const label_t* weights, data_size_t start_index, data_size_t len) {
  if (!weights) {
    Log::Fatal("Passed null weights");
  }
  if (num_weights_ <= 0) {
    Log::Fatal("Inserting weight data into dataset with no weights");
  }
  if (start_index + len > num_weights_) {
    Log::Fatal("Inserted weight data is too large for dataset");
  }
  if (weights_.empty()) { weights_.resize(num_weights_); }

  memcpy(weights_.data() + start_index, weights, sizeof(label_t) * len);

  weight_load_from_file_ = false;
  // CUDA is handled after all insertions are complete
}

Guolin Ke's avatar
Guolin Ke committed
498
void Metadata::SetQuery(const data_size_t* query, data_size_t len) {
499
  std::lock_guard<std::mutex> lock(mutex_);
500
  // save to nullptr
Guolin Ke's avatar
Guolin Ke committed
501
  if (query == nullptr || len == 0) {
502
503
504
505
    query_boundaries_.clear();
    num_queries_ = 0;
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
506
  data_size_t sum = 0;
507
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static) reduction(+:sum)
Guolin Ke's avatar
Guolin Ke committed
508
  for (data_size_t i = 0; i < len; ++i) {
Guolin Ke's avatar
Guolin Ke committed
509
    sum += query[i];
Guolin Ke's avatar
Guolin Ke committed
510
511
  }
  if (num_data_ != sum) {
512
    Log::Fatal("Sum of query counts is not same with #data");
Guolin Ke's avatar
Guolin Ke committed
513
514
  }
  num_queries_ = len;
515
  query_boundaries_.resize(num_queries_ + 1);
Guolin Ke's avatar
Guolin Ke committed
516
  query_boundaries_[0] = 0;
Guolin Ke's avatar
Guolin Ke committed
517
  for (data_size_t i = 0; i < num_queries_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
518
    query_boundaries_[i + 1] = query_boundaries_[i] + query[i];
Guolin Ke's avatar
Guolin Ke committed
519
  }
520
  CalculateQueryWeights();
521
  query_load_from_file_ = false;
522
  #ifdef USE_CUDA
523
524
525
526
527
528
529
530
  if (cuda_metadata_ != nullptr) {
    if (query_weights_.size() > 0) {
      CHECK_EQ(query_weights_.size(), static_cast<size_t>(num_queries_));
      cuda_metadata_->SetQuery(query_boundaries_.data(), query_weights_.data(), num_queries_);
    } else {
      cuda_metadata_->SetQuery(query_boundaries_.data(), nullptr, num_queries_);
    }
  }
531
  #endif  // USE_CUDA
532
}
Guolin Ke's avatar
Guolin Ke committed
533

534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
void Metadata::SetPosition(const data_size_t* positions, data_size_t len) {
  std::lock_guard<std::mutex> lock(mutex_);
  // save to nullptr
  if (positions == nullptr || len == 0) {
    positions_.clear();
    num_positions_ = 0;
    return;
  }
  #ifdef USE_CUDA
  Log::Fatal("Positions in learning to rank is not supported in CUDA version yet.");
  #endif  // USE_CUDA
  if (num_data_ != len) {
    Log::Fatal("Positions size (%i) doesn't match data size (%i)", len, num_data_);
  }
  if (positions_.empty()) {
    positions_.resize(num_data_);
  } else {
    Log::Warning("Overwritting positions in dataset.");
  }
  num_positions_ = num_data_;

  position_load_from_file_ = false;

  position_ids_.clear();
  std::unordered_map<data_size_t, int> map_id2pos;
  for (data_size_t i = 0; i < num_positions_; ++i) {
    if (map_id2pos.count(positions[i]) == 0) {
      int pos = static_cast<int>(map_id2pos.size());
      map_id2pos[positions[i]] = pos;
      position_ids_.push_back(std::to_string(positions[i]));
    }
  }

  Log::Debug("number of unique positions found = %ld", position_ids_.size());

569
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static, 512) if (num_positions_ >= 1024)
570
571
572
573
574
  for (data_size_t i = 0; i < num_positions_; ++i) {
    positions_[i] = map_id2pos.at(positions[i]);
  }
}

575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
void Metadata::InsertQueries(const data_size_t* queries, data_size_t start_index, data_size_t len) {
  if (!queries) {
    Log::Fatal("Passed null queries");
  }
  if (queries_.size() <= 0) {
    Log::Fatal("Inserting query data into dataset with no queries");
  }
  if (static_cast<size_t>(start_index + len) > queries_.size()) {
    Log::Fatal("Inserted query data is too large for dataset");
  }

  memcpy(queries_.data() + start_index, queries, sizeof(data_size_t) * len);

  query_load_from_file_ = false;
  // CUDA is handled after all insertions are complete
}

Guolin Ke's avatar
Guolin Ke committed
592
593
594
595
596
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
597
  TextReader<size_t> reader(weight_filename.c_str(), false);
Guolin Ke's avatar
Guolin Ke committed
598
  reader.ReadAllLines();
Guolin Ke's avatar
Guolin Ke committed
599
  if (reader.Lines().empty()) {
Guolin Ke's avatar
Guolin Ke committed
600
601
    return;
  }
602
  Log::Info("Loading weights...");
Guolin Ke's avatar
Guolin Ke committed
603
  num_weights_ = static_cast<data_size_t>(reader.Lines().size());
604
  weights_ = std::vector<label_t>(num_weights_);
605
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
Guolin Ke's avatar
Guolin Ke committed
606
  for (data_size_t i = 0; i < num_weights_; ++i) {
607
    double tmp_weight = 0.0f;
Guolin Ke's avatar
Guolin Ke committed
608
    Common::Atof(reader.Lines()[i].c_str(), &tmp_weight);
609
    weights_[i] = Common::AvoidInf(static_cast<label_t>(tmp_weight));
Guolin Ke's avatar
Guolin Ke committed
610
  }
611
  weight_load_from_file_ = true;
Guolin Ke's avatar
Guolin Ke committed
612
613
}

614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
void Metadata::LoadPositions() {
  num_positions_ = 0;
  std::string position_filename(data_filename_);
  // default position file name
  position_filename.append(".position");
  TextReader<size_t> reader(position_filename.c_str(), false);
  reader.ReadAllLines();
  if (reader.Lines().empty()) {
    return;
  }
  Log::Info("Loading positions from %s ...", position_filename.c_str());
  num_positions_ = static_cast<data_size_t>(reader.Lines().size());
  positions_ = std::vector<data_size_t>(num_positions_);
  position_ids_ = std::vector<std::string>();
  std::unordered_map<std::string, data_size_t> map_id2pos;
  for (data_size_t i = 0; i < num_positions_; ++i) {
    std::string& line = reader.Lines()[i];
    if (map_id2pos.count(line) == 0) {
      map_id2pos[line] = static_cast<data_size_t>(position_ids_.size());
      position_ids_.push_back(line);
    }
    positions_[i] = map_id2pos.at(line);
  }
  position_load_from_file_ = true;
}

640
void Metadata::LoadInitialScore(const std::string& data_filename) {
Guolin Ke's avatar
Guolin Ke committed
641
  num_init_score_ = 0;
642
643
  std::string init_score_filename(data_filename);
  init_score_filename = std::string(data_filename);
644
645
  // default init_score file name
  init_score_filename.append(".init");
Guolin Ke's avatar
Guolin Ke committed
646
  TextReader<size_t> reader(init_score_filename.c_str(), false);
Guolin Ke's avatar
Guolin Ke committed
647
  reader.ReadAllLines();
Guolin Ke's avatar
Guolin Ke committed
648
  if (reader.Lines().empty()) {
Guolin Ke's avatar
Guolin Ke committed
649
650
    return;
  }
651
652
  Log::Info("Loading initial scores...");

653
654
655
  // use first line to count number class
  int num_class = static_cast<int>(Common::Split(reader.Lines()[0].c_str(), '\t').size());
  data_size_t num_line = static_cast<data_size_t>(reader.Lines().size());
Guolin Ke's avatar
Guolin Ke committed
656
  num_init_score_ = static_cast<int64_t>(num_line) * num_class;
657

Guolin Ke's avatar
Guolin Ke committed
658
  init_score_ = std::vector<double>(num_init_score_);
659
  if (num_class == 1) {
660
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
661
    for (data_size_t i = 0; i < num_line; ++i) {
Guolin Ke's avatar
Guolin Ke committed
662
      double tmp = 0.0f;
663
      Common::Atof(reader.Lines()[i].c_str(), &tmp);
664
      init_score_[i] = Common::AvoidInf(static_cast<double>(tmp));
665
    }
666
  } else {
667
    std::vector<std::string> oneline_init_score;
668
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
669
    for (data_size_t i = 0; i < num_line; ++i) {
Guolin Ke's avatar
Guolin Ke committed
670
      double tmp = 0.0f;
671
672
      oneline_init_score = Common::Split(reader.Lines()[i].c_str(), '\t');
      if (static_cast<int>(oneline_init_score.size()) != num_class) {
673
        Log::Fatal("Invalid initial score file. Redundant or insufficient columns");
674
675
676
      }
      for (int k = 0; k < num_class; ++k) {
        Common::Atof(oneline_init_score[k].c_str(), &tmp);
677
        init_score_[static_cast<size_t>(k) * num_line + i] = Common::AvoidInf(static_cast<double>(tmp));
678
      }
679
    }
Guolin Ke's avatar
Guolin Ke committed
680
  }
681
  init_score_load_from_file_ = true;
Guolin Ke's avatar
Guolin Ke committed
682
683
684
685
686
687
688
}

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
689
  TextReader<size_t> reader(query_filename.c_str(), false);
Guolin Ke's avatar
Guolin Ke committed
690
  reader.ReadAllLines();
Guolin Ke's avatar
Guolin Ke committed
691
  if (reader.Lines().empty()) {
Guolin Ke's avatar
Guolin Ke committed
692
693
    return;
  }
694
  Log::Info("Calculating query boundaries...");
Guolin Ke's avatar
Guolin Ke committed
695
  query_boundaries_ = std::vector<data_size_t>(reader.Lines().size() + 1);
Guolin Ke's avatar
Guolin Ke committed
696
697
698
699
700
701
702
  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);
  }
703
  query_load_from_file_ = true;
Guolin Ke's avatar
Guolin Ke committed
704
705
}

706
void Metadata::CalculateQueryWeights() {
Guolin Ke's avatar
Guolin Ke committed
707
  if (weights_.size() == 0 || query_boundaries_.size() == 0) {
Guolin Ke's avatar
Guolin Ke committed
708
709
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
710
  query_weights_.clear();
711
  Log::Info("Calculating query weights...");
712
  query_weights_ = std::vector<label_t>(num_queries_);
Guolin Ke's avatar
Guolin Ke committed
713
714
715
716
717
718
719
720
721
  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]);
  }
}

722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
void Metadata::InsertAt(data_size_t start_index,
  data_size_t count,
  const float* labels,
  const float* weights,
  const double* init_scores,
  const int32_t* queries) {
  if (num_data_ < count + start_index) {
    Log::Fatal("Length of metadata is too long to append #data");
  }
  InsertLabels(labels, start_index, count);
  if (weights) {
    InsertWeights(weights, start_index, count);
  }
  if (init_scores) {
    InsertInitScores(init_scores, start_index, count, count);
  }
  if (queries) {
    InsertQueries(queries, start_index, count);
  }
}

void Metadata::FinishLoad() {
  CalculateQueryBoundaries();
}

747
#ifdef USE_CUDA
748
749
750
751
void Metadata::CreateCUDAMetadata(const int gpu_device_id) {
  cuda_metadata_.reset(new CUDAMetadata(gpu_device_id));
  cuda_metadata_->Init(label_, weights_, query_boundaries_, query_weights_, init_score_);
}
752
#endif  // USE_CUDA
753

Guolin Ke's avatar
Guolin Ke committed
754
755
756
757
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));
758
  mem_ptr += VirtualFileWriter::AlignedSize(sizeof(num_data_));
Guolin Ke's avatar
Guolin Ke committed
759
  num_weights_ = *(reinterpret_cast<const data_size_t*>(mem_ptr));
760
  mem_ptr += VirtualFileWriter::AlignedSize(sizeof(num_weights_));
Guolin Ke's avatar
Guolin Ke committed
761
  num_queries_ = *(reinterpret_cast<const data_size_t*>(mem_ptr));
762
  mem_ptr += VirtualFileWriter::AlignedSize(sizeof(num_queries_));
Guolin Ke's avatar
Guolin Ke committed
763

Guolin Ke's avatar
Guolin Ke committed
764
  if (!label_.empty()) { label_.clear(); }
765
  label_ = std::vector<label_t>(num_data_);
766
  std::memcpy(label_.data(), mem_ptr, sizeof(label_t) * num_data_);
767
  mem_ptr += VirtualFileWriter::AlignedSize(sizeof(label_t) * num_data_);
Guolin Ke's avatar
Guolin Ke committed
768
769

  if (num_weights_ > 0) {
Guolin Ke's avatar
Guolin Ke committed
770
    if (!weights_.empty()) { weights_.clear(); }
771
    weights_ = std::vector<label_t>(num_weights_);
772
    std::memcpy(weights_.data(), mem_ptr, sizeof(label_t) * num_weights_);
773
    mem_ptr += VirtualFileWriter::AlignedSize(sizeof(label_t) * num_weights_);
774
    weight_load_from_file_ = true;
Guolin Ke's avatar
Guolin Ke committed
775
776
  }
  if (num_queries_ > 0) {
Guolin Ke's avatar
Guolin Ke committed
777
    if (!query_boundaries_.empty()) { query_boundaries_.clear(); }
Guolin Ke's avatar
Guolin Ke committed
778
    query_boundaries_ = std::vector<data_size_t>(num_queries_ + 1);
779
    std::memcpy(query_boundaries_.data(), mem_ptr, sizeof(data_size_t) * (num_queries_ + 1));
780
781
    mem_ptr += VirtualFileWriter::AlignedSize(sizeof(data_size_t) *
                                              (num_queries_ + 1));
782
    query_load_from_file_ = true;
Guolin Ke's avatar
Guolin Ke committed
783
  }
784
  CalculateQueryWeights();
Guolin Ke's avatar
Guolin Ke committed
785
786
}

787
void Metadata::SaveBinaryToFile(BinaryWriter* writer) const {
788
789
790
791
  writer->AlignedWrite(&num_data_, sizeof(num_data_));
  writer->AlignedWrite(&num_weights_, sizeof(num_weights_));
  writer->AlignedWrite(&num_queries_, sizeof(num_queries_));
  writer->AlignedWrite(label_.data(), sizeof(label_t) * num_data_);
Guolin Ke's avatar
Guolin Ke committed
792
  if (!weights_.empty()) {
793
    writer->AlignedWrite(weights_.data(), sizeof(label_t) * num_weights_);
Guolin Ke's avatar
Guolin Ke committed
794
  }
Guolin Ke's avatar
Guolin Ke committed
795
  if (!query_boundaries_.empty()) {
796
797
    writer->AlignedWrite(query_boundaries_.data(),
                         sizeof(data_size_t) * (num_queries_ + 1));
Guolin Ke's avatar
Guolin Ke committed
798
  }
799
800
801
802
  if (num_init_score_ > 0) {
    Log::Warning("Please note that `init_score` is not saved in binary file.\n"
      "If you need it, please set it again after loading Dataset.");
  }
Guolin Ke's avatar
Guolin Ke committed
803
804
}

805
size_t Metadata::SizesInByte() const {
806
807
808
809
  size_t size = VirtualFileWriter::AlignedSize(sizeof(num_data_)) +
                VirtualFileWriter::AlignedSize(sizeof(num_weights_)) +
                VirtualFileWriter::AlignedSize(sizeof(num_queries_));
  size += VirtualFileWriter::AlignedSize(sizeof(label_t) * num_data_);
Guolin Ke's avatar
Guolin Ke committed
810
  if (!weights_.empty()) {
811
    size += VirtualFileWriter::AlignedSize(sizeof(label_t) * num_weights_);
Guolin Ke's avatar
Guolin Ke committed
812
  }
Guolin Ke's avatar
Guolin Ke committed
813
  if (!query_boundaries_.empty()) {
814
815
    size += VirtualFileWriter::AlignedSize(sizeof(data_size_t) *
                                           (num_queries_ + 1));
Guolin Ke's avatar
Guolin Ke committed
816
817
818
819
820
821
  }
  return size;
}


}  // namespace LightGBM