dataset.h 33.7 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.
 */
Guolin Ke's avatar
Guolin Ke committed
5
6
#ifndef LIGHTGBM_DATASET_H_
#define LIGHTGBM_DATASET_H_
Guolin Ke's avatar
Guolin Ke committed
7

8
9
10
#include <LightGBM/config.h>
#include <LightGBM/feature_group.h>
#include <LightGBM/meta.h>
11
#include <LightGBM/train_share_states.h>
12
13
14
15
#include <LightGBM/utils/openmp_wrapper.h>
#include <LightGBM/utils/random.h>
#include <LightGBM/utils/text_reader.h>

Guolin Ke's avatar
Guolin Ke committed
16
#include <string>
17
#include <functional>
18
#include <map>
19
#include <memory>
20
#include <mutex>
21
22
23
#include <unordered_set>
#include <utility>
#include <vector>
Guolin Ke's avatar
Guolin Ke committed
24

25
26
27
#include <LightGBM/cuda/cuda_column_data.hpp>
#include <LightGBM/cuda/cuda_metadata.hpp>

Guolin Ke's avatar
Guolin Ke committed
28
29
30
namespace LightGBM {

/*! \brief forward declaration */
Guolin Ke's avatar
Guolin Ke committed
31
class DatasetLoader;
Guolin Ke's avatar
Guolin Ke committed
32
/*!
Hui Xue's avatar
Hui Xue committed
33
* \brief This class is used to store some meta(non-feature) data for training data,
Andrew Ziem's avatar
Andrew Ziem committed
34
*        e.g. labels, weights, initial scores, query level information.
Guolin Ke's avatar
Guolin Ke committed
35
*
Qiwei Ye's avatar
Qiwei Ye committed
36
*        Some details:
37
*        1. Label, used for training.
Qiwei Ye's avatar
Qiwei Ye committed
38
*        2. Weights, weighs of records, optional
Andrew Ziem's avatar
Andrew Ziem committed
39
*        3. Query Boundaries, necessary for LambdaRank.
40
41
42
43
*           The documents of i-th query is in [ query_boundaries[i], query_boundaries[i+1] )
*        4. Query Weights, auto calculate by weights and query_boundaries(if both of them are existed)
*           the weight for i-th query is sum(query_boundaries[i] , .., query_boundaries[i+1]) / (query_boundaries[i + 1] -  query_boundaries[i+1])
*        5. Initial score. optional. if existing, the model will boost from this score, otherwise will start from 0.
Guolin Ke's avatar
Guolin Ke committed
44
45
*/
class Metadata {
Nikita Titov's avatar
Nikita Titov committed
46
 public:
47
  /*!
48
  * \brief Null constructor
Guolin Ke's avatar
Guolin Ke committed
49
50
51
  */
  Metadata();
  /*!
Andrew Ziem's avatar
Andrew Ziem committed
52
  * \brief Initialization will load query level information, since it is need for sampling data
Guolin Ke's avatar
Guolin Ke committed
53
54
  * \param data_filename Filename of data
  */
55
  void Init(const char* data_filename);
Guolin Ke's avatar
Guolin Ke committed
56
  /*!
Guolin Ke's avatar
Guolin Ke committed
57
58
  * \brief init as subset
  * \param metadata Filename of data
59
  * \param used_indices
Guolin Ke's avatar
Guolin Ke committed
60
61
62
63
  * \param num_used_indices
  */
  void Init(const Metadata& metadata, const data_size_t* used_indices, data_size_t num_used_indices);
  /*!
Guolin Ke's avatar
Guolin Ke committed
64
65
66
67
68
69
70
71
  * \brief Initial with binary memory
  * \param memory Pointer to memory
  */
  void LoadFromMemory(const void* memory);
  /*! \brief Destructor */
  ~Metadata();

  /*!
72
  * \brief Initial work, will allocate space for label, weight (if exists) and query (if exists)
Guolin Ke's avatar
Guolin Ke committed
73
  * \param num_data Number of training data
Guolin Ke's avatar
Guolin Ke committed
74
75
  * \param weight_idx Index of weight column, < 0 means doesn't exists
  * \param query_idx Index of query id column, < 0 means doesn't exists
Guolin Ke's avatar
Guolin Ke committed
76
  */
77
  void Init(data_size_t num_data, int weight_idx, int query_idx);
Guolin Ke's avatar
Guolin Ke committed
78

79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
  /*!
  * \brief Allocate space for label, weight (if exists), initial score (if exists) and query (if exists)
  * \param num_data Number of data
  * \param reference Reference metadata
  */
  void InitByReference(data_size_t num_data, const Metadata* reference);

  /*!
  * \brief Allocate space for label, weight (if exists), initial score (if exists) and query (if exists)
  * \param num_data Number of data rows
  * \param has_weights Whether the metadata has weights
  * \param has_init_scores Whether the metadata has initial scores
  * \param has_queries Whether the metadata has queries
  * \param nclasses Number of classes for initial scores
  */
  void Init(data_size_t num_data, int32_t has_weights, int32_t has_init_scores, int32_t has_queries, int32_t nclasses);

Guolin Ke's avatar
Guolin Ke committed
96
97
  /*!
  * \brief Partition label by used indices
98
  * \param used_indices Indices of local used
Guolin Ke's avatar
Guolin Ke committed
99
100
101
102
103
  */
  void PartitionLabel(const std::vector<data_size_t>& used_indices);

