c_api.cpp 104 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
#include <LightGBM/c_api.h>
Guolin Ke's avatar
Guolin Ke committed
6

Guolin Ke's avatar
Guolin Ke committed
7
8
#include <LightGBM/boosting.h>
#include <LightGBM/config.h>
9
10
11
#include <LightGBM/dataset.h>
#include <LightGBM/dataset_loader.h>
#include <LightGBM/metric.h>
12
#include <LightGBM/network.h>
13
14
15
16
17
18
19
#include <LightGBM/objective_function.h>
#include <LightGBM/prediction_early_stop.h>
#include <LightGBM/utils/common.h>
#include <LightGBM/utils/log.h>
#include <LightGBM/utils/openmp_wrapper.h>
#include <LightGBM/utils/random.h>
#include <LightGBM/utils/threading.h>
Guolin Ke's avatar
Guolin Ke committed
20

21
22
23
24
25
26
27
28
#include <string>
#include <cstdio>
#include <functional>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <vector>

29
#include "application/predictor.hpp"
30
31
#include <LightGBM/utils/yamc/alternate_shared_mutex.hpp>
#include <LightGBM/utils/yamc/yamc_shared_lock.hpp>
Guolin Ke's avatar
Guolin Ke committed
32

Guolin Ke's avatar
Guolin Ke committed
33
34
namespace LightGBM {

Guolin Ke's avatar
Guolin Ke committed
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
inline int LGBM_APIHandleException(const std::exception& ex) {
  LGBM_SetLastError(ex.what());
  return -1;
}
inline int LGBM_APIHandleException(const std::string& ex) {
  LGBM_SetLastError(ex.c_str());
  return -1;
}

#define API_BEGIN() try {
#define API_END() } \
catch(std::exception& ex) { return LGBM_APIHandleException(ex); } \
catch(std::string& ex) { return LGBM_APIHandleException(ex); } \
catch(...) { return LGBM_APIHandleException("unknown exception"); } \
return 0;

51
52
53
54
55
56
#define UNIQUE_LOCK(mtx) \
std::unique_lock<yamc::alternate::shared_mutex> lock(mtx);

#define SHARED_LOCK(mtx) \
yamc::shared_lock<yamc::alternate::shared_mutex> lock(&mtx);

57
58
59
60
61
62
63
64
const int PREDICTOR_TYPES = 4;

// Single row predictor to abstract away caching logic
class SingleRowPredictor {
 public:
  PredictFunction predict_function;
  int64_t num_pred_in_one_row;

65
  SingleRowPredictor(int predict_type, Boosting* boosting, const Config& config, int start_iter, int num_iter) {
66
67
68
69
70
71
72
73
74
75
76
77
78
    bool is_predict_leaf = false;
    bool is_raw_score = false;
    bool predict_contrib = false;
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
      is_predict_leaf = true;
    } else if (predict_type == C_API_PREDICT_RAW_SCORE) {
      is_raw_score = true;
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
      predict_contrib = true;
    }
    early_stop_ = config.pred_early_stop;
    early_stop_freq_ = config.pred_early_stop_freq;
    early_stop_margin_ = config.pred_early_stop_margin;
79
80
    iter_ = num_iter;
    predictor_.reset(new Predictor(boosting, start_iter, iter_, is_raw_score, is_predict_leaf, predict_contrib,
81
                                   early_stop_, early_stop_freq_, early_stop_margin_));
82
    num_pred_in_one_row = boosting->NumPredictOneRow(start_iter, iter_, is_predict_leaf, predict_contrib);
83
    predict_function = predictor_->GetPredictFunction();
Guolin Ke's avatar
Guolin Ke committed
84
    num_total_model_ = boosting->NumberOfTotalModel();
85
  }
86

87
  ~SingleRowPredictor() {}
88

Guolin Ke's avatar
Guolin Ke committed
89
  bool IsPredictorEqual(const Config& config, int iter, Boosting* boosting) {
90
91
92
93
94
    return early_stop_ == config.pred_early_stop &&
      early_stop_freq_ == config.pred_early_stop_freq &&
      early_stop_margin_ == config.pred_early_stop_margin &&
      iter_ == iter &&
      num_total_model_ == boosting->NumberOfTotalModel();
95
  }
Guolin Ke's avatar
Guolin Ke committed
96

97
98
99
100
101
102
103
104
105
 private:
  std::unique_ptr<Predictor> predictor_;
  bool early_stop_;
  int early_stop_freq_;
  double early_stop_margin_;
  int iter_;
  int num_total_model_;
};

Guolin Ke's avatar
Guolin Ke committed
106
class Booster {
Nikita Titov's avatar
Nikita Titov committed
107
 public:
Guolin Ke's avatar
Guolin Ke committed
108
  explicit Booster(const char* filename) {
109
    boosting_.reset(Boosting::CreateBoosting("gbdt", filename));
110
111
  }

Guolin Ke's avatar
Guolin Ke committed
112
  Booster(const Dataset* train_data,
113
          const char* parameters) {
Guolin Ke's avatar
Guolin Ke committed
114
    auto param = Config::Str2Map(parameters);
wxchan's avatar
wxchan committed
115
    config_.Set(param);
116
    OMP_SET_NUM_THREADS(config_.num_threads);
Guolin Ke's avatar
Guolin Ke committed
117
    // create boosting
Guolin Ke's avatar
Guolin Ke committed
118
    if (config_.input_model.size() > 0) {
119
120
      Log::Warning("Continued train from model is not supported for c_api,\n"
                   "please use continued train with input score");
Guolin Ke's avatar
Guolin Ke committed
121
    }
Guolin Ke's avatar
Guolin Ke committed
122

Guolin Ke's avatar
Guolin Ke committed
123
    boosting_.reset(Boosting::CreateBoosting(config_.boosting, nullptr));
Guolin Ke's avatar
Guolin Ke committed
124

125
126
    train_data_ = train_data;
    CreateObjectiveAndMetrics();
Guolin Ke's avatar
Guolin Ke committed
127
    // initialize the boosting
Guolin Ke's avatar
Guolin Ke committed
128
    if (config_.tree_learner == std::string("feature")) {
129
      Log::Fatal("Do not support feature parallel in c api");
130
    }
Guolin Ke's avatar
Guolin Ke committed
131
    if (Network::num_machines() == 1 && config_.tree_learner != std::string("serial")) {
132
      Log::Warning("Only find one worker, will switch to serial tree learner");
Guolin Ke's avatar
Guolin Ke committed
133
      config_.tree_learner = "serial";
134
    }
Guolin Ke's avatar
Guolin Ke committed
135
    boosting_->Init(&config_, train_data_, objective_fun_.get(),
136
                    Common::ConstPtrInVectorWrapper<Metric>(train_metric_));
wxchan's avatar
wxchan committed
137
138
139
  }

  void MergeFrom(const Booster* other) {
140
    UNIQUE_LOCK(mutex_)
wxchan's avatar
wxchan committed
141
    boosting_->MergeFrom(other->boosting_.get());
Guolin Ke's avatar
Guolin Ke committed
142
143
144
145
  }

  ~Booster() {
  }
146

147
  void CreateObjectiveAndMetrics() {
Guolin Ke's avatar
Guolin Ke committed
148
    // create objective function
Guolin Ke's avatar
Guolin Ke committed
149
150
    objective_fun_.reset(ObjectiveFunction::CreateObjectiveFunction(config_.objective,
                                                                    config_));
Guolin Ke's avatar
Guolin Ke committed
151
152
153
154
155
156
157
158
159
160
    if (objective_fun_ == nullptr) {
      Log::Warning("Using self-defined objective function");
    }
    // initialize the objective function
    if (objective_fun_ != nullptr) {
      objective_fun_->Init(train_data_->metadata(), train_data_->num_data());
    }

    // create training metric
    train_metric_.clear();
Guolin Ke's avatar
Guolin Ke committed
161
    for (auto metric_type : config_.metric) {
Guolin Ke's avatar
Guolin Ke committed
162
      auto metric = std::unique_ptr<Metric>(
Guolin Ke's avatar
Guolin Ke committed
163
        Metric::CreateMetric(metric_type, config_));
Guolin Ke's avatar
Guolin Ke committed
164
165
166
167
168
      if (metric == nullptr) { continue; }
      metric->Init(train_data_->metadata(), train_data_->num_data());
      train_metric_.push_back(std::move(metric));
    }
    train_metric_.shrink_to_fit();
169
170
171
172
  }

  void ResetTrainingData(const Dataset* train_data) {
    if (train_data != train_data_) {
173
      UNIQUE_LOCK(mutex_)
174
175
176
177
178
179
      train_data_ = train_data;
      CreateObjectiveAndMetrics();
      // reset the boosting
      boosting_->ResetTrainingData(train_data_,
                                   objective_fun_.get(), Common::ConstPtrInVectorWrapper<Metric>(train_metric_));
    }
wxchan's avatar
wxchan committed
180
181
  }

182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
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
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
  static void CheckDatasetResetConfig(
      const Config& old_config,
      const std::unordered_map<std::string, std::string>& new_param) {
    Config new_config;
    new_config.Set(new_param);
    if (new_param.count("data_random_seed") &&
        new_config.data_random_seed != old_config.data_random_seed) {
      Log::Fatal("Cannot change data_random_seed after constructed Dataset handle.");
    }
    if (new_param.count("max_bin") &&
        new_config.max_bin != old_config.max_bin) {
      Log::Fatal("Cannot change max_bin after constructed Dataset handle.");
    }
    if (new_param.count("max_bin_by_feature") &&
        new_config.max_bin_by_feature != old_config.max_bin_by_feature) {
      Log::Fatal(
          "Cannot change max_bin_by_feature after constructed Dataset handle.");
    }
    if (new_param.count("bin_construct_sample_cnt") &&
        new_config.bin_construct_sample_cnt !=
            old_config.bin_construct_sample_cnt) {
      Log::Fatal(
          "Cannot change bin_construct_sample_cnt after constructed Dataset "
          "handle.");
    }
    if (new_param.count("min_data_in_bin") &&
        new_config.min_data_in_bin != old_config.min_data_in_bin) {
      Log::Fatal(
          "Cannot change min_data_in_bin after constructed Dataset handle.");
    }
    if (new_param.count("use_missing") &&
        new_config.use_missing != old_config.use_missing) {
      Log::Fatal("Cannot change use_missing after constructed Dataset handle.");
    }
    if (new_param.count("zero_as_missing") &&
        new_config.zero_as_missing != old_config.zero_as_missing) {
      Log::Fatal(
          "Cannot change zero_as_missing after constructed Dataset handle.");
    }
    if (new_param.count("categorical_feature") &&
        new_config.categorical_feature != old_config.categorical_feature) {
      Log::Fatal(
          "Cannot change categorical_feature after constructed Dataset "
          "handle.");
    }
    if (new_param.count("feature_pre_filter") &&
        new_config.feature_pre_filter != old_config.feature_pre_filter) {
      Log::Fatal(
          "Cannot change feature_pre_filter after constructed Dataset handle.");
    }
    if (new_param.count("is_enable_sparse") &&
        new_config.is_enable_sparse != old_config.is_enable_sparse) {
      Log::Fatal(
          "Cannot change is_enable_sparse after constructed Dataset handle.");
    }
    if (new_param.count("pre_partition") &&
        new_config.pre_partition != old_config.pre_partition) {
      Log::Fatal(
          "Cannot change pre_partition after constructed Dataset handle.");
    }
    if (new_param.count("enable_bundle") &&
        new_config.enable_bundle != old_config.enable_bundle) {
      Log::Fatal(
          "Cannot change enable_bundle after constructed Dataset handle.");
    }
    if (new_param.count("header") && new_config.header != old_config.header) {
      Log::Fatal("Cannot change header after constructed Dataset handle.");
    }
    if (new_param.count("two_round") &&
        new_config.two_round != old_config.two_round) {
      Log::Fatal("Cannot change two_round after constructed Dataset handle.");
    }
    if (new_param.count("label_column") &&
        new_config.label_column != old_config.label_column) {
      Log::Fatal(
          "Cannot change label_column after constructed Dataset handle.");
    }
    if (new_param.count("weight_column") &&
        new_config.weight_column != old_config.weight_column) {
      Log::Fatal(
          "Cannot change weight_column after constructed Dataset handle.");
    }
    if (new_param.count("group_column") &&
        new_config.group_column != old_config.group_column) {
      Log::Fatal(
          "Cannot change group_column after constructed Dataset handle.");
    }
    if (new_param.count("ignore_column") &&
        new_config.ignore_column != old_config.ignore_column) {
      Log::Fatal(
          "Cannot change ignore_column after constructed Dataset handle.");
    }
    if (new_param.count("forcedbins_filename")) {
      Log::Fatal("Cannot change forced bins after constructed Dataset handle.");
    }
    if (new_param.count("min_data_in_leaf") &&
        new_config.min_data_in_leaf < old_config.min_data_in_leaf &&
        old_config.feature_pre_filter) {
      Log::Fatal(
          "Reducing `min_data_in_leaf` with `feature_pre_filter=true` may "
          "cause unexpected behaviour "
          "for features that were pre-filtered by the larger "
          "`min_data_in_leaf`.\n"
          "You need to set `feature_pre_filter=false` to dynamically change "
          "the `min_data_in_leaf`.");
    }
Nikita Titov's avatar
Nikita Titov committed
288
    if (new_param.count("linear_tree") && new_config.linear_tree != old_config.linear_tree) {
289
      Log::Fatal("Cannot change linear_tree after constructed Dataset handle.");
290
    }
Nikita Titov's avatar
Nikita Titov committed
291
292
293
294
    if (new_param.count("precise_float_parser") &&
        new_config.precise_float_parser != old_config.precise_float_parser) {
      Log::Fatal("Cannot change precise_float_parser after constructed Dataset handle.");
    }
295
296
  }

wxchan's avatar
wxchan committed
297
  void ResetConfig(const char* parameters) {
298
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
299
    auto param = Config::Str2Map(parameters);
300
301
302
    Config new_config;
    new_config.Set(param);
    if (param.count("num_class") && new_config.num_class != config_.num_class) {
303
      Log::Fatal("Cannot change num_class during training");
wxchan's avatar
wxchan committed
304
    }
305
    if (param.count("boosting") && new_config.boosting != config_.boosting) {
Guolin Ke's avatar
Guolin Ke committed
306
      Log::Fatal("Cannot change boosting during training");
wxchan's avatar
wxchan committed
307
    }
308
    if (param.count("metric") && new_config.metric != config_.metric) {
309
      Log::Fatal("Cannot change metric during training");
Guolin Ke's avatar
Guolin Ke committed
310
    }
311
312
    CheckDatasetResetConfig(config_, param);

Guolin Ke's avatar
Guolin Ke committed
313
    config_.Set(param);
314

315
    OMP_SET_NUM_THREADS(config_.num_threads);
Guolin Ke's avatar
Guolin Ke committed
316
317
318

    if (param.count("objective")) {
      // create objective function
Guolin Ke's avatar
Guolin Ke committed
319
320
      objective_fun_.reset(ObjectiveFunction::CreateObjectiveFunction(config_.objective,
                                                                      config_));
Guolin Ke's avatar
Guolin Ke committed
321
322
323
324
325
326
327
      if (objective_fun_ == nullptr) {
        Log::Warning("Using self-defined objective function");
      }
      // initialize the objective function
      if (objective_fun_ != nullptr) {
        objective_fun_->Init(train_data_->metadata(), train_data_->num_data());
      }
328
329
      boosting_->ResetTrainingData(train_data_,
                                   objective_fun_.get(), Common::ConstPtrInVectorWrapper<Metric>(train_metric_));
wxchan's avatar
wxchan committed
330
    }
Guolin Ke's avatar
Guolin Ke committed
331

Guolin Ke's avatar
Guolin Ke committed
332
    boosting_->ResetConfig(&config_);
wxchan's avatar
wxchan committed
333
334
335
  }

