c_api.cpp 96.2 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
int LGBM_RegisterLogCallback(void (*callback)(const char*)) {
  API_BEGIN();
  Log::ResetCallBack(callback);
  API_END();
}

897
898
899
900
901
902
903
904
905
906
907
908
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
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
942
int LGBM_DatasetCreateFromFile(const char* filename,
943
944
945
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
946
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
947
948
  auto param = Config::Str2Map(parameters);
  Config config;
949
  config.Set(param);
950
  OMP_SET_NUM_THREADS(config.num_threads);
951
  DatasetLoader loader(config, nullptr, 1, filename);
Guolin Ke's avatar
Guolin Ke committed
952
  if (reference == nullptr) {
953
    if (Network::num_machines() == 1) {
954
      *out = loader.LoadFromFile(filename);
955
    } else {
956
      *out = loader.LoadFromFile(filename, Network::rank(), Network::num_machines());
957
    }
Guolin Ke's avatar
Guolin Ke committed
958
  } else {
959
    *out = loader.LoadFromFileAlignWithOtherDataset(filename,
960
                                                    reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
961
  }
962
  API_END();
Guolin Ke's avatar
Guolin Ke committed
963
964
}

965

Guolin Ke's avatar
Guolin Ke committed
966
int LGBM_DatasetCreateFromSampledColumn(double** sample_data,
967
968
969
970
971
972
973
                                        int** sample_indices,
                                        int32_t ncol,
                                        const int* num_per_col,
                                        int32_t num_sample_row,
                                        int32_t num_total_row,
                                        const char* parameters,
                                        DatasetHandle* out) {
974
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
975
976
  auto param = Config::Str2Map(parameters);
  Config config;
977
  config.Set(param);
978
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
979
  DatasetLoader loader(config, nullptr, 1, nullptr);
980
981
982
  *out = loader.ConstructFromSampleData(sample_data, sample_indices, ncol, num_per_col,
                                        num_sample_row,
                                        static_cast<data_size_t>(num_total_row));
983
  API_END();
Guolin Ke's avatar
Guolin Ke committed
984
985
}

986

Guolin Ke's avatar
Guolin Ke committed
987
int LGBM_DatasetCreateByReference(const DatasetHandle reference,
988
989
                                  int64_t num_total_row,
                                  DatasetHandle* out) {
Guolin Ke's avatar
Guolin Ke committed
990
991
992
993
994
995
996
997
  API_BEGIN();
  std::unique_ptr<Dataset> ret;
  ret.reset(new Dataset(static_cast<data_size_t>(num_total_row)));
  ret->CreateValid(reinterpret_cast<const Dataset*>(reference));
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
998
int LGBM_DatasetPushRows(DatasetHandle dataset,
999
1000
1001
1002
1003
                         const void* data,
                         int data_type,
                         int32_t nrow,
                         int32_t ncol,
                         int32_t start_row) {
Guolin Ke's avatar
Guolin Ke committed
1004
1005
1006
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto get_row_fun = RowFunctionFromDenseMatric(data, nrow, ncol, data_type, 1);
1007
1008
1009
  if (p_dataset->has_raw()) {
    p_dataset->ResizeRaw(p_dataset->num_numeric_features() + nrow);
  }
1010
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1011
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1012
  for (int i = 0; i < nrow; ++i) {
1013
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1014
1015
1016
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(tid, start_row + i, one_row);
1017
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1018
  }
1019
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1020
1021
1022
1023
1024
1025
  if (start_row + nrow == p_dataset->num_data()) {
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1026
int LGBM_DatasetPushRowsByCSR(DatasetHandle dataset,
1027
1028
1029
1030
1031
1032
1033
1034
1035
                              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
1036
1037
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
1038
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
1039
  int32_t nrow = static_cast<int32_t>(nindptr - 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
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
1049
    p_dataset->PushOneRow(tid, static_cast<data_size_t>(start_row + i), one_row);
1050
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1051
  }
1052
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1053
1054
1055
1056
1057
1058
  if (start_row + nrow == static_cast<int64_t>(p_dataset->num_data())) {
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1059
int LGBM_DatasetCreateFromMat(const void* data,
1060
1061
1062
1063
1064
1065
1066
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
  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) {
1087
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1088
1089
  auto param = Config::Str2Map(parameters);
  Config config;
1090
  config.Set(param);
1091
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1092
  std::unique_ptr<Dataset> ret;
1093
1094
1095
1096
1097
1098
1099
1100
1101
  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));
  }
1102

Guolin Ke's avatar
Guolin Ke committed
1103
1104
  if (reference == nullptr) {
    // sample data first
1105
1106
    auto sample_indices = CreateSampleIndices(total_nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
1107
    std::vector<std::vector<double>> sample_values(ncol);
Guolin Ke's avatar
Guolin Ke committed
1108
    std::vector<std::vector<int>> sample_idx(ncol);
1109
1110
1111

    int offset = 0;
    int j = 0;
Guolin Ke's avatar
Guolin Ke committed
1112
    for (size_t i = 0; i < sample_indices.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1113
      auto idx = sample_indices[i];
1114
1115
1116
1117
      while ((idx - offset) >= nrow[j]) {
        offset += nrow[j];
        ++j;
      }
1118

1119
1120
1121
1122
1123
      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
1124
        }
Guolin Ke's avatar
Guolin Ke committed
1125
1126
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1127
    DatasetLoader loader(config, nullptr, 1, nullptr);
1128
1129
1130
1131
1132
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             ncol,
                                             VectorSize<double>(sample_values).data(),
                                             sample_cnt, total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1133
  } else {
1134
    ret.reset(new Dataset(total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1135
    ret->CreateValid(
1136
      reinterpret_cast<const Dataset*>(reference));
1137
1138
1139
    if (ret->has_raw()) {
      ret->ResizeRaw(total_nrow);
    }
Guolin Ke's avatar
Guolin Ke committed
1140
  }
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
  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
1155
1156
  }
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1157
  *out = ret.release();
1158
  API_END();
1159
1160
}

Guolin Ke's avatar
Guolin Ke committed
1161
int LGBM_DatasetCreateFromCSR(const void* indptr,
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
                              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) {
1172
  API_BEGIN();
1173
1174
1175
1176
1177
  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
1178
1179
  auto param = Config::Str2Map(parameters);
  Config config;
1180
  config.Set(param);
1181
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1182
  std::unique_ptr<Dataset> ret;
1183
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1184
1185
1186
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
  if (reference == nullptr) {
    // sample data first
1187
1188
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1189
1190
    std::vector<std::vector<double>> sample_values(num_col);
    std::vector<std::vector<int>> sample_idx(num_col);
1191
1192
1193
1194
    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
1195
        CHECK_LT(inner_data.first, num_col);
Guolin Ke's avatar
Guolin Ke committed
1196
        if (std::fabs(inner_data.second) > kZeroThreshold || std::isnan(inner_data.second)) {
Guolin Ke's avatar
Guolin Ke committed
1197
1198
          sample_values[inner_data.first].emplace_back(inner_data.second);
          sample_idx[inner_data.first].emplace_back(static_cast<int>(i));
1199
1200
1201
        }
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1202
    DatasetLoader loader(config, nullptr, 1, nullptr);
1203
1204
1205
1206
1207
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             static_cast<int>(num_col),
                                             VectorSize<double>(sample_values).data(),
                                             sample_cnt, nrow));
1208
  } else {
1209
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1210
    ret->CreateValid(
1211
      reinterpret_cast<const Dataset*>(reference));
1212
1213
1214
    if (ret->has_raw()) {
      ret->ResizeRaw(nrow);
    }
1215
  }
1216
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1217
  #pragma omp parallel for schedule(static)
1218
  for (int i = 0; i < nindptr - 1; ++i) {
1219
    OMP_LOOP_EX_BEGIN();
1220
1221
1222
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    ret->PushOneRow(tid, i, one_row);
1223
    OMP_LOOP_EX_END();
1224
  }
1225
  OMP_THROW_EX();
1226
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1227
  *out = ret.release();
1228
  API_END();
1229
1230
}

1231
int LGBM_DatasetCreateFromCSRFunc(void* get_row_funptr,
1232
1233
1234
1235
1236
                                  int num_rows,
                                  int64_t num_col,
                                  const char* parameters,
                                  const DatasetHandle reference,
                                  DatasetHandle* out) {
1237
  API_BEGIN();
1238
1239
1240
1241
1242
  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.");
  }
1243
1244
1245
1246
  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);
1247
  OMP_SET_NUM_THREADS(config.num_threads);
1248
1249
1250
1251
  std::unique_ptr<Dataset> ret;
  int32_t nrow = num_rows;
  if (reference == nullptr) {
    // sample data first
1252
1253
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
1254
1255
1256
1257
1258
1259
1260
1261
    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
1262
        CHECK_LT(inner_data.first, num_col);
1263
1264
1265
1266
1267
1268
1269
        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);
1270
1271
1272
1273
1274
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             static_cast<int>(num_col),
                                             VectorSize<double>(sample_values).data(),
                                             sample_cnt, nrow));
