dataset.h 36.9 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
#ifndef LIGHTGBM_INCLUDE_LIGHTGBM_DATASET_H_
#define LIGHTGBM_INCLUDE_LIGHTGBM_DATASET_H_
Guolin Ke's avatar
Guolin Ke committed
7

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

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

28
29
30
#include <LightGBM/cuda/cuda_column_data.hpp>
#include <LightGBM/cuda/cuda_metadata.hpp>

Guolin Ke's avatar
Guolin Ke committed
31
32
33
namespace LightGBM {

/*! \brief forward declaration */
Guolin Ke's avatar
Guolin Ke committed
34
class DatasetLoader;
Guolin Ke's avatar
Guolin Ke committed
35
/*!
Hui Xue's avatar
Hui Xue committed
36
* \brief This class is used to store some meta(non-feature) data for training data,
Andrew Ziem's avatar
Andrew Ziem committed
37
*        e.g. labels, weights, initial scores, query level information.
Guolin Ke's avatar
Guolin Ke committed
38
*
Qiwei Ye's avatar
Qiwei Ye committed
39
*        Some details:
40
*        1. Label, used for training.
Qiwei Ye's avatar
Qiwei Ye committed
41
*        2. Weights, weighs of records, optional
Andrew Ziem's avatar
Andrew Ziem committed
42
*        3. Query Boundaries, necessary for LambdaRank.
43
44
45
46
*           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
47
48
*/
class Metadata {
Nikita Titov's avatar
Nikita Titov committed
49
 public:
50
  /*!
51
  * \brief Null constructor
Guolin Ke's avatar
Guolin Ke committed
52
53
54
  */
  Metadata();
  /*!
Andrew Ziem's avatar
Andrew Ziem committed
55
  * \brief Initialization will load query level information, since it is need for sampling data
Guolin Ke's avatar
Guolin Ke committed
56
57
  * \param data_filename Filename of data
  */
58
  void Init(const char* data_filename);
Guolin Ke's avatar
Guolin Ke committed
59
  /*!
Guolin Ke's avatar
Guolin Ke committed
60
61
  * \brief init as subset
  * \param metadata Filename of data
62
  * \param used_indices
Guolin Ke's avatar
Guolin Ke committed
63
64
65
66
  * \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
67
68
69
70
71
72
73
74
  * \brief Initial with binary memory
  * \param memory Pointer to memory
  */
  void LoadFromMemory(const void* memory);
  /*! \brief Destructor */
  ~Metadata();

  /*!
75
  * \brief Initial work, will allocate space for label, weight (if exists) and query (if exists)
Guolin Ke's avatar
Guolin Ke committed
76
  * \param num_data Number of training data
Guolin Ke's avatar
Guolin Ke committed
77
78
  * \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
79
  */
80
  void Init(data_size_t num_data, int weight_idx, int query_idx);
Guolin Ke's avatar
Guolin Ke committed
81

82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
  /*!
  * \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
99
100
  /*!
  * \brief Partition label by used indices
101
  * \param used_indices Indices of local used
Guolin Ke's avatar
Guolin Ke committed
102
103
104
105
106
  */
  void PartitionLabel(const std::vector<data_size_t>& used_indices);

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

113
  void SetLabel(const label_t* label, data_size_t len);
114
  void SetLabel(const ArrowChunkedArray& array);
Guolin Ke's avatar
Guolin Ke committed
115

116
  void SetWeights(const label_t* weights, data_size_t len);
117
  void SetWeights(const ArrowChunkedArray& array);
Guolin Ke's avatar
Guolin Ke committed
118

Guolin Ke's avatar
Guolin Ke committed
119
  void SetQuery(const data_size_t* query, data_size_t len);
120
  void SetQuery(const ArrowChunkedArray& array);
Guolin Ke's avatar
Guolin Ke committed
121

122
123
  void SetPosition(const data_size_t* position, data_size_t len);

Guolin Ke's avatar
Guolin Ke committed
124
125
126
127
  /*!
  * \brief Set initial scores
  * \param init_score Initial scores, this class will manage memory for init_score.
  */
128
  void SetInitScore(const double* init_score, data_size_t len);
129
  void SetInitScore(const ArrowChunkedArray& array);
Guolin Ke's avatar
Guolin Ke committed
130
131
132
133
134
135