  void AddValidData(const Dataset* valid_data) {
336
    UNIQUE_LOCK(mutex_)
wxchan's avatar
wxchan committed
337
    valid_metrics_.emplace_back();
Guolin Ke's avatar
Guolin Ke committed
338
339
    for (auto metric_type : config_.metric) {
      auto metric = std::unique_ptr<Metric>(Metric::CreateMetric(metric_type, config_));
wxchan's avatar
wxchan committed
340
341
342
343
344
345
      if (metric == nullptr) { continue; }
      metric->Init(valid_data->metadata(), valid_data->num_data());
      valid_metrics_.back().push_back(std::move(metric));
    }
    valid_metrics_.back().shrink_to_fit();
    boosting_->AddValidDataset(valid_data,
346
                               Common::ConstPtrInVectorWrapper<Metric>(valid_metrics_.back()));
wxchan's avatar
wxchan committed
347
  }
Guolin Ke's avatar
Guolin Ke committed
348

349
  bool TrainOneIter() {
350
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
351
    return boosting_->TrainOneIter(nullptr, nullptr);
352
353
  }

Guolin Ke's avatar
Guolin Ke committed
354
  void Refit(const int32_t* leaf_preds, int32_t nrow, int32_t ncol) {
355
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
356
357
358
    std::vector<std::vector<int32_t>> v_leaf_preds(nrow, std::vector<int32_t>(ncol, 0));
    for (int i = 0; i < nrow; ++i) {
      for (int j = 0; j < ncol; ++j) {
359
        v_leaf_preds[i][j] = leaf_preds[static_cast<size_t>(i) * static_cast<size_t>(ncol) + static_cast<size_t>(j)];
Guolin Ke's avatar
Guolin Ke committed
360
361
362
363
364
      }
    }
    boosting_->RefitTree(v_leaf_preds);
  }

365
  bool TrainOneIter(const score_t* gradients, const score_t* hessians) {
366
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
367
    return boosting_->TrainOneIter(gradients, hessians);
368
369
  }

wxchan's avatar
wxchan committed
370
  void RollbackOneIter() {
371
    UNIQUE_LOCK(mutex_)
wxchan's avatar
wxchan committed
372
373
374
    boosting_->RollbackOneIter();
  }

375
  void SetSingleRowPredictor(int start_iteration, int num_iteration, int predict_type, const Config& config) {
376
377
378
379
      UNIQUE_LOCK(mutex_)
      if (single_row_predictor_[predict_type].get() == nullptr ||
          !single_row_predictor_[predict_type]->IsPredictorEqual(config, num_iteration, boosting_.get())) {
        single_row_predictor_[predict_type].reset(new SingleRowPredictor(predict_type, boosting_.get(),
380
                                                                         config, start_iteration, num_iteration));
381
382
383
384
      }
  }

  void PredictSingleRow(int predict_type, int ncol,
385
386
               std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun,
               const Config& config,
387
               double* out_result, int64_t* out_len) const {
388
389
390
    if (!config.predict_disable_shape_check && ncol != boosting_->MaxFeatureIdx() + 1) {
      Log::Fatal("The number of features in data (%d) is not the same as it was in training data (%d).\n"\
                 "You can set ``predict_disable_shape_check=true`` to discard this error, but please be aware what you are doing.", ncol, boosting_->MaxFeatureIdx() + 1);
391
    }
392
    UNIQUE_LOCK(mutex_)
393
    const auto& single_row_predictor = single_row_predictor_[predict_type];
394
395
    auto one_row = get_row_fun(0);
    auto pred_wrt_ptr = out_result;
396
    single_row_predictor->predict_function(one_row, pred_wrt_ptr);
397

398
    *out_len = single_row_predictor->num_pred_in_one_row;
399
400
  }

401
  Predictor CreatePredictor(int start_iteration, int num_iteration, int predict_type, int ncol, const Config& config) const {
402
403
404
    if (!config.predict_disable_shape_check && ncol != boosting_->MaxFeatureIdx() + 1) {
      Log::Fatal("The number of features in data (%d) is not the same as it was in training data (%d).\n" \
                 "You can set ``predict_disable_shape_check=true`` to discard this error, but please be aware what you are doing.", ncol, boosting_->MaxFeatureIdx() + 1);
405
    }
Guolin Ke's avatar
Guolin Ke committed
406
407
    bool is_predict_leaf = false;
    bool is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
408
    bool predict_contrib = false;
Guolin Ke's avatar
Guolin Ke committed
409
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
Guolin Ke's avatar
Guolin Ke committed
410
      is_predict_leaf = true;
Guolin Ke's avatar
Guolin Ke committed
411
    } else if (predict_type == C_API_PREDICT_RAW_SCORE) {
Guolin Ke's avatar
Guolin Ke committed
412
      is_raw_score = true;
413
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
Guolin Ke's avatar
Guolin Ke committed
414
      predict_contrib = true;
Guolin Ke's avatar
Guolin Ke committed
415
416
    } else {
      is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
417
    }
Guolin Ke's avatar
Guolin Ke committed
418

419
    return Predictor(boosting_.get(), start_iteration, num_iteration, is_raw_score, is_predict_leaf, predict_contrib,
420
                        config.pred_early_stop, config.pred_early_stop_freq, config.pred_early_stop_margin);
421
422
  }

423
  void Predict(int start_iteration, int num_iteration, int predict_type, int nrow, int ncol,
424
425
               std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun,
               const Config& config,
426
427
               double* out_result, int64_t* out_len) const {
    SHARED_LOCK(mutex_);
428
    auto predictor = CreatePredictor(start_iteration, num_iteration, predict_type, ncol, config);
429
430
431
432
433
434
435
    bool is_predict_leaf = false;
    bool predict_contrib = false;
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
      is_predict_leaf = true;
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
      predict_contrib = true;
    }
436
    int64_t num_pred_in_one_row = boosting_->NumPredictOneRow(start_iteration, num_iteration, is_predict_leaf, predict_contrib);
Guolin Ke's avatar
Guolin Ke committed
437
    auto pred_fun = predictor.GetPredictFunction();
438
439
    OMP_INIT_EX();
    #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
440
    for (int i = 0; i < nrow; ++i) {
441
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
442
      auto one_row = get_row_fun(i);
Tony-Y's avatar
Tony-Y committed
443
      auto pred_wrt_ptr = out_result + static_cast<size_t>(num_pred_in_one_row) * i;
Guolin Ke's avatar
Guolin Ke committed
444
      pred_fun(one_row, pred_wrt_ptr);
445
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
446
    }
447
    OMP_THROW_EX();
448
    *out_len = num_pred_in_one_row * nrow;
Guolin Ke's avatar
Guolin Ke committed
449
450
  }

451
  void PredictSparse(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
452
453
454
455
                     std::function<std::vector<std::pair<int, double>>(int64_t row_idx)> get_row_fun,
                     const Config& config, int64_t* out_elements_size,
                     std::vector<std::vector<std::unordered_map<int, double>>>* agg_ptr,
                     int32_t** out_indices, void** out_data, int data_type,
456
                     bool* is_data_float32_ptr, int num_matrices) const {
457
    auto predictor = CreatePredictor(start_iteration, num_iteration, predict_type, ncol, config);
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
    auto pred_sparse_fun = predictor.GetPredictSparseFunction();
    std::vector<std::vector<std::unordered_map<int, double>>>& agg = *agg_ptr;
    OMP_INIT_EX();
    #pragma omp parallel for schedule(static)
    for (int64_t i = 0; i < nrow; ++i) {
      OMP_LOOP_EX_BEGIN();
      auto one_row = get_row_fun(i);
      agg[i] = std::vector<std::unordered_map<int, double>>(num_matrices);
      pred_sparse_fun(one_row, &agg[i]);
      OMP_LOOP_EX_END();
    }
    OMP_THROW_EX();
    // calculate the nonzero data and indices size
    int64_t elements_size = 0;
    for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
      auto row_vector = agg[i];
      for (int j = 0; j < static_cast<int>(row_vector.size()); ++j) {
        elements_size += static_cast<int64_t>(row_vector[j].size());
      }
    }
    *out_elements_size = elements_size;
    *is_data_float32_ptr = false;
    // allocate data and indices arrays
    if (data_type == C_API_DTYPE_FLOAT32) {
      *out_data = new float[elements_size];
      *is_data_float32_ptr = true;
    } else if (data_type == C_API_DTYPE_FLOAT64) {
      *out_data = new double[elements_size];
    } else {
      Log::Fatal("Unknown data type in PredictSparse");
      return;
    }
    *out_indices = new int32_t[elements_size];
  }

493
  void PredictSparseCSR(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
494
495
496
                        std::function<std::vector<std::pair<int, double>>(int64_t row_idx)> get_row_fun,
                        const Config& config,
                        int64_t* out_len, void** out_indptr, int indptr_type,
497
498
                        int32_t** out_indices, void** out_data, int data_type) const {
    SHARED_LOCK(mutex_);
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
    // Get the number of trees per iteration (for multiclass scenario we output multiple sparse matrices)
    int num_matrices = boosting_->NumModelPerIteration();
    bool is_indptr_int32 = false;
    bool is_data_float32 = false;
    int64_t indptr_size = (nrow + 1) * num_matrices;
    if (indptr_type == C_API_DTYPE_INT32) {
      *out_indptr = new int32_t[indptr_size];
      is_indptr_int32 = true;
    } else if (indptr_type == C_API_DTYPE_INT64) {
      *out_indptr = new int64_t[indptr_size];
    } else {
      Log::Fatal("Unknown indptr type in PredictSparseCSR");
      return;
    }
    // aggregated per row feature contribution results
    std::vector<std::vector<std::unordered_map<int, double>>> agg(nrow);
    int64_t elements_size = 0;
516
    PredictSparse(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
517
518
519
                  out_indices, out_data, data_type, &is_data_float32, num_matrices);
    std::vector<int> row_sizes(num_matrices * nrow);
    std::vector<int64_t> row_matrix_offsets(num_matrices * nrow);
520
    std::vector<int64_t> matrix_offsets(num_matrices);
521
522
523
524
525
526
527
528
529
530
531
532
533
534
    int64_t row_vector_cnt = 0;
    for (int m = 0; m < num_matrices; ++m) {
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        auto row_vector = agg[i];
        auto row_vector_size = row_vector[m].size();
        // keep track of the row_vector sizes for parallelization
        row_sizes[row_vector_cnt] = static_cast<int>(row_vector_size);
        if (i == 0) {
          row_matrix_offsets[row_vector_cnt] = 0;
        } else {
          row_matrix_offsets[row_vector_cnt] = static_cast<int64_t>(row_sizes[row_vector_cnt - 1] + row_matrix_offsets[row_vector_cnt - 1]);
        }
        row_vector_cnt++;
      }
535
536
537
538
539
540
      if (m == 0) {
        matrix_offsets[m] = 0;
      }
      if (m + 1 < num_matrices) {
        matrix_offsets[m + 1] = static_cast<int64_t>(matrix_offsets[m] + row_matrix_offsets[row_vector_cnt - 1] + row_sizes[row_vector_cnt - 1]);
      }
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
    }
    // copy vector results to output for each row
    int64_t indptr_index = 0;
    for (int m = 0; m < num_matrices; ++m) {
      if (is_indptr_int32) {
        (reinterpret_cast<int32_t*>(*out_indptr))[indptr_index] = 0;
      } else {
        (reinterpret_cast<int64_t*>(*out_indptr))[indptr_index] = 0;
      }
      indptr_index++;
      int64_t matrix_start_index = m * static_cast<int64_t>(agg.size());
      OMP_INIT_EX();
      #pragma omp parallel for schedule(static)
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        OMP_LOOP_EX_BEGIN();
        auto row_vector = agg[i];
        int64_t row_start_index = matrix_start_index + i;
558
        int64_t element_index = row_matrix_offsets[row_start_index] + matrix_offsets[m];
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
        int64_t indptr_loop_index = indptr_index + i;
        for (auto it = row_vector[m].begin(); it != row_vector[m].end(); ++it) {
          (*out_indices)[element_index] = it->first;
          if (is_data_float32) {
            (reinterpret_cast<float*>(*out_data))[element_index] = static_cast<float>(it->second);
          } else {
            (reinterpret_cast<double*>(*out_data))[element_index] = it->second;
          }
          element_index++;
        }
        int64_t indptr_value = row_matrix_offsets[row_start_index] + row_sizes[row_start_index];
        if (is_indptr_int32) {
          (reinterpret_cast<int32_t*>(*out_indptr))[indptr_loop_index] = static_cast<int32_t>(indptr_value);
        } else {
          (reinterpret_cast<int64_t*>(*out_indptr))[indptr_loop_index] = indptr_value;
        }
        OMP_LOOP_EX_END();
      }
      OMP_THROW_EX();
      indptr_index += static_cast<int64_t>(agg.size());
    }
    out_len[0] = elements_size;
    out_len[1] = indptr_size;
  }