  /*!
  * \brief Partition meta data according to local used indices if need
104
  * \param num_all_data Number of total training data, including other machines' data on distributed learning
Guolin Ke's avatar
Guolin Ke committed
105
106
107
  * \param used_data_indices Indices of local used training data
  */
  void CheckOrPartition(data_size_t num_all_data,
108
                        const std::vector<data_size_t>& used_data_indices);
Guolin Ke's avatar
Guolin Ke committed
109

110
  void SetLabel(const label_t* label, data_size_t len);
Guolin Ke's avatar
Guolin Ke committed
111

112
  void SetWeights(const label_t* weights, data_size_t len);
Guolin Ke's avatar
Guolin Ke committed
113

Guolin Ke's avatar
Guolin Ke committed
114
  void SetQuery(const data_size_t* query, data_size_t len);
Guolin Ke's avatar
Guolin Ke committed
115

Guolin Ke's avatar
Guolin Ke committed
116
117
118
119
  /*!
  * \brief Set initial scores
  * \param init_score Initial scores, this class will manage memory for init_score.
  */
120
  void SetInitScore(const double* init_score, data_size_t len);
Guolin Ke's avatar
Guolin Ke committed
121
122
123
124
125
126


  /*!
  * \brief Save binary data to file
  * \param file File want to write
  */
127
  void SaveBinaryToFile(const VirtualFileWriter* writer) const;
Guolin Ke's avatar
Guolin Ke committed
128
129
130
131
132
133
134
135
136
137

  /*!
  * \brief Get sizes in byte of this object
  */
  size_t SizesInByte() const;

  /*!
  * \brief Get pointer of label
  * \return Pointer of label
  */
138
  inline const label_t* label() const { return label_.data(); }
Guolin Ke's avatar
Guolin Ke committed
139
140
141
142
143
144

  /*!
  * \brief Set label for one record
  * \param idx Index of this record
  * \param value Label value of this record
  */
145
  inline void SetLabelAt(data_size_t idx, label_t value) {
146
    label_[idx] = value;
Guolin Ke's avatar
Guolin Ke committed
147
148
  }

Guolin Ke's avatar
Guolin Ke committed
149
150
151
152
153
  /*!
  * \brief Set Weight for one record
  * \param idx Index of this record
  * \param value Weight value of this record
  */
154
  inline void SetWeightAt(data_size_t idx, label_t value) {
155
    weights_[idx] = value;
Guolin Ke's avatar
Guolin Ke committed
156
157
  }

158
159
160
161
162
163
  /*!
  * \brief Set initial scores for one record.  Note that init_score might have multiple columns and is stored in column format.
  * \param idx Index of this record
  * \param values Initial score values for this record, one per class
  */
  inline void SetInitScoreAt(data_size_t idx, const double* values) {
164
    const auto nclasses = num_init_score_classes();
165
166
167
168
169
170
    const double* val_ptr = values;
    for (int i = idx; i < nclasses * num_data_; i += num_data_, ++val_ptr) {
      init_score_[i] = *val_ptr;
    }
  }

Guolin Ke's avatar
Guolin Ke committed
171
172
173
174
175
  /*!
  * \brief Set Query Id for one record
  * \param idx Index of this record
  * \param value Query Id value of this record
  */
176
  inline void SetQueryAt(data_size_t idx, data_size_t value) {
Guolin Ke's avatar
Guolin Ke committed
177
178
179
    queries_[idx] = static_cast<data_size_t>(value);
  }

180
181
182
  /*! \brief Load initial scores from file */
  void LoadInitialScore(const std::string& data_filename);

183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
  /*!
  * \brief Insert data from a given data to the current data at a specified index
  * \param start_index The target index to begin the insertion
  * \param count Number of records to insert
  * \param labels Pointer to label data
  * \param weights Pointer to weight data, or null
  * \param init_scores Pointer to init-score data, or null
  * \param queries Pointer to query data, or null
  */
  void InsertAt(data_size_t start_index,
    data_size_t count,
    const float* labels,
    const float* weights,
    const double* init_scores,
    const int32_t* queries);

  /*!
  * \brief Perform any extra operations after all data has been loaded
  */
  void FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
203
  /*!
Hui Xue's avatar
Hui Xue committed
204
  * \brief Get weights, if not exists, will return nullptr
Guolin Ke's avatar
Guolin Ke committed
205
206
  * \return Pointer of weights
  */
207
  inline const label_t* weights() const {
Guolin Ke's avatar
Guolin Ke committed
208
    if (!weights_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
209
210
211
212
213
      return weights_.data();
    } else {
      return nullptr;
    }
  }
Guolin Ke's avatar
Guolin Ke committed
214
215

  /*!
Hui Xue's avatar
Hui Xue committed
216
  * \brief Get data boundaries on queries, if not exists, will return nullptr
217
  *        we assume data will order by query,
Guolin Ke's avatar
Guolin Ke committed
218
219
220
221
  *        the interval of [query_boundaris[i], query_boundaris[i+1])
  *        is the data indices for query i.
  * \return Pointer of data boundaries on queries
  */
222
  inline const data_size_t* query_boundaries() const {
Guolin Ke's avatar
Guolin Ke committed
223
    if (!query_boundaries_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
224
225
226
227
228
      return query_boundaries_.data();
    } else {
      return nullptr;
    }
  }
Guolin Ke's avatar
Guolin Ke committed
229
230
231
232
233

  /*!
  * \brief Get Number of queries
  * \return Number of queries
  */
234
  inline data_size_t num_queries() const { return num_queries_; }
Guolin Ke's avatar
Guolin Ke committed
235
236

  /*!
Hui Xue's avatar
Hui Xue committed
237
  * \brief Get weights for queries, if not exists, will return nullptr
Guolin Ke's avatar
Guolin Ke committed
238
239
  * \return Pointer of weights for queries
  */
240
  inline const label_t* query_weights() const {
Guolin Ke's avatar
Guolin Ke committed
241
    if (!query_weights_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
242
243
244
245
246
      return query_weights_.data();
    } else {
      return nullptr;
    }
  }
Guolin Ke's avatar
Guolin Ke committed
247
248

