c_api.cpp 95.8 KB
Newer Older
1
2
3
4
/*!
 * Copyright (c) 2016 Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See LICENSE file in the project root for license information.
 */
Guolin Ke's avatar
Guolin Ke committed
5
#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
117
118
    if (config_.num_threads > 0) {
      omp_set_num_threads(config_.num_threads);
    }
Guolin Ke's avatar
Guolin Ke committed
119
    // create boosting
Guolin Ke's avatar
Guolin Ke committed
120
    if (config_.input_model.size() > 0) {
121
122
      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
123
    }
Guolin Ke's avatar
Guolin Ke committed
124

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

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

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

  ~Booster() {
  }
148

149
  void CreateObjectiveAndMetrics() {
Guolin Ke's avatar
Guolin Ke committed
150
    // create objective function
Guolin Ke's avatar
Guolin Ke committed
151
152
    objective_fun_.reset(ObjectiveFunction::CreateObjectiveFunction(config_.objective,
                                                                    config_));
Guolin Ke's avatar
Guolin Ke committed
153
154
155
156
157
158
159
160
161
162
    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
163
    for (auto metric_type : config_.metric) {
Guolin Ke's avatar
Guolin Ke committed
164
      auto metric = std::unique_ptr<Metric>(
Guolin Ke's avatar
Guolin Ke committed
165
        Metric::CreateMetric(metric_type, config_));
Guolin Ke's avatar
Guolin Ke committed
166
167
168
169
170
      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();
171
172
173
174
  }

  void ResetTrainingData(const Dataset* train_data) {
    if (train_data != train_data_) {
175
      UNIQUE_LOCK(mutex_)
176
177
178
179
180
181
      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
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
288
289
  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`.");
    }
290
    if (new_param.count("linear_tree") && (new_config.linear_tree != old_config.linear_tree)) {
291
      Log::Fatal("Cannot change linear_tree after constructed Dataset handle.");
292
    }
293
294
  }

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

Guolin Ke's avatar
Guolin Ke committed
311
    config_.Set(param);
312

313
314
315
    if (config_.num_threads > 0) {
      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;
712
    predictor.Predict(data_filename, result_filename, bool_data_has_header, config.predict_disable_shape_check);
Guolin Ke's avatar
Guolin Ke committed
713
714
  }

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

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

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

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

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

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

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

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

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

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

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

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

778
  int GetEvalNames(char** out_strs, const int len, const size_t buffer_len, size_t *out_buffer_len) const {
779
    SHARED_LOCK(mutex_)
780
    *out_buffer_len = 0;
wxchan's avatar
wxchan committed
781
782
783
    int idx = 0;
    for (const auto& metric : train_metric_) {
      for (const auto& name : metric->GetName()) {
784
785
786
787
788
        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
789
790
791
792
793
794
        ++idx;
      }
    }
    return idx;
  }

795
  int GetFeatureNames(char** out_strs, const int len, const size_t buffer_len, size_t *out_buffer_len) const {
796
    SHARED_LOCK(mutex_)
797
    *out_buffer_len = 0;
wxchan's avatar
wxchan committed
798
799
    int idx = 0;
    for (const auto& name : boosting_->FeatureNames()) {
800
801
802
803
804
      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
805
806
807
808
809
      ++idx;
    }
    return idx;
  }

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

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

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

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

831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
// 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
848

Guolin Ke's avatar
Guolin Ke committed
849
850
851
852
853
854
855
856
// 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);

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

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

// Row iterator of on column for CSC matrix
class CSC_RowIterator {
Nikita Titov's avatar
Nikita Titov committed
867
 public:
Guolin Ke's avatar
Guolin Ke committed
868
  CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
869
                  const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx);
Guolin Ke's avatar
Guolin Ke committed
870
871
872
873
874
  ~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
875
876

 private:
Guolin Ke's avatar
Guolin Ke committed
877
878
879
880
881
882
883
884
885
  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
886
const char* LGBM_GetLastError() {
wxchan's avatar
wxchan committed
887
  return LastErrorMsg();
Guolin Ke's avatar
Guolin Ke committed
888
889
}

890
891
892
893
894
895
int LGBM_RegisterLogCallback(void (*callback)(const char*)) {
  API_BEGIN();
  Log::ResetCallBack(callback);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
896
int LGBM_DatasetCreateFromFile(const char* filename,
897
898
899
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
900
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
901
902
  auto param = Config::Str2Map(parameters);
  Config config;
903
904
905
906
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
907
  DatasetLoader loader(config, nullptr, 1, filename);
Guolin Ke's avatar
Guolin Ke committed
908
  if (reference == nullptr) {
909
    if (Network::num_machines() == 1) {
910
      *out = loader.LoadFromFile(filename);
911
    } else {
912
      *out = loader.LoadFromFile(filename, Network::rank(), Network::num_machines());
913
    }
Guolin Ke's avatar
Guolin Ke committed
914
  } else {
915
    *out = loader.LoadFromFileAlignWithOtherDataset(filename,
916
                                                    reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
917
  }
918
  API_END();
Guolin Ke's avatar
Guolin Ke committed
919
920
}

921

Guolin Ke's avatar
Guolin Ke committed
922
int LGBM_DatasetCreateFromSampledColumn(double** sample_data,
923
924
925
926
927
928
929
                                        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) {
930
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
931
932
  auto param = Config::Str2Map(parameters);
  Config config;
933
934
935
936
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
937
  DatasetLoader loader(config, nullptr, 1, nullptr);
938
939
940
  *out = loader.ConstructFromSampleData(sample_data, sample_indices, ncol, num_per_col,
                                        num_sample_row,
                                        static_cast<data_size_t>(num_total_row));
941
  API_END();
Guolin Ke's avatar
Guolin Ke committed
942
943
}

944

Guolin Ke's avatar
Guolin Ke committed
945
int LGBM_DatasetCreateByReference(const DatasetHandle reference,
946
947
                                  int64_t num_total_row,
                                  DatasetHandle* out) {
Guolin Ke's avatar
Guolin Ke committed
948
949
950
951
952
953
954
955
  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
956
int LGBM_DatasetPushRows(DatasetHandle dataset,
957
958
959
960
961
                         const void* data,
                         int data_type,
                         int32_t nrow,
                         int32_t ncol,
                         int32_t start_row) {
Guolin Ke's avatar
Guolin Ke committed
962
963
964
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto get_row_fun = RowFunctionFromDenseMatric(data, nrow, ncol, data_type, 1);
965
966
967
  if (p_dataset->has_raw()) {
    p_dataset->ResizeRaw(p_dataset->num_numeric_features() + nrow);
  }
968
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
969
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
970
  for (int i = 0; i < nrow; ++i) {
971
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
972
973
974
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(tid, start_row + i, one_row);
975
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
976
  }
977
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
978
979
980
981
982
983
  if (start_row + nrow == p_dataset->num_data()) {
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
984
int LGBM_DatasetPushRowsByCSR(DatasetHandle dataset,
985
986
987
988
989
990
991
992
993
                              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
994
995
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
996
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
997
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
998
999
1000
  if (p_dataset->has_raw()) {
    p_dataset->ResizeRaw(p_dataset->num_numeric_features() + nrow);
  }
1001
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1002
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1003
  for (int i = 0; i < nrow; ++i) {
1004
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1005
1006
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
1007
    p_dataset->PushOneRow(tid, static_cast<data_size_t>(start_row + i), one_row);
1008
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1009
  }
1010
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1011
1012
1013
1014
1015
1016
  if (start_row + nrow == static_cast<int64_t>(p_dataset->num_data())) {
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1017
int LGBM_DatasetCreateFromMat(const void* data,
1018
1019
1020
1021
1022
1023
1024
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
  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) {
1046
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1047
1048
  auto param = Config::Str2Map(parameters);
  Config config;
1049
1050
1051
1052
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1053
  std::unique_ptr<Dataset> ret;
1054
1055
1056
1057
1058
1059
1060
1061
1062
  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));
  }
1063

Guolin Ke's avatar
Guolin Ke committed
1064
1065
  if (reference == nullptr) {
    // sample data first
Guolin Ke's avatar
Guolin Ke committed
1066
    Random rand(config.data_random_seed);
1067
1068
    int sample_cnt = static_cast<int>(total_nrow < config.bin_construct_sample_cnt ? total_nrow : config.bin_construct_sample_cnt);
    auto sample_indices = rand.Sample(total_nrow, sample_cnt);
1069
    sample_cnt = static_cast<int>(sample_indices.size());
1070
    std::vector<std::vector<double>> sample_values(ncol);
Guolin Ke's avatar
Guolin Ke committed
1071
    std::vector<std::vector<int>> sample_idx(ncol);
1072
1073
1074

    int offset = 0;
    int j = 0;
Guolin Ke's avatar
Guolin Ke committed
1075
    for (size_t i = 0; i < sample_indices.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1076
      auto idx = sample_indices[i];
1077
1078
1079
1080
      while ((idx - offset) >= nrow[j]) {
        offset += nrow[j];
        ++j;
      }
1081

1082
1083
1084
1085
1086
      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
1087
        }
Guolin Ke's avatar
Guolin Ke committed
1088
1089
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1090
    DatasetLoader loader(config, nullptr, 1, nullptr);
1091
1092
1093
1094
1095
    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
1096
  } else {
1097
    ret.reset(new Dataset(total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1098
    ret->CreateValid(
1099
      reinterpret_cast<const Dataset*>(reference));
1100
1101
1102
    if (ret->has_raw()) {
      ret->ResizeRaw(total_nrow);
    }
Guolin Ke's avatar
Guolin Ke committed
1103
  }
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
  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
1118
1119
  }
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1120
  *out = ret.release();
1121
  API_END();
1122
1123
}

Guolin Ke's avatar
Guolin Ke committed
1124
int LGBM_DatasetCreateFromCSR(const void* indptr,
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
                              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) {
1135
  API_BEGIN();
1136
1137
1138
1139
1140
  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
1141
1142
  auto param = Config::Str2Map(parameters);
  Config config;
1143
1144
1145
1146
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1147
  std::unique_ptr<Dataset> ret;
1148
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1149
1150
1151
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
  if (reference == nullptr) {
    // sample data first
Guolin Ke's avatar
Guolin Ke committed
1152
1153
    Random rand(config.data_random_seed);
    int sample_cnt = static_cast<int>(nrow < config.bin_construct_sample_cnt ? nrow : config.bin_construct_sample_cnt);
1154
    auto sample_indices = rand.Sample(nrow, sample_cnt);
1155
    sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1156
1157
    std::vector<std::vector<double>> sample_values(num_col);
    std::vector<std::vector<int>> sample_idx(num_col);
1158
1159
1160
1161
    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
1162
        CHECK_LT(inner_data.first, num_col);
Guolin Ke's avatar
Guolin Ke committed
1163
        if (std::fabs(inner_data.second) > kZeroThreshold || std::isnan(inner_data.second)) {
Guolin Ke's avatar
Guolin Ke committed
1164
1165
          sample_values[inner_data.first].emplace_back(inner_data.second);
          sample_idx[inner_data.first].emplace_back(static_cast<int>(i));
1166
1167
1168
        }
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1169
    DatasetLoader loader(config, nullptr, 1, nullptr);
1170
1171
1172
1173
1174
    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));
1175
  } else {
1176
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1177
    ret->CreateValid(
1178
      reinterpret_cast<const Dataset*>(reference));
1179
1180
1181
    if (ret->has_raw()) {
      ret->ResizeRaw(nrow);
    }
1182
  }
1183
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1184
  #pragma omp parallel for schedule(static)
1185
  for (int i = 0; i < nindptr - 1; ++i) {
1186
    OMP_LOOP_EX_BEGIN();
1187
1188
1189
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    ret->PushOneRow(tid, i, one_row);
1190
    OMP_LOOP_EX_END();
1191
  }
1192
  OMP_THROW_EX();
1193
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1194
  *out = ret.release();
1195
  API_END();
1196
1197
}

1198
int LGBM_DatasetCreateFromCSRFunc(void* get_row_funptr,
1199
1200
1201
1202
1203
                                  int num_rows,
                                  int64_t num_col,
                                  const char* parameters,
                                  const DatasetHandle reference,
                                  DatasetHandle* out) {
1204
  API_BEGIN();
1205
1206
1207
1208
1209
  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.");
  }
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
  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);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  std::unique_ptr<Dataset> ret;
  int32_t nrow = num_rows;
  if (reference == nullptr) {
    // sample data first
    Random rand(config.data_random_seed);
    int sample_cnt = static_cast<int>(nrow < config.bin_construct_sample_cnt ? nrow : config.bin_construct_sample_cnt);
    auto sample_indices = rand.Sample(nrow, sample_cnt);
    sample_cnt = static_cast<int>(sample_indices.size());
    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
1233
        CHECK_LT(inner_data.first, num_col);
1234
1235
1236
1237
1238
1239
1240
        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);
1241
1242
1243
1244
1245
    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));
1246
1247
1248
1249
  } else {
    ret.reset(new Dataset(nrow));
    ret->CreateValid(
      reinterpret_cast<const Dataset*>(reference));
1250
1251
1252
    if (ret->has_raw()) {
      ret->ResizeRaw(nrow);
    }
1253
  }
1254

1255
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1256
1257
  std::vector<std::pair<int, double>> thread_buffer;
  #pragma omp parallel for schedule(static) private(thread_buffer)
1258
1259
1260
  for (int i = 0; i < num_rows; ++i) {
    OMP_LOOP_EX_BEGIN();
    {
1261
      const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1262
1263
      get_row_fun(i, thread_buffer);
      ret->PushOneRow(tid, i, thread_buffer);
1264
1265
1266
1267
1268
1269
1270
1271
1272
    }
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();
  ret->FinishLoad();
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1273
int LGBM_DatasetCreateFromCSC(const void* col_ptr,
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
                              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) {
1284
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1285
1286
  auto param = Config::Str2Map(parameters);
  Config config;
1287
1288
1289
1290
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1291
  std::unique_ptr<Dataset> ret;
Guolin Ke's avatar
Guolin Ke committed
1292
1293
1294
  int32_t nrow = static_cast<int32_t>(num_row);
  if (reference == nullptr) {
    // sample data first
Guolin Ke's avatar
Guolin Ke committed
1295
1296
    Random rand(config.data_random_seed);
    int sample_cnt = static_cast<int>(nrow < config.bin_construct_sample_cnt ? nrow : config.bin_construct_sample_cnt);
Guolin Ke's avatar
Guolin Ke committed
1297
    auto sample_indices = rand.Sample(nrow, sample_cnt);
1298
    sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1299
    std::vector<std::vector<double>> sample_values(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
1300
    std::vector<std::vector<int>> sample_idx(ncol_ptr - 1);
1301
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1302
    #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1303
    for (int i = 0; i < static_cast<int>(sample_values.size()); ++i) {
1304
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1305
1306
1307
      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
1308
        if (std::fabs(val) > kZeroThreshold || std::isnan(val)) {
Guolin Ke's avatar
Guolin Ke committed
1309
1310
          sample_values[i].emplace_back(val);
          sample_idx[i].emplace_back(j);
Guolin Ke's avatar
Guolin Ke committed
1311
1312
        }
      }
1313
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1314
    }
1315
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1316
    DatasetLoader loader(config, nullptr, 1, nullptr);
1317
1318
1319
1320
1321
    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
1322
  } else {
1323
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1324
    ret->CreateValid(
1325
      reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
1326
  }
1327
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1328
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1329
  for (int i = 0; i < ncol_ptr - 1; ++i) {
1330
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1331
    const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1332
    int feature_idx = ret->InnerFeatureIndex(i);
Guolin Ke's avatar
Guolin Ke committed
1333
    if (feature_idx < 0) { continue; }
Guolin Ke's avatar
Guolin Ke committed
1334
1335
    int group = ret->Feature2Group(feature_idx);
    int sub_feature = ret->Feture2SubFeature(feature_idx);
Guolin Ke's avatar
Guolin Ke committed
1336
    CSC_RowIterator col_it(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, i);
Guolin Ke's avatar
Guolin Ke committed
1337
1338
1339
1340
1341
1342
1343
1344
    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; }
1345
        ret->PushOneData(tid, row_idx, group, feature_idx, sub_feature, pair.second);
Guolin Ke's avatar
Guolin Ke committed
1346
1347
1348
1349
      }
    } else {
      for (int row_idx = 0; row_idx < nrow; ++row_idx) {
        auto val = col_it.Get(row_idx);
1350
        ret->PushOneData(tid, row_idx, group, feature_idx, sub_feature, val);
Guolin Ke's avatar
Guolin Ke committed
1351
      }
Guolin Ke's avatar
Guolin Ke committed
1352
    }
1353
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1354
  }
1355
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1356
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1357
  *out = ret.release();
1358
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1359
1360
}

Guolin Ke's avatar
Guolin Ke committed
1361
int LGBM_DatasetGetSubset(
1362
  const DatasetHandle handle,
wxchan's avatar
wxchan committed
1363
1364
1365
  const int32_t* used_row_indices,
  int32_t num_used_row_indices,
  const char* parameters,
Guolin Ke's avatar
typo  
Guolin Ke committed
1366
  DatasetHandle* out) {
wxchan's avatar
wxchan committed
1367
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1368
1369
  auto param = Config::Str2Map(parameters);
  Config config;
1370
1371
1372
1373
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
1374
  auto full_dataset = reinterpret_cast<const Dataset*>(handle);
1375
  CHECK_GT(num_used_row_indices, 0);
1376
1377
  const int32_t lower = 0;
  const int32_t upper = full_dataset->num_data() - 1;
1378
  CheckElementsIntervalClosed(used_row_indices, lower, upper, num_used_row_indices, "Used indices of subset");
1379
1380
1381
  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
1382
  auto ret = std::unique_ptr<Dataset>(new Dataset(num_used_row_indices));
1383
  ret->CopyFeatureMapperFrom(full_dataset);
1384
  ret->CopySubrow(full_dataset, used_row_indices, num_used_row_indices, true);
wxchan's avatar
wxchan committed
1385
1386
1387
1388
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1389
int LGBM_DatasetSetFeatureNames(
Guolin Ke's avatar
typo  
Guolin Ke committed
1390
  DatasetHandle handle,
Guolin Ke's avatar
Guolin Ke committed
1391
  const char** feature_names,
Guolin Ke's avatar
Guolin Ke committed
1392
  int num_feature_names) {
Guolin Ke's avatar
Guolin Ke committed
1393
1394
1395
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  std::vector<std::string> feature_names_str;
Guolin Ke's avatar
Guolin Ke committed
1396
  for (int i = 0; i < num_feature_names; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1397
1398
1399
1400
1401
1402
    feature_names_str.emplace_back(feature_names[i]);
  }
  dataset->set_feature_names(feature_names_str);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1403
int LGBM_DatasetGetFeatureNames(
1404
1405
1406
1407
1408
1409
    DatasetHandle handle,
    const int len,
    int* num_feature_names,
    const size_t buffer_len,
    size_t* out_buffer_len,
    char** feature_names) {
1410
  API_BEGIN();
1411
  *out_buffer_len = 0;
1412
1413
  auto dataset = reinterpret_cast<Dataset*>(handle);
  auto inside_feature_name = dataset->feature_names();
Guolin Ke's avatar
Guolin Ke committed
1414
1415
  *num_feature_names = static_cast<int>(inside_feature_name.size());
  for (int i = 0; i < *num_feature_names; ++i) {
1416
1417
1418
1419
1420
    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);
1421
1422
1423
1424
  }
  API_END();
}

1425
1426
1427
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1428
int LGBM_DatasetFree(DatasetHandle handle) {
1429
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1430
  delete reinterpret_cast<Dataset*>(handle);
1431
  API_END();
1432
1433
}

Guolin Ke's avatar
Guolin Ke committed
1434
int LGBM_DatasetSaveBinary(DatasetHandle handle,
1435
                           const char* filename) {
1436
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1437
1438
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->SaveBinaryFile(filename);
1439
  API_END();
1440
1441
}

1442
1443
1444
1445
1446
1447
1448
1449
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
1450
int LGBM_DatasetSetField(DatasetHandle handle,
1451
1452
1453
1454
                         const char* field_name,
                         const void* field_data,
                         int num_element,
                         int type) {
1455
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1456
  auto dataset = reinterpret_cast<Dataset*>(handle);
1457
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1458
  if (type == C_API_DTYPE_FLOAT32) {
Guolin Ke's avatar
Guolin Ke committed
1459
    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
1460
  } else if (type == C_API_DTYPE_INT32) {
Guolin Ke's avatar
Guolin Ke committed
1461
    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
1462
1463
  } 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));
1464
  }
1465
  if (!is_success) { Log::Fatal("Input data type error or field not found"); }
1466
  API_END();
1467
1468
}

Guolin Ke's avatar
Guolin Ke committed
1469
int LGBM_DatasetGetField(DatasetHandle handle,
1470
1471
1472
1473
                         const char* field_name,
                         int* out_len,
                         const void** out_ptr,
                         int* out_type) {
1474
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1475
  auto dataset = reinterpret_cast<Dataset*>(handle);
1476
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1477
  if (dataset->GetFloatField(field_name, out_len, reinterpret_cast<const float**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1478
    *out_type = C_API_DTYPE_FLOAT32;
1479
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1480
  } else if (dataset->GetIntField(field_name, out_len, reinterpret_cast<const int**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1481
    *out_type = C_API_DTYPE_INT32;
1482
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1483
1484
1485
  } 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
1486
  }
1487
  if (!is_success) { Log::Fatal("Field not found"); }
wxchan's avatar
wxchan committed
1488
  if (*out_ptr == nullptr) { *out_len = 0; }
1489
  API_END();
1490
1491
}

1492
int LGBM_DatasetUpdateParamChecking(const char* old_parameters, const char* new_parameters) {
1493
  API_BEGIN();
1494
1495
1496
1497
1498
  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);
1499
1500
1501
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1502
int LGBM_DatasetGetNumData(DatasetHandle handle,
1503
                           int* out) {
1504
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1505
1506
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_data();
1507
  API_END();
1508
1509
}

Guolin Ke's avatar
Guolin Ke committed
1510
int LGBM_DatasetGetNumFeature(DatasetHandle handle,
1511
                              int* out) {
1512
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1513
1514
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_total_features();
1515
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1516
}
1517

1518
1519
1520
1521
1522
int LGBM_DatasetAddFeaturesFrom(DatasetHandle target,
                                DatasetHandle source) {
  API_BEGIN();
  auto target_d = reinterpret_cast<Dataset*>(target);
  auto source_d = reinterpret_cast<Dataset*>(source);
1523
  target_d->AddFeaturesFrom(source_d);
1524
1525
1526
  API_END();
}

1527
1528
// ---- start of booster

Guolin Ke's avatar
Guolin Ke committed
1529
int LGBM_BoosterCreate(const DatasetHandle train_data,
1530
1531
                       const char* parameters,
                       BoosterHandle* out) {
1532
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1533
  const Dataset* p_train_data = reinterpret_cast<const Dataset*>(train_data);
wxchan's avatar
wxchan committed
1534
1535
  auto ret = std::unique_ptr<Booster>(new Booster(p_train_data, parameters));
  *out = ret.release();
1536
  API_END();
1537
1538
}

Guolin Ke's avatar
Guolin Ke committed
1539
int LGBM_BoosterCreateFromModelfile(
1540
  const char* filename,
Guolin Ke's avatar
Guolin Ke committed
1541
  int* out_num_iterations,
1542
  BoosterHandle* out) {
1543
  API_BEGIN();
wxchan's avatar
wxchan committed
1544
  auto ret = std::unique_ptr<Booster>(new Booster(filename));
Guolin Ke's avatar
Guolin Ke committed
1545
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
wxchan's avatar
wxchan committed
1546
  *out = ret.release();
1547
  API_END();
1548
1549
}

Guolin Ke's avatar
Guolin Ke committed
1550
int LGBM_BoosterLoadModelFromString(
1551
1552
1553
1554
  const char* model_str,
  int* out_num_iterations,
  BoosterHandle* out) {
  API_BEGIN();
wxchan's avatar
wxchan committed
1555
  auto ret = std::unique_ptr<Booster>(new Booster(nullptr));
1556
1557
1558
1559
1560
1561
  ret->LoadModelFromString(model_str);
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
  *out = ret.release();
  API_END();
}

1562
1563
1564
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1565
int LGBM_BoosterFree(BoosterHandle handle) {
1566
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1567
  delete reinterpret_cast<Booster*>(handle);
1568
  API_END();
1569
1570
}

1571
int LGBM_BoosterShuffleModels(BoosterHandle handle, int start_iter, int end_iter) {
1572
1573
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1574
  ref_booster->ShuffleModels(start_iter, end_iter);
1575
1576
1577
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1578
int LGBM_BoosterMerge(BoosterHandle handle,
1579
                      BoosterHandle other_handle) {
wxchan's avatar
wxchan committed
1580
1581
1582
1583
1584
1585
1586
  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
1587
int LGBM_BoosterAddValidData(BoosterHandle handle,
1588
                             const DatasetHandle valid_data) {
wxchan's avatar
wxchan committed
1589
1590
1591
1592
1593
1594
1595
  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
1596
int LGBM_BoosterResetTrainingData(BoosterHandle handle,
1597
                                  const DatasetHandle train_data) {
wxchan's avatar
wxchan committed
1598
1599
1600
1601
1602
1603
1604
  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
1605
int LGBM_BoosterResetParameter(BoosterHandle handle, const char* parameters) {
wxchan's avatar
wxchan committed
1606
1607
1608
1609
1610
1611
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->ResetConfig(parameters);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1612
int LGBM_BoosterGetNumClasses(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1613
1614
1615
1616
1617
1618
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->NumberOfClasses();
  API_END();
}

1619
1620
1621
1622
1623
1624
1625
int LGBM_BoosterGetLinear(BoosterHandle handle, bool* out) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out = ref_booster->GetBoosting()->IsLinear();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1626
1627
1628
1629
1630
1631
1632
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
1633
int LGBM_BoosterUpdateOneIter(BoosterHandle handle, int* is_finished) {
1634
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1635
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1636
1637
1638
1639
1640
  if (ref_booster->TrainOneIter()) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1641
  API_END();
1642
1643
}

Guolin Ke's avatar
Guolin Ke committed
1644
int LGBM_BoosterUpdateOneIterCustom(BoosterHandle handle,
1645
1646
1647
                                    const float* grad,
                                    const float* hess,
                                    int* is_finished) {
1648
  API_BEGIN();
1649
  #ifdef SCORE_T_USE_DOUBLE
1650
1651
1652
1653
  (void) handle;       // UNUSED VARIABLE
  (void) grad;         // UNUSED VARIABLE
  (void) hess;         // UNUSED VARIABLE
  (void) is_finished;  // UNUSED VARIABLE
1654
  Log::Fatal("Don't support custom loss function when SCORE_T_USE_DOUBLE is enabled");
1655
  #else
1656
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1657
1658
1659
1660
1661
  if (ref_booster->TrainOneIter(grad, hess)) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1662
  #endif
1663
  API_END();
1664
1665
}

Guolin Ke's avatar
Guolin Ke committed
1666
int LGBM_BoosterRollbackOneIter(BoosterHandle handle) {
wxchan's avatar
wxchan committed
1667
1668
1669
1670
1671
1672
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->RollbackOneIter();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1673
int LGBM_BoosterGetCurrentIteration(BoosterHandle handle, int* out_iteration) {
wxchan's avatar
wxchan committed
1674
1675
1676
1677
1678
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_iteration = ref_booster->GetBoosting()->GetCurrentIteration();
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
1679

1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
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
1694
int LGBM_BoosterGetEvalCounts(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1695
1696
1697
1698
1699
1700
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetEvalCounts();
  API_END();
}

1701
1702
1703
1704
1705
1706
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
1707
1708
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1709
  *out_len = ref_booster->GetEvalNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1710
1711
1712
  API_END();
}

1713
1714
1715
1716
1717
1718
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
1719
1720
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1721
  *out_len = ref_booster->GetFeatureNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1722
1723
1724
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1725
int LGBM_BoosterGetNumFeature(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1726
1727
1728
1729
1730
1731
  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
1732
int LGBM_BoosterGetEval(BoosterHandle handle,
1733
1734
1735
                        int data_idx,
                        int* out_len,
                        double* out_results) {
1736
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1737
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1738
  auto boosting = ref_booster->GetBoosting();
wxchan's avatar
wxchan committed
1739
  auto result_buf = boosting->GetEvalAt(data_idx);
Guolin Ke's avatar
Guolin Ke committed
1740
  *out_len = static_cast<int>(result_buf.size());
1741
  for (size_t i = 0; i < result_buf.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1742
    (out_results)[i] = static_cast<double>(result_buf[i]);
1743
  }
1744
  API_END();
1745
1746
}

Guolin Ke's avatar
Guolin Ke committed
1747
int LGBM_BoosterGetNumPredict(BoosterHandle handle,
1748
1749
                              int data_idx,
                              int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1750
1751
1752
1753
1754
1755
  API_BEGIN();
  auto boosting = reinterpret_cast<Booster*>(handle)->GetBoosting();
  *out_len = boosting->GetNumPredictAt(data_idx);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1756
int LGBM_BoosterGetPredict(BoosterHandle handle,
1757
1758
1759
                           int data_idx,
                           int64_t* out_len,
                           double* out_result) {
1760
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1761
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1762
  ref_booster->GetPredictAt(data_idx, out_result, out_len);
1763
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1764
1765
}

Guolin Ke's avatar
Guolin Ke committed
1766
int LGBM_BoosterPredictForFile(BoosterHandle handle,
1767
1768
1769
                               const char* data_filename,
                               int data_has_header,
                               int predict_type,
1770
                               int start_iteration,
1771
                               int num_iteration,
1772
                               const char* parameter,
1773
                               const char* result_filename) {
1774
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1775
1776
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1777
1778
1779
1780
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1781
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1782
  ref_booster->Predict(start_iteration, num_iteration, predict_type, data_filename, data_has_header,
Guolin Ke's avatar
Guolin Ke committed
1783
                       config, result_filename);
1784
  API_END();
1785
1786
}

Guolin Ke's avatar
Guolin Ke committed
1787
int LGBM_BoosterCalcNumPredict(BoosterHandle handle,
1788
1789
                               int num_row,
                               int predict_type,
1790
                               int start_iteration,
1791
1792
                               int num_iteration,
                               int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1793
1794
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1795
  *out_len = static_cast<int64_t>(num_row) * ref_booster->GetBoosting()->NumPredictOneRow(start_iteration,
1796
    num_iteration, predict_type == C_API_PREDICT_LEAF_INDEX, predict_type == C_API_PREDICT_CONTRIB);
Guolin Ke's avatar
Guolin Ke committed
1797
1798
1799
  API_END();
}

1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
/*!
 * \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,
1813
             const int predict_type_,
1814
             const int data_type_,
1815
             const int32_t num_cols) : booster(booster_ptr), predict_type(predict_type_), data_type(data_type_), ncol(num_cols) {
1816
1817
1818
1819
1820
    config.Set(Config::Str2Map(parameter));
  }

  Booster* const booster;
  Config config;
1821
  const int predict_type;
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
  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
1832
int LGBM_BoosterPredictForCSR(BoosterHandle handle,
1833
1834
1835
1836
1837
1838
1839
                              const void* indptr,
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
1840
                              int64_t num_col,
1841
                              int predict_type,
1842
                              int start_iteration,
1843
                              int num_iteration,
1844
                              const char* parameter,
1845
1846
                              int64_t* out_len,
                              double* out_result) {
1847
  API_BEGIN();
1848
1849
1850
1851
1852
  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
1853
1854
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1855
1856
1857
1858
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1859
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1860
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
1861
  int nrow = static_cast<int>(nindptr - 1);
1862
  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
1863
                       config, out_result, out_len);
1864
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1865
}
1866

1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
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,
1877
                                    int start_iteration,
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
                                    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);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  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;
1901
    ref_booster->PredictSparseCSR(start_iteration, num_iteration, predict_type, nrow, static_cast<int>(num_col_or_row), get_row_fun,
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
                                  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;
    };
1925
    ref_booster->PredictSparseCSC(start_iteration, num_iteration, predict_type, num_col_or_row, ncol, get_row_fun, config,
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
                                  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();
}

1953
int LGBM_BoosterPredictForCSRSingleRow(BoosterHandle handle,
1954
1955
1956
1957
1958
1959
1960
                                       const void* indptr,
                                       int indptr_type,
                                       const int32_t* indices,
                                       const void* data,
                                       int data_type,
                                       int64_t nindptr,
                                       int64_t nelem,
1961
                                       int64_t num_col,
1962
                                       int predict_type,
1963
                                       int start_iteration,
1964
1965
1966
1967
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
1968
  API_BEGIN();
1969
1970
1971
1972
1973
  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.");
  }
1974
1975
1976
1977
1978
1979
1980
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1981
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1982
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
1983
  ref_booster->PredictSingleRow(predict_type, static_cast<int32_t>(num_col), get_row_fun, config, out_result, out_len);
1984
1985
1986
  API_END();
}

1987
int LGBM_BoosterPredictForCSRSingleRowFastInit(BoosterHandle handle,
1988
                                               const int predict_type,
1989
                                               const int start_iteration,
1990
                                               const int num_iteration,
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
                                               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,
2005
    predict_type,
2006
2007
2008
2009
2010
2011
2012
    data_type,
    static_cast<int32_t>(num_col)));

  if (fastConfig_ptr->config.num_threads > 0) {
    omp_set_num_threads(fastConfig_ptr->config.num_threads);
  }

2013
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2014

2015
2016
2017
2018
2019
2020
  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

int LGBM_BoosterPredictForCSRSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* indptr,
2021
                                           const int indptr_type,
2022
2023
                                           const int32_t* indices,
                                           const void* data,
2024
2025
                                           const int64_t nindptr,
                                           const int64_t nelem,
2026
2027
2028
2029
2030
                                           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);
2031
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
2032
2033
2034
2035
                                        get_row_fun, fastConfig->config, out_result, out_len);
  API_END();
}

2036

Guolin Ke's avatar
Guolin Ke committed
2037
int LGBM_BoosterPredictForCSC(BoosterHandle handle,
2038
2039
2040
2041
2042
2043
2044
2045
2046
                              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,
2047
                              int start_iteration,
2048
                              int num_iteration,
2049
                              const char* parameter,
2050
2051
                              int64_t* out_len,
                              double* out_result) {
Guolin Ke's avatar
Guolin Ke committed
2052
2053
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2054
2055
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2056
2057
2058
2059
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
2060
  int num_threads = OMP_NUM_THREADS();
Guolin Ke's avatar
Guolin Ke committed
2061
  int ncol = static_cast<int>(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
2062
2063
2064
2065
2066
  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
2067
2068
  }
  std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun =
Guolin Ke's avatar
Guolin Ke committed
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
      [&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;
      };
2081
  ref_booster->Predict(start_iteration, num_iteration, predict_type, static_cast<int>(num_row), ncol, get_row_fun, config,
cbecker's avatar
cbecker committed
2082
                       out_result, out_len);
Guolin Ke's avatar
Guolin Ke committed
2083
2084
2085
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2086
int LGBM_BoosterPredictForMat(BoosterHandle handle,
2087
2088
2089
2090
2091
2092
                              const void* data,
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              int predict_type,
2093
                              int start_iteration,
2094
                              int num_iteration,
2095
                              const char* parameter,
2096
2097
                              int64_t* out_len,
                              double* out_result) {
2098
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2099
2100
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2101
2102
2103
2104
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
2105
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2106
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, nrow, ncol, data_type, is_row_major);
2107
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
2108
                       config, out_result, out_len);
2109
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2110
}
2111

2112
int LGBM_BoosterPredictForMatSingleRow(BoosterHandle handle,
2113
2114
2115
2116
2117
                                       const void* data,
                                       int data_type,
                                       int32_t ncol,
                                       int is_row_major,
                                       int predict_type,
2118
                                       int start_iteration,
2119
2120
2121
2122
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2123
2124
2125
2126
2127
2128
2129
2130
2131
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, ncol, data_type, is_row_major);
2132
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
2133
  ref_booster->PredictSingleRow(predict_type, ncol, get_row_fun, config, out_result, out_len);
2134
2135
2136
  API_END();
}

2137
int LGBM_BoosterPredictForMatSingleRowFastInit(BoosterHandle handle,
2138
                                               const int predict_type,
2139
                                               const int start_iteration,
2140
                                               const int num_iteration,
2141
2142
2143
2144
2145
2146
2147
2148
                                               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,
2149
    predict_type,
2150
2151
2152
2153
2154
2155
2156
    data_type,
    ncol));

  if (fastConfig_ptr->config.num_threads > 0) {
    omp_set_num_threads(fastConfig_ptr->config.num_threads);
  }

2157
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2158

2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
  *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);
2171
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
2172
2173
2174
2175
2176
                                        get_row_fun, fastConfig->config,
                                        out_result, out_len);
  API_END();
}

2177

2178
2179
2180
2181
2182
2183
int LGBM_BoosterPredictForMats(BoosterHandle handle,
                               const void** data,
                               int data_type,
                               int32_t nrow,
                               int32_t ncol,
                               int predict_type,
2184
                               int start_iteration,
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
                               int num_iteration,
                               const char* parameter,
                               int64_t* out_len,
                               double* out_result) {
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseRows(data, ncol, data_type);
2198
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, out_result, out_len);
2199
2200
2201
  API_END();
}

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

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

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

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

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

2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
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();
}

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

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

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

Guolin Ke's avatar
Guolin Ke committed
2333
// ---- start of some help functions
2334

2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359

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

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

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

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

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

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

2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464


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

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