584
  void PredictSparseCSC(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
585
586
587
                        std::function<std::vector<std::pair<int, double>>(int64_t row_idx)> get_row_fun,
                        const Config& config,
                        int64_t* out_len, void** out_col_ptr, int col_ptr_type,
588
589
                        int32_t** out_indices, void** out_data, int data_type) const {
    SHARED_LOCK(mutex_);
590
591
    // Get the number of trees per iteration (for multiclass scenario we output multiple sparse matrices)
    int num_matrices = boosting_->NumModelPerIteration();
592
    auto predictor = CreatePredictor(start_iteration, num_iteration, predict_type, ncol, config);
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
    auto pred_sparse_fun = predictor.GetPredictSparseFunction();
    bool is_col_ptr_int32 = false;
    bool is_data_float32 = false;
    int num_output_cols = ncol + 1;
    int col_ptr_size = (num_output_cols + 1) * num_matrices;
    if (col_ptr_type == C_API_DTYPE_INT32) {
      *out_col_ptr = new int32_t[col_ptr_size];
      is_col_ptr_int32 = true;
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
      *out_col_ptr = new int64_t[col_ptr_size];
    } else {
      Log::Fatal("Unknown col_ptr type in PredictSparseCSC");
      return;
    }
    // aggregated per row feature contribution results
    std::vector<std::vector<std::unordered_map<int, double>>> agg(nrow);
    int64_t elements_size = 0;
610
    PredictSparse(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
                  out_indices, out_data, data_type, &is_data_float32, num_matrices);
    // calculate number of elements per column to construct
    // the CSC matrix with random access
    std::vector<std::vector<int64_t>> column_sizes(num_matrices);
    for (int m = 0; m < num_matrices; ++m) {
      column_sizes[m] = std::vector<int64_t>(num_output_cols, 0);
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        auto row_vector = agg[i];
        for (auto it = row_vector[m].begin(); it != row_vector[m].end(); ++it) {
          column_sizes[m][it->first] += 1;
        }
      }
    }
    // keep track of column counts
    std::vector<std::vector<int64_t>> column_counts(num_matrices);
    // keep track of beginning index for each column
    std::vector<std::vector<int64_t>> column_start_indices(num_matrices);
    // keep track of beginning index for each matrix
    std::vector<int64_t> matrix_start_indices(num_matrices, 0);
    int col_ptr_index = 0;
    for (int m = 0; m < num_matrices; ++m) {
      int64_t col_ptr_value = 0;
      column_start_indices[m] = std::vector<int64_t>(num_output_cols, 0);
      column_counts[m] = std::vector<int64_t>(num_output_cols, 0);
      if (is_col_ptr_int32) {
        (reinterpret_cast<int32_t*>(*out_col_ptr))[col_ptr_index] = static_cast<int32_t>(col_ptr_value);
      } else {
        (reinterpret_cast<int64_t*>(*out_col_ptr))[col_ptr_index] = col_ptr_value;
      }
      col_ptr_index++;
      for (int64_t i = 1; i < static_cast<int64_t>(column_sizes[m].size()); ++i) {
        column_start_indices[m][i] = column_sizes[m][i - 1] + column_start_indices[m][i - 1];
        if (is_col_ptr_int32) {
          (reinterpret_cast<int32_t*>(*out_col_ptr))[col_ptr_index] = static_cast<int32_t>(column_start_indices[m][i]);
        } else {
          (reinterpret_cast<int64_t*>(*out_col_ptr))[col_ptr_index] = column_start_indices[m][i];
        }
        col_ptr_index++;
      }
      int64_t last_elem_index = static_cast<int64_t>(column_sizes[m].size()) - 1;
      int64_t last_column_start_index = column_start_indices[m][last_elem_index];
      int64_t last_column_size = column_sizes[m][last_elem_index];
      if (is_col_ptr_int32) {
        (reinterpret_cast<int32_t*>(*out_col_ptr))[col_ptr_index] = static_cast<int32_t>(last_column_start_index + last_column_size);
      } else {
        (reinterpret_cast<int64_t*>(*out_col_ptr))[col_ptr_index] = last_column_start_index + last_column_size;
      }
658
659
      if (m + 1 < num_matrices) {
        matrix_start_indices[m + 1] = matrix_start_indices[m] + last_column_start_index + last_column_size;
660
      }
661
      col_ptr_index++;
662
    }
663
664
665
    // Note: we parallelize across matrices instead of rows because of the column_counts[m][col_idx] increment inside the loop
    OMP_INIT_EX();
    #pragma omp parallel for schedule(static)
666
    for (int m = 0; m < num_matrices; ++m) {
667
      OMP_LOOP_EX_BEGIN();
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        auto row_vector = agg[i];
        for (auto it = row_vector[m].begin(); it != row_vector[m].end(); ++it) {
          int64_t col_idx = it->first;
          int64_t element_index = column_start_indices[m][col_idx] +
            matrix_start_indices[m] +
            column_counts[m][col_idx];
          // store the row index
          (*out_indices)[element_index] = static_cast<int32_t>(i);
          // update column count
          column_counts[m][col_idx]++;
          if (is_data_float32) {
            (reinterpret_cast<float*>(*out_data))[element_index] = static_cast<float>(it->second);
          } else {
            (reinterpret_cast<double*>(*out_data))[element_index] = it->second;
          }
        }
      }
686
      OMP_LOOP_EX_END();
687
    }
688
    OMP_THROW_EX();
689
690
691
692
    out_len[0] = elements_size;
    out_len[1] = col_ptr_size;
  }

693
  void Predict(int start_iteration, int num_iteration, int predict_type, const char* data_filename,
Guolin Ke's avatar
Guolin Ke committed
694
               int data_has_header, const Config& config,
695
696
               const char* result_filename) const {
    SHARED_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
697
698
    bool is_predict_leaf = false;
    bool is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
699
    bool predict_contrib = false;
Guolin Ke's avatar
Guolin Ke committed
700
701
702
703
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
      is_predict_leaf = true;
    } else if (predict_type == C_API_PREDICT_RAW_SCORE) {
      is_raw_score = true;
704
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
Guolin Ke's avatar
Guolin Ke committed
705
      predict_contrib = true;
Guolin Ke's avatar
Guolin Ke committed
706
707
708
    } else {
      is_raw_score = false;
    }
709
    Predictor predictor(boosting_.get(), start_iteration, num_iteration, is_raw_score, is_predict_leaf, predict_contrib,
710
                        config.pred_early_stop, config.pred_early_stop_freq, config.pred_early_stop_margin);
Guolin Ke's avatar
Guolin Ke committed
711
    bool bool_data_has_header = data_has_header > 0 ? true : false;
Chen Yufei's avatar
Chen Yufei committed
712
713
    predictor.Predict(data_filename, result_filename, bool_data_has_header, config.predict_disable_shape_check,
                      config.precise_float_parser);
Guolin Ke's avatar
Guolin Ke committed
714
715
  }

716
  void GetPredictAt(int data_idx, double* out_result, int64_t* out_len) const {
wxchan's avatar
wxchan committed
717
718
719
    boosting_->GetPredictAt(data_idx, out_result, out_len);
  }

720
  void SaveModelToFile(int start_iteration, int num_iteration, int feature_importance_type, const char* filename) const {
721
    boosting_->SaveModelToFile(start_iteration, num_iteration, feature_importance_type, filename);
Guolin Ke's avatar
Guolin Ke committed
722
  }
723

724
  void LoadModelFromString(const char* model_str) {
725
726
    size_t len = std::strlen(model_str);
    boosting_->LoadModelFromString(model_str, len);
727
728
  }

729
  std::string SaveModelToString(int start_iteration, int num_iteration,
730
                                int feature_importance_type) const {
731
732
    return boosting_->SaveModelToString(start_iteration,
                                        num_iteration, feature_importance_type);
733
734
  }

735
  std::string DumpModel(int start_iteration, int num_iteration,
736
                        int feature_importance_type) const {
737
738
    return boosting_->DumpModel(start_iteration, num_iteration,
                                feature_importance_type);
wxchan's avatar
wxchan committed
739
  }
740

741
  std::vector<double> FeatureImportance(int num_iteration, int importance_type) const {
742
743
744
    return boosting_->FeatureImportance(num_iteration, importance_type);
  }

745
  double UpperBoundValue() const {
746
    SHARED_LOCK(mutex_)
747
748
749
750
    return boosting_->GetUpperBoundValue();
  }

  double LowerBoundValue() const {
751
    SHARED_LOCK(mutex_)
752
753
754
    return boosting_->GetLowerBoundValue();
  }

Guolin Ke's avatar
Guolin Ke committed
755
  double GetLeafValue(int tree_idx, int leaf_idx) const {
756
    SHARED_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
757
    return dynamic_cast<GBDTBase*>(boosting_.get())->GetLeafValue(tree_idx, leaf_idx);
Guolin Ke's avatar
Guolin Ke committed
758
759
760
  }

  void SetLeafValue(int tree_idx, int leaf_idx, double val) {
761
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
762
    dynamic_cast<GBDTBase*>(boosting_.get())->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
763
764
  }

765
  void ShuffleModels(int start_iter, int end_iter) {
766
    UNIQUE_LOCK(mutex_)
767
    boosting_->ShuffleModels(start_iter, end_iter);
768
769
  }

wxchan's avatar
wxchan committed
770
  int GetEvalCounts() const {
771
    SHARED_LOCK(mutex_)
wxchan's avatar
wxchan committed
772
773
774
775
776
777
    int ret = 0;
    for (const auto& metric : train_metric_) {
      ret += static_cast<int>(metric->GetName().size());
    }
    return ret;
  }
778

779
  int GetEvalNames(char** out_strs, const int len, const size_t buffer_len, size_t *out_buffer_len) const {
780
    SHARED_LOCK(mutex_)
781
    *out_buffer_len = 0;
wxchan's avatar
wxchan committed
782
783
784
    int idx = 0;
    for (const auto& metric : train_metric_) {
      for (const auto& name : metric->GetName()) {
785
786
787
788
789
        if (idx < len) {
          std::memcpy(out_strs[idx], name.c_str(), std::min(name.size() + 1, buffer_len));
          out_strs[idx][buffer_len - 1] = '\0';
        }
        *out_buffer_len = std::max(name.size() + 1, *out_buffer_len);
wxchan's avatar
wxchan committed
790
791
792
793
794
795
        ++idx;
      }
    }
    return idx;
  }

796
  int GetFeatureNames(char** out_strs, const int len, const size_t buffer_len, size_t *out_buffer_len) const {
797
    SHARED_LOCK(mutex_)
798
    *out_buffer_len = 0;
wxchan's avatar
wxchan committed
799
800
    int idx = 0;
    for (const auto& name : boosting_->FeatureNames()) {
801
802
803
804
805
      if (idx < len) {
        std::memcpy(out_strs[idx], name.c_str(), std::min(name.size() + 1, buffer_len));
        out_strs[idx][buffer_len - 1] = '\0';
      }
      *out_buffer_len = std::max(name.size() + 1, *out_buffer_len);
wxchan's avatar
wxchan committed
806
807
808
809
810
      ++idx;
    }
    return idx;
  }

wxchan's avatar
wxchan committed
811
  const Boosting* GetBoosting() const { return boosting_.get(); }
Guolin Ke's avatar
Guolin Ke committed
812

Nikita Titov's avatar
Nikita Titov committed
813
 private:
wxchan's avatar
wxchan committed
814
  const Dataset* train_data_;
Guolin Ke's avatar
Guolin Ke committed
815
  std::unique_ptr<Boosting> boosting_;
816
  std::unique_ptr<SingleRowPredictor> single_row_predictor_[PREDICTOR_TYPES];
817

Guolin Ke's avatar
Guolin Ke committed
818
  /*! \brief All configs */
Guolin Ke's avatar
Guolin Ke committed
819
  Config config_;
Guolin Ke's avatar
Guolin Ke committed
820
  /*! \brief Metric for training data */
Guolin Ke's avatar
Guolin Ke committed
821
  std::vector<std::unique_ptr<Metric>> train_metric_;
Guolin Ke's avatar
Guolin Ke committed
822
  /*! \brief Metrics for validation data */
Guolin Ke's avatar
Guolin Ke committed
823
  std::vector<std::vector<std::unique_ptr<Metric>>> valid_metrics_;
Guolin Ke's avatar
Guolin Ke committed
824
  /*! \brief Training objective function */
Guolin Ke's avatar
Guolin Ke committed
825
  std::unique_ptr<ObjectiveFunction> objective_fun_;
wxchan's avatar
wxchan committed
826
  /*! \brief mutex for threading safe call */
827
  mutable yamc::alternate::shared_mutex mutex_;
Guolin Ke's avatar
Guolin Ke committed
828
829
};

830
}  // namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
831

832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
// explicitly declare symbols from LightGBM namespace
using LightGBM::AllgatherFunction;
using LightGBM::Booster;
using LightGBM::Common::CheckElementsIntervalClosed;
using LightGBM::Common::RemoveQuotationSymbol;
using LightGBM::Common::Vector2Ptr;
using LightGBM::Common::VectorSize;
using LightGBM::Config;
using LightGBM::data_size_t;
using LightGBM::Dataset;
using LightGBM::DatasetLoader;
using LightGBM::kZeroThreshold;
using LightGBM::LGBM_APIHandleException;
using LightGBM::Log;
using LightGBM::Network;
using LightGBM::Random;
using LightGBM::ReduceScatterFunction;
Guolin Ke's avatar
Guolin Ke committed
849

Guolin Ke's avatar
Guolin Ke committed
850
851
852
853
854
855
856
857
// some help functions used to convert data

std::function<std::vector<double>(int row_idx)>
RowFunctionFromDenseMatric(const void* data, int num_row, int num_col, int data_type, int is_row_major);

std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseMatric(const void* data, int num_row, int num_col, int data_type, int is_row_major);

858
859
860
std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseRows(const void** data, int num_col, int data_type);

861
862
template<typename T>
std::function<std::vector<std::pair<int, double>>(T idx)>
Guolin Ke's avatar
Guolin Ke committed
863
RowFunctionFromCSR(const void* indptr, int indptr_type, const int32_t* indices,
864
                   const void* data, int data_type, int64_t nindptr, int64_t nelem);
Guolin Ke's avatar
Guolin Ke committed
865
866
867

// Row iterator of on column for CSC matrix
class CSC_RowIterator {
Nikita Titov's avatar
Nikita Titov committed
868
 public:
Guolin Ke's avatar
Guolin Ke committed
869
  CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
870
                  const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx);