  /*!
Hui Xue's avatar
Hui Xue committed
249
  * \brief Get initial scores, if not exists, will return nullptr
Guolin Ke's avatar
Guolin Ke committed
250
251
  * \return Pointer of initial scores
  */
252
  inline const double* init_score() const {
Guolin Ke's avatar
Guolin Ke committed
253
    if (!init_score_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
254
255
256
257
258
      return init_score_.data();
    } else {
      return nullptr;
    }
  }
Guolin Ke's avatar
Guolin Ke committed
259

260
261
262
  /*!
  * \brief Get size of initial scores
  */
Guolin Ke's avatar
Guolin Ke committed
263
  inline int64_t num_init_score() const { return num_init_score_; }
264

265
266
267
  /*!
  * \brief Get number of classes
  */
268
  inline int32_t num_init_score_classes() const {
269
270
271
272
273
274
    if (num_data_ && num_init_score_) {
      return static_cast<int>(num_init_score_ / num_data_);
    }
    return 1;
  }

Guolin Ke's avatar
Guolin Ke committed
275
276
277
278
  /*! \brief Disable copy */
  Metadata& operator=(const Metadata&) = delete;
  /*! \brief Disable copy */
  Metadata(const Metadata&) = delete;
Guolin Ke's avatar
Guolin Ke committed
279

280
  #ifdef USE_CUDA
281
282
283
284
285

  CUDAMetadata* cuda_metadata() const { return cuda_metadata_.get(); }

  void CreateCUDAMetadata(const int gpu_device_id);

286
  #endif  // USE_CUDA
287

Nikita Titov's avatar
Nikita Titov committed
288
 private:
Guolin Ke's avatar
Guolin Ke committed
289
290
291
292
  /*! \brief Load wights from file */
  void LoadWeights();
  /*! \brief Load query boundaries from file */
  void LoadQueryBoundaries();
293
294
295
296
297
298
299
300
301
302
303
304
  /*! \brief Calculate query weights from queries */
  void CalculateQueryWeights();
  /*! \brief Calculate query boundaries from queries */
  void CalculateQueryBoundaries();
  /*! \brief Insert labels at the given index */
  void InsertLabels(const label_t* labels, data_size_t start_index, data_size_t len);
  /*! \brief Insert weights at the given index */
  void InsertWeights(const label_t* weights, data_size_t start_index, data_size_t len);
  /*! \brief Insert initial scores at the given index */
  void InsertInitScores(const double* init_scores, data_size_t start_index, data_size_t len, data_size_t source_size);
  /*! \brief Insert queries at the given index */
  void InsertQueries(const data_size_t* queries, data_size_t start_index, data_size_t len);
Guolin Ke's avatar
Guolin Ke committed
305
  /*! \brief Filename of current data */
Guolin Ke's avatar
Guolin Ke committed
306
  std::string data_filename_;
Guolin Ke's avatar
Guolin Ke committed
307
308
309
310
311
  /*! \brief Number of data */
  data_size_t num_data_;
  /*! \brief Number of weights, used to check correct weight file */
  data_size_t num_weights_;
  /*! \brief Label data */
312
  std::vector<label_t> label_;
Guolin Ke's avatar
Guolin Ke committed
313
  /*! \brief Weights data */
314
  std::vector<label_t> weights_;
Guolin Ke's avatar
Guolin Ke committed
315
  /*! \brief Query boundaries */
Guolin Ke's avatar
Guolin Ke committed
316
  std::vector<data_size_t> query_boundaries_;
Guolin Ke's avatar
Guolin Ke committed
317
  /*! \brief Query weights */
318
  std::vector<label_t> query_weights_;
Guolin Ke's avatar
Guolin Ke committed
319
320
321
  /*! \brief Number of querys */
  data_size_t num_queries_;
  /*! \brief Number of Initial score, used to check correct weight file */
Guolin Ke's avatar
Guolin Ke committed
322
  int64_t num_init_score_;
Guolin Ke's avatar
Guolin Ke committed
323
  /*! \brief Initial score */
Guolin Ke's avatar
Guolin Ke committed
324
  std::vector<double> init_score_;
Guolin Ke's avatar
Guolin Ke committed
325
  /*! \brief Queries data */
Guolin Ke's avatar
Guolin Ke committed
326
  std::vector<data_size_t> queries_;
327
328
  /*! \brief mutex for threading safe call */
  std::mutex mutex_;
329
330
331
  bool weight_load_from_file_;
  bool query_load_from_file_;
  bool init_score_load_from_file_;
332
  #ifdef USE_CUDA
333
  std::unique_ptr<CUDAMetadata> cuda_metadata_;
334
  #endif  // USE_CUDA
Guolin Ke's avatar
Guolin Ke committed
335
336
337
338
339
};


/*! \brief Interface for Parser */
class Parser {
Nikita Titov's avatar
Nikita Titov committed
340
 public:
Chen Yufei's avatar
Chen Yufei committed
341
342
  typedef const char* (*AtofFunc)(const char* p, double* out);

343
344
345
346
347
348
349
350
  /*! \brief Default constructor */
  Parser() {}

  /*!
  * \brief Constructor for customized parser. The constructor accepts content not path because need to save/load the config along with model string
  */
  explicit Parser(std::string) {}

Guolin Ke's avatar
Guolin Ke committed
351
352
353
354
355
356
  /*! \brief virtual destructor */
  virtual ~Parser() {}

  /*!
  * \brief Parse one line with label
  * \param str One line record, string format, should end with '\0'
Guolin Ke's avatar
Guolin Ke committed
357
358
  * \param out_features Output columns, store in (column_idx, values)
  * \param out_label Label will store to this if exists
Guolin Ke's avatar
Guolin Ke committed
359
360
  */
  virtual void ParseOneLine(const char* str,
361
                            std::vector<std::pair<int, double>>* out_features, double* out_label) const = 0;
Guolin Ke's avatar
Guolin Ke committed
362

363
  virtual int NumFeatures() const = 0;
Guolin Ke's avatar
Guolin Ke committed
364

Guolin Ke's avatar
Guolin Ke committed
365
  /*!
366
  * \brief Create an object of parser, will auto choose the format depend on file
Guolin Ke's avatar
Guolin Ke committed
367
  * \param filename One Filename of data
368
  * \param header whether input file contains header
369
  * \param num_features Pass num_features of this data file if you know, <=0 means don't know
Guolin Ke's avatar
Guolin Ke committed
370
  * \param label_idx index of label column
Chen Yufei's avatar
Chen Yufei committed
371
  * \param precise_float_parser using precise floating point number parsing if true
Guolin Ke's avatar
Guolin Ke committed
372
373
  * \return Object of parser
  */
Chen Yufei's avatar
Chen Yufei committed
374
  static Parser* CreateParser(const char* filename, bool header, int num_features, int label_idx, bool precise_float_parser);
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419

