c_api.cpp 111 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

7
#include <LightGBM/arrow.h>
Guolin Ke's avatar
Guolin Ke committed
8
9
#include <LightGBM/boosting.h>
#include <LightGBM/config.h>
10
11
12
#include <LightGBM/dataset.h>
#include <LightGBM/dataset_loader.h>
#include <LightGBM/metric.h>
13
#include <LightGBM/network.h>
14
15
#include <LightGBM/objective_function.h>
#include <LightGBM/prediction_early_stop.h>
16
#include <LightGBM/utils/byte_buffer.h>
17
18
19
20
21
#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
22

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

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

Guolin Ke's avatar
Guolin Ke committed
35
36
namespace LightGBM {

Guolin Ke's avatar
Guolin Ke committed
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
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;

53
54
55
56
57
58
#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);

59
60
61
62
63
64
65
66
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;

67
  SingleRowPredictor(int predict_type, Boosting* boosting, const Config& config, int start_iter, int num_iter) {
68
69
70
71
72
73
74
75
76
77
78
79
80
    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;
81
82
    iter_ = num_iter;
    predictor_.reset(new Predictor(boosting, start_iter, iter_, is_raw_score, is_predict_leaf, predict_contrib,
83
                                   early_stop_, early_stop_freq_, early_stop_margin_));
84
    num_pred_in_one_row = boosting->NumPredictOneRow(start_iter, iter_, is_predict_leaf, predict_contrib);
85
    predict_function = predictor_->GetPredictFunction();
Guolin Ke's avatar
Guolin Ke committed
86
    num_total_model_ = boosting->NumberOfTotalModel();
87
  }
88

89
  ~SingleRowPredictor() {}
90

Guolin Ke's avatar
Guolin Ke committed
91
  bool IsPredictorEqual(const Config& config, int iter, Boosting* boosting) {
92
93
94
95
96
    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();
97
  }
Guolin Ke's avatar
Guolin Ke committed
98

99
100
101
102
103
104
105
106
107
 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
108
class Booster {
Nikita Titov's avatar
Nikita Titov committed
109
 public:
Guolin Ke's avatar
Guolin Ke committed
110
  explicit Booster(const char* filename) {
111
    boosting_.reset(Boosting::CreateBoosting("gbdt", filename));
112
113
  }

Guolin Ke's avatar
Guolin Ke committed
114
  Booster(const Dataset* train_data,
115
          const char* parameters) {
Guolin Ke's avatar
Guolin Ke committed
116
    auto param = Config::Str2Map(parameters);
wxchan's avatar
wxchan committed
117
    config_.Set(param);
118
    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
    if (objective_fun_ == nullptr) {
154
      Log::Info("Using self-defined objective function");
Guolin Ke's avatar
Guolin Ke committed
155
156
157
158
159
160
161
162
    }
    // 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`.");
    }
Nikita Titov's avatar
Nikita Titov committed
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
    }
Nikita Titov's avatar
Nikita Titov committed
293
294
295
296
    if (new_param.count("precise_float_parser") &&
        new_config.precise_float_parser != old_config.precise_float_parser) {
      Log::Fatal("Cannot change precise_float_parser after constructed Dataset handle.");
    }
297
298
  }

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

Guolin Ke's avatar
Guolin Ke committed
315
    config_.Set(param);
316

317
    OMP_SET_NUM_THREADS(config_.num_threads);
Guolin Ke's avatar
Guolin Ke committed
318
319
320

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

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

  void AddValidData(const Dataset* valid_data) {
338
    UNIQUE_LOCK(mutex_)
wxchan's avatar
wxchan committed
339
    valid_metrics_.emplace_back();
Guolin Ke's avatar
Guolin Ke committed
340
341
    for (auto metric_type : config_.metric) {
      auto metric = std::unique_ptr<Metric>(Metric::CreateMetric(metric_type, config_));
wxchan's avatar
wxchan committed
342
343
344
345
346
347
      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,
348
                               Common::ConstPtrInVectorWrapper<Metric>(valid_metrics_.back()));
wxchan's avatar
wxchan committed
349
  }
Guolin Ke's avatar
Guolin Ke committed
350

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

Guolin Ke's avatar
Guolin Ke committed
356
  void Refit(const int32_t* leaf_preds, int32_t nrow, int32_t ncol) {
357
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
358
359
360
    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) {
361
        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
362
363
364
365
366
      }
    }
    boosting_->RefitTree(v_leaf_preds);
  }

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

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

377
  void SetSingleRowPredictor(int start_iteration, int num_iteration, int predict_type, const Config& config) {
378
379
380
381
      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(),
382
                                                                         config, start_iteration, num_iteration));
383
384
385
386
      }
  }

  void PredictSingleRow(int predict_type, int ncol,
387
388
               std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun,
               const Config& config,
389
               double* out_result, int64_t* out_len) const {
390
391
392
    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);
393
    }
394
    UNIQUE_LOCK(mutex_)
395
    const auto& single_row_predictor = single_row_predictor_[predict_type];
396
397
    auto one_row = get_row_fun(0);
    auto pred_wrt_ptr = out_result;
398
    single_row_predictor->predict_function(one_row, pred_wrt_ptr);
399

400
    *out_len = single_row_predictor->num_pred_in_one_row;
401
402
  }

403
  Predictor CreatePredictor(int start_iteration, int num_iteration, int predict_type, int ncol, const Config& config) const {
404
405
406
    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);
407
    }
Guolin Ke's avatar
Guolin Ke committed
408
409
    bool is_predict_leaf = false;
    bool is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
410
    bool predict_contrib = false;
Guolin Ke's avatar
Guolin Ke committed
411
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
Guolin Ke's avatar
Guolin Ke committed
412
      is_predict_leaf = true;
Guolin Ke's avatar
Guolin Ke committed
413
    } else if (predict_type == C_API_PREDICT_RAW_SCORE) {
Guolin Ke's avatar
Guolin Ke committed
414
      is_raw_score = true;
415
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
Guolin Ke's avatar
Guolin Ke committed
416
      predict_contrib = true;
Guolin Ke's avatar
Guolin Ke committed
417
418
    } else {
      is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
419
    }
Guolin Ke's avatar
Guolin Ke committed
420

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

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

453
  void PredictSparse(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
454
455
456
457
                     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,
458
                     bool* is_data_float32_ptr, int num_matrices) const {
459
    auto predictor = CreatePredictor(start_iteration, num_iteration, predict_type, ncol, config);
460
461
462
    auto pred_sparse_fun = predictor.GetPredictSparseFunction();
    std::vector<std::vector<std::unordered_map<int, double>>>& agg = *agg_ptr;
    OMP_INIT_EX();
463
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
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
493
494
    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];
  }

495
  void PredictSparseCSR(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
496
497
498
                        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,
499
500
                        int32_t** out_indices, void** out_data, int data_type) const {
    SHARED_LOCK(mutex_);
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
    // 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;
518
    PredictSparse(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
519
520
521
                  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);
522
    std::vector<int64_t> matrix_offsets(num_matrices);
523
524
525
526
527
528
529
530
531
532
533
534
535
536
    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++;
      }
537
538
539
540
541
542
      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]);
      }
543
544
545
546
547
548
549
550
551
552
553
554
    }
    // 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();
555
      #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
556
557
558
559
      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;
560
        int64_t element_index = row_matrix_offsets[row_start_index] + matrix_offsets[m];
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
        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;
  }

586
  void PredictSparseCSC(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
587
588
589
                        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,
590
591
                        int32_t** out_indices, void** out_data, int data_type) const {
    SHARED_LOCK(mutex_);
592
593
    // Get the number of trees per iteration (for multiclass scenario we output multiple sparse matrices)
    int num_matrices = boosting_->NumModelPerIteration();
594
    auto predictor = CreatePredictor(start_iteration, num_iteration, predict_type, ncol, config);
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
    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;
612
    PredictSparse(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
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
658
659
                  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;
      }
660
661
      if (m + 1 < num_matrices) {
        matrix_start_indices[m + 1] = matrix_start_indices[m] + last_column_start_index + last_column_size;
662
      }
663
      col_ptr_index++;
664
    }
