dataset.h 36.8 KB
Newer Older
1
2
3
4
/*!
 * Copyright (c) 2016 Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See LICENSE file in the project root for license information.
 */
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
#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>

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

320
  #ifdef USE_CUDA
321
322
323
324
325

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

  void CreateCUDAMetadata(const int gpu_device_id);

326
  #endif  // USE_CUDA
327

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


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

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

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

Guolin Ke's avatar
Guolin Ke committed
426
  /*!
427
  * \brief Create an object of parser, will auto choose the format depend on file
Guolin Ke's avatar
Guolin Ke committed
428
  * \param filename One Filename of data
429
  * \param header whether input file contains header
430
  * \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
431
  * \param label_idx index of label column
Chen Yufei's avatar
Chen Yufei committed
432
  * \param precise_float_parser using precise floating point number parsing if true
Guolin Ke's avatar
Guolin Ke committed
433
434
  * \return Object of parser
  */
Chen Yufei's avatar
Chen Yufei committed
435
  static Parser* CreateParser(const char* filename, bool header, int num_features, int label_idx, bool precise_float_parser);
436
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

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

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

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

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

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

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

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

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

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

Guolin Ke's avatar
Guolin Ke committed
555
556
557
558
559
560
561
562
563
564
  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);
    }
  }

565
566
567
568
569
570
571
572
573
574
575
576
  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);
577
        }
Guolin Ke's avatar
Guolin Ke committed
578
579
580
581
      }
    }
  }

582
583
584
585
586
587
  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]);
    }
  }

588
  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
589
    if (is_finish_load_) { return; }
Guolin Ke's avatar
Guolin Ke committed
590
    std::vector<bool> is_feature_added(num_features_, false);
591
    for (auto& inner_data : feature_values) {
592
      if (inner_data.first >= num_total_features_) { continue; }
593
594
      int feature_idx = used_feature_map_[inner_data.first];
      if (feature_idx >= 0) {
Guolin Ke's avatar
Guolin Ke committed
595
        is_feature_added[feature_idx] = true;
Guolin Ke's avatar
Guolin Ke committed
596
597
598
        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);
599
600
601
        if (has_raw_) {
          int feat_ind = numeric_feature_map_[feature_idx];
          if (feat_ind >= 0) {
602
            raw_data_[feat_ind][row_idx] = static_cast<float>(inner_data.second);
603
604
          }
        }
605
606
      }
    }
Guolin Ke's avatar
Guolin Ke committed
607
    FinishOneRow(tid, row_idx, is_feature_added);
608
609
  }

610
  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
611
    feature_groups_[group]->PushData(tid, sub_feature, row_idx, value);
612
613
614
    if (has_raw_) {
      int feat_ind = numeric_feature_map_[feature_idx];
      if (feat_ind >= 0) {
615
        raw_data_[feat_ind][row_idx] = static_cast<float>(value);
616
617
      }
    }
Guolin Ke's avatar
Guolin Ke committed
618
619
  }

620
621
622
623
624
625
626
627
628
  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
629
630
631
632
633
  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
634
    return used_feature_map_[col_idx];
Guolin Ke's avatar
Guolin Ke committed
635
  }
Guolin Ke's avatar
Guolin Ke committed
636
637
638
639
640
641
  inline int Feature2Group(int feature_idx) const {
    return feature2group_[feature_idx];
  }
  inline int Feture2SubFeature(int feature_idx) const {
    return feature2subfeature_[feature_idx];
  }
642
643
644
  inline uint64_t GroupBinBoundary(int group_idx) const {
    return group_bin_boundaries_[group_idx];
  }
Guolin Ke's avatar
Guolin Ke committed
645
646
647
  inline uint64_t NumTotalBin() const {
    return group_bin_boundaries_.back();
  }
648

649
650
651
652
653
654
655
656
657
  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
658
659
  void ReSize(data_size_t num_data);

660
  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
661

662
  MultiValBin* GetMultiBinFromSparseFeatures(const std::vector<uint32_t>& offsets) const;
663

664
  MultiValBin* GetMultiBinFromAllFeatures(const std::vector<uint32_t>& offsets) const;