  /*!
  * \brief Save binary data to file
  * \param file File want to write
  */
136
  void SaveBinaryToFile(BinaryWriter* writer) const;
Guolin Ke's avatar
Guolin Ke committed
137
138
139
140
141
142
143
144
145
146

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

  /*!
  * \brief Get pointer of label
  * \return Pointer of label
  */
147
  inline const label_t* label() const { return label_.data(); }
Guolin Ke's avatar
Guolin Ke committed
148
149
150
151
152
153

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

Guolin Ke's avatar
Guolin Ke committed
158
159
160
161
162
  /*!
  * \brief Set Weight for one record
  * \param idx Index of this record
  * \param value Weight value of this record
  */
163
  inline void SetWeightAt(data_size_t idx, label_t value) {
164
    weights_[idx] = value;
Guolin Ke's avatar
Guolin Ke committed
165
166
  }

167
168
169
170
171
172
  /*!
  * \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) {
173
    const auto nclasses = num_init_score_classes();
174
175
176
177
178
179
    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
180
181
182
183
184
  /*!
  * \brief Set Query Id for one record
  * \param idx Index of this record
  * \param value Query Id value of this record
  */
185
  inline void SetQueryAt(data_size_t idx, data_size_t value) {
Guolin Ke's avatar
Guolin Ke committed
186
187
188
    queries_[idx] = static_cast<data_size_t>(value);
  }

189
190
191
  /*! \brief Load initial scores from file */
  void LoadInitialScore(const std::string& data_filename);

192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
  /*!
  * \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
212
  /*!
Hui Xue's avatar
Hui Xue committed
213
  * \brief Get weights, if not exists, will return nullptr
Guolin Ke's avatar
Guolin Ke committed
214
215
  * \return Pointer of weights
  */
216
  inline const label_t* weights() const {
Guolin Ke's avatar
Guolin Ke committed
217
    if (!weights_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
218
219
220
221
222
      return weights_.data();
    } else {
      return nullptr;
    }
  }
Guolin Ke's avatar
Guolin Ke committed
223

224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
  /*!
  * \brief Get positions, if does not exist then return nullptr
  * \return Pointer of positions
  */
  inline const data_size_t* positions() const {
    if (!positions_.empty()) {
      return positions_.data();
    } else {
      return nullptr;
    }
  }

  /*!
  * \brief Get position IDs, if does not exist then return nullptr
  * \return Pointer of position IDs
  */
  inline const std::string* position_ids() const {
    if (!position_ids_.empty()) {
      return position_ids_.data();
    } else {
      return nullptr;
    }
  }

  /*!
  * \brief Get Number of different position IDs
  * \return number of different position IDs
  */
  inline size_t num_position_ids() const {
      return position_ids_.size();
  }

Guolin Ke's avatar
Guolin Ke committed
256
  /*!
Hui Xue's avatar
Hui Xue committed
257
  * \brief Get data boundaries on queries, if not exists, will return nullptr
258
  *        we assume data will order by query,
Guolin Ke's avatar
Guolin Ke committed
259
260
261
262
  *        the interval of [query_boundaris[i], query_boundaris[i+1])
  *        is the data indices for query i.
  * \return Pointer of data boundaries on queries
  */
263
  inline const data_size_t* query_boundaries() const {
Guolin Ke's avatar
Guolin Ke committed
264
    if (!query_boundaries_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
265
266
267
268
269
      return query_boundaries_.data();
    } else {
      return nullptr;
    }
  }
Guolin Ke's avatar
Guolin Ke committed
270
271
272
273
274

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

  /*!
Hui Xue's avatar
Hui Xue committed
278
  * \brief Get weights for queries, if not exists, will return nullptr
Guolin Ke's avatar
Guolin Ke committed
279
280
  * \return Pointer of weights for queries
  */
281
  inline const label_t* query_weights() const {
Guolin Ke's avatar
Guolin Ke committed
282
    if (!query_weights_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
283
284
285
286
287
      return query_weights_.data();
    } else {
      return nullptr;
    }
  }
Guolin Ke's avatar
Guolin Ke committed
288
289