665
666
    // Note: we parallelize across matrices instead of rows because of the column_counts[m][col_idx] increment inside the loop
    OMP_INIT_EX();
667
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
668
    for (int m = 0; m < num_matrices; ++m) {
669
      OMP_LOOP_EX_BEGIN();
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
      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;
          }
        }
      }
688
      OMP_LOOP_EX_END();
689
    }
690
    OMP_THROW_EX();
691
692
693
694
    out_len[0] = elements_size;
    out_len[1] = col_ptr_size;
  }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

832
}  // namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
833

834
835
// explicitly declare symbols from LightGBM namespace
using LightGBM::AllgatherFunction;
836
using LightGBM::ArrowChunkedArray;
837
using LightGBM::ArrowTable;
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
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
853

Guolin Ke's avatar
Guolin Ke committed
854
855
856
857
858
859
860
861
// 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);

862
863
864
std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseRows(const void** data, int num_col, int data_type);

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

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

 private:
Guolin Ke's avatar
Guolin Ke committed
882
883
884
885
886
887
888
889
890
  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
891
const char* LGBM_GetLastError() {
wxchan's avatar
wxchan committed
892
  return LastErrorMsg();
Guolin Ke's avatar
Guolin Ke committed
893
894
}

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

907
908
909
910
911
912
int LGBM_RegisterLogCallback(void (*callback)(const char*)) {
  API_BEGIN();
  Log::ResetCallBack(callback);
  API_END();
}

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

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

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

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

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

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

958
959
960
961
962
963
964
965
966
967
968
969
970
int LGBM_ByteBufferGetAt(ByteBufferHandle handle, int32_t index, uint8_t* out_val) {
  API_BEGIN();
  LightGBM::ByteBuffer* byteBuffer = reinterpret_cast<LightGBM::ByteBuffer*>(handle);
  *out_val = byteBuffer->GetAt(index);
  API_END();
}