1275
1276
1277
1278
  } else {
    ret.reset(new Dataset(nrow));
    ret->CreateValid(
      reinterpret_cast<const Dataset*>(reference));
1279
1280
1281
    if (ret->has_raw()) {
      ret->ResizeRaw(nrow);
    }
1282
  }
1283

1284
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1285
1286
  std::vector<std::pair<int, double>> thread_buffer;
  #pragma omp parallel for schedule(static) private(thread_buffer)
1287
1288
1289
  for (int i = 0; i < num_rows; ++i) {
    OMP_LOOP_EX_BEGIN();
    {
1290
      const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1291
1292
      get_row_fun(i, thread_buffer);
      ret->PushOneRow(tid, i, thread_buffer);
1293
1294
1295
1296
1297
1298
1299
1300
1301
    }
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();
  ret->FinishLoad();
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1302
int LGBM_DatasetCreateFromCSC(const void* col_ptr,
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
                              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) {
1313
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1314
1315
  auto param = Config::Str2Map(parameters);
  Config config;
1316
  config.Set(param);
1317
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1318
  std::unique_ptr<Dataset> ret;
Guolin Ke's avatar
Guolin Ke committed
1319
1320
1321
  int32_t nrow = static_cast<int32_t>(num_row);
  if (reference == nullptr) {
    // sample data first
1322
1323
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1324
    std::vector<std::vector<double>> sample_values(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
1325
    std::vector<std::vector<int>> sample_idx(ncol_ptr - 1);
1326
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1327
    #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1328
    for (int i = 0; i < static_cast<int>(sample_values.size()); ++i) {
1329
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1330
1331
1332
      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
1333
        if (std::fabs(val) > kZeroThreshold || std::isnan(val)) {
Guolin Ke's avatar
Guolin Ke committed
1334
1335
          sample_values[i].emplace_back(val);
          sample_idx[i].emplace_back(j);
Guolin Ke's avatar
Guolin Ke committed
1336
1337
        }
      }
1338
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1339
    }
1340
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1341
    DatasetLoader loader(config, nullptr, 1, nullptr);
1342
1343
1344
1345
1346
    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(),
                                             sample_cnt, nrow));
Guolin Ke's avatar
Guolin Ke committed
1347
  } else {
1348
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1349
    ret->CreateValid(
1350
      reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
1351
  }
1352
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1353
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1354
  for (int i = 0; i < ncol_ptr - 1; ++i) {
1355
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1356
    const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1357
    int feature_idx = ret->InnerFeatureIndex(i);
Guolin Ke's avatar
Guolin Ke committed
1358
    if (feature_idx < 0) { continue; }
Guolin Ke's avatar
Guolin Ke committed
1359
1360
    int group = ret->Feature2Group(feature_idx);
    int sub_feature = ret->Feture2SubFeature(feature_idx);
Guolin Ke's avatar
Guolin Ke committed
1361
    CSC_RowIterator col_it(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, i);
Guolin Ke's avatar
Guolin Ke committed
1362
1363
1364
1365
1366
1367
1368
1369
    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; }
1370
        ret->PushOneData(tid, row_idx, group, feature_idx, sub_feature, pair.second);
Guolin Ke's avatar
Guolin Ke committed
1371
1372
1373
1374
      }
    } else {
      for (int row_idx = 0; row_idx < nrow; ++row_idx) {
        auto val = col_it.Get(row_idx);
1375
        ret->PushOneData(tid, row_idx, group, feature_idx, sub_feature, val);
Guolin Ke's avatar
Guolin Ke committed
1376
      }
Guolin Ke's avatar
Guolin Ke committed
1377
    }