  /*!
Hui Xue's avatar
Hui Xue committed
290
  * \brief Get initial scores, if not exists, will return nullptr
Guolin Ke's avatar
Guolin Ke committed
291
292
  * \return Pointer of initial scores
  */
293
  inline const double* init_score() const {
Guolin Ke's avatar
Guolin Ke committed
294
    if (!init_score_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
295
296
297
298
299
      return init_score_.data();
    } else {
      return nullptr;
    }
  }
Guolin Ke's avatar
Guolin Ke committed
300

301
302
303
  /*!
  * \brief Get size of initial scores
  */
Guolin Ke's avatar
Guolin Ke committed
304
  inline int64_t num_init_score() const { return num_init_score_; }
305

306
307
308
  /*!
  * \brief Get number of classes
  */
309
  inline int32_t num_init_score_classes() const {
310
311
312
313
314
315
    if (num_data_ && num_init_score_) {
      return static_cast<int>(num_init_score_ / num_data_);
    }
    return 1;
  }

Guolin Ke's avatar
Guolin Ke committed
316
317
318
319
  /*! \brief Disable copy */
  Metadata& operator=(const Metadata&) = delete;
  /*! \brief Disable copy */
  Metadata(const Metadata&) = delete;
Guolin Ke's avatar
Guolin Ke committed
320

321
  #ifdef USE_CUDA
322
323
324
325
326

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

  void CreateCUDAMetadata(const int gpu_device_id);

327
  #endif  // USE_CUDA
328

Nikita Titov's avatar
Nikita Titov committed
329
 private:
Guolin Ke's avatar
Guolin Ke committed
330
331
  /*! \brief Load wights from file */
  void LoadWeights();
332
333
  /*! \brief Load positions from file */
  void LoadPositions();
Guolin Ke's avatar
Guolin Ke committed
334
335
  /*! \brief Load query boundaries from file */
  void LoadQueryBoundaries();
336
337
338
339
340
341
  /*! \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);
342
343
344
  /*! \brief Set labels from pointers to the first element and the end of an iterator. */
  template <typename It>
  void SetLabelsFromIterator(It first, It last);
345
346
  /*! \brief Insert weights at the given index */
  void InsertWeights(const label_t* weights, data_size_t start_index, data_size_t len);
347
348
349
  /*! \brief Set weights from pointers to the first element and the end of an iterator. */
  template <typename It>
  void SetWeightsFromIterator(It first, It last);
350
351
  /*! \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);
352
353
354
  /*! \brief Set init scores from pointers to the first element and the end of an iterator. */
  template <typename It>
  void SetInitScoresFromIterator(It first, It last);
355
356
  /*! \brief Insert queries at the given index */
  void InsertQueries(const data_size_t* queries, data_size_t start_index, data_size_t len);
357
358
359
  /*! \brief Set queries from pointers to the first element and the end of an iterator. */
  template <typename It>
  void SetQueriesFromIterator(It first, It last);
Guolin Ke's avatar
Guolin Ke committed
360
  /*! \brief Filename of current data */
Guolin Ke's avatar
Guolin Ke committed
361
  std::string data_filename_;
Guolin Ke's avatar
Guolin Ke committed
362
363
364
365
  /*! \brief Number of data */
  data_size_t num_data_;
  /*! \brief Number of weights, used to check correct weight file */
  data_size_t num_weights_;
366
367
  /*! \brief Number of positions, used to check correct position file */
  data_size_t num_positions_;
Guolin Ke's avatar
Guolin Ke committed
368
  /*! \brief Label data */
369
  std::vector<label_t> label_;
Guolin Ke's avatar
Guolin Ke committed
370
  /*! \brief Weights data */
371
  std::vector<label_t> weights_;
372
373
374
375
  /*! \brief Positions data */
  std::vector<data_size_t> positions_;
  /*! \brief Position identifiers */
  std::vector<std::string> position_ids_;
Guolin Ke's avatar
Guolin Ke committed
376
  /*! \brief Query boundaries */
Guolin Ke's avatar
Guolin Ke committed
377
  std::vector<data_size_t> query_boundaries_;
Guolin Ke's avatar
Guolin Ke committed
378
  /*! \brief Query weights */
379
  std::vector<label_t> query_weights_;
380
  /*! \brief Number of queries */
Guolin Ke's avatar
Guolin Ke committed
381
382
  data_size_t num_queries_;
  /*! \brief Number of Initial score, used to check correct weight file */
Guolin Ke's avatar
Guolin Ke committed
383
  int64_t num_init_score_;
Guolin Ke's avatar
Guolin Ke committed
384
  /*! \brief Initial score */
Guolin Ke's avatar
Guolin Ke committed
385
  std::vector<double> init_score_;
Guolin Ke's avatar
Guolin Ke committed
386
  /*! \brief Queries data */
Guolin Ke's avatar
Guolin Ke committed
387
  std::vector<data_size_t> queries_;
388
389
  /*! \brief mutex for threading safe call */
  std::mutex mutex_;
390
  bool weight_load_from_file_;
391
  bool position_load_from_file_;
392
393
  bool query_load_from_file_;
  bool init_score_load_from_file_;
394
  #ifdef USE_CUDA
395
  std::unique_ptr<CUDAMetadata> cuda_metadata_;
396
  #endif  // USE_CUDA
Guolin Ke's avatar
Guolin Ke committed
397
398
399
400
401
};


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

405
406
407
408
409
410
411
412
  /*! \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
413
414
415
416
417
418
  /*! \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
419
420
  * \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
421
422
  */
  virtual void ParseOneLine(const char* str,
423
                            std::vector<std::pair<int, double>>* out_features, double* out_label) const = 0;
Guolin Ke's avatar
Guolin Ke committed
424

425
  virtual int NumFeatures() const = 0;
Guolin Ke's avatar
Guolin Ke committed
426

Guolin Ke's avatar
Guolin Ke committed
427
  /*!
428
  * \brief Create an object of parser, will auto choose the format depend on file
Guolin Ke's avatar
Guolin Ke committed
429
  * \param filename One Filename of data
430
  * \param header whether input file contains header
431
  * \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
432
  * \param label_idx index of label column
Chen Yufei's avatar
Chen Yufei committed
433
  * \param precise_float_parser using precise floating point number parsing if true
Guolin Ke's avatar
Guolin Ke committed
434
435
  * \return Object of parser
  */
Chen Yufei's avatar
Chen Yufei committed
436
  static Parser* CreateParser(const char* filename, bool header, int num_features, int label_idx, bool precise_float_parser);
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481

