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
int LGBM_BoosterGetLinear(BoosterHandle handle, int* out) {
1643
1644
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1645
1646
1647
1648
1649
  if (ref_booster->GetBoosting()->IsLinear()) {
    *out = 1;
  } else {
    *out = 0;
  }
1650
1651
1652
  API_END();
}

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

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

Guolin Ke's avatar
Guolin Ke committed
1693
int LGBM_BoosterRollbackOneIter(BoosterHandle handle) {
wxchan's avatar
wxchan committed
1694
1695
1696
1697
1698
1699
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->RollbackOneIter();
  API_END();
}

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

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

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

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

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

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

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

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

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

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

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

1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
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,
1900
                                    int start_iteration,
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
                                    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);
1913
  OMP_SET_NUM_THREADS(config.num_threads);
1914
1915
1916
1917
1918
1919
1920
1921
  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;
1922
    ref_booster->PredictSparseCSR(start_iteration, num_iteration, predict_type, nrow, static_cast<int>(num_col_or_row), get_row_fun,
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
                                  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;
    };
1946
    ref_booster->PredictSparseCSC(start_iteration, num_iteration, predict_type, num_col_or_row, ncol, get_row_fun, config,
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
                                  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();
}

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

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

2028
  OMP_SET_NUM_THREADS(fastConfig_ptr->config.num_threads);
2029

2030
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2031

2032
2033
2034
2035
2036
2037
  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

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

2053

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

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

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

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

2164
  OMP_SET_NUM_THREADS(fastConfig_ptr->config.num_threads);
2165

2166
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2167

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

2186

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

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

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

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

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

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

2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
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();
}

2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
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();
}

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

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

Guolin Ke's avatar
Guolin Ke committed
2340
// ---- start of some help functions
2341

2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366

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

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

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

2397
2398
2399
2400
2401
2402
2403
// 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
2404
    ret.reserve(raw_values.size());
2405
2406
2407
2408
2409
2410
2411
2412
2413
    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;
  };
}

2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
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;
  };
}

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

2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471


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

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