  /*!
  * \brief Create an object of parser, could use customized parser, or auto choose the format depend on file
  * \param filename One Filename of data
  * \param header whether input file contains header
  * \param num_features Pass num_features of this data file if you know, <=0 means don't know
  * \param label_idx index of label column
  * \param precise_float_parser using precise floating point number parsing if true
  * \param parser_config_str Customized parser config content
  * \return Object of parser
  */
  static Parser* CreateParser(const char* filename, bool header, int num_features, int label_idx, bool precise_float_parser,
                              std::string parser_config_str);

  /*!
  * \brief Generate parser config str used for custom parser initialization, may save values of label id and header
  * \param filename One Filename of data
  * \param parser_config_filename One Filename of parser config
  * \param header whether input file contains header
  * \param label_idx index of label column
  * \return Parser config str
  */
  static std::string GenerateParserConfigStr(const char* filename, const char* parser_config_filename, bool header, int label_idx);
};

/*! \brief Interface for parser factory, used by customized parser */
class ParserFactory {
 private:
  ParserFactory() {}
  std::map<std::string, std::function<Parser*(std::string)>> object_map_;

 public:
  ~ParserFactory() {}
  static ParserFactory& getInstance();
  void Register(std::string class_name, std::function<Parser*(std::string)> objc);
  Parser* getObject(std::string class_name, std::string config_str);
};

/*! \brief Interface for parser reflector, used by customized parser */
class ParserReflector {
 public:
  ParserReflector(std::string class_name, std::function<Parser*(std::string)> objc) {
    ParserFactory::getInstance().Register(class_name, objc);
  }
  virtual ~ParserReflector() {}
Guolin Ke's avatar
Guolin Ke committed
420
421
422
};

/*! \brief The main class of data set,
423
*          which are used to training or validation
Guolin Ke's avatar
Guolin Ke committed
424
425
*/
class Dataset {
Nikita Titov's avatar
Nikita Titov committed
426
 public:
Guolin Ke's avatar
Guolin Ke committed
427
  friend DatasetLoader;
Guolin Ke's avatar
Guolin Ke committed
428

429
  LIGHTGBM_EXPORT Dataset();
Guolin Ke's avatar
Guolin Ke committed
430

431
  LIGHTGBM_EXPORT Dataset(data_size_t num_data);
Guolin Ke's avatar
Guolin Ke committed
432

Guolin Ke's avatar
Guolin Ke committed
433
  void Construct(
Guolin Ke's avatar
Guolin Ke committed
434
    std::vector<std::unique_ptr<BinMapper>>* bin_mappers,
435
    int num_total_features,
436
    const std::vector<std::vector<double>>& forced_bins,
437
    int** sample_non_zero_indices,
Guolin Ke's avatar
Guolin Ke committed
438
    double** sample_values,
439
    const int* num_per_col,
440
    int num_sample_col,
Guolin Ke's avatar
Guolin Ke committed
441
    size_t total_sample_cnt,
Guolin Ke's avatar
Guolin Ke committed
442
    const Config& io_config);
Guolin Ke's avatar
Guolin Ke committed
443

Guolin Ke's avatar
Guolin Ke committed
444
  /*! \brief Destructor */
445
  LIGHTGBM_EXPORT ~Dataset();
Guolin Ke's avatar
Guolin Ke committed
446

447
448
449
450
451
452
453
454
455
456
457
458
459
460
  /*!
  * \brief Initialize from the given reference
  * \param num_data Number of data
  * \param reference Reference dataset
  */
  LIGHTGBM_EXPORT void InitByReference(data_size_t num_data, const Dataset* reference) {
    metadata_.InitByReference(num_data, &reference->metadata());
  }