  /*!
  * \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
482
483
484
};

/*! \brief The main class of data set,
485
*          which are used to training or validation
Guolin Ke's avatar
Guolin Ke committed
486
487
*/
class Dataset {
Nikita Titov's avatar
Nikita Titov committed
488
 public:
Guolin Ke's avatar
Guolin Ke committed
489
  friend DatasetLoader;
Guolin Ke's avatar
Guolin Ke committed
490

491
  LIGHTGBM_EXPORT Dataset();
Guolin Ke's avatar
Guolin Ke committed
492

493
  LIGHTGBM_EXPORT Dataset(data_size_t num_data);
Guolin Ke's avatar
Guolin Ke committed
494

Guolin Ke's avatar
Guolin Ke committed
495
  void Construct(
Guolin Ke's avatar
Guolin Ke committed
496
    std::vector<std::unique_ptr<BinMapper>>* bin_mappers,
497
    int num_total_features,
498
    const std::vector<std::vector<double>>& forced_bins,
499
    int** sample_non_zero_indices,
Guolin Ke's avatar
Guolin Ke committed
500
    double** sample_values,
501
    const int* num_per_col,
502
    int num_sample_col,
Guolin Ke's avatar
Guolin Ke committed
503
    size_t total_sample_cnt,
Guolin Ke's avatar
Guolin Ke committed
504
    const Config& io_config);
Guolin Ke's avatar
Guolin Ke committed
505

Guolin Ke's avatar
Guolin Ke committed
506
  /*! \brief Destructor */
507
  LIGHTGBM_EXPORT ~Dataset();
Guolin Ke's avatar
Guolin Ke committed
508

509
510
511
512
513
514
515
516
517
518
519
520
521
522
  /*!
  * \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,
523
524
525
526
527
528
529
530
531
                                     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();
    }

532
533
    metadata_.Init(num_data, has_weights, has_init_scores, has_queries, nclasses);
    for (int i = 0; i < num_groups_; ++i) {
534
      feature_groups_[i]->InitStreaming(nthreads, omp_max_threads_);
535
536
537
    }
  }

538
  LIGHTGBM_EXPORT bool CheckAlign(const Dataset& other) const {
539
540
541
542
543
544
545
546
547
548
    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
549
      if (!FeatureBinMapper(i)->CheckAlign(*(other.FeatureBinMapper(i)))) {
550
551
552
553
554
555
        return false;
      }
    }
    return true;
  }

Guolin Ke's avatar
Guolin Ke committed
556
  inline void FinishOneRow(int tid, data_size_t row_idx, const std::vector<bool>& is_feature_added) {
557
558
559
    if (is_finish_load_) {
      return;
    }
Guolin Ke's avatar
Guolin Ke committed
560
    for (auto fidx : feature_need_push_zeros_) {
561
562
563
      if (is_feature_added[fidx]) {
        continue;
      }
Guolin Ke's avatar
Guolin Ke committed
564
565
566
567
568
569
      const int group = feature2group_[fidx];
      const int sub_feature = feature2subfeature_[fidx];
      feature_groups_[group]->PushData(tid, sub_feature, row_idx, 0.0f);
    }
  }

570
571
572
573
574
575
576
577
578
579
580
581
  inline void PushOneValue(int tid, data_size_t row_idx, size_t col_idx, double value) {
    if (this->is_finish_load_)
      return;
    auto feature_idx = this->used_feature_map_[col_idx];
    if (feature_idx >= 0) {
      auto group = this->feature2group_[feature_idx];
      auto sub_feature = this->feature2subfeature_[feature_idx];
      this->feature_groups_[group]->PushData(tid, sub_feature, row_idx, value);
      if (this->has_raw_) {
        auto feat_ind = numeric_feature_map_[feature_idx];
        if (feat_ind >= 0) {
          raw_data_[feat_ind][row_idx] = static_cast<float>(value);
582
        }
Guolin Ke's avatar
Guolin Ke committed
583
584
585
586
      }
    }
  }

587
588
589
590
591
592
  inline void PushOneRow(int tid, data_size_t row_idx, const std::vector<double>& feature_values) {
    for (size_t i = 0; i < feature_values.size() && i < static_cast<size_t>(num_total_features_); ++i) {
      this->PushOneValue(tid, row_idx, i, feature_values[i]);
    }
  }

593
  inline void PushOneRow(int tid, data_size_t row_idx, const std::vector<std::pair<int, double>>& feature_values) {
594
595
596
    if (is_finish_load_) {
      return;
    }
Guolin Ke's avatar
Guolin Ke committed
597
    std::vector<bool> is_feature_added(num_features_, false);
598
    for (auto& inner_data : feature_values) {
599
600
601
      if (inner_data.first >= num_total_features_) {
        continue;
      }
602
603
      int feature_idx = used_feature_map_[inner_data.first];
      if (feature_idx >= 0) {
Guolin Ke's avatar
Guolin Ke committed
604
        is_feature_added[feature_idx] = true;
Guolin Ke's avatar
Guolin Ke committed
605
606
607
        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);
608
609
610
        if (has_raw_) {
          int feat_ind = numeric_feature_map_[feature_idx];
          if (feat_ind >= 0) {
611
            raw_data_[feat_ind][row_idx] = static_cast<float>(inner_data.second);
612
613
          }
        }
614
615
      }
    }
Guolin Ke's avatar
Guolin Ke committed
616
    FinishOneRow(tid, row_idx, is_feature_added);
617
618
  }

619
  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
620
    feature_groups_[group]->PushData(tid, sub_feature, row_idx, value);
621
622
623
    if (has_raw_) {
      int feat_ind = numeric_feature_map_[feature_idx];
      if (feat_ind >= 0) {
624
        raw_data_[feat_ind][row_idx] = static_cast<float>(value);
625
626
      }
    }
Guolin Ke's avatar
Guolin Ke committed
627
628
  }

629
630
631
632
633
634
635
636
637
  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
638
639
640
641
642
  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
643
    return used_feature_map_[col_idx];
Guolin Ke's avatar
Guolin Ke committed
644
  }
Guolin Ke's avatar
Guolin Ke committed
645
646
647
648
649
650
  inline int Feature2Group(int feature_idx) const {
    return feature2group_[feature_idx];
  }
  inline int Feture2SubFeature(int feature_idx) const {
    return feature2subfeature_[feature_idx];
  }
651
652
653
  inline uint64_t GroupBinBoundary(int group_idx) const {
    return group_bin_boundaries_[group_idx];
  }
Guolin Ke's avatar
Guolin Ke committed
654
655
656
  inline uint64_t NumTotalBin() const {
    return group_bin_boundaries_.back();
  }
657

658
659
660
661
662
663
664
665
666
  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
667
668
  void ReSize(data_size_t num_data);

669
  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
670

671
  MultiValBin* GetMultiBinFromSparseFeatures(const std::vector<uint32_t>& offsets) const;
672

673
  MultiValBin* GetMultiBinFromAllFeatures(const std::vector<uint32_t>& offsets) const;
674

675
  template <bool USE_QUANT_GRAD, int HIST_BITS>
676
677
678
  TrainingShareStates* GetShareStates(
      score_t* gradients, score_t* hessians,
      const std::vector<int8_t>& is_feature_used, bool is_constant_hessian,
679
      bool force_col_wise, bool force_row_wise, const int num_grad_quant_bins) const;
680

681
  LIGHTGBM_EXPORT void FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
682

683
684
  bool SetFieldFromArrow(const char* field_name, const ArrowChunkedArray& ca);

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

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

689
  LIGHTGBM_EXPORT bool SetIntField(const char* field_name, const int* field_data, data_size_t num_element);
690

691
  LIGHTGBM_EXPORT bool GetFloatField(const char* field_name, data_size_t* out_len, const float** out_ptr);
692

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

695
  LIGHTGBM_EXPORT bool GetIntField(const char* field_name, data_size_t* out_len, const int** out_ptr);
696

Guolin Ke's avatar
Guolin Ke committed
697
698
699
  /*!
  * \brief Save current dataset into binary file, will save to "filename.bin"
  */
700
  LIGHTGBM_EXPORT void SaveBinaryFile(const char* bin_filename);
Guolin Ke's avatar
Guolin Ke committed
701

702
703
704
705
706
  /*!
   * \brief Serialize the overall Dataset definition/schema to a binary buffer (i.e., without data)
   */
  LIGHTGBM_EXPORT void SerializeReference(ByteBuffer* out);

707
708
  LIGHTGBM_EXPORT void DumpTextFile(const char* text_filename);

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

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

713
  void InitTrain(const std::vector<int8_t>& is_feature_used,
714
                 TrainingShareStates* share_state) const;
715

716
  template <bool USE_INDICES, bool USE_HESSIAN, bool USE_QUANT_GRAD, int HIST_BITS>
Guolin Ke's avatar
Guolin Ke committed
717
718
719
720
721
722
723
724
725
  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;

726
  template <bool USE_INDICES, bool ORDERED, bool USE_QUANT_GRAD, int HIST_BITS>
727
728
729
730
  void ConstructHistogramsMultiVal(const data_size_t* data_indices,
                                   data_size_t num_data,
                                   const score_t* gradients,
                                   const score_t* hessians,
731
                                   TrainingShareStates* share_state,
Guolin Ke's avatar
Guolin Ke committed
732
733
                                   hist_t* hist_data) const;

734
  template <bool USE_QUANT_GRAD, int HIST_BITS>
Guolin Ke's avatar
Guolin Ke committed
735
736
737
738
739
740
741
742
743
744
745
746
  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) {
747
        ConstructHistogramsInner<true, false, USE_QUANT_GRAD, HIST_BITS>(
Guolin Ke's avatar
Guolin Ke committed
748
749
750
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      } else {
751
        ConstructHistogramsInner<false, false, USE_QUANT_GRAD, HIST_BITS>(
Guolin Ke's avatar
Guolin Ke committed
752
753
754
755
756
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      }
    } else {
      if (use_indices) {
757
        ConstructHistogramsInner<true, true, USE_QUANT_GRAD, HIST_BITS>(
Guolin Ke's avatar
Guolin Ke committed
758
759
760
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      } else {
761
        ConstructHistogramsInner<false, true, USE_QUANT_GRAD, HIST_BITS>(
Guolin Ke's avatar
Guolin Ke committed
762
763
764
765
766
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      }
    }
  }
Guolin Ke's avatar
Guolin Ke committed
767

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

770
771
772
  template <typename PACKED_HIST_BIN_T, typename PACKED_HIST_ACC_T, int HIST_BITS_BIN, int HIST_BITS_ACC>
  void FixHistogramInt(int feature_idx, int64_t sum_gradient_and_hessian, hist_t* data) const;

773
774
  inline data_size_t Split(int feature, const uint32_t* threshold,
                           int num_threshold, bool default_left,
775
                           const data_size_t* data_indices,
776
777
                           data_size_t cnt, data_size_t* lte_indices,
                           data_size_t* gt_indices) const {
Guolin Ke's avatar
Guolin Ke committed
778
779
    const int group = feature2group_[feature];
    const int sub_feature = feature2subfeature_[feature];
780
781
782
    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
783
784
785
786
787
788
789
790
791
792
793
794
795
796
  }

  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];
797
    return feature_groups_[group]->bin_mappers_[sub_feature]->num_bin();
Guolin Ke's avatar
Guolin Ke committed
798
  }
Guolin Ke's avatar
Guolin Ke committed
799

800
801
802
  inline int FeatureGroupNumBin(int group) const {
    return feature_groups_[group]->num_total_bin_;
  }
803

Guolin Ke's avatar
Guolin Ke committed
804
805
806
807
808
809
  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();
  }