int LGBM_ByteBufferFree(ByteBufferHandle handle) {
  API_BEGIN();
  delete reinterpret_cast<LightGBM::ByteBuffer*>(handle);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
971
int LGBM_DatasetCreateFromFile(const char* filename,
972
973
974
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
975
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
976
977
  auto param = Config::Str2Map(parameters);
  Config config;
978
  config.Set(param);
979
  OMP_SET_NUM_THREADS(config.num_threads);
980
  DatasetLoader loader(config, nullptr, 1, filename);
Guolin Ke's avatar
Guolin Ke committed
981
  if (reference == nullptr) {
982
    if (Network::num_machines() == 1) {
983
      *out = loader.LoadFromFile(filename);
984
    } else {
985
      *out = loader.LoadFromFile(filename, Network::rank(), Network::num_machines());
986
    }
Guolin Ke's avatar
Guolin Ke committed
987
  } else {
988
    *out = loader.LoadFromFileAlignWithOtherDataset(filename,
989
                                                    reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
990
  }
991
  API_END();
Guolin Ke's avatar
Guolin Ke committed
992
993
}

Guolin Ke's avatar
Guolin Ke committed
994
int LGBM_DatasetCreateFromSampledColumn(double** sample_data,
995
996
997
998
                                        int** sample_indices,
                                        int32_t ncol,
                                        const int* num_per_col,
                                        int32_t num_sample_row,
999
1000
                                        int32_t num_local_row,
                                        int64_t num_dist_row,
1001
1002
                                        const char* parameters,
                                        DatasetHandle* out) {
1003
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1004
1005
  auto param = Config::Str2Map(parameters);
  Config config;
1006
  config.Set(param);
1007
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1008
  DatasetLoader loader(config, nullptr, 1, nullptr);
1009
1010
1011
1012
  *out = loader.ConstructFromSampleData(sample_data,
                                        sample_indices,
                                        ncol,
                                        num_per_col,
1013
                                        num_sample_row,
1014
1015
                                        static_cast<data_size_t>(num_local_row),
                                        num_dist_row);
1016
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1017
1018
}

Guolin Ke's avatar
Guolin Ke committed
1019
int LGBM_DatasetCreateByReference(const DatasetHandle reference,
1020
1021
                                  int64_t num_total_row,
                                  DatasetHandle* out) {
Guolin Ke's avatar
Guolin Ke committed
1022
1023
  API_BEGIN();
  std::unique_ptr<Dataset> ret;
1024
1025
1026
1027
1028
  data_size_t nrows = static_cast<data_size_t>(num_total_row);
  ret.reset(new Dataset(nrows));
  const Dataset* reference_dataset = reinterpret_cast<const Dataset*>(reference);
  ret->CreateValid(reference_dataset);
  ret->InitByReference(nrows, reference_dataset);
Guolin Ke's avatar
Guolin Ke committed
1029
1030
1031
1032
  *out = ret.release();
  API_END();
}

1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
int LGBM_DatasetCreateFromSerializedReference(const void* ref_buffer,
                                              int32_t ref_buffer_size,
                                              int64_t num_row,
                                              int32_t num_classes,
                                              const char* parameters,
                                              DatasetHandle* out) {
  API_BEGIN();
  auto param = Config::Str2Map(parameters);
  Config config;
  config.Set(param);
  OMP_SET_NUM_THREADS(config.num_threads);
  DatasetLoader loader(config, nullptr, 1, nullptr);
  *out = loader.LoadFromSerializedReference(static_cast<const char*>(ref_buffer),
    static_cast<size_t>(ref_buffer_size),
    static_cast<data_size_t>(num_row),
    num_classes);
  API_END();
}

1052
1053
1054
1055
1056
int LGBM_DatasetInitStreaming(DatasetHandle dataset,
                              int32_t has_weights,
                              int32_t has_init_scores,
                              int32_t has_queries,
                              int32_t nclasses,
1057
1058
                              int32_t nthreads,
                              int32_t omp_max_threads) {
1059
1060
1061
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto num_data = p_dataset->num_data();
1062
  p_dataset->InitStreaming(num_data, has_weights, has_init_scores, has_queries, nclasses, nthreads, omp_max_threads);
1063
1064
1065
1066
  p_dataset->set_wait_for_manual_finish(true);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1067
int LGBM_DatasetPushRows(DatasetHandle dataset,
1068
1069
1070
1071
1072
                         const void* data,
                         int data_type,
                         int32_t nrow,
                         int32_t ncol,
                         int32_t start_row) {
Guolin Ke's avatar
Guolin Ke committed
1073
1074
1075
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto get_row_fun = RowFunctionFromDenseMatric(data, nrow, ncol, data_type, 1);
1076
1077
1078
  if (p_dataset->has_raw()) {
    p_dataset->ResizeRaw(p_dataset->num_numeric_features() + nrow);
  }
1079
  OMP_INIT_EX();
1080
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1081
  for (int i = 0; i < nrow; ++i) {
1082
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1083
1084
1085
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(tid, start_row + i, one_row);
1086
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1087
  }
1088
  OMP_THROW_EX();
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
  if (!p_dataset->wait_for_manual_finish() && (start_row + nrow == p_dataset->num_data())) {
    p_dataset->FinishLoad();
  }
  API_END();
}

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

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

1121
  OMP_INIT_EX();
1122
#pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
1123
1124
1125
  for (int i = 0; i < nrow; ++i) {
    OMP_LOOP_EX_BEGIN();
    // convert internal thread id to be unique based on external thread id
1126
    const int internal_tid = omp_get_thread_num() + (max_omp_threads * tid);
1127
1128
1129
1130
1131
1132
1133
1134
1135
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(internal_tid, start_row + i, one_row);
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();

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

  if (!p_dataset->wait_for_manual_finish() && (start_row + nrow == p_dataset->num_data())) {
Guolin Ke's avatar
Guolin Ke committed
1136
1137
1138
1139
1140
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1141
int LGBM_DatasetPushRowsByCSR(DatasetHandle dataset,
1142
1143
1144
1145
1146
1147
1148
1149
1150
                              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
1151
1152
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
1153
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
1154
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
1155
1156
1157
  if (p_dataset->has_raw()) {
    p_dataset->ResizeRaw(p_dataset->num_numeric_features() + nrow);
  }
1158
  OMP_INIT_EX();
1159
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1160
  for (int i = 0; i < nrow; ++i) {
1161
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1162
1163
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
1164
    p_dataset->PushOneRow(tid, static_cast<data_size_t>(start_row + i), one_row);
1165
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1166
  }
1167
  OMP_THROW_EX();
1168
  if (!p_dataset->wait_for_manual_finish() && (start_row + nrow == static_cast<int64_t>(p_dataset->num_data()))) {
Guolin Ke's avatar
Guolin Ke committed
1169
1170
    p_dataset->FinishLoad();
  }
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
  API_END();
}

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

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

1204
  OMP_INIT_EX();
1205
#pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
1206
1207
1208
  for (int i = 0; i < nrow; ++i) {
    OMP_LOOP_EX_BEGIN();
    // convert internal thread id to be unique based on external thread id
1209
    const int internal_tid = omp_get_thread_num() + (max_omp_threads * tid);
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(internal_tid, static_cast<data_size_t>(start_row + i), one_row);
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();

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

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

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

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

Guolin Ke's avatar
Guolin Ke committed
1238
int LGBM_DatasetCreateFromMat(const void* data,
1239
1240
1241
1242
1243
1244
1245
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
  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) {
1266
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1267
1268
  auto param = Config::Str2Map(parameters);
  Config config;
1269
  config.Set(param);
1270
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1271
  std::unique_ptr<Dataset> ret;
1272
1273
1274
1275
1276
1277
1278
1279
1280
  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));
  }
1281

Guolin Ke's avatar
Guolin Ke committed
1282
1283
  if (reference == nullptr) {
    // sample data first
1284
1285
    auto sample_indices = CreateSampleIndices(total_nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
1286
    std::vector<std::vector<double>> sample_values(ncol);
Guolin Ke's avatar
Guolin Ke committed
1287
    std::vector<std::vector<int>> sample_idx(ncol);
1288
1289
1290

    int offset = 0;
    int j = 0;
Guolin Ke's avatar
Guolin Ke committed
1291
    for (size_t i = 0; i < sample_indices.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1292
      auto idx = sample_indices[i];
1293
1294
1295
1296
      while ((idx - offset) >= nrow[j]) {
        offset += nrow[j];
        ++j;
      }
1297

1298
1299
1300
1301
1302
      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
1303
        }
Guolin Ke's avatar
Guolin Ke committed
1304
1305
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1306
    DatasetLoader loader(config, nullptr, 1, nullptr);
1307
1308
1309
1310
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             ncol,
                                             VectorSize<double>(sample_values).data(),
1311
1312
1313
                                             sample_cnt,
                                             total_nrow,
                                             total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1314
  } else {
1315
    ret.reset(new Dataset(total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1316
    ret->CreateValid(
1317
      reinterpret_cast<const Dataset*>(reference));
1318
1319
1320
    if (ret->has_raw()) {
      ret->ResizeRaw(total_nrow);
    }
Guolin Ke's avatar
Guolin Ke committed
1321
  }
1322
1323
1324
  int32_t start_row = 0;
  for (int j = 0; j < nmat; ++j) {
    OMP_INIT_EX();
1325
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
    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
1336
1337
  }
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1338
  *out = ret.release();
1339
  API_END();
1340
1341
}

Guolin Ke's avatar
Guolin Ke committed
1342
int LGBM_DatasetCreateFromCSR(const void* indptr,
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
                              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) {
1353
  API_BEGIN();
1354
1355
1356
1357
1358
  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
1359
1360
  auto param = Config::Str2Map(parameters);
  Config config;
1361
  config.Set(param);
1362
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1363
  std::unique_ptr<Dataset> ret;
1364
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1365
1366
1367
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
  if (reference == nullptr) {
    // sample data first
1368
1369
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1370
1371
    std::vector<std::vector<double>> sample_values(num_col);
    std::vector<std::vector<int>> sample_idx(num_col);
1372
1373
1374
1375
    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
1376
        CHECK_LT(inner_data.first, num_col);
Guolin Ke's avatar
Guolin Ke committed
1377
        if (std::fabs(inner_data.second) > kZeroThreshold || std::isnan(inner_data.second)) {
Guolin Ke's avatar
Guolin Ke committed
1378
1379
          sample_values[inner_data.first].emplace_back(inner_data.second);
          sample_idx[inner_data.first].emplace_back(static_cast<int>(i));
1380
1381
1382
        }
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1383
    DatasetLoader loader(config, nullptr, 1, nullptr);
1384
1385
1386
1387
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             static_cast<int>(num_col),
                                             VectorSize<double>(sample_values).data(),
1388
1389
1390
                                             sample_cnt,
                                             nrow,
                                             nrow));
1391
  } else {
1392
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1393
    ret->CreateValid(
1394
      reinterpret_cast<const Dataset*>(reference));
1395
1396
1397
    if (ret->has_raw()) {
      ret->ResizeRaw(nrow);
    }
1398
  }
1399
  OMP_INIT_EX();
1400
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
1401
  for (int i = 0; i < nindptr - 1; ++i) {
1402
    OMP_LOOP_EX_BEGIN();
1403
1404
1405
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    ret->PushOneRow(tid, i, one_row);
1406
    OMP_LOOP_EX_END();
1407
  }
1408
  OMP_THROW_EX();
1409
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1410
  *out = ret.release();
1411
  API_END();
1412
1413
}

1414
int LGBM_DatasetCreateFromCSRFunc(void* get_row_funptr,
1415
1416
1417
1418
1419
                                  int num_rows,
                                  int64_t num_col,
                                  const char* parameters,
                                  const DatasetHandle reference,
                                  DatasetHandle* out) {
1420
  API_BEGIN();
1421
1422
1423
1424
1425
  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.");
  }
1426
1427
1428
1429
  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);
1430
  OMP_SET_NUM_THREADS(config.num_threads);
1431
1432
1433
1434
  std::unique_ptr<Dataset> ret;
  int32_t nrow = num_rows;
  if (reference == nullptr) {
    // sample data first
1435
1436
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
1437
1438
1439
1440
1441
1442
1443
1444
    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
1445
        CHECK_LT(inner_data.first, num_col);
1446
1447
1448
1449
1450
1451
1452
        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);
1453
1454
1455
1456
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             static_cast<int>(num_col),
                                             VectorSize<double>(sample_values).data(),
1457
1458
1459
                                             sample_cnt,
                                             nrow,
                                             nrow));
1460
1461
1462
1463
  } else {
    ret.reset(new Dataset(nrow));
    ret->CreateValid(
      reinterpret_cast<const Dataset*>(reference));
1464
1465
1466
    if (ret->has_raw()) {
      ret->ResizeRaw(nrow);
    }
1467
  }
1468

1469
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1470
  std::vector<std::pair<int, double>> thread_buffer;