1378
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1379
  }
1380
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1381
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1382
  *out = ret.release();
1383
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1384
1385
}

Guolin Ke's avatar
Guolin Ke committed
1386
int LGBM_DatasetGetSubset(
1387
  const DatasetHandle handle,
wxchan's avatar
wxchan committed
1388
1389
1390
  const int32_t* used_row_indices,
  int32_t num_used_row_indices,
  const char* parameters,
Guolin Ke's avatar
typo  
Guolin Ke committed
1391
  DatasetHandle* out) {
wxchan's avatar
wxchan committed
1392
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1393
1394
  auto param = Config::Str2Map(parameters);
  Config config;
1395
  config.Set(param);
1396
  OMP_SET_NUM_THREADS(config.num_threads);
1397
  auto full_dataset = reinterpret_cast<const Dataset*>(handle);
1398
  CHECK_GT(num_used_row_indices, 0);
1399
1400
  const int32_t lower = 0;
  const int32_t upper = full_dataset->num_data() - 1;
1401
  CheckElementsIntervalClosed(used_row_indices, lower, upper, num_used_row_indices, "Used indices of subset");
1402
1403
1404
  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
1405
  auto ret = std::unique_ptr<Dataset>(new Dataset(num_used_row_indices));
1406
  ret->CopyFeatureMapperFrom(full_dataset);
1407
  ret->CopySubrow(full_dataset, used_row_indices, num_used_row_indices, true);
wxchan's avatar
wxchan committed
1408
1409
1410
1411
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1412
int LGBM_DatasetSetFeatureNames(
Guolin Ke's avatar
typo  
Guolin Ke committed
1413
  DatasetHandle handle,
Guolin Ke's avatar
Guolin Ke committed
1414
  const char** feature_names,
Guolin Ke's avatar
Guolin Ke committed
1415
  int num_feature_names) {
Guolin Ke's avatar
Guolin Ke committed
1416
1417
1418
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  std::vector<std::string> feature_names_str;
Guolin Ke's avatar
Guolin Ke committed
1419
  for (int i = 0; i < num_feature_names; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1420
1421
1422
1423
1424
1425
    feature_names_str.emplace_back(feature_names[i]);
  }
  dataset->set_feature_names(feature_names_str);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1426
int LGBM_DatasetGetFeatureNames(
1427
1428
1429
1430
1431
1432
    DatasetHandle handle,
    const int len,
    int* num_feature_names,
    const size_t buffer_len,
    size_t* out_buffer_len,
    char** feature_names) {
1433
  API_BEGIN();
1434
  *out_buffer_len = 0;
1435
1436
  auto dataset = reinterpret_cast<Dataset*>(handle);
  auto inside_feature_name = dataset->feature_names();
Guolin Ke's avatar
Guolin Ke committed
1437
1438
  *num_feature_names = static_cast<int>(inside_feature_name.size());
  for (int i = 0; i < *num_feature_names; ++i) {
1439
1440
1441
1442
1443
    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);
1444
1445
1446
1447
  }
  API_END();
}

1448
1449
1450
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1451
int LGBM_DatasetFree(DatasetHandle handle) {
1452
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1453
  delete reinterpret_cast<Dataset*>(handle);
1454
  API_END();
1455
1456
}

Guolin Ke's avatar
Guolin Ke committed
1457
int LGBM_DatasetSaveBinary(DatasetHandle handle,
1458
                           const char* filename) {
1459
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1460
1461
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->SaveBinaryFile(filename);
1462
  API_END();
1463
1464
}

1465
1466
1467
1468
1469
1470
1471
1472
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
1473
int LGBM_DatasetSetField(DatasetHandle handle,
1474
1475
1476
1477
                         const char* field_name,
                         const void* field_data,
                         int num_element,
                         int type) {
1478
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1479
  auto dataset = reinterpret_cast<Dataset*>(handle);
1480
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1481
  if (type == C_API_DTYPE_FLOAT32) {
Guolin Ke's avatar
Guolin Ke committed
1482
    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
1483
  } else if (type == C_API_DTYPE_INT32) {
Guolin Ke's avatar
Guolin Ke committed
1484
    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
1485
1486
  } 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));