810
811
812
813
  inline const Bin* FeatureGroupBin(int group) const {
    return feature_groups_[group]->bin_data_.get();
  }

Guolin Ke's avatar
Guolin Ke committed
814
815
816
  inline BinIterator* FeatureIterator(int i) const {
    const int group = feature2group_[i];
    const int sub_feature = feature2subfeature_[i];
zhangyafeikimi's avatar
zhangyafeikimi committed
817
    return feature_groups_[group]->SubFeatureIterator(sub_feature);
Guolin Ke's avatar
Guolin Ke committed
818
819
  }

820
821
822
  inline BinIterator* FeatureGroupIterator(int group) const {
    return feature_groups_[group]->FeatureGroupIterator();
  }
823

824
825
826
827
  inline bool IsMultiGroup(int i) const {
    return feature_groups_[i]->is_multi_val_;
  }

828
829
830
831
832
833
834
835
  inline size_t FeatureGroupSizesInByte(int group) const {
    return feature_groups_[group]->FeatureGroupSizesInByte();
  }

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

836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
  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
851
852
853
854
855
856
  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);
  }

857
858
859
860
861
862
863
  // 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);
  }

864
865
866
867
868
869
  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
870
871
872
873
874
875
876
877
878
  /*!
  * \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_; }

879
880
881
  /*! \brief Get number of numeric features */
  inline int num_numeric_features() const { return num_numeric_features_; }