1471
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static) private(thread_buffer)
1472
1473
1474
  for (int i = 0; i < num_rows; ++i) {
    OMP_LOOP_EX_BEGIN();
    {
1475
      const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1476
1477
      get_row_fun(i, thread_buffer);
      ret->PushOneRow(tid, i, thread_buffer);
1478
1479
1480
1481
1482
1483
1484
1485
1486
    }
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();
  ret->FinishLoad();
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1487
int LGBM_DatasetCreateFromCSC(const void* col_ptr,
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
                              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) {
1498
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1499
1500
  auto param = Config::Str2Map(parameters);
  Config config;
1501
  config.Set(param);
1502
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1503
  std::unique_ptr<Dataset> ret;
Guolin Ke's avatar
Guolin Ke committed
1504
1505
1506
  int32_t nrow = static_cast<int32_t>(num_row);
  if (reference == nullptr) {
    // sample data first
1507
1508
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1509
    std::vector<std::vector<double>> sample_values(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
1510
    std::vector<std::vector<int>> sample_idx(ncol_ptr - 1);
1511
    OMP_INIT_EX();
1512
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1513
    for (int i = 0; i < static_cast<int>(sample_values.size()); ++i) {
1514
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1515
1516
1517
      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
1518
        if (std::fabs(val) > kZeroThreshold || std::isnan(val)) {
Guolin Ke's avatar
Guolin Ke committed
1519
1520
          sample_values[i].emplace_back(val);
          sample_idx[i].emplace_back(j);
Guolin Ke's avatar
Guolin Ke committed
1521
1522
        }
      }
1523
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1524
    }
1525
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1526
    DatasetLoader loader(config, nullptr, 1, nullptr);
1527
1528
1529
1530
    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(),
1531
1532
1533
                                             sample_cnt,
                                             nrow,
                                             nrow));
Guolin Ke's avatar
Guolin Ke committed
1534
  } else {
1535
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1536
    ret->CreateValid(
1537
      reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
1538
  }
1539
  OMP_INIT_EX();
1540
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1541
  for (int i = 0; i < ncol_ptr - 1; ++i) {
1542
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1543
    const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1544
    int feature_idx = ret->InnerFeatureIndex(i);
Guolin Ke's avatar
Guolin Ke committed
1545
    if (feature_idx < 0) { continue; }
Guolin Ke's avatar
Guolin Ke committed
1546
1547
    int group = ret->Feature2Group(feature_idx);
    int sub_feature = ret->Feture2SubFeature(feature_idx);
Guolin Ke's avatar
Guolin Ke committed
1548
    CSC_RowIterator col_it(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, i);
Guolin Ke's avatar
Guolin Ke committed
1549
1550
1551
1552
1553
1554
1555
1556
    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; }
1557
        ret->PushOneData(tid, row_idx, group, feature_idx, sub_feature, pair.second);
Guolin Ke's avatar
Guolin Ke committed
1558
1559
1560
1561
      }
    } else {
      for (int row_idx = 0; row_idx < nrow; ++row_idx) {
        auto val = col_it.Get(row_idx);
1562
        ret->PushOneData(tid, row_idx, group, feature_idx, sub_feature, val);
Guolin Ke's avatar
Guolin Ke committed
1563
      }
Guolin Ke's avatar
Guolin Ke committed
1564
    }
1565
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1566
  }
1567
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1568
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1569
  *out = ret.release();
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
  API_END();
}

int LGBM_DatasetCreateFromArrow(int64_t n_chunks,
                                const ArrowArray* chunks,
                                const ArrowSchema* schema,
                                const char* parameters,
                                const DatasetHandle reference,
                                DatasetHandle *out) {
  API_BEGIN();

  auto param = Config::Str2Map(parameters);
  Config config;
  config.Set(param);
  OMP_SET_NUM_THREADS(config.num_threads);

  std::unique_ptr<Dataset> ret;

  // Prepare the Arrow data
  ArrowTable table(n_chunks, chunks, schema);

  // Initialize the dataset
  if (reference == nullptr) {
    // If there is no reference dataset, we first sample indices
    auto sample_indices = CreateSampleIndices(static_cast<int32_t>(table.get_num_rows()), config);
    auto sample_count = static_cast<int>(sample_indices.size());
    std::vector<std::vector<double>> sample_values(table.get_num_columns());
    std::vector<std::vector<int>> sample_idx(table.get_num_columns());

    // Then, we obtain sample values by parallelizing across columns
    OMP_INIT_EX();
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
    for (int64_t j = 0; j < table.get_num_columns(); ++j) {
      OMP_LOOP_EX_BEGIN();

      // Values need to be copied from the record batches.
      sample_values[j].reserve(sample_indices.size());
      sample_idx[j].reserve(sample_indices.size());

      // The chunks are iterated over in the inner loop as columns can be treated independently.
      int last_idx = 0;
      int i = 0;
      auto it = table.get_column(j).begin<double>();
      for (auto idx : sample_indices) {
        std::advance(it, idx - last_idx);
        auto v = *it;
        if (std::fabs(v) > kZeroThreshold || std::isnan(v)) {
          sample_values[j].emplace_back(v);
          sample_idx[j].emplace_back(i);
        }
        last_idx = idx;
        i++;
      }
      OMP_LOOP_EX_END();
    }
    OMP_THROW_EX();

    // Finally, we initialize a loader from the sampled values
    DatasetLoader loader(config, nullptr, 1, nullptr);
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             table.get_num_columns(),
                                             VectorSize<double>(sample_values).data(),
                                             sample_count,
                                             table.get_num_rows(),
                                             table.get_num_rows()));
  } else {
    ret.reset(new Dataset(static_cast<data_size_t>(table.get_num_rows())));
    ret->CreateValid(reinterpret_cast<const Dataset*>(reference));
    if (ret->has_raw()) {
      ret->ResizeRaw(static_cast<int>(table.get_num_rows()));
    }
  }

  // After sampling and properly initializing all bins, we can add our data to the dataset. Here,
  // we parallelize across rows.
  OMP_INIT_EX();
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
  for (int64_t j = 0; j < table.get_num_columns(); ++j) {
    OMP_LOOP_EX_BEGIN();
    const int tid = omp_get_thread_num();
    data_size_t idx = 0;
    auto column = table.get_column(j);
    for (auto it = column.begin<double>(), end = column.end<double>(); it != end; ++it) {
      ret->PushOneValue(tid, idx++, j, *it);
    }
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();

  ret->FinishLoad();
  *out = ret.release();
1662
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1663
1664
}