665

666
  template <bool USE_QUANT_GRAD, int HIST_BITS>
667
668
669
  TrainingShareStates* GetShareStates(
      score_t* gradients, score_t* hessians,
      const std::vector<int8_t>& is_feature_used, bool is_constant_hessian,
670
      bool force_col_wise, bool force_row_wise, const int num_grad_quant_bins) const;
671

672
  LIGHTGBM_EXPORT void FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
673

674
675
  bool SetFieldFromArrow(const char* field_name, const ArrowChunkedArray& ca);

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

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

680
  LIGHTGBM_EXPORT bool SetIntField(const char* field_name, const int* field_data, data_size_t num_element);
681

682
  LIGHTGBM_EXPORT bool GetFloatField(const char* field_name, data_size_t* out_len, const float** out_ptr);
683

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

686
  LIGHTGBM_EXPORT bool GetIntField(const char* field_name, data_size_t* out_len, const int** out_ptr);
687

Guolin Ke's avatar
Guolin Ke committed
688
689
690
  /*!
  * \brief Save current dataset into binary file, will save to "filename.bin"
  */
691
  LIGHTGBM_EXPORT void SaveBinaryFile(const char* bin_filename);
Guolin Ke's avatar
Guolin Ke committed
692

693
694
695
696
697
  /*!
   * \brief Serialize the overall Dataset definition/schema to a binary buffer (i.e., without data)
   */
  LIGHTGBM_EXPORT void SerializeReference(ByteBuffer* out);

698
699
  LIGHTGBM_EXPORT void DumpTextFile(const char* text_filename);

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

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

704
  void InitTrain(const std::vector<int8_t>& is_feature_used,
705
                 TrainingShareStates* share_state) const;
706

707
  template <bool USE_INDICES, bool USE_HESSIAN, bool USE_QUANT_GRAD, int HIST_BITS>
Guolin Ke's avatar
Guolin Ke committed
708
709
710
711
712
713
714
715
716
  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;

717
  template <bool USE_INDICES, bool ORDERED, bool USE_QUANT_GRAD, int HIST_BITS>
718
719
720
721
  void ConstructHistogramsMultiVal(const data_size_t* data_indices,
                                   data_size_t num_data,
                                   const score_t* gradients,
                                   const score_t* hessians,
722
                                   TrainingShareStates* share_state,
Guolin Ke's avatar
Guolin Ke committed
723
724
                                   hist_t* hist_data) const;

725
  template <bool USE_QUANT_GRAD, int HIST_BITS>
Guolin Ke's avatar
Guolin Ke committed
726
727
728
729
730
731
732
733
734
735
736
737
  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) {
738
        ConstructHistogramsInner<true, false, USE_QUANT_GRAD, HIST_BITS>(
Guolin Ke's avatar
Guolin Ke committed
739
740
741
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      } else {
742
        ConstructHistogramsInner<false, false, USE_QUANT_GRAD, HIST_BITS>(
Guolin Ke's avatar
Guolin Ke committed
743
744
745
746
747
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      }
    } else {
      if (use_indices) {
748
        ConstructHistogramsInner<true, true, USE_QUANT_GRAD, HIST_BITS>(
Guolin Ke's avatar
Guolin Ke committed
749
750
751
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      } else {
752
        ConstructHistogramsInner<false, true, USE_QUANT_GRAD, HIST_BITS>(
Guolin Ke's avatar
Guolin Ke committed
753
754
755
756
757
            is_feature_used, data_indices, num_data, gradients, hessians,
            ordered_gradients, ordered_hessians, share_state, hist_data);
      }
    }
  }
Guolin Ke's avatar
Guolin Ke committed
758

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

761
762
763
  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;

764
765
  inline data_size_t Split(int feature, const uint32_t* threshold,
                           int num_threshold, bool default_left,
766
                           const data_size_t* data_indices,
767
768
                           data_size_t cnt, data_size_t* lte_indices,
                           data_size_t* gt_indices) const {
Guolin Ke's avatar
Guolin Ke committed
769
770
    const int group = feature2group_[feature];
    const int sub_feature = feature2subfeature_[feature];
771
772
773
    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
774
775
776
777
778
779
780
781
782
783
784
785
786
787
  }

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

791
792
793
  inline int FeatureGroupNumBin(int group) const {
    return feature_groups_[group]->num_total_bin_;
  }
794

Guolin Ke's avatar
Guolin Ke committed
795
796
797
798
799
800
  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();
  }