1487
  }
1488
  if (!is_success) { Log::Fatal("Input data type error or field not found"); }
1489
  API_END();
1490
1491
}

Guolin Ke's avatar
Guolin Ke committed
1492
int LGBM_DatasetGetField(DatasetHandle handle,
1493
1494
1495
1496
                         const char* field_name,
                         int* out_len,
                         const void** out_ptr,
                         int* out_type) {
1497
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1498
  auto dataset = reinterpret_cast<Dataset*>(handle);
1499
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1500
  if (dataset->GetFloatField(field_name, out_len, reinterpret_cast<const float**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1501
    *out_type = C_API_DTYPE_FLOAT32;
1502
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1503
  } else if (dataset->GetIntField(field_name, out_len, reinterpret_cast<const int**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1504
    *out_type = C_API_DTYPE_INT32;
1505
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1506
1507
1508
  } 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
1509
  }
1510
  if (!is_success) { Log::Fatal("Field not found"); }
wxchan's avatar
wxchan committed
1511
  if (*out_ptr == nullptr) { *out_len = 0; }
1512
  API_END();
1513
1514
}

1515
int LGBM_DatasetUpdateParamChecking(const char* old_parameters, const char* new_parameters) {
1516
  API_BEGIN();
1517
1518
1519
1520
1521
  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);
1522
1523
1524
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1525
int LGBM_DatasetGetNumData(DatasetHandle handle,
1526
                           int* out) {
1527
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1528
1529
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_data();
1530
  API_END();
1531
1532
}

Guolin Ke's avatar
Guolin Ke committed
1533
int LGBM_DatasetGetNumFeature(DatasetHandle handle,
1534
                              int* out) {
1535
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1536
1537
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_total_features();
1538
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1539
}
1540

1541
1542
1543
1544
1545
int LGBM_DatasetAddFeaturesFrom(DatasetHandle target,
                                DatasetHandle source) {
  API_BEGIN();
  auto target_d = reinterpret_cast<Dataset*>(target);
  auto source_d = reinterpret_cast<Dataset*>(source);
1546
  target_d->AddFeaturesFrom(source_d);
1547
1548
1549
  API_END();
}

1550
1551
// ---- start of booster

Guolin Ke's avatar
Guolin Ke committed
1552
int LGBM_BoosterCreate(const DatasetHandle train_data,
1553
1554
                       const char* parameters,
                       BoosterHandle* out) {
1555
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1556
  const Dataset* p_train_data = reinterpret_cast<const Dataset*>(train_data);
wxchan's avatar
wxchan committed
1557
1558
  auto ret = std::unique_ptr<Booster>(new Booster(p_train_data, parameters));
  *out = ret.release();
1559
  API_END();
1560
1561
}

Guolin Ke's avatar
Guolin Ke committed
1562
int LGBM_BoosterCreateFromModelfile(
1563
  const char* filename,
Guolin Ke's avatar
Guolin Ke committed
1564
  int* out_num_iterations,
1565
  BoosterHandle* out) {
1566
  API_BEGIN();
wxchan's avatar
wxchan committed
1567
  auto ret = std::unique_ptr<Booster>(new Booster(filename));
Guolin Ke's avatar
Guolin Ke committed
1568
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
wxchan's avatar
wxchan committed
1569
  *out = ret.release();
1570
  API_END();
1571
1572
}

Guolin Ke's avatar
Guolin Ke committed
1573
int LGBM_BoosterLoadModelFromString(
1574
1575
1576
1577
  const char* model_str,
  int* out_num_iterations,
  BoosterHandle* out) {
  API_BEGIN();
wxchan's avatar
wxchan committed
1578
  auto ret = std::unique_ptr<Booster>(new Booster(nullptr));
1579
1580
1581
1582
1583
1584
  ret->LoadModelFromString(model_str);
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
  *out = ret.release();
  API_END();
}

1585
1586
1587
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1588
int LGBM_BoosterFree(BoosterHandle handle) {
1589
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1590
  delete reinterpret_cast<Booster*>(handle);
1591
  API_END();
1592
1593
}

1594
int LGBM_BoosterShuffleModels(BoosterHandle handle, int start_iter, int end_iter) {
1595
1596
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1597
  ref_booster->ShuffleModels(start_iter, end_iter);
1598
1599
1600
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1601
int LGBM_BoosterMerge(BoosterHandle handle,
1602
                      BoosterHandle other_handle) {
wxchan's avatar
wxchan committed
1603
1604
1605
1606
1607
1608
1609
  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
1610
int LGBM_BoosterAddValidData(BoosterHandle handle,
1611
                             const DatasetHandle valid_data) {
wxchan's avatar
wxchan committed
1612
1613
1614
1615
1616
1617
1618
  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
1619
int LGBM_BoosterResetTrainingData(BoosterHandle handle,
1620
                                  const DatasetHandle train_data) {
wxchan's avatar
wxchan committed
1621
1622
1623
1624
1625
1626
1627
  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
1628
int LGBM_BoosterResetParameter(BoosterHandle handle, const char* parameters) {
wxchan's avatar
wxchan committed
1629
1630
1631
1632
1633
1634
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->ResetConfig(parameters);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1635
int LGBM_BoosterGetNumClasses(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1636
1637
1638
1639
1640
1641
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->NumberOfClasses();
  API_END();
}

1642
1643
1644
1645
1646
1647
1648
int LGBM_BoosterGetLinear(BoosterHandle handle, bool* out) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out = ref_booster->GetBoosting()->IsLinear();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1649
1650
1651
1652
1653
1654
1655
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
1656
int LGBM_BoosterUpdateOneIter(BoosterHandle handle, int* is_finished) {
1657
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1658
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1659
1660
1661
1662
1663
  if (ref_booster->TrainOneIter()) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1664
  API_END();
1665
1666
}

Guolin Ke's avatar
Guolin Ke committed
1667
int LGBM_BoosterUpdateOneIterCustom(BoosterHandle handle,
1668
1669
1670
                                    const float* grad,
                                    const float* hess,
                                    int* is_finished) {
1671
  API_BEGIN();
1672
  #ifdef SCORE_T_USE_DOUBLE
1673
1674
1675
1676
  (void) handle;       // UNUSED VARIABLE
  (void) grad;         // UNUSED VARIABLE
  (void) hess;         // UNUSED VARIABLE
  (void) is_finished;  // UNUSED VARIABLE
1677
  Log::Fatal("Don't support custom loss function when SCORE_T_USE_DOUBLE is enabled");
1678
  #else
1679
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1680
1681
1682
1683
1684
  if (ref_booster->TrainOneIter(grad, hess)) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1685
  #endif
1686
  API_END();
1687
1688
}

Guolin Ke's avatar
Guolin Ke committed
1689
int LGBM_BoosterRollbackOneIter(BoosterHandle handle) {
wxchan's avatar
wxchan committed
1690
1691
1692
1693
1694
1695
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->RollbackOneIter();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1696
int LGBM_BoosterGetCurrentIteration(BoosterHandle handle, int* out_iteration) {
wxchan's avatar
wxchan committed
1697
1698
1699
1700
1701
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_iteration = ref_booster->GetBoosting()->GetCurrentIteration();
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
1702

1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
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
1717
int LGBM_BoosterGetEvalCounts(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1718
1719
1720
1721
1722
1723
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetEvalCounts();
  API_END();
}

1724
1725
1726
1727
1728
1729
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
1730
1731
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1732
  *out_len = ref_booster->GetEvalNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1733
1734
1735
  API_END();
}

1736
1737
1738
1739
1740
1741
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
1742
1743
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1744
  *out_len = ref_booster->GetFeatureNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1745
1746
1747
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1748
int LGBM_BoosterGetNumFeature(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1749
1750
1751
1752
1753
1754
  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
1755
int LGBM_BoosterGetEval(BoosterHandle handle,
1756
1757
1758
                        int data_idx,
                        int* out_len,
                        double* out_results) {
1759
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1760
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1761
  auto boosting = ref_booster->GetBoosting();
wxchan's avatar
wxchan committed
1762
  auto result_buf = boosting->GetEvalAt(data_idx);
Guolin Ke's avatar
Guolin Ke committed
1763
  *out_len = static_cast<int>(result_buf.size());
1764
  for (size_t i = 0; i < result_buf.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1765
    (out_results)[i] = static_cast<double>(result_buf[i]);
1766
  }
1767
  API_END();
1768
1769
}

Guolin Ke's avatar
Guolin Ke committed
1770
int LGBM_BoosterGetNumPredict(BoosterHandle handle,
1771
1772
                              int data_idx,
                              int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1773
1774
1775
1776
1777
1778
  API_BEGIN();
  auto boosting = reinterpret_cast<Booster*>(handle)->GetBoosting();
  *out_len = boosting->GetNumPredictAt(data_idx);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1779
int LGBM_BoosterGetPredict(BoosterHandle handle,
1780
1781
1782
                           int data_idx,
                           int64_t* out_len,
                           double* out_result) {
1783
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1784
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1785
  ref_booster->GetPredictAt(data_idx, out_result, out_len);
1786
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1787
1788
}

Guolin Ke's avatar
Guolin Ke committed
1789
int LGBM_BoosterPredictForFile(BoosterHandle handle,
1790
1791
1792
                               const char* data_filename,
                               int data_has_header,
                               int predict_type,
1793
                               int start_iteration,
1794
                               int num_iteration,
1795
                               const char* parameter,
1796
                               const char* result_filename) {
1797
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1798
1799
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1800
  config.Set(param);
1801
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1802
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1803
  ref_booster->Predict(start_iteration, num_iteration, predict_type, data_filename, data_has_header,
Guolin Ke's avatar
Guolin Ke committed
1804
                       config, result_filename);
1805
  API_END();
1806
1807
}

Guolin Ke's avatar
Guolin Ke committed
1808
int LGBM_BoosterCalcNumPredict(BoosterHandle handle,
1809
1810
                               int num_row,
                               int predict_type,
1811
                               int start_iteration,
1812
1813
                               int num_iteration,
                               int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1814
1815
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1816
  *out_len = static_cast<int64_t>(num_row) * ref_booster->GetBoosting()->NumPredictOneRow(start_iteration,
1817
    num_iteration, predict_type == C_API_PREDICT_LEAF_INDEX, predict_type == C_API_PREDICT_CONTRIB);
Guolin Ke's avatar
Guolin Ke committed
1818
1819
1820
  API_END();
}

1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
/*!
 * \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,
1834
             const int predict_type_,
1835
             const int data_type_,
1836
             const int32_t num_cols) : booster(booster_ptr), predict_type(predict_type_), data_type(data_type_), ncol(num_cols) {
1837
1838
1839
1840
1841
    config.Set(Config::Str2Map(parameter));
  }

  Booster* const booster;
  Config config;
1842
  const int predict_type;
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
  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
1853
int LGBM_BoosterPredictForCSR(BoosterHandle handle,
1854
1855
1856
1857
1858
1859
1860
                              const void* indptr,
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
1861
                              int64_t num_col,
1862
                              int predict_type,
1863
                              int start_iteration,
1864
                              int num_iteration,
1865
                              const char* parameter,
1866
1867
                              int64_t* out_len,
                              double* out_result) {
1868
  API_BEGIN();
1869
1870
1871
1872
1873
  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
1874
1875
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1876
  config.Set(param);
1877
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1878
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1879
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
1880
  int nrow = static_cast<int>(nindptr - 1);
1881
  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
1882
                       config, out_result, out_len);
1883
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1884
}
1885

1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
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,
1896
                                    int start_iteration,
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
                                    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);
1909
  OMP_SET_NUM_THREADS(config.num_threads);
1910
1911
1912
1913
1914
1915
1916
1917
  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;
1918
    ref_booster->PredictSparseCSR(start_iteration, num_iteration, predict_type, nrow, static_cast<int>(num_col_or_row), get_row_fun,
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
                                  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;
    };
1942
    ref_booster->PredictSparseCSC(start_iteration, num_iteration, predict_type, num_col_or_row, ncol, get_row_fun, config,
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
                                  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) {
    delete reinterpret_cast<int32_t*>(indptr);
  } else if (indptr_type == C_API_DTYPE_INT64) {
    delete reinterpret_cast<int64_t*>(indptr);
  } else {
    Log::Fatal("Unknown indptr type in LGBM_BoosterFreePredictSparse");
  }
  delete indices;
  if (data_type == C_API_DTYPE_FLOAT32) {
    delete reinterpret_cast<float*>(data);
  } else if (data_type == C_API_DTYPE_FLOAT64) {
    delete reinterpret_cast<double*>(data);
  } else {
    Log::Fatal("Unknown data type in LGBM_BoosterFreePredictSparse");
  }
  API_END();
}

1970
int LGBM_BoosterPredictForCSRSingleRow(BoosterHandle handle,
1971
1972
1973
1974
1975
1976
1977
                                       const void* indptr,
                                       int indptr_type,
                                       const int32_t* indices,
                                       const void* data,
                                       int data_type,
                                       int64_t nindptr,
                                       int64_t nelem,
1978
                                       int64_t num_col,
1979
                                       int predict_type,
1980
                                       int start_iteration,
1981
1982
1983
1984
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
1985
  API_BEGIN();
1986
1987
1988
1989
1990
  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.");
  }
1991
1992
1993
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
1994
  OMP_SET_NUM_THREADS(config.num_threads);
1995
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1996
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1997
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
1998
  ref_booster->PredictSingleRow(predict_type, static_cast<int32_t>(num_col), get_row_fun, config, out_result, out_len);
1999
2000
2001
  API_END();
}

2002
int LGBM_BoosterPredictForCSRSingleRowFastInit(BoosterHandle handle,
2003
                                               const int predict_type,
2004
                                               const int start_iteration,
2005
                                               const int num_iteration,
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
                                               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,
2020
    predict_type,
2021
2022
2023
    data_type,
    static_cast<int32_t>(num_col)));

2024
  OMP_SET_NUM_THREADS(fastConfig_ptr->config.num_threads);
2025

2026
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2027

2028
2029
2030
2031
2032
2033
  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

int LGBM_BoosterPredictForCSRSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* indptr,
2034
                                           const int indptr_type,
2035
2036
                                           const int32_t* indices,
                                           const void* data,
2037
2038
                                           const int64_t nindptr,
                                           const int64_t nelem,
2039
2040
2041
2042
2043
                                           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);
2044
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
2045
2046
2047
2048
                                        get_row_fun, fastConfig->config, out_result, out_len);
  API_END();
}

2049

Guolin Ke's avatar
Guolin Ke committed
2050
int LGBM_BoosterPredictForCSC(BoosterHandle handle,
2051
2052
2053
2054
2055
2056
2057
2058
2059
                              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,
2060
                              int start_iteration,
2061
                              int num_iteration,
2062
                              const char* parameter,
2063
2064
                              int64_t* out_len,
                              double* out_result) {
Guolin Ke's avatar
Guolin Ke committed
2065
2066
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2067
2068
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2069
  config.Set(param);
2070
  OMP_SET_NUM_THREADS(config.num_threads);
2071
  int num_threads = OMP_NUM_THREADS();
Guolin Ke's avatar
Guolin Ke committed
2072
  int ncol = static_cast<int>(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
2073
2074
2075
2076
2077
  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
2078
2079
  }
  std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun =
Guolin Ke's avatar
Guolin Ke committed
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
      [&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;
      };
2092
  ref_booster->Predict(start_iteration, num_iteration, predict_type, static_cast<int>(num_row), ncol, get_row_fun, config,
cbecker's avatar
cbecker committed
2093
                       out_result, out_len);
Guolin Ke's avatar
Guolin Ke committed
2094
2095
2096
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2097
int LGBM_BoosterPredictForMat(BoosterHandle handle,
2098
2099
2100
2101
2102
2103
                              const void* data,
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              int predict_type,
2104
                              int start_iteration,
2105
                              int num_iteration,
2106
                              const char* parameter,
2107
2108
                              int64_t* out_len,
                              double* out_result) {
2109
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2110
2111
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2112
  config.Set(param);
2113
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
2114
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2115
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, nrow, ncol, data_type, is_row_major);
2116
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
2117
                       config, out_result, out_len);
2118
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2119
}
2120

2121
int LGBM_BoosterPredictForMatSingleRow(BoosterHandle handle,
2122
2123
2124
2125
2126
                                       const void* data,
                                       int data_type,
                                       int32_t ncol,
                                       int is_row_major,
                                       int predict_type,
2127
                                       int start_iteration,
2128
2129
2130
2131
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2132
2133
2134
2135
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
2136
  OMP_SET_NUM_THREADS(config.num_threads);
2137
2138
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, ncol, data_type, is_row_major);
2139
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
2140
  ref_booster->PredictSingleRow(predict_type, ncol, get_row_fun, config, out_result, out_len);
2141
2142
2143
  API_END();
}

2144
int LGBM_BoosterPredictForMatSingleRowFastInit(BoosterHandle handle,
2145
                                               const int predict_type,
2146
                                               const int start_iteration,
2147
                                               const int num_iteration,
2148
2149
2150
2151
2152
2153
2154
2155
                                               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,
2156
    predict_type,
2157
2158
2159
    data_type,
    ncol));

2160
  OMP_SET_NUM_THREADS(fastConfig_ptr->config.num_threads);
2161

2162
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2163

2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
  *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);
2176
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
2177
2178
2179
2180
2181
                                        get_row_fun, fastConfig->config,
                                        out_result, out_len);
  API_END();
}

2182

2183
2184
2185
2186
2187
2188
int LGBM_BoosterPredictForMats(BoosterHandle handle,
                               const void** data,
                               int data_type,
                               int32_t nrow,
                               int32_t ncol,
                               int predict_type,
2189
                               int start_iteration,
2190
2191
2192
2193
2194
2195
2196
2197
                               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);
2198
  OMP_SET_NUM_THREADS(config.num_threads);
2199
2200
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseRows(data, ncol, data_type);
2201
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, out_result, out_len);
2202
2203
2204
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2205
int LGBM_BoosterSaveModel(BoosterHandle handle,
2206
                          int start_iteration,
2207
                          int num_iteration,
2208
                          int feature_importance_type,
2209
                          const char* filename) {
2210
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2211
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2212
2213
  ref_booster->SaveModelToFile(start_iteration, num_iteration,
                               feature_importance_type, filename);
wxchan's avatar
wxchan committed
2214
2215
2216
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2217
int LGBM_BoosterSaveModelToString(BoosterHandle handle,
2218
                                  int start_iteration,
2219
                                  int num_iteration,
2220
                                  int feature_importance_type,
2221
                                  int64_t buffer_len,
2222
                                  int64_t* out_len,
2223
                                  char* out_str) {
2224
2225
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2226
2227
  std::string model = ref_booster->SaveModelToString(
      start_iteration, num_iteration, feature_importance_type);
2228
  *out_len = static_cast<int64_t>(model.size()) + 1;
2229
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2230
    std::memcpy(out_str, model.c_str(), *out_len);
2231
2232
2233
2234
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2235
int LGBM_BoosterDumpModel(BoosterHandle handle,
2236
                          int start_iteration,
2237
                          int num_iteration,
2238
                          int feature_importance_type,
2239
2240
                          int64_t buffer_len,
                          int64_t* out_len,
2241
                          char* out_str) {
wxchan's avatar
wxchan committed
2242
2243
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2244
2245
  std::string model = ref_booster->DumpModel(start_iteration, num_iteration,
                                             feature_importance_type);
2246
  *out_len = static_cast<int64_t>(model.size()) + 1;
wxchan's avatar
wxchan committed
2247
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2248
    std::memcpy(out_str, model.c_str(), *out_len);
wxchan's avatar
wxchan committed
2249
  }
2250
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2251
}
2252

Guolin Ke's avatar
Guolin Ke committed
2253
int LGBM_BoosterGetLeafValue(BoosterHandle handle,
2254
2255
2256
                             int tree_idx,
                             int leaf_idx,
                             double* out_val) {
Guolin Ke's avatar
Guolin Ke committed
2257
2258
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2259
  *out_val = static_cast<double>(ref_booster->GetLeafValue(tree_idx, leaf_idx));
Guolin Ke's avatar
Guolin Ke committed
2260
2261
2262
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2263
int LGBM_BoosterSetLeafValue(BoosterHandle handle,
2264
2265
2266
                             int tree_idx,
                             int leaf_idx,
                             double val) {
Guolin Ke's avatar
Guolin Ke committed
2267
2268
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2269
  ref_booster->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
2270
2271
2272
  API_END();
}

2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
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();
}

2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
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();
}

2304
2305
2306
2307
2308
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
2309
  Config config;
2310
  config.machines = RemoveQuotationSymbol(std::string(machines));
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
  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();
}

2326
2327
2328
int LGBM_NetworkInitWithFunctions(int num_machines, int rank,
                                  void* reduce_scatter_ext_fun,
                                  void* allgather_ext_fun) {
ww's avatar
ww committed
2329
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2330
  if (num_machines > 1) {
2331
    Network::Init(num_machines, rank, (ReduceScatterFunction)reduce_scatter_ext_fun, (AllgatherFunction)allgather_ext_fun);
ww's avatar
ww committed
2332
2333
2334
  }
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
2335

Guolin Ke's avatar
Guolin Ke committed
2336
// ---- start of some help functions
2337

2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362

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;
    };
  }
}

2363
2364
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
2365
  if (data_type == C_API_DTYPE_FLOAT32) {
2366
    return RowFunctionFromDenseMatric_helper<float>(data, num_row, num_col, is_row_major);
Guolin Ke's avatar
Guolin Ke committed
2367
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2368
    return RowFunctionFromDenseMatric_helper<double>(data, num_row, num_col, is_row_major);
2369
  }
2370
  Log::Fatal("Unknown data type in RowFunctionFromDenseMatric");
2371
  return nullptr;
2372
2373
2374
2375
}

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
2376
2377
  auto inner_function = RowFunctionFromDenseMatric(data, num_row, num_col, data_type, is_row_major);
  if (inner_function != nullptr) {
2378
    return [inner_function] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2379
2380
      auto raw_values = inner_function(row_idx);
      std::vector<std::pair<int, double>> ret;
Guolin Ke's avatar
Guolin Ke committed
2381
      ret.reserve(raw_values.size());
Guolin Ke's avatar
Guolin Ke committed
2382
      for (int i = 0; i < static_cast<int>(raw_values.size()); ++i) {
Guolin Ke's avatar
Guolin Ke committed
2383
        if (std::fabs(raw_values[i]) > kZeroThreshold || std::isnan(raw_values[i])) {
Guolin Ke's avatar
Guolin Ke committed
2384
          ret.emplace_back(i, raw_values[i]);
2385
        }
Guolin Ke's avatar
Guolin Ke committed
2386
2387
2388
      }
      return ret;
    };
2389
  }
Guolin Ke's avatar
Guolin Ke committed
2390
  return nullptr;
2391
2392
}

2393
2394
2395
2396
2397
2398
2399
// 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
2400
    ret.reserve(raw_values.size());
2401
2402
2403
2404
2405
2406
2407
2408
2409
    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;
  };
}

2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
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;
  };
}

2429
2430
template<typename T>
std::function<std::vector<std::pair<int, double>>(T idx)>
2431
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
2432
2433
  if (data_type == C_API_DTYPE_FLOAT32) {
    if (indptr_type == C_API_DTYPE_INT32) {
2434
     return RowFunctionFromCSR_helper<T, float, int32_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2435
    } else if (indptr_type == C_API_DTYPE_INT64) {
2436
     return RowFunctionFromCSR_helper<T, float, int64_t>(indptr, indices, data);
2437
    }
Guolin Ke's avatar
Guolin Ke committed
2438
2439
  } else if (data_type == C_API_DTYPE_FLOAT64) {
    if (indptr_type == C_API_DTYPE_INT32) {
2440
     return RowFunctionFromCSR_helper<T, double, int32_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2441
    } else if (indptr_type == C_API_DTYPE_INT64) {
2442
     return RowFunctionFromCSR_helper<T, double, int64_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2443
2444
    }
  }
2445
  Log::Fatal("Unknown data type in RowFunctionFromCSR");
2446
  return nullptr;
2447
2448
}

2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467


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
2468
std::function<std::pair<int, double>(int idx)>
2469
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
2470
  CHECK(col_idx < ncol_ptr && col_idx >= 0);
Guolin Ke's avatar
Guolin Ke committed
2471
2472
  if (data_type == C_API_DTYPE_FLOAT32) {
    if (col_ptr_type == C_API_DTYPE_INT32) {
2473
      return IterateFunctionFromCSC_helper<float, int32_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2474
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2475
      return IterateFunctionFromCSC_helper<float, int64_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2476
    }
Guolin Ke's avatar
Guolin Ke committed
2477
2478
  } else if (data_type == C_API_DTYPE_FLOAT64) {
    if (col_ptr_type == C_API_DTYPE_INT32) {
2479
      return IterateFunctionFromCSC_helper<double, int32_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2480
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2481
      return IterateFunctionFromCSC_helper<double, int64_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2482
2483
    }
  }
2484
  Log::Fatal("Unknown data type in CSC matrix");
2485
  return nullptr;
2486
2487
}

Guolin Ke's avatar
Guolin Ke committed
2488
CSC_RowIterator::CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
2489
                                 const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx) {
Guolin Ke's avatar
Guolin Ke committed
2490
2491
2492
2493
2494
2495
2496
2497
2498
  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;
2499
    }
Guolin Ke's avatar
Guolin Ke committed
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
    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;
2517
    }
Guolin Ke's avatar
Guolin Ke committed
2518
2519
2520
    return ret;
  } else {
    return std::make_pair(-1, 0.0);
2521
  }
Guolin Ke's avatar
Guolin Ke committed
2522
}