Guolin Ke's avatar
Guolin Ke committed
1665
int LGBM_DatasetGetSubset(
1666
  const DatasetHandle handle,
wxchan's avatar
wxchan committed
1667
1668
1669
  const int32_t* used_row_indices,
  int32_t num_used_row_indices,
  const char* parameters,
Guolin Ke's avatar
typo  
Guolin Ke committed
1670
  DatasetHandle* out) {
wxchan's avatar
wxchan committed
1671
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1672
1673
  auto param = Config::Str2Map(parameters);
  Config config;
1674
  config.Set(param);
1675
  OMP_SET_NUM_THREADS(config.num_threads);
1676
  auto full_dataset = reinterpret_cast<const Dataset*>(handle);
1677
  CHECK_GT(num_used_row_indices, 0);
1678
1679
  const int32_t lower = 0;
  const int32_t upper = full_dataset->num_data() - 1;
1680
  CheckElementsIntervalClosed(used_row_indices, lower, upper, num_used_row_indices, "Used indices of subset");
1681
1682
1683
  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
1684
  auto ret = std::unique_ptr<Dataset>(new Dataset(num_used_row_indices));
1685
  ret->CopyFeatureMapperFrom(full_dataset);
1686
  ret->CopySubrow(full_dataset, used_row_indices, num_used_row_indices, true);
wxchan's avatar
wxchan committed
1687
1688
1689
1690
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1691
int LGBM_DatasetSetFeatureNames(
Guolin Ke's avatar
typo  
Guolin Ke committed
1692
  DatasetHandle handle,
Guolin Ke's avatar
Guolin Ke committed
1693
  const char** feature_names,
Guolin Ke's avatar
Guolin Ke committed
1694
  int num_feature_names) {
Guolin Ke's avatar
Guolin Ke committed
1695
1696
1697
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  std::vector<std::string> feature_names_str;
Guolin Ke's avatar
Guolin Ke committed
1698
  for (int i = 0; i < num_feature_names; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1699
1700
1701
1702
1703
1704
    feature_names_str.emplace_back(feature_names[i]);
  }
  dataset->set_feature_names(feature_names_str);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1705
int LGBM_DatasetGetFeatureNames(
1706
1707
1708
1709
1710
1711
    DatasetHandle handle,
    const int len,
    int* num_feature_names,
    const size_t buffer_len,
    size_t* out_buffer_len,
    char** feature_names) {
1712
  API_BEGIN();
1713
  *out_buffer_len = 0;
1714
1715
  auto dataset = reinterpret_cast<Dataset*>(handle);
  auto inside_feature_name = dataset->feature_names();
Guolin Ke's avatar
Guolin Ke committed
1716
1717
  *num_feature_names = static_cast<int>(inside_feature_name.size());
  for (int i = 0; i < *num_feature_names; ++i) {
1718
1719
1720
1721
1722
    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);
1723
1724
1725
1726
  }
  API_END();
}

1727
1728
1729
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1730
int LGBM_DatasetFree(DatasetHandle handle) {
1731
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1732
  delete reinterpret_cast<Dataset*>(handle);
1733
  API_END();
1734
1735
}

Guolin Ke's avatar
Guolin Ke committed
1736
int LGBM_DatasetSaveBinary(DatasetHandle handle,
1737
                           const char* filename) {
1738
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1739
1740
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->SaveBinaryFile(filename);
1741
  API_END();
1742
1743
}

1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
int LGBM_DatasetSerializeReferenceToBinary(DatasetHandle handle,
                                           ByteBufferHandle* out,
                                           int32_t* out_len) {
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  std::unique_ptr<LightGBM::ByteBuffer> ret;
  ret.reset(new LightGBM::ByteBuffer());
  dataset->SerializeReference(ret.get());
  *out_len = static_cast<int32_t>(ret->GetSize());
  *out = ret.release();
  API_END();
}

1757
1758
1759
1760
1761
1762
1763
1764
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
1765
int LGBM_DatasetSetField(DatasetHandle handle,
1766
1767
1768
1769
                         const char* field_name,
                         const void* field_data,
                         int num_element,
                         int type) {
1770
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1771
  auto dataset = reinterpret_cast<Dataset*>(handle);
1772
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1773
  if (type == C_API_DTYPE_FLOAT32) {
Guolin Ke's avatar
Guolin Ke committed
1774
    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
1775
  } else if (type == C_API_DTYPE_INT32) {
Guolin Ke's avatar
Guolin Ke committed
1776
    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
1777
1778
  } 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));
1779
  }
1780
  if (!is_success) { Log::Fatal("Input data type error or field not found"); }
1781
  API_END();
1782
1783
}