  LIGHTGBM_EXPORT void InitStreaming(data_size_t num_data,
                                     int32_t has_weights,
                                     int32_t has_init_scores,
                                     int32_t has_queries,
                                     int32_t nclasses,
461
462
463
464
465
466
467
468
469
                                     int32_t nthreads,
                                     int32_t omp_max_threads) {
    // Initialize optional max thread count with either parameter or OMP setting
    if (omp_max_threads > 0) {
      omp_max_threads_ = omp_max_threads;
    } else if (omp_max_threads_ <= 0) {
      omp_max_threads_ = OMP_NUM_THREADS();
    }

470
471
    metadata_.Init(num_data, has_weights, has_init_scores, has_queries, nclasses);
    for (int i = 0; i < num_groups_; ++i) {
472
      feature_groups_[i]->InitStreaming(nthreads, omp_max_threads_);
473
474
475
    }
  }

476
  LIGHTGBM_EXPORT bool CheckAlign(const Dataset& other) const {
477
478
479
480
481
482
483
484
485
486
    if (num_features_ != other.num_features_) {
      return false;
    }
    if (num_total_features_ != other.num_total_features_) {
      return false;
    }
    if (label_idx_ != other.label_idx_) {
      return false;
    }
    for (int i = 0; i < num_features_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
487
      if (!FeatureBinMapper(i)->CheckAlign(*(other.FeatureBinMapper(i)))) {
488
489
490
491
492
493
        return false;
      }
    }
    return true;
  }

Guolin Ke's avatar
Guolin Ke committed
494
495
496
497
498
499
500
501
502
503
  inline void FinishOneRow(int tid, data_size_t row_idx, const std::vector<bool>& is_feature_added) {
    if (is_finish_load_) { return; }
    for (auto fidx : feature_need_push_zeros_) {
      if (is_feature_added[fidx]) { continue; }
      const int group = feature2group_[fidx];
      const int sub_feature = feature2subfeature_[fidx];
      feature_groups_[group]->PushData(tid, sub_feature, row_idx, 0.0f);
    }
  }

Guolin Ke's avatar
Guolin Ke committed
504
  inline void PushOneRow(int tid, data_size_t row_idx, const std::vector<double>& feature_values) {
Guolin Ke's avatar
Guolin Ke committed
505
    if (is_finish_load_) { return; }
Guolin Ke's avatar
Guolin Ke committed
506
    for (size_t i = 0; i < feature_values.size() && i < static_cast<size_t>(num_total_features_); ++i) {
Guolin Ke's avatar
Guolin Ke committed
507
508
      int feature_idx = used_feature_map_[i];
      if (feature_idx >= 0) {
Guolin Ke's avatar
Guolin Ke committed
509
510
511
        const int group = feature2group_[feature_idx];
        const int sub_feature = feature2subfeature_[feature_idx];
        feature_groups_[group]->PushData(tid, sub_feature, row_idx, feature_values[i]);
512
513
514
        if (has_raw_) {
          int feat_ind = numeric_feature_map_[feature_idx];
          if (feat_ind >= 0) {
sisco0's avatar
sisco0 committed
515
            raw_data_[feat_ind][row_idx] = static_cast<float>(feature_values[i]);
516
517
          }
        }
Guolin Ke's avatar
Guolin Ke committed
518
519
520
521
      }
    }
  }

522
  inline void PushOneRow(int tid, data_size_t row_idx, const std::vector<std::pair<int, double>>& feature_values) {
Guolin Ke's avatar
Guolin Ke committed
523
    if (is_finish_load_) { return; }
Guolin Ke's avatar
Guolin Ke committed
524
    std::vector<bool> is_feature_added(num_features_, false);
525
    for (auto& inner_data : feature_values) {
526
      if (inner_data.first >= num_total_features_) { continue; }
527
528
      int feature_idx = used_feature_map_[inner_data.first];
      if (feature_idx >= 0) {
Guolin Ke's avatar
Guolin Ke committed
529
        is_feature_added[feature_idx] = true;
Guolin Ke's avatar
Guolin Ke committed
530
531
532
        const int group = feature2group_[feature_idx];
        const int sub_feature = feature2subfeature_[feature_idx];
        feature_groups_[group]->PushData(tid, sub_feature, row_idx, inner_data.second);
533
534
535
        if (has_raw_) {
          int feat_ind = numeric_feature_map_[feature_idx];
          if (feat_ind >= 0) {
536
            raw_data_[feat_ind][row_idx] = static_cast<float>(inner_data.second);
537
538
          }
        }
539
540
      }
    }
Guolin Ke's avatar
Guolin Ke committed
541
    FinishOneRow(tid, row_idx, is_feature_added);
542
543
  }

544
  inline void PushOneData(int tid, data_size_t row_idx, int group, int feature_idx, int sub_feature, double value) {
Guolin Ke's avatar
Guolin Ke committed
545
    feature_groups_[group]->PushData(tid, sub_feature, row_idx, value);
546
547
548
    if (has_raw_) {
      int feat_ind = numeric_feature_map_[feature_idx];
      if (feat_ind >= 0) {
549
        raw_data_[feat_ind][row_idx] = static_cast<float>(value);
550
551
      }
    }
Guolin Ke's avatar
Guolin Ke committed
552
553
  }

554
555
556
557
558
559
560
561
562
  inline void InsertMetadataAt(data_size_t start_index,
    data_size_t count,
    const label_t* labels,
    const label_t* weights,
    const double* init_scores,
    const data_size_t* queries) {
    metadata_.InsertAt(start_index, count, labels, weights, init_scores, queries);
  }

Guolin Ke's avatar
Guolin Ke committed
563
564
565
566
567
  inline int RealFeatureIndex(int fidx) const {
    return real_feature_idx_[fidx];
  }