Guolin Ke's avatar
Guolin Ke committed
871
872
873
874
875
  ~CSC_RowIterator() {}
  // return value at idx, only can access by ascent order
  double Get(int idx);
  // return next non-zero pair, if index < 0, means no more data
  std::pair<int, double> NextNonZero();
Nikita Titov's avatar
Nikita Titov committed
876
877

 private:
Guolin Ke's avatar
Guolin Ke committed
878
879
880
881
882
883
884
885
886
  int nonzero_idx_ = 0;
  int cur_idx_ = -1;
  double cur_val_ = 0.0f;
  bool is_end_ = false;
  std::function<std::pair<int, double>(int idx)> iter_fun_;
};

// start of c_api functions

Guolin Ke's avatar
Guolin Ke committed
887
const char* LGBM_GetLastError() {
wxchan's avatar
wxchan committed
888
  return LastErrorMsg();
Guolin Ke's avatar
Guolin Ke committed
889
890
}

891
892
893
894
895
896
897
898
899
900
901
902
int LGBM_DumpParamAliases(int64_t buffer_len,
                          int64_t* out_len,
                          char* out_str) {
  API_BEGIN();
  std::string aliases = Config::DumpAliases();
  *out_len = static_cast<int64_t>(aliases.size()) + 1;
  if (*out_len <= buffer_len) {
    std::memcpy(out_str, aliases.c_str(), *out_len);
  }
  API_END();
}

903
904
905
906
907
908
int LGBM_RegisterLogCallback(void (*callback)(const char*)) {
  API_BEGIN();
  Log::ResetCallBack(callback);
  API_END();
}

909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
static inline int SampleCount(int32_t total_nrow, const Config& config) {
  return static_cast<int>(total_nrow < config.bin_construct_sample_cnt ? total_nrow : config.bin_construct_sample_cnt);
}

static inline std::vector<int32_t> CreateSampleIndices(int32_t total_nrow, const Config& config) {
  Random rand(config.data_random_seed);
  int sample_cnt = SampleCount(total_nrow, config);
  return rand.Sample(total_nrow, sample_cnt);
}

int LGBM_GetSampleCount(int32_t num_total_row,
                        const char* parameters,
                        int* out) {
  API_BEGIN();
  if (out == nullptr) {
    Log::Fatal("LGBM_GetSampleCount output is nullptr");
  }
  auto param = Config::Str2Map(parameters);
  Config config;
  config.Set(param);

  *out = SampleCount(num_total_row, config);
  API_END();
}