882
883
884
  /*! \brief Get Number of feature groups */
  inline int num_feature_groups() const { return num_groups_;}

885
886
887
  /*! \brief Get Number of total features */
  inline int num_total_features() const { return num_total_features_; }

Guolin Ke's avatar
Guolin Ke committed
888
889
890
891
  /*! \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
892
893
  inline const std::vector<std::string>& feature_names() const { return feature_names_; }

894
895
896
  /*! \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
897
898
  inline void set_feature_names(const std::vector<std::string>& feature_names) {
    if (feature_names.size() != static_cast<size_t>(num_total_features_)) {
899
      Log::Fatal("Size of feature_names error, should equal with total number of features");
Guolin Ke's avatar
Guolin Ke committed
900
901
    }
    feature_names_ = std::vector<std::string>(feature_names);
Guolin Ke's avatar
Guolin Ke committed
902
    std::unordered_set<std::string> feature_name_set;
903
904
    // replace ' ' in feature_names with '_'
    bool spaceInFeatureName = false;
905
    for (auto& feature_name : feature_names_) {
Andrew Ziem's avatar
Andrew Ziem committed
906
      // check JSON
907
908
      if (!Common::CheckAllowedJSON(feature_name)) {
        Log::Fatal("Do not support special JSON characters in feature name.");
909
      }
910
      if (feature_name.find(' ') != std::string::npos) {
911
912
913
        spaceInFeatureName = true;
        std::replace(feature_name.begin(), feature_name.end(), ' ', '_');
      }
Guolin Ke's avatar
Guolin Ke committed
914
915
916
917
      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);
918
    }
919
    if (spaceInFeatureName) {
Andrew Ziem's avatar
Andrew Ziem committed
920
      Log::Warning("Found whitespace in feature_names, replace with underlines");
921
    }
Guolin Ke's avatar
Guolin Ke committed
922
  }
Guolin Ke's avatar
Guolin Ke committed
923

Guolin Ke's avatar
Guolin Ke committed
924
925
  inline std::vector<std::string> feature_infos() const {
    std::vector<std::string> bufs;
926
    for (int i = 0; i < num_total_features_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
927
      int fidx = used_feature_map_[i];
928
      if (fidx < 0) {
Guolin Ke's avatar
Guolin Ke committed
929
930
931
        bufs.push_back("none");
      } else {
        const auto bin_mapper = FeatureBinMapper(fidx);
932
        bufs.push_back(bin_mapper->bin_info_string());
Guolin Ke's avatar
Guolin Ke committed
933
934
935
936
937
      }
    }
    return bufs;
  }

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

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

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

947
948
949
950
951
952
953
954
955
  /*! \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
956
957
958
959
960
  /*! \brief Disable copy */
  Dataset& operator=(const Dataset&) = delete;
  /*! \brief Disable copy */
  Dataset(const Dataset&) = delete;