  inline int InnerFeatureIndex(int col_idx) const {
Guolin Ke's avatar
Guolin Ke committed
568
    return used_feature_map_[col_idx];
Guolin Ke's avatar
Guolin Ke committed
569
  }
Guolin Ke's avatar
Guolin Ke committed
570
571
572
573
574
575
  inline int Feature2Group(int feature_idx) const {
    return feature2group_[feature_idx];
  }
  inline int Feture2SubFeature(int feature_idx) const {
    return feature2subfeature_[feature_idx];
  }
576
577
578
  inline uint64_t GroupBinBoundary(int group_idx) const {
    return group_bin_boundaries_[group_idx];
  }
Guolin Ke's avatar
Guolin Ke committed
579
580
581
  inline uint64_t NumTotalBin() const {
    return group_bin_boundaries_.back();
  }
582

583
584
585
586
587
588
589
590
591
  inline std::vector<int> ValidFeatureIndices() const {
    std::vector<int> ret;
    for (int i = 0; i < num_total_features_; ++i) {
      if (used_feature_map_[i] >= 0) {
        ret.push_back(i);
      }
    }
    return ret;
  }
Guolin Ke's avatar
Guolin Ke committed
592
593
  void ReSize(data_size_t num_data);

594
  void CopySubrow(const Dataset* fullset, const data_size_t* used_indices, data_size_t num_used_indices, bool need_meta_data);
Guolin Ke's avatar
Guolin Ke committed
595

596
  MultiValBin* GetMultiBinFromSparseFeatures(const std::vector<uint32_t>& offsets) const;
597

598
  MultiValBin* GetMultiBinFromAllFeatures(const std::vector<uint32_t>& offsets) const;
599

600
601
602
  TrainingShareStates* GetShareStates(
      score_t* gradients, score_t* hessians,
      const std::vector<int8_t>& is_feature_used, bool is_constant_hessian,
603
      bool force_col_wise, bool force_row_wise) const;
604

605
  LIGHTGBM_EXPORT void FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
606

607
  LIGHTGBM_EXPORT bool SetFloatField(const char* field_name, const float* field_data, data_size_t num_element);
Guolin Ke's avatar
Guolin Ke committed
608

609
  LIGHTGBM_EXPORT bool SetDoubleField(const char* field_name, const double* field_data, data_size_t num_element);
Guolin Ke's avatar
Guolin Ke committed
610

611
  LIGHTGBM_EXPORT bool SetIntField(const char* field_name, const int* field_data, data_size_t num_element);
612

613
  LIGHTGBM_EXPORT bool GetFloatField(const char* field_name, data_size_t* out_len, const float** out_ptr);
614

615
  LIGHTGBM_EXPORT bool GetDoubleField(const char* field_name, data_size_t* out_len, const double** out_ptr);
Guolin Ke's avatar
Guolin Ke committed
616

617
  LIGHTGBM_EXPORT bool GetIntField(const char* field_name, data_size_t* out_len, const int** out_ptr);
618

Guolin Ke's avatar
Guolin Ke committed
619
620
621
  /*!
  * \brief Save current dataset into binary file, will save to "filename.bin"
  */
622
  LIGHTGBM_EXPORT void SaveBinaryFile(const char* bin_filename);
Guolin Ke's avatar
Guolin Ke committed
623

624
625
  LIGHTGBM_EXPORT void DumpTextFile(const char* text_filename);

626
  LIGHTGBM_EXPORT void CopyFeatureMapperFrom(const Dataset* dataset);
Guolin Ke's avatar
Guolin Ke committed
627

Guolin Ke's avatar
Guolin Ke committed
628
629
  LIGHTGBM_EXPORT void CreateValid(const Dataset* dataset);

630
  void InitTrain(const std::vector<int8_t>& is_feature_used,
631
                 TrainingShareStates* share_state) const;
632

Guolin Ke's avatar
Guolin Ke committed
633
634
635
636
637
638
639
640
641
642
643
  template <bool USE_INDICES, bool USE_HESSIAN>
  void ConstructHistogramsInner(const std::vector<int8_t>& is_feature_used,
                                const data_size_t* data_indices,
                                data_size_t num_data, const score_t* gradients,
                                const score_t* hessians,
                                score_t* ordered_gradients,
                                score_t* ordered_hessians,
                                TrainingShareStates* share_state,
                                hist_t* hist_data) const;

  template <bool USE_INDICES, bool ORDERED>
644
645
646
647
  void ConstructHistogramsMultiVal(const data_size_t* data_indices,
                                   data_size_t num_data,
                                   const score_t* gradients,
                                   const score_t* hessians,
648
                                   TrainingShareStates* share_state,
Guolin Ke's avatar
Guolin Ke committed
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
                                   hist_t* hist_data) const;