1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
int LGBM_DatasetSetFieldFromArrow(DatasetHandle handle,
                                  const char* field_name,
                                  int64_t n_chunks,
                                  const ArrowArray* chunks,
                                  const ArrowSchema* schema) {
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  ArrowChunkedArray ca(n_chunks, chunks, schema);
  auto is_success = dataset->SetFieldFromArrow(field_name, ca);
  if (!is_success) {
    Log::Fatal("Input field is not supported");
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1799
int LGBM_DatasetGetField(DatasetHandle handle,
1800
1801
1802
1803
                         const char* field_name,
                         int* out_len,
                         const void** out_ptr,
                         int* out_type) {
1804
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1805
  auto dataset = reinterpret_cast<Dataset*>(handle);
1806
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1807
  if (dataset->GetFloatField(field_name, out_len, reinterpret_cast<const float**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1808
    *out_type = C_API_DTYPE_FLOAT32;
1809
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1810
  } else if (dataset->GetIntField(field_name, out_len, reinterpret_cast<const int**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1811
    *out_type = C_API_DTYPE_INT32;
1812
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1813
1814
1815
  } 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
1816
  }
1817
  if (!is_success) { Log::Fatal("Field not found"); }
wxchan's avatar
wxchan committed
1818
  if (*out_ptr == nullptr) { *out_len = 0; }
1819
  API_END();
1820
1821
}

1822
int LGBM_DatasetUpdateParamChecking(const char* old_parameters, const char* new_parameters) {
1823
  API_BEGIN();
1824
1825
1826
1827
1828
  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);
1829
1830
1831
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1832
int LGBM_DatasetGetNumData(DatasetHandle handle,
1833
                           int* out) {
1834
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1835
1836
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_data();
1837
  API_END();
1838
1839
}

Guolin Ke's avatar
Guolin Ke committed
1840
int LGBM_DatasetGetNumFeature(DatasetHandle handle,
1841
                              int* out) {
1842
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1843
1844
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_total_features();
1845
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1846
}
1847

1848
1849
1850
1851
1852
int LGBM_DatasetGetFeatureNumBin(DatasetHandle handle,
                                 int feature,
                                 int* out) {
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
1853
1854
1855
1856
1857
  int num_features = dataset->num_total_features();
  if (feature < 0 || feature >= num_features) {
    Log::Fatal("Tried to retrieve number of bins for feature index %d, "
               "but the valid feature indices are [0, %d].", feature, num_features - 1);
  }
1858
1859
1860
1861
1862
1863
1864
1865
1866
  int inner_idx = dataset->InnerFeatureIndex(feature);
  if (inner_idx >= 0) {
    *out = dataset->FeatureNumBin(inner_idx);
  } else {
    *out = 0;
  }
  API_END();
}

1867
1868
1869
1870
1871
int LGBM_DatasetAddFeaturesFrom(DatasetHandle target,
                                DatasetHandle source) {
  API_BEGIN();
  auto target_d = reinterpret_cast<Dataset*>(target);
  auto source_d = reinterpret_cast<Dataset*>(source);
1872
  target_d->AddFeaturesFrom(source_d);
1873
1874
1875
  API_END();
}

1876
1877
// ---- start of booster

Guolin Ke's avatar
Guolin Ke committed
1878
int LGBM_BoosterCreate(const DatasetHandle train_data,
1879
1880
                       const char* parameters,
                       BoosterHandle* out) {
1881
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1882
  const Dataset* p_train_data = reinterpret_cast<const Dataset*>(train_data);
wxchan's avatar
wxchan committed
1883
1884
  auto ret = std::unique_ptr<Booster>(new Booster(p_train_data, parameters));
  *out = ret.release();
1885
  API_END();
1886
1887
}

Guolin Ke's avatar
Guolin Ke committed
1888
int LGBM_BoosterCreateFromModelfile(
1889
  const char* filename,
Guolin Ke's avatar
Guolin Ke committed
1890
  int* out_num_iterations,
1891
  BoosterHandle* out) {
1892
  API_BEGIN();
wxchan's avatar
wxchan committed
1893
  auto ret = std::unique_ptr<Booster>(new Booster(filename));
Guolin Ke's avatar
Guolin Ke committed
1894
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
wxchan's avatar
wxchan committed
1895
  *out = ret.release();
1896
  API_END();
1897
1898
}

Guolin Ke's avatar
Guolin Ke committed
1899
int LGBM_BoosterLoadModelFromString(
1900
1901
1902
1903
  const char* model_str,
  int* out_num_iterations,
  BoosterHandle* out) {
  API_BEGIN();
wxchan's avatar
wxchan committed
1904
  auto ret = std::unique_ptr<Booster>(new Booster(nullptr));
1905
1906
1907
1908
1909
1910
  ret->LoadModelFromString(model_str);
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
  *out = ret.release();
  API_END();
}

1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
int LGBM_BoosterGetLoadedParam(
  BoosterHandle handle,
  int64_t buffer_len,
  int64_t* out_len,
  char* out_str) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  std::string params = ref_booster->GetBoosting()->GetLoadedParam();
  *out_len = static_cast<int64_t>(params.size()) + 1;
  if (*out_len <= buffer_len) {
    std::memcpy(out_str, params.c_str(), *out_len);
  }
  API_END();
}

1926
1927
1928
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1929
int LGBM_BoosterFree(BoosterHandle handle) {
1930
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1931
  delete reinterpret_cast<Booster*>(handle);
1932
  API_END();
1933
1934
}

1935
int LGBM_BoosterShuffleModels(BoosterHandle handle, int start_iter, int end_iter) {
1936
1937
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1938
  ref_booster->ShuffleModels(start_iter, end_iter);
1939
1940
1941
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1942
int LGBM_BoosterMerge(BoosterHandle handle,
1943
                      BoosterHandle other_handle) {
wxchan's avatar
wxchan committed
1944
1945
1946
1947
1948
1949
1950
  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
1951
int LGBM_BoosterAddValidData(BoosterHandle handle,
1952
                             const DatasetHandle valid_data) {
wxchan's avatar
wxchan committed
1953
1954
1955
1956
1957
1958
1959
  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
1960
int LGBM_BoosterResetTrainingData(BoosterHandle handle,
1961
                                  const DatasetHandle train_data) {
wxchan's avatar
wxchan committed
1962
1963
1964
1965
1966
1967
1968
  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
1969
int LGBM_BoosterResetParameter(BoosterHandle handle, const char* parameters) {
wxchan's avatar
wxchan committed
1970
1971
1972
1973
1974
1975
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->ResetConfig(parameters);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1976
int LGBM_BoosterGetNumClasses(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1977
1978
1979
1980
1981
1982
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->NumberOfClasses();
  API_END();
}

1983
int LGBM_BoosterGetLinear(BoosterHandle handle, int* out) {
1984
1985
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1986
1987
1988
1989
1990
  if (ref_booster->GetBoosting()->IsLinear()) {
    *out = 1;
  } else {
    *out = 0;
  }
1991
1992
1993
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1994
1995
1996
1997
1998
1999
2000
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
2001
int LGBM_BoosterUpdateOneIter(BoosterHandle handle, int* is_finished) {
2002
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2003
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2004
2005
2006
2007
2008
  if (ref_booster->TrainOneIter()) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
2009
  API_END();
2010
2011
}

Guolin Ke's avatar
Guolin Ke committed
2012
int LGBM_BoosterUpdateOneIterCustom(BoosterHandle handle,
2013
2014
2015
                                    const float* grad,
                                    const float* hess,
                                    int* is_finished) {
2016
  API_BEGIN();
2017
  #ifdef SCORE_T_USE_DOUBLE
2018
2019
2020
2021
  (void) handle;       // UNUSED VARIABLE
  (void) grad;         // UNUSED VARIABLE
  (void) hess;         // UNUSED VARIABLE
  (void) is_finished;  // UNUSED VARIABLE
2022
  Log::Fatal("Don't support custom loss function when SCORE_T_USE_DOUBLE is enabled");
2023
  #else
2024
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2025
2026
2027
2028
2029
  if (ref_booster->TrainOneIter(grad, hess)) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
2030
  #endif
2031
  API_END();
2032
2033
}

Guolin Ke's avatar
Guolin Ke committed
2034
int LGBM_BoosterRollbackOneIter(BoosterHandle handle) {
wxchan's avatar
wxchan committed
2035
2036
2037
2038
2039
2040
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->RollbackOneIter();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2041
int LGBM_BoosterGetCurrentIteration(BoosterHandle handle, int* out_iteration) {
wxchan's avatar
wxchan committed
2042
2043
2044
2045
2046
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_iteration = ref_booster->GetBoosting()->GetCurrentIteration();
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
2047

2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
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
2062
int LGBM_BoosterGetEvalCounts(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
2063
2064
2065
2066
2067
2068
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetEvalCounts();
  API_END();
}

2069
2070
2071
2072
2073
2074
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
2075
2076
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2077
  *out_len = ref_booster->GetEvalNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
2078
2079
2080
  API_END();
}

2081
2082
2083
2084
2085
2086
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
2087
2088
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2089
  *out_len = ref_booster->GetFeatureNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
2090
2091
2092
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2093
int LGBM_BoosterGetNumFeature(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
2094
2095
2096
2097
2098
2099
  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
2100
int LGBM_BoosterGetEval(BoosterHandle handle,
2101
2102
2103
                        int data_idx,
                        int* out_len,
                        double* out_results) {
2104
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2105
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2106
  auto boosting = ref_booster->GetBoosting();
wxchan's avatar
wxchan committed
2107
  auto result_buf = boosting->GetEvalAt(data_idx);
Guolin Ke's avatar
Guolin Ke committed
2108
  *out_len = static_cast<int>(result_buf.size());
2109
  for (size_t i = 0; i < result_buf.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
2110
    (out_results)[i] = static_cast<double>(result_buf[i]);
2111
  }
2112
  API_END();
2113
2114
}

Guolin Ke's avatar
Guolin Ke committed
2115
int LGBM_BoosterGetNumPredict(BoosterHandle handle,
2116
2117
                              int data_idx,
                              int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
2118
2119
2120
2121
2122
2123
  API_BEGIN();
  auto boosting = reinterpret_cast<Booster*>(handle)->GetBoosting();
  *out_len = boosting->GetNumPredictAt(data_idx);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2124
int LGBM_BoosterGetPredict(BoosterHandle handle,
2125
2126
2127
                           int data_idx,
                           int64_t* out_len,
                           double* out_result) {
2128
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2129
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2130
  ref_booster->GetPredictAt(data_idx, out_result, out_len);
2131
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2132
2133
}

Guolin Ke's avatar
Guolin Ke committed
2134
int LGBM_BoosterPredictForFile(BoosterHandle handle,
2135
2136
2137
                               const char* data_filename,
                               int data_has_header,
                               int predict_type,
2138
                               int start_iteration,
2139
                               int num_iteration,
2140
                               const char* parameter,
2141
                               const char* result_filename) {
2142
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2143
2144
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2145
  config.Set(param);
2146
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
2147
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2148
  ref_booster->Predict(start_iteration, num_iteration, predict_type, data_filename, data_has_header,
Guolin Ke's avatar
Guolin Ke committed
2149
                       config, result_filename);
2150
  API_END();
2151
2152
}

Guolin Ke's avatar
Guolin Ke committed
2153
int LGBM_BoosterCalcNumPredict(BoosterHandle handle,
2154
2155
                               int num_row,
                               int predict_type,
2156
                               int start_iteration,
2157
2158
                               int num_iteration,
                               int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
2159
2160
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2161
  *out_len = static_cast<int64_t>(num_row) * ref_booster->GetBoosting()->NumPredictOneRow(start_iteration,
2162
    num_iteration, predict_type == C_API_PREDICT_LEAF_INDEX, predict_type == C_API_PREDICT_CONTRIB);
Guolin Ke's avatar
Guolin Ke committed
2163
2164
2165
  API_END();
}

2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
/*!
 * \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,
2179
             const int predict_type_,
2180
             const int data_type_,
2181
             const int32_t num_cols) : booster(booster_ptr), predict_type(predict_type_), data_type(data_type_), ncol(num_cols) {
2182
2183
2184
2185
2186
    config.Set(Config::Str2Map(parameter));
  }

  Booster* const booster;
  Config config;
2187
  const int predict_type;
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
  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
2198
int LGBM_BoosterPredictForCSR(BoosterHandle handle,
2199
2200
2201
2202
2203
2204
2205
                              const void* indptr,
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
2206
                              int64_t num_col,
2207
                              int predict_type,
2208
                              int start_iteration,
2209
                              int num_iteration,
2210
                              const char* parameter,
2211
2212
                              int64_t* out_len,
                              double* out_result) {
2213
  API_BEGIN();
2214
2215
2216
2217
2218
  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
2219
2220
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2221
  config.Set(param);
2222
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
2223
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2224
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
2225
  int nrow = static_cast<int>(nindptr - 1);
2226
  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
2227
                       config, out_result, out_len);
2228
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2229
}
2230

2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
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,
2241
                                    int start_iteration,
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
                                    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);
2254
  OMP_SET_NUM_THREADS(config.num_threads);
2255
2256
2257
2258
2259
2260
2261
2262
  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;
2263
    ref_booster->PredictSparseCSR(start_iteration, num_iteration, predict_type, nrow, static_cast<int>(num_col_or_row), get_row_fun,
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
                                  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;
    };
2287
    ref_booster->PredictSparseCSC(start_iteration, num_iteration, predict_type, num_col_or_row, ncol, get_row_fun, config,
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
                                  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) {
2298
    delete[] reinterpret_cast<int32_t*>(indptr);
2299
  } else if (indptr_type == C_API_DTYPE_INT64) {
2300
    delete[] reinterpret_cast<int64_t*>(indptr);
2301
2302
2303
  } else {
    Log::Fatal("Unknown indptr type in LGBM_BoosterFreePredictSparse");
  }
2304
  delete[] indices;
2305
  if (data_type == C_API_DTYPE_FLOAT32) {
2306
    delete[] reinterpret_cast<float*>(data);
2307
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2308
    delete[] reinterpret_cast<double*>(data);
2309
2310
2311
2312
2313
2314
  } else {
    Log::Fatal("Unknown data type in LGBM_BoosterFreePredictSparse");
  }
  API_END();
}

2315
int LGBM_BoosterPredictForCSRSingleRow(BoosterHandle handle,
2316
2317
2318
2319
2320
2321
2322
                                       const void* indptr,
                                       int indptr_type,
                                       const int32_t* indices,
                                       const void* data,
                                       int data_type,
                                       int64_t nindptr,
                                       int64_t nelem,
2323
                                       int64_t num_col,
2324
                                       int predict_type,
2325
                                       int start_iteration,
2326
2327
2328
2329
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2330
  API_BEGIN();
2331
2332
2333
2334
2335
  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.");
  }
2336
2337
2338
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
2339
  OMP_SET_NUM_THREADS(config.num_threads);
2340
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2341
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
2342
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
2343
  ref_booster->PredictSingleRow(predict_type, static_cast<int32_t>(num_col), get_row_fun, config, out_result, out_len);
2344
2345
2346
  API_END();
}

2347
int LGBM_BoosterPredictForCSRSingleRowFastInit(BoosterHandle handle,
2348
                                               const int predict_type,
2349
                                               const int start_iteration,
2350
                                               const int num_iteration,
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
                                               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,
2365
    predict_type,
2366
2367
2368
    data_type,
    static_cast<int32_t>(num_col)));

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

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

2373
2374
2375
2376
2377
2378
  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

int LGBM_BoosterPredictForCSRSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* indptr,
2379
                                           const int indptr_type,
2380
2381
                                           const int32_t* indices,
                                           const void* data,
2382
2383
                                           const int64_t nindptr,
                                           const int64_t nelem,
2384
2385
2386
2387
2388
                                           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);
2389
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
2390
2391
2392
2393
                                        get_row_fun, fastConfig->config, out_result, out_len);
  API_END();
}

2394

Guolin Ke's avatar
Guolin Ke committed
2395
int LGBM_BoosterPredictForCSC(BoosterHandle handle,
2396
2397
2398
2399
2400
2401
2402
2403
2404
                              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,
2405
                              int start_iteration,
2406
                              int num_iteration,
2407
                              const char* parameter,
2408
2409
                              int64_t* out_len,
                              double* out_result) {
Guolin Ke's avatar
Guolin Ke committed
2410
2411
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2412
2413
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2414
  config.Set(param);
2415
  OMP_SET_NUM_THREADS(config.num_threads);
2416
  int num_threads = OMP_NUM_THREADS();
Guolin Ke's avatar
Guolin Ke committed
2417
  int ncol = static_cast<int>(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
2418
2419
2420
2421
2422
  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
2423
2424
  }
  std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun =
Guolin Ke's avatar
Guolin Ke committed
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
      [&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;
      };
2437
  ref_booster->Predict(start_iteration, num_iteration, predict_type, static_cast<int>(num_row), ncol, get_row_fun, config,
cbecker's avatar
cbecker committed
2438
                       out_result, out_len);
Guolin Ke's avatar
Guolin Ke committed
2439
2440
2441
  API_END();
}

2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
int LGBM_BoosterValidateFeatureNames(BoosterHandle handle,
                                     const char** data_names,
                                     int data_num_features) {
  API_BEGIN();
  int booster_num_features;
  size_t out_buffer_len;
  LGBM_BoosterGetFeatureNames(handle, 0, &booster_num_features, 0, &out_buffer_len, nullptr);
  if (booster_num_features != data_num_features) {
    Log::Fatal("Model was trained on %d features, but got %d input features to predict.", booster_num_features, data_num_features);
  }
  std::vector<std::vector<char>> tmp_names(booster_num_features, std::vector<char>(out_buffer_len));
  std::vector<char*> booster_names = Vector2Ptr(&tmp_names);
  LGBM_BoosterGetFeatureNames(handle, data_num_features, &booster_num_features, out_buffer_len, &out_buffer_len, booster_names.data());
  for (int i = 0; i < booster_num_features; ++i) {
    if (strcmp(data_names[i], booster_names[i]) != 0) {
      Log::Fatal("Expected '%s' at position %d but found '%s'", booster_names[i], i, data_names[i]);
    }
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2463
int LGBM_BoosterPredictForMat(BoosterHandle handle,
2464
2465
2466
2467
2468
2469
                              const void* data,
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              int predict_type,
2470
                              int start_iteration,
2471
                              int num_iteration,
2472
                              const char* parameter,
2473
2474
                              int64_t* out_len,
                              double* out_result) {
2475
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2476
2477
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2478
  config.Set(param);
2479
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
2480
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2481
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, nrow, ncol, data_type, is_row_major);
2482
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
2483
                       config, out_result, out_len);
2484
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2485
}
2486

2487
int LGBM_BoosterPredictForMatSingleRow(BoosterHandle handle,
2488
2489
2490
2491
2492
                                       const void* data,
                                       int data_type,
                                       int32_t ncol,
                                       int is_row_major,
                                       int predict_type,
2493
                                       int start_iteration,
2494
2495
2496
2497
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2498
2499
2500
2501
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
2502
  OMP_SET_NUM_THREADS(config.num_threads);
2503
2504
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, ncol, data_type, is_row_major);
2505
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
2506
  ref_booster->PredictSingleRow(predict_type, ncol, get_row_fun, config, out_result, out_len);
2507
2508
2509
  API_END();
}

2510
int LGBM_BoosterPredictForMatSingleRowFastInit(BoosterHandle handle,
2511
                                               const int predict_type,
2512
                                               const int start_iteration,
2513
                                               const int num_iteration,
2514
2515
2516
2517
2518
2519
2520
2521
                                               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,
2522
    predict_type,
2523
2524
2525
    data_type,
    ncol));

2526
  OMP_SET_NUM_THREADS(fastConfig_ptr->config.num_threads);
2527

2528
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2529

2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
  *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);
2542
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
2543
2544
2545
2546
2547
                                        get_row_fun, fastConfig->config,
                                        out_result, out_len);
  API_END();
}

2548

2549
2550
2551
2552
2553
2554
int LGBM_BoosterPredictForMats(BoosterHandle handle,
                               const void** data,
                               int data_type,
                               int32_t nrow,
                               int32_t ncol,
                               int predict_type,
2555
                               int start_iteration,
2556
2557
2558
2559
2560
2561
2562
2563
                               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);
2564
  OMP_SET_NUM_THREADS(config.num_threads);
2565
2566
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseRows(data, ncol, data_type);
2567
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, out_result, out_len);
2568
2569
2570
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2571
int LGBM_BoosterSaveModel(BoosterHandle handle,
2572
                          int start_iteration,
2573
                          int num_iteration,
2574
                          int feature_importance_type,
2575
                          const char* filename) {
2576
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2577
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2578
2579
  ref_booster->SaveModelToFile(start_iteration, num_iteration,
                               feature_importance_type, filename);
wxchan's avatar
wxchan committed
2580
2581
2582
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2583
int LGBM_BoosterSaveModelToString(BoosterHandle handle,
2584
                                  int start_iteration,
2585
                                  int num_iteration,
2586
                                  int feature_importance_type,
2587
                                  int64_t buffer_len,
2588
                                  int64_t* out_len,
2589
                                  char* out_str) {
2590
2591
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2592
2593
  std::string model = ref_booster->SaveModelToString(
      start_iteration, num_iteration, feature_importance_type);
2594
  *out_len = static_cast<int64_t>(model.size()) + 1;
2595
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2596
    std::memcpy(out_str, model.c_str(), *out_len);
2597
2598
2599
2600
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2601
int LGBM_BoosterDumpModel(BoosterHandle handle,
2602
                          int start_iteration,
2603
                          int num_iteration,
2604
                          int feature_importance_type,
2605
2606
                          int64_t buffer_len,
                          int64_t* out_len,
2607
                          char* out_str) {
wxchan's avatar
wxchan committed
2608
2609
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2610
2611
  std::string model = ref_booster->DumpModel(start_iteration, num_iteration,
                                             feature_importance_type);
2612
  *out_len = static_cast<int64_t>(model.size()) + 1;
wxchan's avatar
wxchan committed
2613
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2614
    std::memcpy(out_str, model.c_str(), *out_len);
wxchan's avatar
wxchan committed
2615
  }
2616
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2617
}
2618

Guolin Ke's avatar
Guolin Ke committed
2619
int LGBM_BoosterGetLeafValue(BoosterHandle handle,
2620
2621
2622
                             int tree_idx,
                             int leaf_idx,
                             double* out_val) {
Guolin Ke's avatar
Guolin Ke committed
2623
2624
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2625
  *out_val = static_cast<double>(ref_booster->GetLeafValue(tree_idx, leaf_idx));
Guolin Ke's avatar
Guolin Ke committed
2626
2627
2628
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2629
int LGBM_BoosterSetLeafValue(BoosterHandle handle,
2630
2631
2632
                             int tree_idx,
                             int leaf_idx,
                             double val) {
Guolin Ke's avatar
Guolin Ke committed
2633
2634
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2635
  ref_booster->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
2636
2637
2638
  API_END();
}

2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
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();
}

2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
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();
}

2670
2671
2672
2673
2674
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
2675
  Config config;
2676
  config.machines = RemoveQuotationSymbol(std::string(machines));
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
  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();
}

2692
2693
2694
int LGBM_NetworkInitWithFunctions(int num_machines, int rank,
                                  void* reduce_scatter_ext_fun,
                                  void* allgather_ext_fun) {
ww's avatar
ww committed
2695
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2696
  if (num_machines > 1) {
2697
    Network::Init(num_machines, rank, (ReduceScatterFunction)reduce_scatter_ext_fun, (AllgatherFunction)allgather_ext_fun);
ww's avatar
ww committed
2698
2699
2700
  }
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
2701

2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
2714
2715
2716
2717
2718
int LGBM_SetMaxThreads(int num_threads) {
  API_BEGIN();
  if (num_threads <= 0) {
    LGBM_MAX_NUM_THREADS = -1;
  } else {
    LGBM_MAX_NUM_THREADS = num_threads;
  }
  API_END();
}

int LGBM_GetMaxThreads(int* out) {
  API_BEGIN();
  *out = LGBM_MAX_NUM_THREADS;
  API_END();
}


Guolin Ke's avatar
Guolin Ke committed
2719
// ---- start of some help functions
2720

2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
2742
2743
2744
2745

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

2746
2747
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
2748
  if (data_type == C_API_DTYPE_FLOAT32) {
2749
    return RowFunctionFromDenseMatric_helper<float>(data, num_row, num_col, is_row_major);
Guolin Ke's avatar
Guolin Ke committed
2750
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2751
    return RowFunctionFromDenseMatric_helper<double>(data, num_row, num_col, is_row_major);
2752
  }
2753
  Log::Fatal("Unknown data type in RowFunctionFromDenseMatric");
2754
  return nullptr;
2755
2756
2757
2758
}

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
2759
2760
  auto inner_function = RowFunctionFromDenseMatric(data, num_row, num_col, data_type, is_row_major);
  if (inner_function != nullptr) {
2761
    return [inner_function] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2762
2763
      auto raw_values = inner_function(row_idx);
      std::vector<std::pair<int, double>> ret;
Guolin Ke's avatar
Guolin Ke committed
2764
      ret.reserve(raw_values.size());
Guolin Ke's avatar
Guolin Ke committed
2765
      for (int i = 0; i < static_cast<int>(raw_values.size()); ++i) {
Guolin Ke's avatar
Guolin Ke committed
2766
        if (std::fabs(raw_values[i]) > kZeroThreshold || std::isnan(raw_values[i])) {
Guolin Ke's avatar
Guolin Ke committed
2767
          ret.emplace_back(i, raw_values[i]);
2768
        }
Guolin Ke's avatar
Guolin Ke committed
2769
2770
2771
      }
      return ret;
    };
2772
  }
Guolin Ke's avatar
Guolin Ke committed
2773
  return nullptr;
2774
2775
}

2776
2777
2778
2779
2780
2781
2782
// 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
2783
    ret.reserve(raw_values.size());
2784
2785
2786
2787
2788
2789
2790
2791
2792
    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;
  };
}

2793
2794
2795
2796
2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
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;
  };
}

2812
2813
template<typename T>
std::function<std::vector<std::pair<int, double>>(T idx)>
2814
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
2815
2816
  if (data_type == C_API_DTYPE_FLOAT32) {
    if (indptr_type == C_API_DTYPE_INT32) {
2817
     return RowFunctionFromCSR_helper<T, float, int32_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2818
    } else if (indptr_type == C_API_DTYPE_INT64) {
2819
     return RowFunctionFromCSR_helper<T, float, int64_t>(indptr, indices, data);
2820
    }
Guolin Ke's avatar
Guolin Ke committed
2821
2822
  } else if (data_type == C_API_DTYPE_FLOAT64) {
    if (indptr_type == C_API_DTYPE_INT32) {
2823
     return RowFunctionFromCSR_helper<T, double, int32_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2824
    } else if (indptr_type == C_API_DTYPE_INT64) {
2825
     return RowFunctionFromCSR_helper<T, double, int64_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2826
2827
    }
  }
2828
  Log::Fatal("Unknown data type in RowFunctionFromCSR");
2829
  return nullptr;
2830
2831
}

2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850


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
2851
std::function<std::pair<int, double>(int idx)>
2852
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
2853
  CHECK(col_idx < ncol_ptr && col_idx >= 0);
Guolin Ke's avatar
Guolin Ke committed
2854
2855
  if (data_type == C_API_DTYPE_FLOAT32) {
    if (col_ptr_type == C_API_DTYPE_INT32) {
2856
      return IterateFunctionFromCSC_helper<float, int32_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2857
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2858
      return IterateFunctionFromCSC_helper<float, int64_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2859
    }
Guolin Ke's avatar
Guolin Ke committed
2860
2861
  } else if (data_type == C_API_DTYPE_FLOAT64) {
    if (col_ptr_type == C_API_DTYPE_INT32) {
2862
      return IterateFunctionFromCSC_helper<double, int32_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2863
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2864
      return IterateFunctionFromCSC_helper<double, int64_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2865
2866
    }
  }
2867
  Log::Fatal("Unknown data type in CSC matrix");
2868
  return nullptr;
2869
2870
}

Guolin Ke's avatar
Guolin Ke committed
2871
CSC_RowIterator::CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
2872
                                 const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx) {
Guolin Ke's avatar
Guolin Ke committed
2873
2874
2875
2876
2877
2878
2879
2880
2881
  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;
2882
    }
Guolin Ke's avatar
Guolin Ke committed
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
    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;
2900
    }
Guolin Ke's avatar
Guolin Ke committed
2901
2902
2903
    return ret;
  } else {
    return std::make_pair(-1, 0.0);
2904
  }
Guolin Ke's avatar
Guolin Ke committed
2905
}