961
  void AddFeaturesFrom(Dataset* other);
962

963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
  /*! \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();
  }

988
989
990
991
992
993
994
995
996
997
998
999
  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);
  }

1000
  #ifdef USE_CUDA
1001
1002
1003
1004
1005

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

1006
  #endif  // USE_CUDA
1007

Nikita Titov's avatar
Nikita Titov committed
1008
 private:
1009
1010
1011
1012
  void SerializeHeader(BinaryWriter* serializer);

  size_t GetSerializedHeaderSize();

1013
1014
  void CreateCUDAColumnData();

Guolin Ke's avatar
Guolin Ke committed
1015
  std::string data_filename_;
Guolin Ke's avatar
Guolin Ke committed
1016
  /*! \brief Store used features */
Guolin Ke's avatar
Guolin Ke committed
1017
  std::vector<std::unique_ptr<FeatureGroup>> feature_groups_;
Guolin Ke's avatar
Guolin Ke committed
1018
1019
1020
1021
  /*! \brief Mapper from real feature index to used index*/
  std::vector<int> used_feature_map_;
  /*! \brief Number of used features*/
  int num_features_;
1022
1023
  /*! \brief Number of total features*/
  int num_total_features_;
Guolin Ke's avatar
Guolin Ke committed
1024
1025
1026
1027
  /*! \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
1028
1029
1030
1031
  /*! \brief index of label column */
  int label_idx_ = 0;
  /*! \brief store feature names */
  std::vector<std::string> feature_names_;