  inline void ConstructHistograms(
      const std::vector<int8_t>& is_feature_used,
      const data_size_t* data_indices, data_size_t num_data,
      const score_t* gradients, const score_t* hessians,
      score_t* ordered_gradients, score_t* ordered_hessians,
      TrainingShareStates* share_state, hist_t* hist_data) const {
    if (num_data <= 0) {
      return;
    }
    bool use_indices = data_indices != nullptr && (num_data < num_data_);
    if (share_state->is_constant_hessian) {
      if (use_indices) {
        ConstructHistogramsInner<true, false>(
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      } else {
        ConstructHistogramsInner<false, false>(
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      }
    } else {
      if (use_indices) {
        ConstructHistogramsInner<true, true>(
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      } else {
        ConstructHistogramsInner<false, true>(
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      }
    }
  }
Guolin Ke's avatar
Guolin Ke committed
683

684
  void FixHistogram(int feature_idx, double sum_gradient, double sum_hessian, hist_t* data) const;
Guolin Ke's avatar
Guolin Ke committed
685

686
687
  inline data_size_t Split(int feature, const uint32_t* threshold,
                           int num_threshold, bool default_left,
688
                           const data_size_t* data_indices,
689
690
                           data_size_t cnt, data_size_t* lte_indices,
                           data_size_t* gt_indices) const {
Guolin Ke's avatar
Guolin Ke committed
691
692
    const int group = feature2group_[feature];
    const int sub_feature = feature2subfeature_[feature];
693
694
695
    return feature_groups_[group]->Split(
        sub_feature, threshold, num_threshold, default_left, data_indices,
        cnt, lte_indices, gt_indices);
Guolin Ke's avatar
Guolin Ke committed
696
697
698
699
700
701
702
703
704
705
706
707
708
709
  }

  inline int SubFeatureBinOffset(int i) const {
    const int sub_feature = feature2subfeature_[i];
    if (sub_feature == 0) {
      return 1;
    } else {
      return 0;
    }
  }

  inline int FeatureNumBin(int i) const {
    const int group = feature2group_[i];
    const int sub_feature = feature2subfeature_[i];
710
    return feature_groups_[group]->bin_mappers_[sub_feature]->num_bin();
Guolin Ke's avatar
Guolin Ke committed
711
  }
Guolin Ke's avatar
Guolin Ke committed
712

713
714
715
  inline int FeatureGroupNumBin(int group) const {
    return feature_groups_[group]->num_total_bin_;
  }
716

Guolin Ke's avatar
Guolin Ke committed
717
718
719
720
721
722
  inline const BinMapper* FeatureBinMapper(int i) const {
    const int group = feature2group_[i];
    const int sub_feature = feature2subfeature_[i];
    return feature_groups_[group]->bin_mappers_[sub_feature].get();
  }

723
724
725
726
  inline const Bin* FeatureGroupBin(int group) const {
    return feature_groups_[group]->bin_data_.get();
  }

Guolin Ke's avatar
Guolin Ke committed
727
728
729
  inline BinIterator* FeatureIterator(int i) const {
    const int group = feature2group_[i];
    const int sub_feature = feature2subfeature_[i];
zhangyafeikimi's avatar
zhangyafeikimi committed
730
    return feature_groups_[group]->SubFeatureIterator(sub_feature);
Guolin Ke's avatar
Guolin Ke committed
731
732
  }

733
734
735
  inline BinIterator* FeatureGroupIterator(int group) const {
    return feature_groups_[group]->FeatureGroupIterator();
  }
736

737
738
739
740
  inline bool IsMultiGroup(int i) const {
    return feature_groups_[i]->is_multi_val_;
  }

741
742
743
744
745
746
747
748
  inline size_t FeatureGroupSizesInByte(int group) const {
    return feature_groups_[group]->FeatureGroupSizesInByte();
  }

  inline void* FeatureGroupData(int group) const {
    return feature_groups_[group]->FeatureGroupData();
  }

749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
  const void* GetColWiseData(
    const int feature_group_index,
    const int sub_feature_index,
    uint8_t* bit_type,
    bool* is_sparse,
    std::vector<BinIterator*>* bin_iterator,
    const int num_threads) const;

  const void* GetColWiseData(
    const int feature_group_index,
    const int sub_feature_index,
    uint8_t* bit_type,
    bool* is_sparse,
    BinIterator** bin_iterator) const;

Guolin Ke's avatar
Guolin Ke committed
764
765
766
767
768
769
  inline double RealThreshold(int i, uint32_t threshold) const {
    const int group = feature2group_[i];
    const int sub_feature = feature2subfeature_[i];
    return feature_groups_[group]->bin_mappers_[sub_feature]->BinToValue(threshold);
  }

770
771
772
773
774
775
776
  // given a real threshold, find the closest threshold bin
  inline uint32_t BinThreshold(int i, double threshold_double) const {
    const int group = feature2group_[i];
    const int sub_feature = feature2subfeature_[i];
    return feature_groups_[group]->bin_mappers_[sub_feature]->ValueToBin(threshold_double);
  }

777
778
779
780
781
782
  inline int MaxRealCatValue(int i) const {
    const int group = feature2group_[i];
    const int sub_feature = feature2subfeature_[i];
    return feature_groups_[group]->bin_mappers_[sub_feature]->MaxCatValue();
  }

Guolin Ke's avatar
Guolin Ke committed
783
784
785
786
787
788
789
790
791
  /*!
  * \brief Get meta data pointer
  * \return Pointer of meta data
  */
  inline const Metadata& metadata() const { return metadata_; }

  /*! \brief Get Number of used features */
  inline int num_features() const { return num_features_; }

792
793
794
  /*! \brief Get number of numeric features */
  inline int num_numeric_features() const { return num_numeric_features_; }

795
796
797
  /*! \brief Get Number of feature groups */
  inline int num_feature_groups() const { return num_groups_;}

798
799
800
  /*! \brief Get Number of total features */
  inline int num_total_features() const { return num_total_features_; }

Guolin Ke's avatar
Guolin Ke committed
801
802
803
804
  /*! \brief Get the index of label column */
  inline int label_idx() const { return label_idx_; }

  /*! \brief Get names of current data set */
Guolin Ke's avatar
Guolin Ke committed
805
806
  inline const std::vector<std::string>& feature_names() const { return feature_names_; }

807
808
809
  /*! \brief Get content of parser config file */
  inline const std::string parser_config_str() const { return parser_config_str_; }

Guolin Ke's avatar
Guolin Ke committed
810
811
  inline void set_feature_names(const std::vector<std::string>& feature_names) {
    if (feature_names.size() != static_cast<size_t>(num_total_features_)) {
812
      Log::Fatal("Size of feature_names error, should equal with total number of features");
Guolin Ke's avatar
Guolin Ke committed
813
814
    }
    feature_names_ = std::vector<std::string>(feature_names);
Guolin Ke's avatar
Guolin Ke committed
815
    std::unordered_set<std::string> feature_name_set;
816
817
    // replace ' ' in feature_names with '_'
    bool spaceInFeatureName = false;
818
    for (auto& feature_name : feature_names_) {
Andrew Ziem's avatar
Andrew Ziem committed
819
      // check JSON
820
821
      if (!Common::CheckAllowedJSON(feature_name)) {
        Log::Fatal("Do not support special JSON characters in feature name.");
822
      }
823
      if (feature_name.find(' ') != std::string::npos) {
824
825
826
        spaceInFeatureName = true;
        std::replace(feature_name.begin(), feature_name.end(), ' ', '_');
      }
Guolin Ke's avatar
Guolin Ke committed
827
828
829
830
      if (feature_name_set.count(feature_name) > 0) {
        Log::Fatal("Feature (%s) appears more than one time.", feature_name.c_str());
      }
      feature_name_set.insert(feature_name);
831
    }
832
    if (spaceInFeatureName) {
Andrew Ziem's avatar
Andrew Ziem committed
833
      Log::Warning("Found whitespace in feature_names, replace with underlines");
834
    }
Guolin Ke's avatar
Guolin Ke committed
835
  }
Guolin Ke's avatar
Guolin Ke committed
836

Guolin Ke's avatar
Guolin Ke committed
837
838
  inline std::vector<std::string> feature_infos() const {
    std::vector<std::string> bufs;
839
    for (int i = 0; i < num_total_features_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
840
      int fidx = used_feature_map_[i];
841
      if (fidx < 0) {
Guolin Ke's avatar
Guolin Ke committed
842
843
844
        bufs.push_back("none");
      } else {
        const auto bin_mapper = FeatureBinMapper(fidx);
845
        bufs.push_back(bin_mapper->bin_info_string());
Guolin Ke's avatar
Guolin Ke committed
846
847
848
849
850
      }
    }
    return bufs;
  }

Guolin Ke's avatar
Guolin Ke committed
851
852
853
  /*! \brief Get Number of data */
  inline data_size_t num_data() const { return num_data_; }

854
855
856
  /*! \brief Get whether FinishLoad is automatically called when pushing last row. */
  inline bool wait_for_manual_finish() const { return wait_for_manual_finish_; }

857
858
859
  /*! \brief Get the maximum number of OpenMP threads to allocate for. */
  inline int omp_max_threads() const { return omp_max_threads_; }

860
861
862
863
864
865
866
867
868
  /*! \brief Set whether the Dataset is finished automatically when last row is pushed or with a manual
   *         MarkFinished API call.  Set to true for thread-safe streaming and/or if will be coalesced later.
   *         FinishLoad should not be called on any Dataset that will be coalesced.
   */
  inline void set_wait_for_manual_finish(bool value) {
    std::lock_guard<std::mutex> lock(mutex_);
    wait_for_manual_finish_ = value;
  }

Guolin Ke's avatar
Guolin Ke committed
869
870
871
872
873
  /*! \brief Disable copy */
  Dataset& operator=(const Dataset&) = delete;
  /*! \brief Disable copy */
  Dataset(const Dataset&) = delete;

874
  void AddFeaturesFrom(Dataset* other);
875

876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
  /*! \brief Get has_raw_ */
  inline bool has_raw() const { return has_raw_; }

  /*! \brief Set has_raw_ */
  inline void SetHasRaw(bool has_raw) { has_raw_ = has_raw; }

  /*! \brief Resize raw_data_ */
  inline void ResizeRaw(int num_rows) {
    if (static_cast<int>(raw_data_.size()) > num_numeric_features_) {
      raw_data_.resize(num_numeric_features_);
    }
    for (size_t i = 0; i < raw_data_.size(); ++i) {
      raw_data_[i].resize(num_rows);
    }
    int curr_size = static_cast<int>(raw_data_.size());
    for (int i = curr_size; i < num_numeric_features_; ++i) {
      raw_data_.push_back(std::vector<float>(num_rows, 0));
    }
  }

  /*! \brief Get pointer to raw_data_ feature */
  inline const float* raw_index(int feat_ind) const {
    return raw_data_[numeric_feature_map_[feat_ind]].data();
  }

901
902
903
904
905
906
907
908
909
910
911
912
  inline uint32_t feature_max_bin(const int inner_feature_index) const {
    const int feature_group_index = Feature2Group(inner_feature_index);
    const int sub_feature_index = feature2subfeature_[inner_feature_index];
    return feature_groups_[feature_group_index]->feature_max_bin(sub_feature_index);
  }

  inline uint32_t feature_min_bin(const int inner_feature_index) const {
    const int feature_group_index = Feature2Group(inner_feature_index);
    const int sub_feature_index = feature2subfeature_[inner_feature_index];
    return feature_groups_[feature_group_index]->feature_min_bin(sub_feature_index);
  }

913
  #ifdef USE_CUDA
914
915
916
917
918

  const CUDAColumnData* cuda_column_data() const {
    return cuda_column_data_.get();
  }

919
  #endif  // USE_CUDA
920

Nikita Titov's avatar
Nikita Titov committed
921
 private:
922
923
  void CreateCUDAColumnData();

Guolin Ke's avatar
Guolin Ke committed
924
  std::string data_filename_;
Guolin Ke's avatar
Guolin Ke committed
925
  /*! \brief Store used features */
Guolin Ke's avatar
Guolin Ke committed
926
  std::vector<std::unique_ptr<FeatureGroup>> feature_groups_;
Guolin Ke's avatar
Guolin Ke committed
927
928
929
930
  /*! \brief Mapper from real feature index to used index*/
  std::vector<int> used_feature_map_;
  /*! \brief Number of used features*/
  int num_features_;
931
932
  /*! \brief Number of total features*/
  int num_total_features_;
Guolin Ke's avatar
Guolin Ke committed
933
934
935
936
  /*! \brief Number of total data*/
  data_size_t num_data_;
  /*! \brief Store some label level data*/
  Metadata metadata_;
Guolin Ke's avatar
Guolin Ke committed
937
938
939
940
  /*! \brief index of label column */
  int label_idx_ = 0;
  /*! \brief store feature names */
  std::vector<std::string> feature_names_;
941
942
  /*! \brief store feature names */
  static const char* binary_file_token;
Guolin Ke's avatar
Guolin Ke committed
943
944
945
946
947
948
949
  int num_groups_;
  std::vector<int> real_feature_idx_;
  std::vector<int> feature2group_;
  std::vector<int> feature2subfeature_;
  std::vector<uint64_t> group_bin_boundaries_;
  std::vector<int> group_feature_start_;
  std::vector<int> group_feature_cnt_;
Guolin Ke's avatar
Guolin Ke committed
950
  bool is_finish_load_;
951
  int max_bin_;
Belinda Trotta's avatar
Belinda Trotta committed
952
  std::vector<int32_t> max_bin_by_feature_;
953
  std::vector<std::vector<double>> forced_bin_bounds_;
954
955
956
957
  int bin_construct_sample_cnt_;
  int min_data_in_bin_;
  bool use_missing_;
  bool zero_as_missing_;
Guolin Ke's avatar
Guolin Ke committed
958
  std::vector<int> feature_need_push_zeros_;
959
  std::vector<std::vector<float>> raw_data_;
960
  bool wait_for_manual_finish_;
961
  int omp_max_threads_ = -1;
962
963
964
965
  bool has_raw_;
  /*! map feature (inner index) to its index in the list of numeric (non-categorical) features */
  std::vector<int> numeric_feature_map_;
  int num_numeric_features_;
966
967
  std::string device_type_;
  int gpu_device_id_;
968
969
  /*! \brief mutex for threading safe call */
  std::mutex mutex_;
970

971
  #ifdef USE_CUDA
972
  std::unique_ptr<CUDAColumnData> cuda_column_data_;
973
  #endif  // USE_CUDA
974

975
  std::string parser_config_str_;
Guolin Ke's avatar
Guolin Ke committed
976
977
978
979
};

}  // namespace LightGBM

Guolin Ke's avatar
Guolin Ke committed
980
#endif   // LightGBM_DATA_H_