int LGBM_SampleIndices(int32_t num_total_row,
                       const char* parameters,
                       void* out,
                       int32_t* out_len) {
  // This API is to keep python binding's behavior the same with C++ implementation.
  // Sample count, random seed etc. should be provided in parameters.
  API_BEGIN();
  if (out == nullptr) {
    Log::Fatal("LGBM_SampleIndices output is nullptr");
  }
  auto param = Config::Str2Map(parameters);
  Config config;
  config.Set(param);

  auto sample_indices = CreateSampleIndices(num_total_row, config);
  memcpy(out, sample_indices.data(), sizeof(int32_t) * sample_indices.size());
  *out_len = static_cast<int32_t>(sample_indices.size());
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
954
int LGBM_DatasetCreateFromFile(const char* filename,
955
956
957
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
958
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
959
960
  auto param = Config::Str2Map(parameters);
  Config config;
961
  config.Set(param);
962
  OMP_SET_NUM_THREADS(config.num_threads);
963
  DatasetLoader loader(config, nullptr, 1, filename);
Guolin Ke's avatar
Guolin Ke committed
964
  if (reference == nullptr) {
965
    if (Network::num_machines() == 1) {
966
      *out = loader.LoadFromFile(filename);
967
    } else {
968
      *out = loader.LoadFromFile(filename, Network::rank(), Network::num_machines());
969
    }
Guolin Ke's avatar
Guolin Ke committed
970
  } else {
971
    *out = loader.LoadFromFileAlignWithOtherDataset(filename,
972
                                                    reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
973
  }
974
  API_END();
Guolin Ke's avatar
Guolin Ke committed
975
976
}

Guolin Ke's avatar
Guolin Ke committed
977
int LGBM_DatasetCreateFromSampledColumn(double** sample_data,
978
979
980
981
                                        int** sample_indices,
                                        int32_t ncol,
                                        const int* num_per_col,
                                        int32_t num_sample_row,
982
983
                                        int32_t num_local_row,
                                        int64_t num_dist_row,
984
985
                                        const char* parameters,
                                        DatasetHandle* out) {
986
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
987
988
  auto param = Config::Str2Map(parameters);
  Config config;
989
  config.Set(param);
990
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
991
  DatasetLoader loader(config, nullptr, 1, nullptr);
992
993
994
995
  *out = loader.ConstructFromSampleData(sample_data,
                                        sample_indices,
                                        ncol,
                                        num_per_col,
996
                                        num_sample_row,
997
998
                                        static_cast<data_size_t>(num_local_row),
                                        num_dist_row);
999
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1000
1001
}

Guolin Ke's avatar
Guolin Ke committed
1002
int LGBM_DatasetCreateByReference(const DatasetHandle reference,
1003
1004
                                  int64_t num_total_row,
                                  DatasetHandle* out) {
Guolin Ke's avatar
Guolin Ke committed
1005
1006
  API_BEGIN();
  std::unique_ptr<Dataset> ret;
1007
1008
1009
1010
1011
  data_size_t nrows = static_cast<data_size_t>(num_total_row);
  ret.reset(new Dataset(nrows));
  const Dataset* reference_dataset = reinterpret_cast<const Dataset*>(reference);
  ret->CreateValid(reference_dataset);
  ret->InitByReference(nrows, reference_dataset);
Guolin Ke's avatar
Guolin Ke committed
1012
1013
1014
1015
  *out = ret.release();
  API_END();
}

1016
1017
1018
1019
1020
int LGBM_DatasetInitStreaming(DatasetHandle dataset,
                              int32_t has_weights,
                              int32_t has_init_scores,
                              int32_t has_queries,
                              int32_t nclasses,
1021
1022
                              int32_t nthreads,
                              int32_t omp_max_threads) {
1023
1024
1025
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto num_data = p_dataset->num_data();
1026
  p_dataset->InitStreaming(num_data, has_weights, has_init_scores, has_queries, nclasses, nthreads, omp_max_threads);
1027
1028
1029
1030
  p_dataset->set_wait_for_manual_finish(true);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1031
int LGBM_DatasetPushRows(DatasetHandle dataset,
1032
1033
1034
1035
1036
                         const void* data,
                         int data_type,
                         int32_t nrow,
                         int32_t ncol,
                         int32_t start_row) {
Guolin Ke's avatar
Guolin Ke committed
1037
1038
1039
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto get_row_fun = RowFunctionFromDenseMatric(data, nrow, ncol, data_type, 1);
1040
1041
1042
  if (p_dataset->has_raw()) {
    p_dataset->ResizeRaw(p_dataset->num_numeric_features() + nrow);
  }
1043
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1044
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1045
  for (int i = 0; i < nrow; ++i) {
1046
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1047
1048
1049
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(tid, start_row + i, one_row);
1050
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1051
  }
1052
  OMP_THROW_EX();
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
  if (!p_dataset->wait_for_manual_finish() && (start_row + nrow == p_dataset->num_data())) {
    p_dataset->FinishLoad();
  }
  API_END();
}

int LGBM_DatasetPushRowsWithMetadata(DatasetHandle dataset,
                                     const void* data,
                                     int data_type,
                                     int32_t nrow,
                                     int32_t ncol,
                                     int32_t start_row,
                                     const float* labels,
                                     const float* weights,
                                     const double* init_scores,
                                     const int32_t* queries,
                                     int32_t tid) {
  API_BEGIN();
#ifdef LABEL_T_USE_DOUBLE
  Log::Fatal("Don't support LABEL_T_USE_DOUBLE");
#endif
  if (!data) {
    Log::Fatal("data cannot be null.");
  }
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto get_row_fun = RowFunctionFromDenseMatric(data, nrow, ncol, data_type, 1);
  if (p_dataset->has_raw()) {
    p_dataset->ResizeRaw(p_dataset->num_numeric_features() + nrow);
  }

1083
1084
  const int max_omp_threads = p_dataset->omp_max_threads() > 0 ? p_dataset->omp_max_threads() : OMP_NUM_THREADS();

1085
1086
1087
1088
1089
  OMP_INIT_EX();
#pragma omp parallel for schedule(static)
  for (int i = 0; i < nrow; ++i) {
    OMP_LOOP_EX_BEGIN();
    // convert internal thread id to be unique based on external thread id
1090
    const int internal_tid = omp_get_thread_num() + (max_omp_threads * tid);
1091
1092
1093
1094
1095
1096
1097
1098
1099
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(internal_tid, start_row + i, one_row);
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();

  p_dataset->InsertMetadataAt(start_row, nrow, labels, weights, init_scores, queries);

  if (!p_dataset->wait_for_manual_finish() && (start_row + nrow == p_dataset->num_data())) {
Guolin Ke's avatar
Guolin Ke committed
1100
1101
1102
1103
1104
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1105
int LGBM_DatasetPushRowsByCSR(DatasetHandle dataset,
1106
1107
1108
1109
1110
1111
1112
1113
1114
                              const void* indptr,
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
                              int64_t,
                              int64_t start_row) {
Guolin Ke's avatar
Guolin Ke committed
1115
1116
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
1117
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
1118
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
1119
1120
1121
  if (p_dataset->has_raw()) {
    p_dataset->ResizeRaw(p_dataset->num_numeric_features() + nrow);
  }
1122
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1123
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1124
  for (int i = 0; i < nrow; ++i) {
1125
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1126
1127
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
1128
    p_dataset->PushOneRow(tid, static_cast<data_size_t>(start_row + i), one_row);
1129
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1130
  }
1131
  OMP_THROW_EX();
1132
  if (!p_dataset->wait_for_manual_finish() && (start_row + nrow == static_cast<int64_t>(p_dataset->num_data()))) {
Guolin Ke's avatar
Guolin Ke committed
1133
1134
    p_dataset->FinishLoad();
  }
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
  API_END();
}

int LGBM_DatasetPushRowsByCSRWithMetadata(DatasetHandle dataset,
                                          const void* indptr,
                                          int indptr_type,
                                          const int32_t* indices,
                                          const void* data,
                                          int data_type,
                                          int64_t nindptr,
                                          int64_t nelem,
                                          int64_t start_row,
                                          const float* labels,
                                          const float* weights,
                                          const double* init_scores,
                                          const int32_t* queries,
                                          int32_t tid) {
  API_BEGIN();
#ifdef LABEL_T_USE_DOUBLE
  Log::Fatal("Don't support LABEL_T_USE_DOUBLE");
#endif
  if (!data) {
    Log::Fatal("data cannot be null.");
  }
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
  if (p_dataset->has_raw()) {
    p_dataset->ResizeRaw(p_dataset->num_numeric_features() + nrow);
  }
1165
1166
1167

  const int max_omp_threads = p_dataset->omp_max_threads() > 0 ? p_dataset->omp_max_threads() : OMP_NUM_THREADS();

1168
1169
1170
1171
1172
  OMP_INIT_EX();
#pragma omp parallel for schedule(static)
  for (int i = 0; i < nrow; ++i) {
    OMP_LOOP_EX_BEGIN();
    // convert internal thread id to be unique based on external thread id
1173
    const int internal_tid = omp_get_thread_num() + (max_omp_threads * tid);
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(internal_tid, static_cast<data_size_t>(start_row + i), one_row);
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();

  p_dataset->InsertMetadataAt(static_cast<int32_t>(start_row), nrow, labels, weights, init_scores, queries);

  if (!p_dataset->wait_for_manual_finish() && (start_row + nrow == static_cast<int64_t>(p_dataset->num_data()))) {
    p_dataset->FinishLoad();
  }
  API_END();
}

int LGBM_DatasetSetWaitForManualFinish(DatasetHandle dataset, int wait) {
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  p_dataset->set_wait_for_manual_finish(wait);
  API_END();
}

int LGBM_DatasetMarkFinished(DatasetHandle dataset) {
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  p_dataset->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1199
1200
1201
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1202
int LGBM_DatasetCreateFromMat(const void* data,
1203
1204
1205
1206
1207
1208
1209
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
  return LGBM_DatasetCreateFromMats(1,
                                    &data,
                                    data_type,
                                    &nrow,
                                    ncol,
                                    is_row_major,
                                    parameters,
                                    reference,
                                    out);
}

int LGBM_DatasetCreateFromMats(int32_t nmat,
                               const void** data,
                               int data_type,
                               int32_t* nrow,
                               int32_t ncol,
                               int is_row_major,
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
1230
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1231
1232
  auto param = Config::Str2Map(parameters);
  Config config;
1233
  config.Set(param);
1234
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1235
  std::unique_ptr<Dataset> ret;
1236
1237
1238
1239
1240
1241
1242
1243
1244
  int32_t total_nrow = 0;
  for (int j = 0; j < nmat; ++j) {
    total_nrow += nrow[j];
  }

  std::vector<std::function<std::vector<double>(int row_idx)>> get_row_fun;
  for (int j = 0; j < nmat; ++j) {
    get_row_fun.push_back(RowFunctionFromDenseMatric(data[j], nrow[j], ncol, data_type, is_row_major));
  }
1245

Guolin Ke's avatar
Guolin Ke committed
1246
1247
  if (reference == nullptr) {
    // sample data first
1248
1249
    auto sample_indices = CreateSampleIndices(total_nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
1250
    std::vector<std::vector<double>> sample_values(ncol);
Guolin Ke's avatar
Guolin Ke committed
1251
    std::vector<std::vector<int>> sample_idx(ncol);
1252
1253
1254

    int offset = 0;
    int j = 0;
Guolin Ke's avatar
Guolin Ke committed
1255
    for (size_t i = 0; i < sample_indices.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1256
      auto idx = sample_indices[i];
1257
1258
1259
1260
      while ((idx - offset) >= nrow[j]) {
        offset += nrow[j];
        ++j;
      }
1261

1262
1263
1264
1265
1266
      auto row = get_row_fun[j](static_cast<int>(idx - offset));
      for (size_t k = 0; k < row.size(); ++k) {
        if (std::fabs(row[k]) > kZeroThreshold || std::isnan(row[k])) {
          sample_values[k].emplace_back(row[k]);
          sample_idx[k].emplace_back(static_cast<int>(i));
Guolin Ke's avatar
Guolin Ke committed
1267
        }
Guolin Ke's avatar
Guolin Ke committed
1268
1269
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1270
    DatasetLoader loader(config, nullptr, 1, nullptr);
1271
1272
1273
1274
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             ncol,
                                             VectorSize<double>(sample_values).data(),
1275
1276
1277
                                             sample_cnt,
                                             total_nrow,
                                             total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1278
  } else {
1279
    ret.reset(new Dataset(total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1280
    ret->CreateValid(
1281
      reinterpret_cast<const Dataset*>(reference));
1282
1283
1284
    if (ret->has_raw()) {
      ret->ResizeRaw(total_nrow);
    }
Guolin Ke's avatar
Guolin Ke committed
1285
  }
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
  int32_t start_row = 0;
  for (int j = 0; j < nmat; ++j) {
    OMP_INIT_EX();
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < nrow[j]; ++i) {
      OMP_LOOP_EX_BEGIN();
      const int tid = omp_get_thread_num();
      auto one_row = get_row_fun[j](i);
      ret->PushOneRow(tid, start_row + i, one_row);
      OMP_LOOP_EX_END();
    }
    OMP_THROW_EX();

    start_row += nrow[j];
Guolin Ke's avatar
Guolin Ke committed
1300
1301
  }
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1302
  *out = ret.release();
1303
  API_END();
1304
1305
}

Guolin Ke's avatar
Guolin Ke committed
1306
int LGBM_DatasetCreateFromCSR(const void* indptr,
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
                              int64_t num_col,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
1317
  API_BEGIN();
1318
1319
1320
1321
1322
  if (num_col <= 0) {
    Log::Fatal("The number of columns should be greater than zero.");
  } else if (num_col >= INT32_MAX) {
    Log::Fatal("The number of columns should be smaller than INT32_MAX.");
  }
Guolin Ke's avatar
Guolin Ke committed
1323
1324
  auto param = Config::Str2Map(parameters);
  Config config;
1325
  config.Set(param);
1326
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1327
  std::unique_ptr<Dataset> ret;
1328
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1329
1330
1331
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
  if (reference == nullptr) {
    // sample data first
1332
1333
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1334
1335
    std::vector<std::vector<double>> sample_values(num_col);
    std::vector<std::vector<int>> sample_idx(num_col);
1336
1337
1338
1339
    for (size_t i = 0; i < sample_indices.size(); ++i) {
      auto idx = sample_indices[i];
      auto row = get_row_fun(static_cast<int>(idx));
      for (std::pair<int, double>& inner_data : row) {
Nikita Titov's avatar
Nikita Titov committed
1340
        CHECK_LT(inner_data.first, num_col);
Guolin Ke's avatar
Guolin Ke committed
1341
        if (std::fabs(inner_data.second) > kZeroThreshold || std::isnan(inner_data.second)) {
Guolin Ke's avatar
Guolin Ke committed
1342
1343
          sample_values[inner_data.first].emplace_back(inner_data.second);
          sample_idx[inner_data.first].emplace_back(static_cast<int>(i));
1344
1345
1346
        }
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1347
    DatasetLoader loader(config, nullptr, 1, nullptr);
1348
1349
1350
1351
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             static_cast<int>(num_col),
                                             VectorSize<double>(sample_values).data(),
1352
1353
1354
                                             sample_cnt,
                                             nrow,
                                             nrow));
1355
  } else {
1356
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1357
    ret->CreateValid(
1358
      reinterpret_cast<const Dataset*>(reference));
1359
1360
1361
    if (ret->has_raw()) {
      ret->ResizeRaw(nrow);
    }
1362
  }
1363
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1364
  #pragma omp parallel for schedule(static)
1365
  for (int i = 0; i < nindptr - 1; ++i) {
1366
    OMP_LOOP_EX_BEGIN();
1367
1368
1369
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    ret->PushOneRow(tid, i, one_row);
1370
    OMP_LOOP_EX_END();
1371
  }
1372
  OMP_THROW_EX();
1373
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1374
  *out = ret.release();
1375
  API_END();
1376
1377
}

1378
int LGBM_DatasetCreateFromCSRFunc(void* get_row_funptr,
1379
1380
1381
1382
1383
                                  int num_rows,
                                  int64_t num_col,
                                  const char* parameters,
                                  const DatasetHandle reference,
                                  DatasetHandle* out) {
1384
  API_BEGIN();
1385
1386
1387
1388
1389
  if (num_col <= 0) {
    Log::Fatal("The number of columns should be greater than zero.");
  } else if (num_col >= INT32_MAX) {
    Log::Fatal("The number of columns should be smaller than INT32_MAX.");
  }
1390
1391
1392
1393
  auto get_row_fun = *static_cast<std::function<void(int idx, std::vector<std::pair<int, double>>&)>*>(get_row_funptr);
  auto param = Config::Str2Map(parameters);
  Config config;
  config.Set(param);
1394
  OMP_SET_NUM_THREADS(config.num_threads);
1395
1396
1397
1398
  std::unique_ptr<Dataset> ret;
  int32_t nrow = num_rows;
  if (reference == nullptr) {
    // sample data first
1399
1400
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
1401
1402
1403
1404
1405
1406
1407
1408
    std::vector<std::vector<double>> sample_values(num_col);
    std::vector<std::vector<int>> sample_idx(num_col);
    // local buffer to re-use memory
    std::vector<std::pair<int, double>> buffer;
    for (size_t i = 0; i < sample_indices.size(); ++i) {
      auto idx = sample_indices[i];
      get_row_fun(static_cast<int>(idx), buffer);
      for (std::pair<int, double>& inner_data : buffer) {
Nikita Titov's avatar
Nikita Titov committed
1409
        CHECK_LT(inner_data.first, num_col);
1410
1411
1412
1413
1414
1415
1416
        if (std::fabs(inner_data.second) > kZeroThreshold || std::isnan(inner_data.second)) {
          sample_values[inner_data.first].emplace_back(inner_data.second);
          sample_idx[inner_data.first].emplace_back(static_cast<int>(i));
        }
      }
    }
    DatasetLoader loader(config, nullptr, 1, nullptr);
1417
1418
1419
1420
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             static_cast<int>(num_col),
                                             VectorSize<double>(sample_values).data(),
1421
1422
1423
                                             sample_cnt,
                                             nrow,
                                             nrow));
1424
1425
1426
1427
  } else {
    ret.reset(new Dataset(nrow));
    ret->CreateValid(
      reinterpret_cast<const Dataset*>(reference));
1428
1429
1430
    if (ret->has_raw()) {
      ret->ResizeRaw(nrow);
    }
1431
  }
1432

1433
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1434
1435
  std::vector<std::pair<int, double>> thread_buffer;
  #pragma omp parallel for schedule(static) private(thread_buffer)
1436
1437
1438
  for (int i = 0; i < num_rows; ++i) {
    OMP_LOOP_EX_BEGIN();
    {
1439
      const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1440
1441
      get_row_fun(i, thread_buffer);
      ret->PushOneRow(tid, i, thread_buffer);
1442
1443
1444
1445
1446
1447
1448
1449
1450
    }
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();
  ret->FinishLoad();
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1451
int LGBM_DatasetCreateFromCSC(const void* col_ptr,
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
                              int col_ptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t ncol_ptr,
                              int64_t nelem,
                              int64_t num_row,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
1462
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1463
1464
  auto param = Config::Str2Map(parameters);
  Config config;
1465
  config.Set(param);
1466
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1467
  std::unique_ptr<Dataset> ret;
Guolin Ke's avatar
Guolin Ke committed
1468
1469
1470
  int32_t nrow = static_cast<int32_t>(num_row);
  if (reference == nullptr) {
    // sample data first
1471
1472
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1473
    std::vector<std::vector<double>> sample_values(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
1474
    std::vector<std::vector<int>> sample_idx(ncol_ptr - 1);
1475
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1476
    #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1477
    for (int i = 0; i < static_cast<int>(sample_values.size()); ++i) {
1478
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1479
1480
1481
      CSC_RowIterator col_it(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, i);
      for (int j = 0; j < sample_cnt; j++) {
        auto val = col_it.Get(sample_indices[j]);
Guolin Ke's avatar
Guolin Ke committed
1482
        if (std::fabs(val) > kZeroThreshold || std::isnan(val)) {
Guolin Ke's avatar
Guolin Ke committed
1483
1484
          sample_values[i].emplace_back(val);
          sample_idx[i].emplace_back(j);
Guolin Ke's avatar
Guolin Ke committed
1485
1486
        }
      }
1487
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1488
    }
1489
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1490
    DatasetLoader loader(config, nullptr, 1, nullptr);
1491
1492
1493
1494
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             static_cast<int>(sample_values.size()),
                                             VectorSize<double>(sample_values).data(),
1495
1496
1497
                                             sample_cnt,
                                             nrow,
                                             nrow));
Guolin Ke's avatar
Guolin Ke committed
1498
  } else {
1499
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1500
    ret->CreateValid(
1501
      reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
1502
  }
1503
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1504
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1505
  for (int i = 0; i < ncol_ptr - 1; ++i) {
1506
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1507
    const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1508
    int feature_idx = ret->InnerFeatureIndex(i);
Guolin Ke's avatar
Guolin Ke committed
1509
    if (feature_idx < 0) { continue; }
Guolin Ke's avatar
Guolin Ke committed
1510
1511
    int group = ret->Feature2Group(feature_idx);
    int sub_feature = ret->Feture2SubFeature(feature_idx);
Guolin Ke's avatar
Guolin Ke committed
1512
    CSC_RowIterator col_it(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, i);
Guolin Ke's avatar
Guolin Ke committed
1513
1514
1515
1516
1517
1518
1519
1520
    auto bin_mapper = ret->FeatureBinMapper(feature_idx);
    if (bin_mapper->GetDefaultBin() == bin_mapper->GetMostFreqBin()) {
      int row_idx = 0;
      while (row_idx < nrow) {
        auto pair = col_it.NextNonZero();
        row_idx = pair.first;
        // no more data
        if (row_idx < 0) { break; }
1521
        ret->PushOneData(tid, row_idx, group, feature_idx, sub_feature, pair.second);
Guolin Ke's avatar
Guolin Ke committed
1522
1523
1524
1525
      }
    } else {
      for (int row_idx = 0; row_idx < nrow; ++row_idx) {
        auto val = col_it.Get(row_idx);
1526
        ret->PushOneData(tid, row_idx, group, feature_idx, sub_feature, val);
Guolin Ke's avatar
Guolin Ke committed
1527
      }
Guolin Ke's avatar
Guolin Ke committed
1528
    }
1529
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1530
  }
1531
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1532
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1533
  *out = ret.release();
1534
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1535
1536
}

Guolin Ke's avatar
Guolin Ke committed
1537
int LGBM_DatasetGetSubset(
1538
  const DatasetHandle handle,
wxchan's avatar
wxchan committed
1539
1540
1541
  const int32_t* used_row_indices,
  int32_t num_used_row_indices,
  const char* parameters,
Guolin Ke's avatar
typo  
Guolin Ke committed
1542
  DatasetHandle* out) {
wxchan's avatar
wxchan committed
1543
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1544
1545
  auto param = Config::Str2Map(parameters);
  Config config;
1546
  config.Set(param);
1547
  OMP_SET_NUM_THREADS(config.num_threads);
1548
  auto full_dataset = reinterpret_cast<const Dataset*>(handle);
1549
  CHECK_GT(num_used_row_indices, 0);
1550
1551
  const int32_t lower = 0;
  const int32_t upper = full_dataset->num_data() - 1;
1552
  CheckElementsIntervalClosed(used_row_indices, lower, upper, num_used_row_indices, "Used indices of subset");
1553
1554
1555
  if (!std::is_sorted(used_row_indices, used_row_indices + num_used_row_indices)) {
    Log::Fatal("used_row_indices should be sorted in Subset");
  }
Guolin Ke's avatar
Guolin Ke committed
1556
  auto ret = std::unique_ptr<Dataset>(new Dataset(num_used_row_indices));
1557
  ret->CopyFeatureMapperFrom(full_dataset);
1558
  ret->CopySubrow(full_dataset, used_row_indices, num_used_row_indices, true);
wxchan's avatar
wxchan committed
1559
1560
1561
1562
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1563
int LGBM_DatasetSetFeatureNames(
Guolin Ke's avatar
typo  
Guolin Ke committed
1564
  DatasetHandle handle,
Guolin Ke's avatar
Guolin Ke committed
1565
  const char** feature_names,
Guolin Ke's avatar
Guolin Ke committed
1566
  int num_feature_names) {
Guolin Ke's avatar
Guolin Ke committed
1567
1568
1569
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  std::vector<std::string> feature_names_str;
Guolin Ke's avatar
Guolin Ke committed
1570
  for (int i = 0; i < num_feature_names; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1571
1572
1573
1574
1575
1576
    feature_names_str.emplace_back(feature_names[i]);
  }
  dataset->set_feature_names(feature_names_str);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1577
int LGBM_DatasetGetFeatureNames(
1578
1579
1580
1581
1582
1583
    DatasetHandle handle,
    const int len,
    int* num_feature_names,
    const size_t buffer_len,
    size_t* out_buffer_len,
    char** feature_names) {
1584
  API_BEGIN();
1585
  *out_buffer_len = 0;
1586
1587
  auto dataset = reinterpret_cast<Dataset*>(handle);
  auto inside_feature_name = dataset->feature_names();
Guolin Ke's avatar
Guolin Ke committed
1588
1589
  *num_feature_names = static_cast<int>(inside_feature_name.size());
  for (int i = 0; i < *num_feature_names; ++i) {
1590
1591
1592
1593
1594
    if (i < len) {
      std::memcpy(feature_names[i], inside_feature_name[i].c_str(), std::min(inside_feature_name[i].size() + 1, buffer_len));
      feature_names[i][buffer_len - 1] = '\0';
    }
    *out_buffer_len = std::max(inside_feature_name[i].size() + 1, *out_buffer_len);
1595
1596
1597
1598
  }
  API_END();
}

1599
1600
1601
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1602
int LGBM_DatasetFree(DatasetHandle handle) {
1603
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1604
  delete reinterpret_cast<Dataset*>(handle);
1605
  API_END();
1606
1607
}

Guolin Ke's avatar
Guolin Ke committed
1608
int LGBM_DatasetSaveBinary(DatasetHandle handle,
1609
                           const char* filename) {
1610
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1611
1612
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->SaveBinaryFile(filename);
1613
  API_END();
1614
1615
}

1616
1617
1618
1619
1620
1621
1622
1623
int LGBM_DatasetDumpText(DatasetHandle handle,
                         const char* filename) {
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->DumpTextFile(filename);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1624
int LGBM_DatasetSetField(DatasetHandle handle,
1625
1626
1627
1628
                         const char* field_name,
                         const void* field_data,
                         int num_element,
                         int type) {
1629
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1630
  auto dataset = reinterpret_cast<Dataset*>(handle);
1631
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1632
  if (type == C_API_DTYPE_FLOAT32) {
Guolin Ke's avatar
Guolin Ke committed
1633
    is_success = dataset->SetFloatField(field_name, reinterpret_cast<const float*>(field_data), static_cast<int32_t>(num_element));
Guolin Ke's avatar
Guolin Ke committed
1634
  } else if (type == C_API_DTYPE_INT32) {
Guolin Ke's avatar
Guolin Ke committed
1635
    is_success = dataset->SetIntField(field_name, reinterpret_cast<const int*>(field_data), static_cast<int32_t>(num_element));
Guolin Ke's avatar
Guolin Ke committed
1636
1637
  } else if (type == C_API_DTYPE_FLOAT64) {
    is_success = dataset->SetDoubleField(field_name, reinterpret_cast<const double*>(field_data), static_cast<int32_t>(num_element));
1638
  }
1639
  if (!is_success) { Log::Fatal("Input data type error or field not found"); }
1640
  API_END();
1641
1642
}

Guolin Ke's avatar
Guolin Ke committed
1643
int LGBM_DatasetGetField(DatasetHandle handle,
1644
1645
1646
1647
                         const char* field_name,
                         int* out_len,
                         const void** out_ptr,
                         int* out_type) {
1648
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1649
  auto dataset = reinterpret_cast<Dataset*>(handle);
1650
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1651
  if (dataset->GetFloatField(field_name, out_len, reinterpret_cast<const float**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1652
    *out_type = C_API_DTYPE_FLOAT32;
1653
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1654
  } else if (dataset->GetIntField(field_name, out_len, reinterpret_cast<const int**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1655
    *out_type = C_API_DTYPE_INT32;
1656
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1657
1658
1659
  } else if (dataset->GetDoubleField(field_name, out_len, reinterpret_cast<const double**>(out_ptr))) {
    *out_type = C_API_DTYPE_FLOAT64;
    is_success = true;
Nikita Titov's avatar
Nikita Titov committed
1660
  }
1661
  if (!is_success) { Log::Fatal("Field not found"); }
wxchan's avatar
wxchan committed
1662
  if (*out_ptr == nullptr) { *out_len = 0; }
1663
  API_END();
1664
1665
}

1666
int LGBM_DatasetUpdateParamChecking(const char* old_parameters, const char* new_parameters) {
1667
  API_BEGIN();
1668
1669
1670
1671
1672
  auto old_param = Config::Str2Map(old_parameters);
  Config old_config;
  old_config.Set(old_param);
  auto new_param = Config::Str2Map(new_parameters);
  Booster::CheckDatasetResetConfig(old_config, new_param);
1673
1674
1675
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1676
int LGBM_DatasetGetNumData(DatasetHandle handle,
1677
                           int* out) {
1678
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1679
1680
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_data();
1681
  API_END();
1682
1683
}

Guolin Ke's avatar
Guolin Ke committed
1684
int LGBM_DatasetGetNumFeature(DatasetHandle handle,
1685
                              int* out) {
1686
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1687
1688
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_total_features();
1689
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1690
}
1691

1692
1693
1694
1695
1696
int LGBM_DatasetGetFeatureNumBin(DatasetHandle handle,
                                 int feature,
                                 int* out) {
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
1697
1698
1699
1700
1701
  int num_features = dataset->num_total_features();
  if (feature < 0 || feature >= num_features) {
    Log::Fatal("Tried to retrieve number of bins for feature index %d, "
               "but the valid feature indices are [0, %d].", feature, num_features - 1);
  }
1702
1703
1704
1705
1706
1707
1708
1709
1710
  int inner_idx = dataset->InnerFeatureIndex(feature);
  if (inner_idx >= 0) {
    *out = dataset->FeatureNumBin(inner_idx);
  } else {
    *out = 0;
  }
  API_END();
}

1711
1712
1713
1714
1715
int LGBM_DatasetAddFeaturesFrom(DatasetHandle target,
                                DatasetHandle source) {
  API_BEGIN();
  auto target_d = reinterpret_cast<Dataset*>(target);
  auto source_d = reinterpret_cast<Dataset*>(source);
1716
  target_d->AddFeaturesFrom(source_d);
1717
1718
1719
  API_END();
}

1720
1721
// ---- start of booster

Guolin Ke's avatar
Guolin Ke committed
1722
int LGBM_BoosterCreate(const DatasetHandle train_data,
1723
1724
                       const char* parameters,
                       BoosterHandle* out) {
1725
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1726
  const Dataset* p_train_data = reinterpret_cast<const Dataset*>(train_data);
wxchan's avatar
wxchan committed
1727
1728
  auto ret = std::unique_ptr<Booster>(new Booster(p_train_data, parameters));
  *out = ret.release();
1729
  API_END();
1730
1731
}

Guolin Ke's avatar
Guolin Ke committed
1732
int LGBM_BoosterCreateFromModelfile(
1733
  const char* filename,
Guolin Ke's avatar
Guolin Ke committed
1734
  int* out_num_iterations,
1735
  BoosterHandle* out) {
1736
  API_BEGIN();
wxchan's avatar
wxchan committed
1737
  auto ret = std::unique_ptr<Booster>(new Booster(filename));
Guolin Ke's avatar
Guolin Ke committed
1738
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
wxchan's avatar
wxchan committed
1739
  *out = ret.release();
1740
  API_END();
1741
1742
}

Guolin Ke's avatar
Guolin Ke committed
1743
int LGBM_BoosterLoadModelFromString(
1744
1745
1746
1747
  const char* model_str,
  int* out_num_iterations,
  BoosterHandle* out) {
  API_BEGIN();
wxchan's avatar
wxchan committed
1748
  auto ret = std::unique_ptr<Booster>(new Booster(nullptr));
1749
1750
1751
1752
1753
1754
  ret->LoadModelFromString(model_str);
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
  *out = ret.release();
  API_END();
}

1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
int LGBM_BoosterGetLoadedParam(
  BoosterHandle handle,
  int64_t buffer_len,
  int64_t* out_len,
  char* out_str) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  std::string params = ref_booster->GetBoosting()->GetLoadedParam();
  *out_len = static_cast<int64_t>(params.size()) + 1;
  if (*out_len <= buffer_len) {
    std::memcpy(out_str, params.c_str(), *out_len);
  }
  API_END();
}

1770
1771
1772
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1773
int LGBM_BoosterFree(BoosterHandle handle) {
1774
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1775
  delete reinterpret_cast<Booster*>(handle);
1776
  API_END();
1777
1778
}

1779
int LGBM_BoosterShuffleModels(BoosterHandle handle, int start_iter, int end_iter) {
1780
1781
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1782
  ref_booster->ShuffleModels(start_iter, end_iter);
1783
1784
1785
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1786
int LGBM_BoosterMerge(BoosterHandle handle,
1787
                      BoosterHandle other_handle) {
wxchan's avatar
wxchan committed
1788
1789
1790
1791
1792
1793
1794
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  Booster* ref_other_booster = reinterpret_cast<Booster*>(other_handle);
  ref_booster->MergeFrom(ref_other_booster);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1795
int LGBM_BoosterAddValidData(BoosterHandle handle,
1796
                             const DatasetHandle valid_data) {
wxchan's avatar
wxchan committed
1797
1798
1799
1800
1801
1802
1803
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  const Dataset* p_dataset = reinterpret_cast<const Dataset*>(valid_data);
  ref_booster->AddValidData(p_dataset);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1804
int LGBM_BoosterResetTrainingData(BoosterHandle handle,
1805
                                  const DatasetHandle train_data) {
wxchan's avatar
wxchan committed
1806
1807
1808
1809
1810
1811
1812
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  const Dataset* p_dataset = reinterpret_cast<const Dataset*>(train_data);
  ref_booster->ResetTrainingData(p_dataset);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1813
int LGBM_BoosterResetParameter(BoosterHandle handle, const char* parameters) {
wxchan's avatar
wxchan committed
1814
1815
1816
1817
1818
1819
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->ResetConfig(parameters);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1820
int LGBM_BoosterGetNumClasses(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1821
1822
1823
1824
1825
1826
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->NumberOfClasses();
  API_END();
}

1827
int LGBM_BoosterGetLinear(BoosterHandle handle, int* out) {
1828
1829
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1830
1831
1832
1833
1834
  if (ref_booster->GetBoosting()->IsLinear()) {
    *out = 1;
  } else {
    *out = 0;
  }
1835
1836
1837
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1838
1839
1840
1841
1842
1843
1844
int LGBM_BoosterRefit(BoosterHandle handle, const int32_t* leaf_preds, int32_t nrow, int32_t ncol) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->Refit(leaf_preds, nrow, ncol);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1845
int LGBM_BoosterUpdateOneIter(BoosterHandle handle, int* is_finished) {
1846
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1847
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1848
1849
1850
1851
1852
  if (ref_booster->TrainOneIter()) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1853
  API_END();
1854
1855
}

Guolin Ke's avatar
Guolin Ke committed
1856
int LGBM_BoosterUpdateOneIterCustom(BoosterHandle handle,
1857
1858
1859
                                    const float* grad,
                                    const float* hess,
                                    int* is_finished) {
1860
  API_BEGIN();
1861
  #ifdef SCORE_T_USE_DOUBLE
1862
1863
1864
1865
  (void) handle;       // UNUSED VARIABLE
  (void) grad;         // UNUSED VARIABLE
  (void) hess;         // UNUSED VARIABLE
  (void) is_finished;  // UNUSED VARIABLE
1866
  Log::Fatal("Don't support custom loss function when SCORE_T_USE_DOUBLE is enabled");
1867
  #else
1868
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1869
1870
1871
1872
1873
  if (ref_booster->TrainOneIter(grad, hess)) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1874
  #endif
1875
  API_END();
1876
1877
}

Guolin Ke's avatar
Guolin Ke committed
1878
int LGBM_BoosterRollbackOneIter(BoosterHandle handle) {
wxchan's avatar
wxchan committed
1879
1880
1881
1882
1883
1884
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->RollbackOneIter();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1885
int LGBM_BoosterGetCurrentIteration(BoosterHandle handle, int* out_iteration) {
wxchan's avatar
wxchan committed
1886
1887
1888
1889
1890
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_iteration = ref_booster->GetBoosting()->GetCurrentIteration();
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
1891

1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
int LGBM_BoosterNumModelPerIteration(BoosterHandle handle, int* out_tree_per_iteration) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_tree_per_iteration = ref_booster->GetBoosting()->NumModelPerIteration();
  API_END();
}

int LGBM_BoosterNumberOfTotalModel(BoosterHandle handle, int* out_models) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_models = ref_booster->GetBoosting()->NumberOfTotalModel();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1906
int LGBM_BoosterGetEvalCounts(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1907
1908
1909
1910
1911
1912
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetEvalCounts();
  API_END();
}

1913
1914
1915
1916
1917
1918
int LGBM_BoosterGetEvalNames(BoosterHandle handle,
                             const int len,
                             int* out_len,
                             const size_t buffer_len,
                             size_t* out_buffer_len,
                             char** out_strs) {
wxchan's avatar
wxchan committed
1919
1920
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1921
  *out_len = ref_booster->GetEvalNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1922
1923
1924
  API_END();
}

1925
1926
1927
1928
1929
1930
int LGBM_BoosterGetFeatureNames(BoosterHandle handle,
                                const int len,
                                int* out_len,
                                const size_t buffer_len,
                                size_t* out_buffer_len,
                                char** out_strs) {
wxchan's avatar
wxchan committed
1931
1932
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1933
  *out_len = ref_booster->GetFeatureNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1934
1935
1936
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1937
int LGBM_BoosterGetNumFeature(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1938
1939
1940
1941
1942
1943
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->MaxFeatureIdx() + 1;
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1944
int LGBM_BoosterGetEval(BoosterHandle handle,
1945
1946
1947
                        int data_idx,
                        int* out_len,
                        double* out_results) {
1948
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1949
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1950
  auto boosting = ref_booster->GetBoosting();
wxchan's avatar
wxchan committed
1951
  auto result_buf = boosting->GetEvalAt(data_idx);
Guolin Ke's avatar
Guolin Ke committed
1952
  *out_len = static_cast<int>(result_buf.size());
1953
  for (size_t i = 0; i < result_buf.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1954
    (out_results)[i] = static_cast<double>(result_buf[i]);
1955
  }
1956
  API_END();
1957
1958
}

Guolin Ke's avatar
Guolin Ke committed
1959
int LGBM_BoosterGetNumPredict(BoosterHandle handle,
1960
1961
                              int data_idx,
                              int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1962
1963
1964
1965
1966
1967
  API_BEGIN();
  auto boosting = reinterpret_cast<Booster*>(handle)->GetBoosting();
  *out_len = boosting->GetNumPredictAt(data_idx);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1968
int LGBM_BoosterGetPredict(BoosterHandle handle,
1969
1970
1971
                           int data_idx,
                           int64_t* out_len,
                           double* out_result) {
1972
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1973
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1974
  ref_booster->GetPredictAt(data_idx, out_result, out_len);
1975
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1976
1977
}

Guolin Ke's avatar
Guolin Ke committed
1978
int LGBM_BoosterPredictForFile(BoosterHandle handle,
1979
1980
1981
                               const char* data_filename,
                               int data_has_header,
                               int predict_type,
1982
                               int start_iteration,
1983
                               int num_iteration,
1984
                               const char* parameter,
1985
                               const char* result_filename) {
1986
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1987
1988
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1989
  config.Set(param);
1990
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1991
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1992
  ref_booster->Predict(start_iteration, num_iteration, predict_type, data_filename, data_has_header,
Guolin Ke's avatar
Guolin Ke committed
1993
                       config, result_filename);
1994
  API_END();
1995
1996
}

Guolin Ke's avatar
Guolin Ke committed
1997
int LGBM_BoosterCalcNumPredict(BoosterHandle handle,
1998
1999
                               int num_row,
                               int predict_type,
2000
                               int start_iteration,
2001
2002
                               int num_iteration,
                               int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
2003
2004
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2005
  *out_len = static_cast<int64_t>(num_row) * ref_booster->GetBoosting()->NumPredictOneRow(start_iteration,
2006
    num_iteration, predict_type == C_API_PREDICT_LEAF_INDEX, predict_type == C_API_PREDICT_CONTRIB);
Guolin Ke's avatar
Guolin Ke committed
2007
2008
2009
  API_END();
}

2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
/*!
 * \brief Object to store resources meant for single-row Fast Predict methods.
 *
 * Meant to be used as a basic struct by the *Fast* predict methods only.
 * It stores the configuration resources for reuse during prediction.
 *
 * Even the row function is stored. We score the instance at the same memory
 * address all the time. One just replaces the feature values at that address
 * and scores again with the *Fast* methods.
 */
struct FastConfig {
  FastConfig(Booster *const booster_ptr,
             const char *parameter,
2023
             const int predict_type_,
2024
             const int data_type_,
2025
             const int32_t num_cols) : booster(booster_ptr), predict_type(predict_type_), data_type(data_type_), ncol(num_cols) {
2026
2027
2028
2029
2030
    config.Set(Config::Str2Map(parameter));
  }

  Booster* const booster;
  Config config;
2031
  const int predict_type;
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
  const int data_type;
  const int32_t ncol;
};

int LGBM_FastConfigFree(FastConfigHandle fastConfig) {
  API_BEGIN();
  delete reinterpret_cast<FastConfig*>(fastConfig);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2042
int LGBM_BoosterPredictForCSR(BoosterHandle handle,
2043
2044
2045
2046
2047
2048
2049
                              const void* indptr,
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
2050
                              int64_t num_col,
2051
                              int predict_type,
2052
                              int start_iteration,
2053
                              int num_iteration,
2054
                              const char* parameter,
2055
2056
                              int64_t* out_len,
                              double* out_result) {
2057
  API_BEGIN();
2058
2059
2060
2061
2062
  if (num_col <= 0) {
    Log::Fatal("The number of columns should be greater than zero.");
  } else if (num_col >= INT32_MAX) {
    Log::Fatal("The number of columns should be smaller than INT32_MAX.");
  }
Guolin Ke's avatar
Guolin Ke committed
2063
2064
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2065
  config.Set(param);
2066
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
2067
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2068
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
2069
  int nrow = static_cast<int>(nindptr - 1);
2070
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, static_cast<int>(num_col), get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
2071
                       config, out_result, out_len);
2072
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2073
}
2074

2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
int LGBM_BoosterPredictSparseOutput(BoosterHandle handle,
                                    const void* indptr,
                                    int indptr_type,
                                    const int32_t* indices,
                                    const void* data,
                                    int data_type,
                                    int64_t nindptr,
                                    int64_t nelem,
                                    int64_t num_col_or_row,
                                    int predict_type,
2085
                                    int start_iteration,
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
                                    int num_iteration,
                                    const char* parameter,
                                    int matrix_type,
                                    int64_t* out_len,
                                    void** out_indptr,
                                    int32_t** out_indices,
                                    void** out_data) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
2098
  OMP_SET_NUM_THREADS(config.num_threads);
2099
2100
2101
2102
2103
2104
2105
2106
  if (matrix_type == C_API_MATRIX_TYPE_CSR) {
    if (num_col_or_row <= 0) {
      Log::Fatal("The number of columns should be greater than zero.");
    } else if (num_col_or_row >= INT32_MAX) {
      Log::Fatal("The number of columns should be smaller than INT32_MAX.");
    }
    auto get_row_fun = RowFunctionFromCSR<int64_t>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
    int64_t nrow = nindptr - 1;
2107
    ref_booster->PredictSparseCSR(start_iteration, num_iteration, predict_type, nrow, static_cast<int>(num_col_or_row), get_row_fun,
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
                                  config, out_len, out_indptr, indptr_type, out_indices, out_data, data_type);
  } else if (matrix_type == C_API_MATRIX_TYPE_CSC) {
    int num_threads = OMP_NUM_THREADS();
    int ncol = static_cast<int>(nindptr - 1);
    std::vector<std::vector<CSC_RowIterator>> iterators(num_threads, std::vector<CSC_RowIterator>());
    for (int i = 0; i < num_threads; ++i) {
      for (int j = 0; j < ncol; ++j) {
        iterators[i].emplace_back(indptr, indptr_type, indices, data, data_type, nindptr, nelem, j);
      }
    }
    std::function<std::vector<std::pair<int, double>>(int64_t row_idx)> get_row_fun =
      [&iterators, ncol](int64_t i) {
      std::vector<std::pair<int, double>> one_row;
      one_row.reserve(ncol);
      const int tid = omp_get_thread_num();
      for (int j = 0; j < ncol; ++j) {
        auto val = iterators[tid][j].Get(static_cast<int>(i));
        if (std::fabs(val) > kZeroThreshold || std::isnan(val)) {
          one_row.emplace_back(j, val);
        }
      }
      return one_row;
    };
2131
    ref_booster->PredictSparseCSC(start_iteration, num_iteration, predict_type, num_col_or_row, ncol, get_row_fun, config,
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
                                  out_len, out_indptr, indptr_type, out_indices, out_data, data_type);
  } else {
    Log::Fatal("Unknown matrix type in LGBM_BoosterPredictSparseOutput");
  }
  API_END();
}

int LGBM_BoosterFreePredictSparse(void* indptr, int32_t* indices, void* data, int indptr_type, int data_type) {
  API_BEGIN();
  if (indptr_type == C_API_DTYPE_INT32) {
2142
    delete[] reinterpret_cast<int32_t*>(indptr);
2143
  } else if (indptr_type == C_API_DTYPE_INT64) {
2144
    delete[] reinterpret_cast<int64_t*>(indptr);
2145
2146
2147
  } else {
    Log::Fatal("Unknown indptr type in LGBM_BoosterFreePredictSparse");
  }
2148
  delete[] indices;
2149
  if (data_type == C_API_DTYPE_FLOAT32) {
2150
    delete[] reinterpret_cast<float*>(data);
2151
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2152
    delete[] reinterpret_cast<double*>(data);
2153
2154
2155
2156
2157
2158
  } else {
    Log::Fatal("Unknown data type in LGBM_BoosterFreePredictSparse");
  }
  API_END();
}

2159
int LGBM_BoosterPredictForCSRSingleRow(BoosterHandle handle,
2160
2161
2162
2163
2164
2165
2166
                                       const void* indptr,
                                       int indptr_type,
                                       const int32_t* indices,
                                       const void* data,
                                       int data_type,
                                       int64_t nindptr,
                                       int64_t nelem,
2167
                                       int64_t num_col,
2168
                                       int predict_type,
2169
                                       int start_iteration,
2170
2171
2172
2173
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2174
  API_BEGIN();
2175
2176
2177
2178
2179
  if (num_col <= 0) {
    Log::Fatal("The number of columns should be greater than zero.");
  } else if (num_col >= INT32_MAX) {
    Log::Fatal("The number of columns should be smaller than INT32_MAX.");
  }
2180
2181
2182
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
2183
  OMP_SET_NUM_THREADS(config.num_threads);
2184
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2185
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
2186
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
2187
  ref_booster->PredictSingleRow(predict_type, static_cast<int32_t>(num_col), get_row_fun, config, out_result, out_len);
2188
2189
2190
  API_END();
}

2191
int LGBM_BoosterPredictForCSRSingleRowFastInit(BoosterHandle handle,
2192
                                               const int predict_type,
2193
                                               const int start_iteration,
2194
                                               const int num_iteration,
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
                                               const int data_type,
                                               const int64_t num_col,
                                               const char* parameter,
                                               FastConfigHandle *out_fastConfig) {
  API_BEGIN();
  if (num_col <= 0) {
    Log::Fatal("The number of columns should be greater than zero.");
  } else if (num_col >= INT32_MAX) {
    Log::Fatal("The number of columns should be smaller than INT32_MAX.");
  }

  auto fastConfig_ptr = std::unique_ptr<FastConfig>(new FastConfig(
    reinterpret_cast<Booster*>(handle),
    parameter,
2209
    predict_type,
2210
2211
2212
    data_type,
    static_cast<int32_t>(num_col)));

2213
  OMP_SET_NUM_THREADS(fastConfig_ptr->config.num_threads);
2214

2215
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2216

2217
2218
2219
2220
2221
2222
  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

int LGBM_BoosterPredictForCSRSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* indptr,
2223
                                           const int indptr_type,
2224
2225
                                           const int32_t* indices,
                                           const void* data,
2226
2227
                                           const int64_t nindptr,
                                           const int64_t nelem,
2228
2229
2230
2231
2232
                                           int64_t* out_len,
                                           double* out_result) {
  API_BEGIN();
  FastConfig *fastConfig = reinterpret_cast<FastConfig*>(fastConfig_handle);
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, fastConfig->data_type, nindptr, nelem);
2233
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
2234
2235
2236
2237
                                        get_row_fun, fastConfig->config, out_result, out_len);
  API_END();
}

2238

Guolin Ke's avatar
Guolin Ke committed
2239
int LGBM_BoosterPredictForCSC(BoosterHandle handle,
2240
2241
2242
2243
2244
2245
2246
2247
2248
                              const void* col_ptr,
                              int col_ptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t ncol_ptr,
                              int64_t nelem,
                              int64_t num_row,
                              int predict_type,
2249
                              int start_iteration,
2250
                              int num_iteration,
2251
                              const char* parameter,
2252
2253
                              int64_t* out_len,
                              double* out_result) {
Guolin Ke's avatar
Guolin Ke committed
2254
2255
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2256
2257
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2258
  config.Set(param);
2259
  OMP_SET_NUM_THREADS(config.num_threads);
2260
  int num_threads = OMP_NUM_THREADS();
Guolin Ke's avatar
Guolin Ke committed
2261
  int ncol = static_cast<int>(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
2262
2263
2264
2265
2266
  std::vector<std::vector<CSC_RowIterator>> iterators(num_threads, std::vector<CSC_RowIterator>());
  for (int i = 0; i < num_threads; ++i) {
    for (int j = 0; j < ncol; ++j) {
      iterators[i].emplace_back(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, j);
    }
Guolin Ke's avatar
Guolin Ke committed
2267
2268
  }
  std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun =
Guolin Ke's avatar
Guolin Ke committed
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
      [&iterators, ncol](int i) {
        std::vector<std::pair<int, double>> one_row;
        one_row.reserve(ncol);
        const int tid = omp_get_thread_num();
        for (int j = 0; j < ncol; ++j) {
          auto val = iterators[tid][j].Get(i);
          if (std::fabs(val) > kZeroThreshold || std::isnan(val)) {
            one_row.emplace_back(j, val);
          }
        }
        return one_row;
      };
2281
  ref_booster->Predict(start_iteration, num_iteration, predict_type, static_cast<int>(num_row), ncol, get_row_fun, config,
cbecker's avatar
cbecker committed
2282
                       out_result, out_len);
Guolin Ke's avatar
Guolin Ke committed
2283
2284
2285
  API_END();
}

2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
int LGBM_BoosterValidateFeatureNames(BoosterHandle handle,
                                     const char** data_names,
                                     int data_num_features) {
  API_BEGIN();
  int booster_num_features;
  size_t out_buffer_len;
  LGBM_BoosterGetFeatureNames(handle, 0, &booster_num_features, 0, &out_buffer_len, nullptr);
  if (booster_num_features != data_num_features) {
    Log::Fatal("Model was trained on %d features, but got %d input features to predict.", booster_num_features, data_num_features);
  }
  std::vector<std::vector<char>> tmp_names(booster_num_features, std::vector<char>(out_buffer_len));
  std::vector<char*> booster_names = Vector2Ptr(&tmp_names);
  LGBM_BoosterGetFeatureNames(handle, data_num_features, &booster_num_features, out_buffer_len, &out_buffer_len, booster_names.data());
  for (int i = 0; i < booster_num_features; ++i) {
    if (strcmp(data_names[i], booster_names[i]) != 0) {
      Log::Fatal("Expected '%s' at position %d but found '%s'", booster_names[i], i, data_names[i]);
    }
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2307
int LGBM_BoosterPredictForMat(BoosterHandle handle,
2308
2309
2310
2311
2312
2313
                              const void* data,
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              int predict_type,
2314
                              int start_iteration,
2315
                              int num_iteration,
2316
                              const char* parameter,
2317
2318
                              int64_t* out_len,
                              double* out_result) {
2319
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2320
2321
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2322
  config.Set(param);
2323
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
2324
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2325
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, nrow, ncol, data_type, is_row_major);
2326
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
2327
                       config, out_result, out_len);
2328
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2329
}
2330

2331
int LGBM_BoosterPredictForMatSingleRow(BoosterHandle handle,
2332
2333
2334
2335
2336
                                       const void* data,
                                       int data_type,
                                       int32_t ncol,
                                       int is_row_major,
                                       int predict_type,
2337
                                       int start_iteration,
2338
2339
2340
2341
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2342
2343
2344
2345
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
2346
  OMP_SET_NUM_THREADS(config.num_threads);
2347
2348
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, ncol, data_type, is_row_major);
2349
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
2350
  ref_booster->PredictSingleRow(predict_type, ncol, get_row_fun, config, out_result, out_len);
2351
2352
2353
  API_END();
}

2354
int LGBM_BoosterPredictForMatSingleRowFastInit(BoosterHandle handle,
2355
                                               const int predict_type,
2356
                                               const int start_iteration,
2357
                                               const int num_iteration,
2358
2359
2360
2361
2362
2363
2364
2365
                                               const int data_type,
                                               const int32_t ncol,
                                               const char* parameter,
                                               FastConfigHandle *out_fastConfig) {
  API_BEGIN();
  auto fastConfig_ptr = std::unique_ptr<FastConfig>(new FastConfig(
    reinterpret_cast<Booster*>(handle),
    parameter,
2366
    predict_type,
2367
2368
2369
    data_type,
    ncol));

2370
  OMP_SET_NUM_THREADS(fastConfig_ptr->config.num_threads);
2371

2372
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2373

2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

int LGBM_BoosterPredictForMatSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* data,
                                           int64_t* out_len,
                                           double* out_result) {
  API_BEGIN();
  FastConfig *fastConfig = reinterpret_cast<FastConfig*>(fastConfig_handle);
  // Single row in row-major format:
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, fastConfig->ncol, fastConfig->data_type, 1);
2386
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
2387
2388
2389
2390
2391
                                        get_row_fun, fastConfig->config,
                                        out_result, out_len);
  API_END();
}

2392

2393
2394
2395
2396
2397
2398
int LGBM_BoosterPredictForMats(BoosterHandle handle,
                               const void** data,
                               int data_type,
                               int32_t nrow,
                               int32_t ncol,
                               int predict_type,
2399
                               int start_iteration,
2400
2401
2402
2403
2404
2405
2406
2407
                               int num_iteration,
                               const char* parameter,
                               int64_t* out_len,
                               double* out_result) {
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
2408
  OMP_SET_NUM_THREADS(config.num_threads);
2409
2410
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseRows(data, ncol, data_type);
2411
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, out_result, out_len);
2412
2413
2414
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2415
int LGBM_BoosterSaveModel(BoosterHandle handle,
2416
                          int start_iteration,
2417
                          int num_iteration,
2418
                          int feature_importance_type,
2419
                          const char* filename) {
2420
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2421
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2422
2423
  ref_booster->SaveModelToFile(start_iteration, num_iteration,
                               feature_importance_type, filename);
wxchan's avatar
wxchan committed
2424
2425
2426
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2427
int LGBM_BoosterSaveModelToString(BoosterHandle handle,
2428
                                  int start_iteration,
2429
                                  int num_iteration,
2430
                                  int feature_importance_type,
2431
                                  int64_t buffer_len,
2432
                                  int64_t* out_len,
2433
                                  char* out_str) {
2434
2435
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2436
2437
  std::string model = ref_booster->SaveModelToString(
      start_iteration, num_iteration, feature_importance_type);
2438
  *out_len = static_cast<int64_t>(model.size()) + 1;
2439
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2440
    std::memcpy(out_str, model.c_str(), *out_len);
2441
2442
2443
2444
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2445
int LGBM_BoosterDumpModel(BoosterHandle handle,
2446
                          int start_iteration,
2447
                          int num_iteration,
2448
                          int feature_importance_type,
2449
2450
                          int64_t buffer_len,
                          int64_t* out_len,
2451
                          char* out_str) {
wxchan's avatar
wxchan committed
2452
2453
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2454
2455
  std::string model = ref_booster->DumpModel(start_iteration, num_iteration,
                                             feature_importance_type);
2456
  *out_len = static_cast<int64_t>(model.size()) + 1;
wxchan's avatar
wxchan committed
2457
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2458
    std::memcpy(out_str, model.c_str(), *out_len);
wxchan's avatar
wxchan committed
2459
  }
2460
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2461
}
2462

Guolin Ke's avatar
Guolin Ke committed
2463
int LGBM_BoosterGetLeafValue(BoosterHandle handle,
2464
2465
2466
                             int tree_idx,
                             int leaf_idx,
                             double* out_val) {
Guolin Ke's avatar
Guolin Ke committed
2467
2468
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2469
  *out_val = static_cast<double>(ref_booster->GetLeafValue(tree_idx, leaf_idx));
Guolin Ke's avatar
Guolin Ke committed
2470
2471
2472
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2473
int LGBM_BoosterSetLeafValue(BoosterHandle handle,
2474
2475
2476
                             int tree_idx,
                             int leaf_idx,
                             double val) {
Guolin Ke's avatar
Guolin Ke committed
2477
2478
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2479
  ref_booster->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
2480
2481
2482
  API_END();
}

2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
int LGBM_BoosterFeatureImportance(BoosterHandle handle,
                                  int num_iteration,
                                  int importance_type,
                                  double* out_results) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  std::vector<double> feature_importances = ref_booster->FeatureImportance(num_iteration, importance_type);
  for (size_t i = 0; i < feature_importances.size(); ++i) {
    (out_results)[i] = feature_importances[i];
  }
  API_END();
}

2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
int LGBM_BoosterGetUpperBoundValue(BoosterHandle handle,
                                   double* out_results) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  double max_value = ref_booster->UpperBoundValue();
  *out_results = max_value;
  API_END();
}

int LGBM_BoosterGetLowerBoundValue(BoosterHandle handle,
                                   double* out_results) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  double min_value = ref_booster->LowerBoundValue();
  *out_results = min_value;
  API_END();
}

2514
2515
2516
2517
2518
int LGBM_NetworkInit(const char* machines,
                     int local_listen_port,
                     int listen_time_out,
                     int num_machines) {
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2519
  Config config;
2520
  config.machines = RemoveQuotationSymbol(std::string(machines));
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
  config.local_listen_port = local_listen_port;
  config.num_machines = num_machines;
  config.time_out = listen_time_out;
  if (num_machines > 1) {
    Network::Init(config);
  }
  API_END();
}

int LGBM_NetworkFree() {
  API_BEGIN();
  Network::Dispose();
  API_END();
}

2536
2537
2538
int LGBM_NetworkInitWithFunctions(int num_machines, int rank,
                                  void* reduce_scatter_ext_fun,
                                  void* allgather_ext_fun) {
ww's avatar
ww committed
2539
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2540
  if (num_machines > 1) {
2541
    Network::Init(num_machines, rank, (ReduceScatterFunction)reduce_scatter_ext_fun, (AllgatherFunction)allgather_ext_fun);
ww's avatar
ww committed
2542
2543
2544
  }
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
2545

Guolin Ke's avatar
Guolin Ke committed
2546
// ---- start of some help functions
2547

2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572

template<typename T>
std::function<std::vector<double>(int row_idx)>
RowFunctionFromDenseMatric_helper(const void* data, int num_row, int num_col, int is_row_major) {
  const T* data_ptr = reinterpret_cast<const T*>(data);
  if (is_row_major) {
    return [=] (int row_idx) {
      std::vector<double> ret(num_col);
      auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
      for (int i = 0; i < num_col; ++i) {
        ret[i] = static_cast<double>(*(tmp_ptr + i));
      }
      return ret;
    };
  } else {
    return [=] (int row_idx) {
      std::vector<double> ret(num_col);
      for (int i = 0; i < num_col; ++i) {
        ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
      }
      return ret;
    };
  }
}

2573
2574
std::function<std::vector<double>(int row_idx)>
RowFunctionFromDenseMatric(const void* data, int num_row, int num_col, int data_type, int is_row_major) {
Guolin Ke's avatar
Guolin Ke committed
2575
  if (data_type == C_API_DTYPE_FLOAT32) {
2576
    return RowFunctionFromDenseMatric_helper<float>(data, num_row, num_col, is_row_major);
Guolin Ke's avatar
Guolin Ke committed
2577
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2578
    return RowFunctionFromDenseMatric_helper<double>(data, num_row, num_col, is_row_major);
2579
  }
2580
  Log::Fatal("Unknown data type in RowFunctionFromDenseMatric");
2581
  return nullptr;
2582
2583
2584
2585
}

std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseMatric(const void* data, int num_row, int num_col, int data_type, int is_row_major) {
Guolin Ke's avatar
Guolin Ke committed
2586
2587
  auto inner_function = RowFunctionFromDenseMatric(data, num_row, num_col, data_type, is_row_major);
  if (inner_function != nullptr) {
2588
    return [inner_function] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2589
2590
      auto raw_values = inner_function(row_idx);
      std::vector<std::pair<int, double>> ret;
Guolin Ke's avatar
Guolin Ke committed
2591
      ret.reserve(raw_values.size());
Guolin Ke's avatar
Guolin Ke committed
2592
      for (int i = 0; i < static_cast<int>(raw_values.size()); ++i) {
Guolin Ke's avatar
Guolin Ke committed
2593
        if (std::fabs(raw_values[i]) > kZeroThreshold || std::isnan(raw_values[i])) {
Guolin Ke's avatar
Guolin Ke committed
2594
          ret.emplace_back(i, raw_values[i]);
2595
        }
Guolin Ke's avatar
Guolin Ke committed
2596
2597
2598
      }
      return ret;
    };
2599
  }
Guolin Ke's avatar
Guolin Ke committed
2600
  return nullptr;
2601
2602
}

2603
2604
2605
2606
2607
2608
2609
// data is array of pointers to individual rows
std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseRows(const void** data, int num_col, int data_type) {
  return [=](int row_idx) {
    auto inner_function = RowFunctionFromDenseMatric(data[row_idx], 1, num_col, data_type, /* is_row_major */ true);
    auto raw_values = inner_function(0);
    std::vector<std::pair<int, double>> ret;
Guolin Ke's avatar
Guolin Ke committed
2610
    ret.reserve(raw_values.size());
2611
2612
2613
2614
2615
2616
2617
2618
2619
    for (int i = 0; i < static_cast<int>(raw_values.size()); ++i) {
      if (std::fabs(raw_values[i]) > kZeroThreshold || std::isnan(raw_values[i])) {
        ret.emplace_back(i, raw_values[i]);
      }
    }
    return ret;
  };
}

2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
template<typename T, typename T1, typename T2>
std::function<std::vector<std::pair<int, double>>(T idx)>
RowFunctionFromCSR_helper(const void* indptr, const int32_t* indices, const void* data) {
  const T1* data_ptr = reinterpret_cast<const T1*>(data);
  const T2* ptr_indptr = reinterpret_cast<const T2*>(indptr);
  return [=] (T idx) {
    std::vector<std::pair<int, double>> ret;
    int64_t start = ptr_indptr[idx];
    int64_t end = ptr_indptr[idx + 1];
    if (end - start > 0)  {
      ret.reserve(end - start);
    }
    for (int64_t i = start; i < end; ++i) {
      ret.emplace_back(indices[i], data_ptr[i]);
    }
    return ret;
  };
}

2639
2640
template<typename T>
std::function<std::vector<std::pair<int, double>>(T idx)>
2641
RowFunctionFromCSR(const void* indptr, int indptr_type, const int32_t* indices, const void* data, int data_type, int64_t , int64_t ) {
Guolin Ke's avatar
Guolin Ke committed
2642
2643
  if (data_type == C_API_DTYPE_FLOAT32) {
    if (indptr_type == C_API_DTYPE_INT32) {
2644
     return RowFunctionFromCSR_helper<T, float, int32_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2645
    } else if (indptr_type == C_API_DTYPE_INT64) {
2646
     return RowFunctionFromCSR_helper<T, float, int64_t>(indptr, indices, data);
2647
    }
Guolin Ke's avatar
Guolin Ke committed
2648
2649
  } else if (data_type == C_API_DTYPE_FLOAT64) {
    if (indptr_type == C_API_DTYPE_INT32) {
2650
     return RowFunctionFromCSR_helper<T, double, int32_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2651
    } else if (indptr_type == C_API_DTYPE_INT64) {
2652
     return RowFunctionFromCSR_helper<T, double, int64_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2653
2654
    }
  }
2655
  Log::Fatal("Unknown data type in RowFunctionFromCSR");
2656
  return nullptr;
2657
2658
}

2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675
2676
2677


template <typename T1, typename T2>
std::function<std::pair<int, double>(int idx)> IterateFunctionFromCSC_helper(const void* col_ptr, const int32_t* indices, const void* data, int col_idx) {
  const T1* data_ptr = reinterpret_cast<const T1*>(data);
  const T2* ptr_col_ptr = reinterpret_cast<const T2*>(col_ptr);
  int64_t start = ptr_col_ptr[col_idx];
  int64_t end = ptr_col_ptr[col_idx + 1];
  return [=] (int offset) {
    int64_t i = static_cast<int64_t>(start + offset);
    if (i >= end) {
      return std::make_pair(-1, 0.0);
    }
    int idx = static_cast<int>(indices[i]);
    double val = static_cast<double>(data_ptr[i]);
    return std::make_pair(idx, val);
  };
}

Guolin Ke's avatar
Guolin Ke committed
2678
std::function<std::pair<int, double>(int idx)>
2679
IterateFunctionFromCSC(const void* col_ptr, int col_ptr_type, const int32_t* indices, const void* data, int data_type, int64_t ncol_ptr, int64_t , int col_idx) {
Guolin Ke's avatar
Guolin Ke committed
2680
  CHECK(col_idx < ncol_ptr && col_idx >= 0);
Guolin Ke's avatar
Guolin Ke committed
2681
2682
  if (data_type == C_API_DTYPE_FLOAT32) {
    if (col_ptr_type == C_API_DTYPE_INT32) {
2683
      return IterateFunctionFromCSC_helper<float, int32_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2684
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2685
      return IterateFunctionFromCSC_helper<float, int64_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2686
    }
Guolin Ke's avatar
Guolin Ke committed
2687
2688
  } else if (data_type == C_API_DTYPE_FLOAT64) {
    if (col_ptr_type == C_API_DTYPE_INT32) {
2689
      return IterateFunctionFromCSC_helper<double, int32_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2690
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2691
      return IterateFunctionFromCSC_helper<double, int64_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2692
2693
    }
  }
2694
  Log::Fatal("Unknown data type in CSC matrix");
2695
  return nullptr;
2696
2697
}

Guolin Ke's avatar
Guolin Ke committed
2698
CSC_RowIterator::CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
2699
                                 const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx) {
Guolin Ke's avatar
Guolin Ke committed
2700
2701
2702
2703
2704
2705
2706
2707
2708
  iter_fun_ = IterateFunctionFromCSC(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, col_idx);
}

double CSC_RowIterator::Get(int idx) {
  while (idx > cur_idx_ && !is_end_) {
    auto ret = iter_fun_(nonzero_idx_);
    if (ret.first < 0) {
      is_end_ = true;
      break;
2709
    }
Guolin Ke's avatar
Guolin Ke committed
2710
2711
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
    cur_idx_ = ret.first;
    cur_val_ = ret.second;
    ++nonzero_idx_;
  }
  if (idx == cur_idx_) {
    return cur_val_;
  } else {
    return 0.0f;
  }
}

std::pair<int, double> CSC_RowIterator::NextNonZero() {
  if (!is_end_) {
    auto ret = iter_fun_(nonzero_idx_);
    ++nonzero_idx_;
    if (ret.first < 0) {
      is_end_ = true;
2727
    }
Guolin Ke's avatar
Guolin Ke committed
2728
2729
2730
    return ret;
  } else {
    return std::make_pair(-1, 0.0);
2731
  }
Guolin Ke's avatar
Guolin Ke committed
2732
}