801
802
803
804
  inline const Bin* FeatureGroupBin(int group) const {
    return feature_groups_[group]->bin_data_.get();
  }

Guolin Ke's avatar
Guolin Ke committed
805
806
807
  inline BinIterator* FeatureIterator(int i) const {
    const int group = feature2group_[i];
    const int sub_feature = feature2subfeature_[i];
zhangyafeikimi's avatar
zhangyafeikimi committed
808
    return feature_groups_[group]->SubFeatureIterator(sub_feature);
Guolin Ke's avatar
Guolin Ke committed
809
810
  }

811
812
813
  inline BinIterator* FeatureGroupIterator(int group) const {
    return feature_groups_[group]->FeatureGroupIterator();
  }
814

815
816
817
818
  inline bool IsMultiGroup(int i) const {
    return feature_groups_[i]->is_multi_val_;
  }

819
820
821
822
823
824
825
826
  inline size_t FeatureGroupSizesInByte(int group) const {
    return feature_groups_[group]->FeatureGroupSizesInByte();
  }

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

827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
  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
842
843
844
845
846
847
  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);
  }

848
849
850
851
852
853
854
  // 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);
  }

855
856
857
858
859
860
  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
861
862
863
864
865
866
867
868
869
  /*!
  * \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_; }

870
871
872
  /*! \brief Get number of numeric features */
  inline int num_numeric_features() const { return num_numeric_features_; }

873
874
875
  /*! \brief Get Number of feature groups */
  inline int num_feature_groups() const { return num_groups_;}

876
877
878
  /*! \brief Get Number of total features */
  inline int num_total_features() const { return num_total_features_; }

Guolin Ke's avatar
Guolin Ke committed
879
880
881
882
  /*! \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
883
884
  inline const std::vector<std::string>& feature_names() const { return feature_names_; }

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

Guolin Ke's avatar
Guolin Ke committed
915
916
  inline std::vector<std::string> feature_infos() const {
    std::vector<std::string> bufs;
917
    for (int i = 0; i < num_total_features_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
918
      int fidx = used_feature_map_[i];
919
      if (fidx < 0) {
Guolin Ke's avatar
Guolin Ke committed
920
921
922
        bufs.push_back("none");
      } else {
        const auto bin_mapper = FeatureBinMapper(fidx);
923
        bufs.push_back(bin_mapper->bin_info_string());
Guolin Ke's avatar
Guolin Ke committed
924
925
926
927
928
      }
    }
    return bufs;
  }

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

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

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

938
939
940
941
942
943
944
945
946
  /*! \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
947
948
949
950
951
  /*! \brief Disable copy */
  Dataset& operator=(const Dataset&) = delete;
  /*! \brief Disable copy */
  Dataset(const Dataset&) = delete;

952
  void AddFeaturesFrom(Dataset* other);
953

954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
  /*! \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();
  }

979
980
981
982
983
984
985
986
987
988
989
990
  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);
  }

991
  #ifdef USE_CUDA
992
993
994
995
996

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

997
  #endif  // USE_CUDA
998

Nikita Titov's avatar
Nikita Titov committed
999
 private:
1000
1001
1002
1003
  void SerializeHeader(BinaryWriter* serializer);

  size_t GetSerializedHeaderSize();

1004
1005
  void CreateCUDAColumnData();

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

1056
  #ifdef USE_CUDA
1057
  std::unique_ptr<CUDAColumnData> cuda_column_data_;
1058
  #endif  // USE_CUDA
1059

1060
  std::string parser_config_str_;
Guolin Ke's avatar
Guolin Ke committed
1061
1062
1063
1064
};

}  // namespace LightGBM

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