1032
1033
1034
  /*! \brief serialized versions */
  static const int kSerializedReferenceVersionLength;
  static const char* serialized_reference_version;
1035
  static const char* binary_file_token;
1036
  static const char* binary_serialized_reference_token;
Guolin Ke's avatar
Guolin Ke committed
1037
1038
1039
1040
1041
1042
1043
  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
1044
  bool is_finish_load_;
1045
  int max_bin_;
Belinda Trotta's avatar
Belinda Trotta committed
1046
  std::vector<int32_t> max_bin_by_feature_;
1047
  std::vector<std::vector<double>> forced_bin_bounds_;
1048
1049
1050
1051
  int bin_construct_sample_cnt_;
  int min_data_in_bin_;
  bool use_missing_;
  bool zero_as_missing_;
Guolin Ke's avatar
Guolin Ke committed
1052
  std::vector<int> feature_need_push_zeros_;
1053
  std::vector<std::vector<float>> raw_data_;
1054
  bool wait_for_manual_finish_;
1055
  int omp_max_threads_ = -1;
1056
1057
1058
1059
  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_;
1060
1061
  std::string device_type_;
  int gpu_device_id_;
1062
1063
  /*! \brief mutex for threading safe call */
  std::mutex mutex_;
1064

1065
  #ifdef USE_CUDA
1066
  std::unique_ptr<CUDAColumnData> cuda_column_data_;
1067
  #endif  // USE_CUDA
1068

1069
  std::string parser_config_str_;
Guolin Ke's avatar
Guolin Ke committed
1070
1071
1072
1073
};

}  // namespace LightGBM

1074
#endif   // LIGHTGBM_INCLUDE_LIGHTGBM_DATASET_H_