c_api.cpp 94.9 KB
Newer Older
1
2
3
4
/*!
 * Copyright (c) 2016 Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See LICENSE file in the project root for license information.
 */
Guolin Ke's avatar
Guolin Ke committed
5
#include <LightGBM/c_api.h>
Guolin Ke's avatar
Guolin Ke committed
6

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

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

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

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

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

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

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

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

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

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

Guolin Ke's avatar
Guolin Ke committed
65
  SingleRowPredictor(int predict_type, Boosting* boosting, const Config& config, int iter) {
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
    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;
    } else {
      is_raw_score = false;
    }
    early_stop_ = config.pred_early_stop;
    early_stop_freq_ = config.pred_early_stop_freq;
    early_stop_margin_ = config.pred_early_stop_margin;
    iter_ = iter;
Guolin Ke's avatar
Guolin Ke committed
82
    predictor_.reset(new Predictor(boosting, iter_, is_raw_score, is_predict_leaf, predict_contrib,
83
                                   early_stop_, early_stop_freq_, early_stop_margin_));
Guolin Ke's avatar
Guolin Ke committed
84
    num_pred_in_one_row = boosting->NumPredictOneRow(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
119
120
    if (config_.num_threads > 0) {
      omp_set_num_threads(config_.num_threads);
    }
Guolin Ke's avatar
Guolin Ke committed
121
    // create boosting
Guolin Ke's avatar
Guolin Ke committed
122
    if (config_.input_model.size() > 0) {
123
124
      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
125
    }
Guolin Ke's avatar
Guolin Ke committed
126

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

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

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

  ~Booster() {
  }
150

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

    // create training metric
    train_metric_.clear();
Guolin Ke's avatar
Guolin Ke committed
165
    for (auto metric_type : config_.metric) {
Guolin Ke's avatar
Guolin Ke committed
166
      auto metric = std::unique_ptr<Metric>(
Guolin Ke's avatar
Guolin Ke committed
167
        Metric::CreateMetric(metric_type, config_));
Guolin Ke's avatar
Guolin Ke committed
168
169
170
171
172
      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();
173
174
175
176
  }

  void ResetTrainingData(const Dataset* train_data) {
    if (train_data != train_data_) {
177
      UNIQUE_LOCK(mutex_)
178
179
180
181
182
183
      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
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
290
291
292
293
  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`.");
    }
  }

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

Guolin Ke's avatar
Guolin Ke committed
308
    config_.Set(param);
309

310
311
312
    if (config_.num_threads > 0) {
      omp_set_num_threads(config_.num_threads);
    }
Guolin Ke's avatar
Guolin Ke committed
313
314
315

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

Guolin Ke's avatar
Guolin Ke committed
329
    boosting_->ResetConfig(&config_);
wxchan's avatar
wxchan committed
330
331
332
  }

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

346
  bool TrainOneIter() {
347
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
348
    return boosting_->TrainOneIter(nullptr, nullptr);
349
350
  }

Guolin Ke's avatar
Guolin Ke committed
351
  void Refit(const int32_t* leaf_preds, int32_t nrow, int32_t ncol) {
352
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
353
354
355
    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) {
356
        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
357
358
359
360
361
      }
    }
    boosting_->RefitTree(v_leaf_preds);
  }

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

wxchan's avatar
wxchan committed
367
  void RollbackOneIter() {
368
    UNIQUE_LOCK(mutex_)
wxchan's avatar
wxchan committed
369
370
371
    boosting_->RollbackOneIter();
  }

372
373
374
375
376
377
378
379
380
381
  void SetSingleRowPredictor(int num_iteration, int predict_type, const Config& config) {
      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(),
                                                                         config, num_iteration));
      }
  }

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

395
    *out_len = single_row_predictor->num_pred_in_one_row;
396
397
  }

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

Guolin Ke's avatar
Guolin Ke committed
416
    Predictor predictor(boosting_.get(), num_iteration, is_raw_score, is_predict_leaf, predict_contrib,
417
                        config.pred_early_stop, config.pred_early_stop_freq, config.pred_early_stop_margin);
418
419
420
421
422
423
    return predictor;
  }

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

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

  void PredictSparseCSR(int num_iteration, int predict_type, int64_t nrow, int ncol,
                        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,
495
496
                        int32_t** out_indices, void** out_data, int data_type) const {
    SHARED_LOCK(mutex_);
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
    // 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;
    PredictSparse(num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
                  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);
    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++;
      }
    }
    // copy vector results to output for each row
    int64_t indptr_index = 0;
    for (int m = 0; m < num_matrices; ++m) {
      if (is_indptr_int32) {
        (reinterpret_cast<int32_t*>(*out_indptr))[indptr_index] = 0;
      } else {
        (reinterpret_cast<int64_t*>(*out_indptr))[indptr_index] = 0;
      }
      indptr_index++;
      int64_t matrix_start_index = m * static_cast<int64_t>(agg.size());
      OMP_INIT_EX();
      #pragma omp parallel for schedule(static)
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        OMP_LOOP_EX_BEGIN();
        auto row_vector = agg[i];
        int64_t row_start_index = matrix_start_index + i;
        int64_t element_index = row_matrix_offsets[row_start_index];
        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;
  }

  void PredictSparseCSC(int num_iteration, int predict_type, int64_t nrow, int ncol,
                        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,
579
580
                        int32_t** out_indices, void** out_data, int data_type) const {
    SHARED_LOCK(mutex_);
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
    // Get the number of trees per iteration (for multiclass scenario we output multiple sparse matrices)
    int num_matrices = boosting_->NumModelPerIteration();
    auto predictor = CreatePredictor(num_iteration, predict_type, ncol, config);
    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;
    PredictSparse(num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
                  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;
      }
      if (m != 0) {
        matrix_start_indices[m] = matrix_start_indices[m - 1] +
          last_column_start_index +
          last_column_size;
      }
    }
    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];
        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;
          }
        }
      }
    }
    out_len[0] = elements_size;
    out_len[1] = col_ptr_size;
  }

Guolin Ke's avatar
Guolin Ke committed
679
  void Predict(int num_iteration, int predict_type, const char* data_filename,
Guolin Ke's avatar
Guolin Ke committed
680
               int data_has_header, const Config& config,
681
682
               const char* result_filename) const {
    SHARED_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
683
684
    bool is_predict_leaf = false;
    bool is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
685
    bool predict_contrib = false;
Guolin Ke's avatar
Guolin Ke committed
686
687
688
689
    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;
690
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
Guolin Ke's avatar
Guolin Ke committed
691
      predict_contrib = true;
Guolin Ke's avatar
Guolin Ke committed
692
693
694
    } else {
      is_raw_score = false;
    }
Guolin Ke's avatar
Guolin Ke committed
695
    Predictor predictor(boosting_.get(), num_iteration, is_raw_score, is_predict_leaf, predict_contrib,
696
                        config.pred_early_stop, config.pred_early_stop_freq, config.pred_early_stop_margin);
Guolin Ke's avatar
Guolin Ke committed
697
    bool bool_data_has_header = data_has_header > 0 ? true : false;
698
    predictor.Predict(data_filename, result_filename, bool_data_has_header, config.predict_disable_shape_check);
Guolin Ke's avatar
Guolin Ke committed
699
700
  }

701
  void GetPredictAt(int data_idx, double* out_result, int64_t* out_len) const {
wxchan's avatar
wxchan committed
702
703
704
    boosting_->GetPredictAt(data_idx, out_result, out_len);
  }

705
  void SaveModelToFile(int start_iteration, int num_iteration, int feature_importance_type, const char* filename) const {
706
    boosting_->SaveModelToFile(start_iteration, num_iteration, feature_importance_type, filename);
Guolin Ke's avatar
Guolin Ke committed
707
  }
708

709
  void LoadModelFromString(const char* model_str) {
710
711
    size_t len = std::strlen(model_str);
    boosting_->LoadModelFromString(model_str, len);
712
713
  }

714
  std::string SaveModelToString(int start_iteration, int num_iteration,
715
                                int feature_importance_type) const {
716
717
    return boosting_->SaveModelToString(start_iteration,
                                        num_iteration, feature_importance_type);
718
719
  }

720
  std::string DumpModel(int start_iteration, int num_iteration,
721
                        int feature_importance_type) const {
722
723
    return boosting_->DumpModel(start_iteration, num_iteration,
                                feature_importance_type);
wxchan's avatar
wxchan committed
724
  }
725

726
  std::vector<double> FeatureImportance(int num_iteration, int importance_type) const {
727
728
729
    return boosting_->FeatureImportance(num_iteration, importance_type);
  }

730
  double UpperBoundValue() const {
731
    SHARED_LOCK(mutex_)
732
733
734
735
    return boosting_->GetUpperBoundValue();
  }

  double LowerBoundValue() const {
736
    SHARED_LOCK(mutex_)
737
738
739
    return boosting_->GetLowerBoundValue();
  }

Guolin Ke's avatar
Guolin Ke committed
740
  double GetLeafValue(int tree_idx, int leaf_idx) const {
741
    SHARED_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
742
    return dynamic_cast<GBDTBase*>(boosting_.get())->GetLeafValue(tree_idx, leaf_idx);
Guolin Ke's avatar
Guolin Ke committed
743
744
745
  }

  void SetLeafValue(int tree_idx, int leaf_idx, double val) {
746
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
747
    dynamic_cast<GBDTBase*>(boosting_.get())->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
748
749
  }

750
  void ShuffleModels(int start_iter, int end_iter) {
751
    UNIQUE_LOCK(mutex_)
752
    boosting_->ShuffleModels(start_iter, end_iter);
753
754
  }

wxchan's avatar
wxchan committed
755
  int GetEvalCounts() const {
756
    SHARED_LOCK(mutex_)
wxchan's avatar
wxchan committed
757
758
759
760
761
762
    int ret = 0;
    for (const auto& metric : train_metric_) {
      ret += static_cast<int>(metric->GetName().size());
    }
    return ret;
  }
763

764
  int GetEvalNames(char** out_strs, const int len, const size_t buffer_len, size_t *out_buffer_len) const {
765
    SHARED_LOCK(mutex_)
766
    *out_buffer_len = 0;
wxchan's avatar
wxchan committed
767
768
769
    int idx = 0;
    for (const auto& metric : train_metric_) {
      for (const auto& name : metric->GetName()) {
770
771
772
773
774
        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
775
776
777
778
779
780
        ++idx;
      }
    }
    return idx;
  }

781
  int GetFeatureNames(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
    int idx = 0;
    for (const auto& name : boosting_->FeatureNames()) {
786
787
788
789
790
      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
791
792
793
794
795
      ++idx;
    }
    return idx;
  }

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

Nikita Titov's avatar
Nikita Titov committed
798
 private:
wxchan's avatar
wxchan committed
799
  const Dataset* train_data_;
Guolin Ke's avatar
Guolin Ke committed
800
  std::unique_ptr<Boosting> boosting_;
801
  std::unique_ptr<SingleRowPredictor> single_row_predictor_[PREDICTOR_TYPES];
802

Guolin Ke's avatar
Guolin Ke committed
803
  /*! \brief All configs */
Guolin Ke's avatar
Guolin Ke committed
804
  Config config_;
Guolin Ke's avatar
Guolin Ke committed
805
  /*! \brief Metric for training data */
Guolin Ke's avatar
Guolin Ke committed
806
  std::vector<std::unique_ptr<Metric>> train_metric_;
Guolin Ke's avatar
Guolin Ke committed
807
  /*! \brief Metrics for validation data */
Guolin Ke's avatar
Guolin Ke committed
808
  std::vector<std::vector<std::unique_ptr<Metric>>> valid_metrics_;
Guolin Ke's avatar
Guolin Ke committed
809
  /*! \brief Training objective function */
Guolin Ke's avatar
Guolin Ke committed
810
  std::unique_ptr<ObjectiveFunction> objective_fun_;
wxchan's avatar
wxchan committed
811
  /*! \brief mutex for threading safe call */
812
  mutable yamc::alternate::shared_mutex mutex_;
Guolin Ke's avatar
Guolin Ke committed
813
814
};

815
}  // namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
816

817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
// explicitly declare symbols from LightGBM namespace
using LightGBM::AllgatherFunction;
using LightGBM::Booster;
using LightGBM::Common::CheckElementsIntervalClosed;
using LightGBM::Common::RemoveQuotationSymbol;
using LightGBM::Common::Vector2Ptr;
using LightGBM::Common::VectorSize;
using LightGBM::Config;
using LightGBM::data_size_t;
using LightGBM::Dataset;
using LightGBM::DatasetLoader;
using LightGBM::kZeroThreshold;
using LightGBM::LGBM_APIHandleException;
using LightGBM::Log;
using LightGBM::Network;
using LightGBM::Random;
using LightGBM::ReduceScatterFunction;
Guolin Ke's avatar
Guolin Ke committed
834

Guolin Ke's avatar
Guolin Ke committed
835
836
837
838
839
840
841
842
// 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);

843
844
845
std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseRows(const void** data, int num_col, int data_type);

846
847
template<typename T>
std::function<std::vector<std::pair<int, double>>(T idx)>
Guolin Ke's avatar
Guolin Ke committed
848
RowFunctionFromCSR(const void* indptr, int indptr_type, const int32_t* indices,
849
                   const void* data, int data_type, int64_t nindptr, int64_t nelem);
Guolin Ke's avatar
Guolin Ke committed
850
851
852

// Row iterator of on column for CSC matrix
class CSC_RowIterator {
Nikita Titov's avatar
Nikita Titov committed
853
 public:
Guolin Ke's avatar
Guolin Ke committed
854
  CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
855
                  const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx);
Guolin Ke's avatar
Guolin Ke committed
856
857
858
859
860
  ~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
861
862

 private:
Guolin Ke's avatar
Guolin Ke committed
863
864
865
866
867
868
869
870
871
  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
872
const char* LGBM_GetLastError() {
wxchan's avatar
wxchan committed
873
  return LastErrorMsg();
Guolin Ke's avatar
Guolin Ke committed
874
875
}

876
877
878
879
880
881
int LGBM_RegisterLogCallback(void (*callback)(const char*)) {
  API_BEGIN();
  Log::ResetCallBack(callback);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
882
int LGBM_DatasetCreateFromFile(const char* filename,
883
884
885
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
886
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
887
888
  auto param = Config::Str2Map(parameters);
  Config config;
889
890
891
892
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
893
  DatasetLoader loader(config, nullptr, 1, filename);
Guolin Ke's avatar
Guolin Ke committed
894
  if (reference == nullptr) {
895
    if (Network::num_machines() == 1) {
896
      *out = loader.LoadFromFile(filename);
897
    } else {
898
      *out = loader.LoadFromFile(filename, Network::rank(), Network::num_machines());
899
    }
Guolin Ke's avatar
Guolin Ke committed
900
  } else {
901
    *out = loader.LoadFromFileAlignWithOtherDataset(filename,
902
                                                    reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
903
  }
904
  API_END();
Guolin Ke's avatar
Guolin Ke committed
905
906
}

907

Guolin Ke's avatar
Guolin Ke committed
908
int LGBM_DatasetCreateFromSampledColumn(double** sample_data,
909
910
911
912
913
914
915
                                        int** sample_indices,
                                        int32_t ncol,
                                        const int* num_per_col,
                                        int32_t num_sample_row,
                                        int32_t num_total_row,
                                        const char* parameters,
                                        DatasetHandle* out) {
916
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
917
918
  auto param = Config::Str2Map(parameters);
  Config config;
919
920
921
922
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
923
  DatasetLoader loader(config, nullptr, 1, nullptr);
924
925
926
927
  *out = loader.CostructFromSampleData(sample_data, sample_indices, ncol, num_per_col,
                                       num_sample_row,
                                       static_cast<data_size_t>(num_total_row));
  API_END();
Guolin Ke's avatar
Guolin Ke committed
928
929
}

930

Guolin Ke's avatar
Guolin Ke committed
931
int LGBM_DatasetCreateByReference(const DatasetHandle reference,
932
933
                                  int64_t num_total_row,
                                  DatasetHandle* out) {
Guolin Ke's avatar
Guolin Ke committed
934
935
936
937
938
939
940
941
  API_BEGIN();
  std::unique_ptr<Dataset> ret;
  ret.reset(new Dataset(static_cast<data_size_t>(num_total_row)));
  ret->CreateValid(reinterpret_cast<const Dataset*>(reference));
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
942
int LGBM_DatasetPushRows(DatasetHandle dataset,
943
944
945
946
947
                         const void* data,
                         int data_type,
                         int32_t nrow,
                         int32_t ncol,
                         int32_t start_row) {
Guolin Ke's avatar
Guolin Ke committed
948
949
950
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto get_row_fun = RowFunctionFromDenseMatric(data, nrow, ncol, data_type, 1);
951
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
952
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
953
  for (int i = 0; i < nrow; ++i) {
954
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
955
956
957
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(tid, start_row + i, one_row);
958
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
959
  }
960
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
961
962
963
964
965
966
  if (start_row + nrow == p_dataset->num_data()) {
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
967
int LGBM_DatasetPushRowsByCSR(DatasetHandle dataset,
968
969
970
971
972
973
974
975
976
                              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
977
978
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
979
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
980
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
981
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
982
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
983
  for (int i = 0; i < nrow; ++i) {
984
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
985
986
987
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(tid,
988
                          static_cast<data_size_t>(start_row + i), one_row);
989
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
990
  }
991
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
992
993
994
995
996
997
  if (start_row + nrow == static_cast<int64_t>(p_dataset->num_data())) {
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
998
int LGBM_DatasetCreateFromMat(const void* data,
999
1000
1001
1002
1003
1004
1005
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
  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) {
1027
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1028
1029
  auto param = Config::Str2Map(parameters);
  Config config;
1030
1031
1032
1033
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1034
  std::unique_ptr<Dataset> ret;
1035
1036
1037
1038
1039
1040
1041
1042
1043
  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));
  }
1044

Guolin Ke's avatar
Guolin Ke committed
1045
1046
  if (reference == nullptr) {
    // sample data first
Guolin Ke's avatar
Guolin Ke committed
1047
    Random rand(config.data_random_seed);
1048
1049
    int sample_cnt = static_cast<int>(total_nrow < config.bin_construct_sample_cnt ? total_nrow : config.bin_construct_sample_cnt);
    auto sample_indices = rand.Sample(total_nrow, sample_cnt);
1050
    sample_cnt = static_cast<int>(sample_indices.size());
1051
    std::vector<std::vector<double>> sample_values(ncol);
Guolin Ke's avatar
Guolin Ke committed
1052
    std::vector<std::vector<int>> sample_idx(ncol);
1053
1054
1055

    int offset = 0;
    int j = 0;
Guolin Ke's avatar
Guolin Ke committed
1056
    for (size_t i = 0; i < sample_indices.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1057
      auto idx = sample_indices[i];
1058
1059
1060
1061
      while ((idx - offset) >= nrow[j]) {
        offset += nrow[j];
        ++j;
      }
1062

1063
1064
1065
1066
1067
      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
1068
        }
Guolin Ke's avatar
Guolin Ke committed
1069
1070
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1071
    DatasetLoader loader(config, nullptr, 1, nullptr);
1072
1073
    ret.reset(loader.CostructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                            Vector2Ptr<int>(&sample_idx).data(),
1074
                                            ncol,
1075
                                            VectorSize<double>(sample_values).data(),
1076
                                            sample_cnt, total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1077
  } else {
1078
    ret.reset(new Dataset(total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1079
    ret->CreateValid(
1080
      reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
1081
  }
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
  int32_t start_row = 0;
  for (int j = 0; j < nmat; ++j) {
    OMP_INIT_EX();
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < nrow[j]; ++i) {
      OMP_LOOP_EX_BEGIN();
      const int tid = omp_get_thread_num();
      auto one_row = get_row_fun[j](i);
      ret->PushOneRow(tid, start_row + i, one_row);
      OMP_LOOP_EX_END();
    }
    OMP_THROW_EX();

    start_row += nrow[j];
Guolin Ke's avatar
Guolin Ke committed
1096
1097
  }
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1098
  *out = ret.release();
1099
  API_END();
1100
1101
}

Guolin Ke's avatar
Guolin Ke committed
1102
int LGBM_DatasetCreateFromCSR(const void* indptr,
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
                              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) {
1113
  API_BEGIN();
1114
1115
1116
1117
1118
  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
1119
1120
  auto param = Config::Str2Map(parameters);
  Config config;
1121
1122
1123
1124
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1125
  std::unique_ptr<Dataset> ret;
1126
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1127
1128
1129
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
  if (reference == nullptr) {
    // sample data first
Guolin Ke's avatar
Guolin Ke committed
1130
1131
    Random rand(config.data_random_seed);
    int sample_cnt = static_cast<int>(nrow < config.bin_construct_sample_cnt ? nrow : config.bin_construct_sample_cnt);
1132
    auto sample_indices = rand.Sample(nrow, sample_cnt);
1133
    sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1134
1135
    std::vector<std::vector<double>> sample_values(num_col);
    std::vector<std::vector<int>> sample_idx(num_col);
1136
1137
1138
1139
    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
1140
        CHECK_LT(inner_data.first, num_col);
Guolin Ke's avatar
Guolin Ke committed
1141
        if (std::fabs(inner_data.second) > kZeroThreshold || std::isnan(inner_data.second)) {
Guolin Ke's avatar
Guolin Ke committed
1142
1143
          sample_values[inner_data.first].emplace_back(inner_data.second);
          sample_idx[inner_data.first].emplace_back(static_cast<int>(i));
1144
1145
1146
        }
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1147
    DatasetLoader loader(config, nullptr, 1, nullptr);
1148
1149
    ret.reset(loader.CostructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                            Vector2Ptr<int>(&sample_idx).data(),
1150
                                            static_cast<int>(num_col),
1151
                                            VectorSize<double>(sample_values).data(),
1152
                                            sample_cnt, nrow));
1153
  } else {
1154
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1155
    ret->CreateValid(
1156
      reinterpret_cast<const Dataset*>(reference));
1157
  }
1158
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1159
  #pragma omp parallel for schedule(static)
1160
  for (int i = 0; i < nindptr - 1; ++i) {
1161
    OMP_LOOP_EX_BEGIN();
1162
1163
1164
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    ret->PushOneRow(tid, i, one_row);
1165
    OMP_LOOP_EX_END();
1166
  }
1167
  OMP_THROW_EX();
1168
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1169
  *out = ret.release();
1170
  API_END();
1171
1172
}

1173
int LGBM_DatasetCreateFromCSRFunc(void* get_row_funptr,
1174
1175
1176
1177
1178
                                  int num_rows,
                                  int64_t num_col,
                                  const char* parameters,
                                  const DatasetHandle reference,
                                  DatasetHandle* out) {
1179
  API_BEGIN();
1180
1181
1182
1183
1184
  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.");
  }
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
  auto get_row_fun = *static_cast<std::function<void(int idx, std::vector<std::pair<int, double>>&)>*>(get_row_funptr);
  auto param = Config::Str2Map(parameters);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  std::unique_ptr<Dataset> ret;
  int32_t nrow = num_rows;
  if (reference == nullptr) {
    // sample data first
    Random rand(config.data_random_seed);
    int sample_cnt = static_cast<int>(nrow < config.bin_construct_sample_cnt ? nrow : config.bin_construct_sample_cnt);
    auto sample_indices = rand.Sample(nrow, sample_cnt);
    sample_cnt = static_cast<int>(sample_indices.size());
    std::vector<std::vector<double>> sample_values(num_col);
    std::vector<std::vector<int>> sample_idx(num_col);
    // local buffer to re-use memory
    std::vector<std::pair<int, double>> buffer;
    for (size_t i = 0; i < sample_indices.size(); ++i) {
      auto idx = sample_indices[i];
      get_row_fun(static_cast<int>(idx), buffer);
      for (std::pair<int, double>& inner_data : buffer) {
Nikita Titov's avatar
Nikita Titov committed
1208
        CHECK_LT(inner_data.first, num_col);
1209
1210
1211
1212
1213
1214
1215
        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);
1216
1217
    ret.reset(loader.CostructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                            Vector2Ptr<int>(&sample_idx).data(),
1218
                                            static_cast<int>(num_col),
1219
                                            VectorSize<double>(sample_values).data(),
1220
1221
1222
1223
1224
1225
                                            sample_cnt, nrow));
  } else {
    ret.reset(new Dataset(nrow));
    ret->CreateValid(
      reinterpret_cast<const Dataset*>(reference));
  }
1226

1227
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1228
1229
  std::vector<std::pair<int, double>> thread_buffer;
  #pragma omp parallel for schedule(static) private(thread_buffer)
1230
1231
1232
  for (int i = 0; i < num_rows; ++i) {
    OMP_LOOP_EX_BEGIN();
    {
1233
      const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1234
1235
      get_row_fun(i, thread_buffer);
      ret->PushOneRow(tid, i, thread_buffer);
1236
1237
1238
1239
1240
1241
1242
1243
1244
    }
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();
  ret->FinishLoad();
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1245
int LGBM_DatasetCreateFromCSC(const void* col_ptr,
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
                              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) {
1256
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1257
1258
  auto param = Config::Str2Map(parameters);
  Config config;
1259
1260
1261
1262
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1263
  std::unique_ptr<Dataset> ret;
Guolin Ke's avatar
Guolin Ke committed
1264
1265
1266
  int32_t nrow = static_cast<int32_t>(num_row);
  if (reference == nullptr) {
    // sample data first
Guolin Ke's avatar
Guolin Ke committed
1267
1268
    Random rand(config.data_random_seed);
    int sample_cnt = static_cast<int>(nrow < config.bin_construct_sample_cnt ? nrow : config.bin_construct_sample_cnt);
Guolin Ke's avatar
Guolin Ke committed
1269
    auto sample_indices = rand.Sample(nrow, sample_cnt);
1270
    sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1271
    std::vector<std::vector<double>> sample_values(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
1272
    std::vector<std::vector<int>> sample_idx(ncol_ptr - 1);
1273
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1274
    #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1275
    for (int i = 0; i < static_cast<int>(sample_values.size()); ++i) {
1276
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1277
1278
1279
      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
1280
        if (std::fabs(val) > kZeroThreshold || std::isnan(val)) {
Guolin Ke's avatar
Guolin Ke committed
1281
1282
          sample_values[i].emplace_back(val);
          sample_idx[i].emplace_back(j);
Guolin Ke's avatar
Guolin Ke committed
1283
1284
        }
      }
1285
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1286
    }
1287
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1288
    DatasetLoader loader(config, nullptr, 1, nullptr);
1289
1290
    ret.reset(loader.CostructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                            Vector2Ptr<int>(&sample_idx).data(),
1291
                                            static_cast<int>(sample_values.size()),
1292
                                            VectorSize<double>(sample_values).data(),
1293
                                            sample_cnt, nrow));
Guolin Ke's avatar
Guolin Ke committed
1294
  } else {
1295
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1296
    ret->CreateValid(
1297
      reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
1298
  }
1299
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1300
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1301
  for (int i = 0; i < ncol_ptr - 1; ++i) {
1302
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1303
    const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1304
    int feature_idx = ret->InnerFeatureIndex(i);
Guolin Ke's avatar
Guolin Ke committed
1305
    if (feature_idx < 0) { continue; }
Guolin Ke's avatar
Guolin Ke committed
1306
1307
    int group = ret->Feature2Group(feature_idx);
    int sub_feature = ret->Feture2SubFeature(feature_idx);
Guolin Ke's avatar
Guolin Ke committed
1308
    CSC_RowIterator col_it(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, i);
Guolin Ke's avatar
Guolin Ke committed
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
    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; }
        ret->PushOneData(tid, row_idx, group, sub_feature, pair.second);
      }
    } else {
      for (int row_idx = 0; row_idx < nrow; ++row_idx) {
        auto val = col_it.Get(row_idx);
        ret->PushOneData(tid, row_idx, group, sub_feature, val);
      }
Guolin Ke's avatar
Guolin Ke committed
1324
    }
1325
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1326
  }
1327
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1328
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1329
  *out = ret.release();
1330
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1331
1332
}

Guolin Ke's avatar
Guolin Ke committed
1333
int LGBM_DatasetGetSubset(
1334
  const DatasetHandle handle,
wxchan's avatar
wxchan committed
1335
1336
1337
  const int32_t* used_row_indices,
  int32_t num_used_row_indices,
  const char* parameters,
Guolin Ke's avatar
typo  
Guolin Ke committed
1338
  DatasetHandle* out) {
wxchan's avatar
wxchan committed
1339
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1340
1341
  auto param = Config::Str2Map(parameters);
  Config config;
1342
1343
1344
1345
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
1346
  auto full_dataset = reinterpret_cast<const Dataset*>(handle);
1347
  CHECK_GT(num_used_row_indices, 0);
1348
1349
  const int32_t lower = 0;
  const int32_t upper = full_dataset->num_data() - 1;
1350
  CheckElementsIntervalClosed(used_row_indices, lower, upper, num_used_row_indices, "Used indices of subset");
1351
1352
1353
  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
1354
  auto ret = std::unique_ptr<Dataset>(new Dataset(num_used_row_indices));
1355
  ret->CopyFeatureMapperFrom(full_dataset);
1356
  ret->CopySubrow(full_dataset, used_row_indices, num_used_row_indices, true);
wxchan's avatar
wxchan committed
1357
1358
1359
1360
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1361
int LGBM_DatasetSetFeatureNames(
Guolin Ke's avatar
typo  
Guolin Ke committed
1362
  DatasetHandle handle,
Guolin Ke's avatar
Guolin Ke committed
1363
  const char** feature_names,
Guolin Ke's avatar
Guolin Ke committed
1364
  int num_feature_names) {
Guolin Ke's avatar
Guolin Ke committed
1365
1366
1367
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  std::vector<std::string> feature_names_str;
Guolin Ke's avatar
Guolin Ke committed
1368
  for (int i = 0; i < num_feature_names; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1369
1370
1371
1372
1373
1374
    feature_names_str.emplace_back(feature_names[i]);
  }
  dataset->set_feature_names(feature_names_str);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1375
int LGBM_DatasetGetFeatureNames(
1376
1377
1378
1379
1380
1381
    DatasetHandle handle,
    const int len,
    int* num_feature_names,
    const size_t buffer_len,
    size_t* out_buffer_len,
    char** feature_names) {
1382
  API_BEGIN();
1383
  *out_buffer_len = 0;
1384
1385
  auto dataset = reinterpret_cast<Dataset*>(handle);
  auto inside_feature_name = dataset->feature_names();
Guolin Ke's avatar
Guolin Ke committed
1386
1387
  *num_feature_names = static_cast<int>(inside_feature_name.size());
  for (int i = 0; i < *num_feature_names; ++i) {
1388
1389
1390
1391
1392
    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);
1393
1394
1395
1396
  }
  API_END();
}

1397
1398
1399
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1400
int LGBM_DatasetFree(DatasetHandle handle) {
1401
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1402
  delete reinterpret_cast<Dataset*>(handle);
1403
  API_END();
1404
1405
}

Guolin Ke's avatar
Guolin Ke committed
1406
int LGBM_DatasetSaveBinary(DatasetHandle handle,
1407
                           const char* filename) {
1408
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1409
1410
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->SaveBinaryFile(filename);
1411
  API_END();
1412
1413
}

1414
1415
1416
1417
1418
1419
1420
1421
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
1422
int LGBM_DatasetSetField(DatasetHandle handle,
1423
1424
1425
1426
                         const char* field_name,
                         const void* field_data,
                         int num_element,
                         int type) {
1427
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1428
  auto dataset = reinterpret_cast<Dataset*>(handle);
1429
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1430
  if (type == C_API_DTYPE_FLOAT32) {
Guolin Ke's avatar
Guolin Ke committed
1431
    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
1432
  } else if (type == C_API_DTYPE_INT32) {
Guolin Ke's avatar
Guolin Ke committed
1433
    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
1434
1435
  } 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));
1436
  }
1437
  if (!is_success) { Log::Fatal("Input data type error or field not found"); }
1438
  API_END();
1439
1440
}

Guolin Ke's avatar
Guolin Ke committed
1441
int LGBM_DatasetGetField(DatasetHandle handle,
1442
1443
1444
1445
                         const char* field_name,
                         int* out_len,
                         const void** out_ptr,
                         int* out_type) {
1446
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1447
  auto dataset = reinterpret_cast<Dataset*>(handle);
1448
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1449
  if (dataset->GetFloatField(field_name, out_len, reinterpret_cast<const float**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1450
    *out_type = C_API_DTYPE_FLOAT32;
1451
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1452
  } else if (dataset->GetIntField(field_name, out_len, reinterpret_cast<const int**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1453
    *out_type = C_API_DTYPE_INT32;
1454
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1455
1456
1457
  } 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
1458
  }
1459
  if (!is_success) { Log::Fatal("Field not found"); }
wxchan's avatar
wxchan committed
1460
  if (*out_ptr == nullptr) { *out_len = 0; }
1461
  API_END();
1462
1463
}

1464
int LGBM_DatasetUpdateParamChecking(const char* old_parameters, const char* new_parameters) {
1465
  API_BEGIN();
1466
1467
1468
1469
1470
  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);
1471
1472
1473
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1474
int LGBM_DatasetGetNumData(DatasetHandle handle,
1475
                           int* out) {
1476
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1477
1478
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_data();
1479
  API_END();
1480
1481
}

Guolin Ke's avatar
Guolin Ke committed
1482
int LGBM_DatasetGetNumFeature(DatasetHandle handle,
1483
                              int* out) {
1484
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1485
1486
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_total_features();
1487
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1488
}
1489

1490
1491
1492
1493
1494
int LGBM_DatasetAddFeaturesFrom(DatasetHandle target,
                                DatasetHandle source) {
  API_BEGIN();
  auto target_d = reinterpret_cast<Dataset*>(target);
  auto source_d = reinterpret_cast<Dataset*>(source);
1495
  target_d->AddFeaturesFrom(source_d);
1496
1497
1498
  API_END();
}

1499
1500
// ---- start of booster

Guolin Ke's avatar
Guolin Ke committed
1501
int LGBM_BoosterCreate(const DatasetHandle train_data,
1502
1503
                       const char* parameters,
                       BoosterHandle* out) {
1504
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1505
  const Dataset* p_train_data = reinterpret_cast<const Dataset*>(train_data);
wxchan's avatar
wxchan committed
1506
1507
  auto ret = std::unique_ptr<Booster>(new Booster(p_train_data, parameters));
  *out = ret.release();
1508
  API_END();
1509
1510
}

Guolin Ke's avatar
Guolin Ke committed
1511
int LGBM_BoosterCreateFromModelfile(
1512
  const char* filename,
Guolin Ke's avatar
Guolin Ke committed
1513
  int* out_num_iterations,
1514
  BoosterHandle* out) {
1515
  API_BEGIN();
wxchan's avatar
wxchan committed
1516
  auto ret = std::unique_ptr<Booster>(new Booster(filename));
Guolin Ke's avatar
Guolin Ke committed
1517
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
wxchan's avatar
wxchan committed
1518
  *out = ret.release();
1519
  API_END();
1520
1521
}

Guolin Ke's avatar
Guolin Ke committed
1522
int LGBM_BoosterLoadModelFromString(
1523
1524
1525
1526
  const char* model_str,
  int* out_num_iterations,
  BoosterHandle* out) {
  API_BEGIN();
wxchan's avatar
wxchan committed
1527
  auto ret = std::unique_ptr<Booster>(new Booster(nullptr));
1528
1529
1530
1531
1532
1533
  ret->LoadModelFromString(model_str);
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
  *out = ret.release();
  API_END();
}

1534
1535
1536
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1537
int LGBM_BoosterFree(BoosterHandle handle) {
1538
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1539
  delete reinterpret_cast<Booster*>(handle);
1540
  API_END();
1541
1542
}

1543
int LGBM_BoosterShuffleModels(BoosterHandle handle, int start_iter, int end_iter) {
1544
1545
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1546
  ref_booster->ShuffleModels(start_iter, end_iter);
1547
1548
1549
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1550
int LGBM_BoosterMerge(BoosterHandle handle,
1551
                      BoosterHandle other_handle) {
wxchan's avatar
wxchan committed
1552
1553
1554
1555
1556
1557
1558
  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
1559
int LGBM_BoosterAddValidData(BoosterHandle handle,
1560
                             const DatasetHandle valid_data) {
wxchan's avatar
wxchan committed
1561
1562
1563
1564
1565
1566
1567
  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
1568
int LGBM_BoosterResetTrainingData(BoosterHandle handle,
1569
                                  const DatasetHandle train_data) {
wxchan's avatar
wxchan committed
1570
1571
1572
1573
1574
1575
1576
  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
1577
int LGBM_BoosterResetParameter(BoosterHandle handle, const char* parameters) {
wxchan's avatar
wxchan committed
1578
1579
1580
1581
1582
1583
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->ResetConfig(parameters);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1584
int LGBM_BoosterGetNumClasses(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1585
1586
1587
1588
1589
1590
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->NumberOfClasses();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1591
1592
1593
1594
1595
1596
1597
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
1598
int LGBM_BoosterUpdateOneIter(BoosterHandle handle, int* is_finished) {
1599
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1600
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1601
1602
1603
1604
1605
  if (ref_booster->TrainOneIter()) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1606
  API_END();
1607
1608
}

Guolin Ke's avatar
Guolin Ke committed
1609
int LGBM_BoosterUpdateOneIterCustom(BoosterHandle handle,
1610
1611
1612
                                    const float* grad,
                                    const float* hess,
                                    int* is_finished) {
1613
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1614
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1615
  #ifdef SCORE_T_USE_DOUBLE
1616
  Log::Fatal("Don't support custom loss function when SCORE_T_USE_DOUBLE is enabled");
1617
  #else
1618
1619
1620
1621
1622
  if (ref_booster->TrainOneIter(grad, hess)) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1623
  #endif
1624
  API_END();
1625
1626
}

Guolin Ke's avatar
Guolin Ke committed
1627
int LGBM_BoosterRollbackOneIter(BoosterHandle handle) {
wxchan's avatar
wxchan committed
1628
1629
1630
1631
1632
1633
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->RollbackOneIter();
  API_END();
}

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

1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
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
1655
int LGBM_BoosterGetEvalCounts(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1656
1657
1658
1659
1660
1661
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetEvalCounts();
  API_END();
}

1662
1663
1664
1665
1666
1667
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
1668
1669
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1670
  *out_len = ref_booster->GetEvalNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1671
1672
1673
  API_END();
}

1674
1675
1676
1677
1678
1679
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
1680
1681
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1682
  *out_len = ref_booster->GetFeatureNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1683
1684
1685
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1686
int LGBM_BoosterGetNumFeature(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1687
1688
1689
1690
1691
1692
  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
1693
int LGBM_BoosterGetEval(BoosterHandle handle,
1694
1695
1696
                        int data_idx,
                        int* out_len,
                        double* out_results) {
1697
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1698
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1699
  auto boosting = ref_booster->GetBoosting();
wxchan's avatar
wxchan committed
1700
  auto result_buf = boosting->GetEvalAt(data_idx);
Guolin Ke's avatar
Guolin Ke committed
1701
  *out_len = static_cast<int>(result_buf.size());
1702
  for (size_t i = 0; i < result_buf.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1703
    (out_results)[i] = static_cast<double>(result_buf[i]);
1704
  }
1705
  API_END();
1706
1707
}

Guolin Ke's avatar
Guolin Ke committed
1708
int LGBM_BoosterGetNumPredict(BoosterHandle handle,
1709
1710
                              int data_idx,
                              int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1711
1712
1713
1714
1715
1716
  API_BEGIN();
  auto boosting = reinterpret_cast<Booster*>(handle)->GetBoosting();
  *out_len = boosting->GetNumPredictAt(data_idx);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1717
int LGBM_BoosterGetPredict(BoosterHandle handle,
1718
1719
1720
                           int data_idx,
                           int64_t* out_len,
                           double* out_result) {
1721
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1722
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1723
  ref_booster->GetPredictAt(data_idx, out_result, out_len);
1724
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1725
1726
}

Guolin Ke's avatar
Guolin Ke committed
1727
int LGBM_BoosterPredictForFile(BoosterHandle handle,
1728
1729
1730
1731
                               const char* data_filename,
                               int data_has_header,
                               int predict_type,
                               int num_iteration,
1732
                               const char* parameter,
1733
                               const char* result_filename) {
1734
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1735
1736
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1737
1738
1739
1740
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1741
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
cbecker's avatar
cbecker committed
1742
  ref_booster->Predict(num_iteration, predict_type, data_filename, data_has_header,
Guolin Ke's avatar
Guolin Ke committed
1743
                       config, result_filename);
1744
  API_END();
1745
1746
}

Guolin Ke's avatar
Guolin Ke committed
1747
int LGBM_BoosterCalcNumPredict(BoosterHandle handle,
1748
1749
1750
1751
                               int num_row,
                               int predict_type,
                               int num_iteration,
                               int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1752
1753
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1754
1755
  *out_len = static_cast<int64_t>(num_row) * ref_booster->GetBoosting()->NumPredictOneRow(
    num_iteration, predict_type == C_API_PREDICT_LEAF_INDEX, predict_type == C_API_PREDICT_CONTRIB);
Guolin Ke's avatar
Guolin Ke committed
1756
1757
1758
  API_END();
}

1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
/*!
 * \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,
             const int data_type_,
             const int32_t num_cols) : booster(booster_ptr), data_type(data_type_), ncol(num_cols) {
    config.Set(Config::Str2Map(parameter));
  }

  Booster* const booster;
  Config config;
  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
1789
int LGBM_BoosterPredictForCSR(BoosterHandle handle,
1790
1791
1792
1793
1794
1795
1796
                              const void* indptr,
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
1797
                              int64_t num_col,
1798
1799
                              int predict_type,
                              int num_iteration,
1800
                              const char* parameter,
1801
1802
                              int64_t* out_len,
                              double* out_result) {
1803
  API_BEGIN();
1804
1805
1806
1807
1808
  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
1809
1810
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1811
1812
1813
1814
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1815
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1816
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
1817
  int nrow = static_cast<int>(nindptr - 1);
1818
  ref_booster->Predict(num_iteration, predict_type, nrow, static_cast<int>(num_col), get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
1819
                       config, out_result, out_len);
1820
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1821
}
1822

1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
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,
                                    int num_iteration,
                                    const char* parameter,
                                    int matrix_type,
                                    int64_t* out_len,
                                    void** out_indptr,
                                    int32_t** out_indices,
                                    void** out_data) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  if (matrix_type == C_API_MATRIX_TYPE_CSR) {
    if (num_col_or_row <= 0) {
      Log::Fatal("The number of columns should be greater than zero.");
    } else if (num_col_or_row >= INT32_MAX) {
      Log::Fatal("The number of columns should be smaller than INT32_MAX.");
    }
    auto get_row_fun = RowFunctionFromCSR<int64_t>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
    int64_t nrow = nindptr - 1;
    ref_booster->PredictSparseCSR(num_iteration, predict_type, nrow, static_cast<int>(num_col_or_row), get_row_fun,
                                  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;
    };
    ref_booster->PredictSparseCSC(num_iteration, predict_type, num_col_or_row, ncol, get_row_fun, config,
                                  out_len, out_indptr, indptr_type, out_indices, out_data, data_type);
  } else {
    Log::Fatal("Unknown matrix type in LGBM_BoosterPredictSparseOutput");
  }
  API_END();
}

int LGBM_BoosterFreePredictSparse(void* indptr, int32_t* indices, void* data, int indptr_type, int data_type) {
  API_BEGIN();
  if (indptr_type == C_API_DTYPE_INT32) {
    delete reinterpret_cast<int32_t*>(indptr);
  } else if (indptr_type == C_API_DTYPE_INT64) {
    delete reinterpret_cast<int64_t*>(indptr);
  } else {
    Log::Fatal("Unknown indptr type in LGBM_BoosterFreePredictSparse");
  }
  delete indices;
  if (data_type == C_API_DTYPE_FLOAT32) {
    delete reinterpret_cast<float*>(data);
  } else if (data_type == C_API_DTYPE_FLOAT64) {
    delete reinterpret_cast<double*>(data);
  } else {
    Log::Fatal("Unknown data type in LGBM_BoosterFreePredictSparse");
  }
  API_END();
}

1908
int LGBM_BoosterPredictForCSRSingleRow(BoosterHandle handle,
1909
1910
1911
1912
1913
1914
1915
                                       const void* indptr,
                                       int indptr_type,
                                       const int32_t* indices,
                                       const void* data,
                                       int data_type,
                                       int64_t nindptr,
                                       int64_t nelem,
1916
                                       int64_t num_col,
1917
1918
1919
1920
1921
                                       int predict_type,
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
1922
  API_BEGIN();
1923
1924
1925
1926
1927
  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.");
  }
1928
1929
1930
1931
1932
1933
1934
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1935
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1936
1937
  ref_booster->SetSingleRowPredictor(num_iteration, predict_type, config);
  ref_booster->PredictSingleRow(predict_type, static_cast<int32_t>(num_col), get_row_fun, config, out_result, out_len);
1938
1939
1940
  API_END();
}

1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
int LGBM_BoosterPredictForCSRSingleRowFastInit(BoosterHandle handle,
                                               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,
    data_type,
    static_cast<int32_t>(num_col)));

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

  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

int LGBM_BoosterPredictForCSRSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* indptr,
                                           int indptr_type,
                                           const int32_t* indices,
                                           const void* data,
                                           int64_t nindptr,
                                           int64_t nelem,
                                           int predict_type,
                                           int num_iteration,
                                           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);
1981
  fastConfig->booster->PredictSingleRow(predict_type, fastConfig->ncol,
1982
1983
1984
1985
                                        get_row_fun, fastConfig->config, out_result, out_len);
  API_END();
}

1986

Guolin Ke's avatar
Guolin Ke committed
1987
int LGBM_BoosterPredictForCSC(BoosterHandle handle,
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
                              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,
                              int num_iteration,
1998
                              const char* parameter,
1999
2000
                              int64_t* out_len,
                              double* out_result) {
Guolin Ke's avatar
Guolin Ke committed
2001
2002
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2003
2004
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2005
2006
2007
2008
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
2009
  int num_threads = OMP_NUM_THREADS();
Guolin Ke's avatar
Guolin Ke committed
2010
  int ncol = static_cast<int>(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
2011
2012
2013
2014
2015
  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
2016
2017
  }
  std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun =
Guolin Ke's avatar
Guolin Ke committed
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
      [&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;
      };
2030
  ref_booster->Predict(num_iteration, predict_type, static_cast<int>(num_row), ncol, get_row_fun, config,
cbecker's avatar
cbecker committed
2031
                       out_result, out_len);
Guolin Ke's avatar
Guolin Ke committed
2032
2033
2034
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2035
int LGBM_BoosterPredictForMat(BoosterHandle handle,
2036
2037
2038
2039
2040
2041
2042
                              const void* data,
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              int predict_type,
                              int num_iteration,
2043
                              const char* parameter,
2044
2045
                              int64_t* out_len,
                              double* out_result) {
2046
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2047
2048
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2049
2050
2051
2052
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
2053
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2054
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, nrow, ncol, data_type, is_row_major);
2055
  ref_booster->Predict(num_iteration, predict_type, nrow, ncol, get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
2056
                       config, out_result, out_len);
2057
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2058
}
2059

2060
int LGBM_BoosterPredictForMatSingleRow(BoosterHandle handle,
2061
2062
2063
2064
2065
2066
2067
2068
2069
                                       const void* data,
                                       int data_type,
                                       int32_t ncol,
                                       int is_row_major,
                                       int predict_type,
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2070
2071
2072
2073
2074
2075
2076
2077
2078
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, ncol, data_type, is_row_major);
2079
2080
  ref_booster->SetSingleRowPredictor(num_iteration, predict_type, config);
  ref_booster->PredictSingleRow(predict_type, ncol, get_row_fun, config, out_result, out_len);
2081
2082
2083
  API_END();
}

2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
int LGBM_BoosterPredictForMatSingleRowFastInit(BoosterHandle handle,
                                               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,
    data_type,
    ncol));

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

  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

int LGBM_BoosterPredictForMatSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* data,
                                           const int predict_type,
                                           const int num_iteration,
                                           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);
2114
  fastConfig->booster->PredictSingleRow(predict_type, fastConfig->ncol,
2115
2116
2117
2118
2119
                                        get_row_fun, fastConfig->config,
                                        out_result, out_len);
  API_END();
}

2120

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

Guolin Ke's avatar
Guolin Ke committed
2144
int LGBM_BoosterSaveModel(BoosterHandle handle,
2145
                          int start_iteration,
2146
                          int num_iteration,
2147
                          int feature_importance_type,
2148
                          const char* filename) {
2149
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2150
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2151
2152
  ref_booster->SaveModelToFile(start_iteration, num_iteration,
                               feature_importance_type, filename);
wxchan's avatar
wxchan committed
2153
2154
2155
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2156
int LGBM_BoosterSaveModelToString(BoosterHandle handle,
2157
                                  int start_iteration,
2158
                                  int num_iteration,
2159
                                  int feature_importance_type,
2160
                                  int64_t buffer_len,
2161
                                  int64_t* out_len,
2162
                                  char* out_str) {
2163
2164
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2165
2166
  std::string model = ref_booster->SaveModelToString(
      start_iteration, num_iteration, feature_importance_type);
2167
  *out_len = static_cast<int64_t>(model.size()) + 1;
2168
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2169
    std::memcpy(out_str, model.c_str(), *out_len);
2170
2171
2172
2173
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2174
int LGBM_BoosterDumpModel(BoosterHandle handle,
2175
                          int start_iteration,
2176
                          int num_iteration,
2177
                          int feature_importance_type,
2178
2179
                          int64_t buffer_len,
                          int64_t* out_len,
2180
                          char* out_str) {
wxchan's avatar
wxchan committed
2181
2182
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2183
2184
  std::string model = ref_booster->DumpModel(start_iteration, num_iteration,
                                             feature_importance_type);
2185
  *out_len = static_cast<int64_t>(model.size()) + 1;
wxchan's avatar
wxchan committed
2186
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2187
    std::memcpy(out_str, model.c_str(), *out_len);
wxchan's avatar
wxchan committed
2188
  }
2189
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2190
}
2191

Guolin Ke's avatar
Guolin Ke committed
2192
int LGBM_BoosterGetLeafValue(BoosterHandle handle,
2193
2194
2195
                             int tree_idx,
                             int leaf_idx,
                             double* out_val) {
Guolin Ke's avatar
Guolin Ke committed
2196
2197
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2198
  *out_val = static_cast<double>(ref_booster->GetLeafValue(tree_idx, leaf_idx));
Guolin Ke's avatar
Guolin Ke committed
2199
2200
2201
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2202
int LGBM_BoosterSetLeafValue(BoosterHandle handle,
2203
2204
2205
                             int tree_idx,
                             int leaf_idx,
                             double val) {
Guolin Ke's avatar
Guolin Ke committed
2206
2207
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2208
  ref_booster->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
2209
2210
2211
  API_END();
}

2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
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();
}

2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
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();
}

2243
2244
2245
2246
2247
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
2248
  Config config;
2249
  config.machines = RemoveQuotationSymbol(std::string(machines));
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
  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();
}

2265
2266
2267
int LGBM_NetworkInitWithFunctions(int num_machines, int rank,
                                  void* reduce_scatter_ext_fun,
                                  void* allgather_ext_fun) {
ww's avatar
ww committed
2268
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2269
  if (num_machines > 1) {
2270
    Network::Init(num_machines, rank, (ReduceScatterFunction)reduce_scatter_ext_fun, (AllgatherFunction)allgather_ext_fun);
ww's avatar
ww committed
2271
2272
2273
  }
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
2274

Guolin Ke's avatar
Guolin Ke committed
2275
// ---- start of some help functions
2276
2277
2278

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
2279
  if (data_type == C_API_DTYPE_FLOAT32) {
2280
2281
    const float* data_ptr = reinterpret_cast<const float*>(data);
    if (is_row_major) {
2282
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2283
        std::vector<double> ret(num_col);
2284
        auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
2285
        for (int i = 0; i < num_col; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2286
          ret[i] = static_cast<double>(*(tmp_ptr + i));
2287
2288
2289
2290
        }
        return ret;
      };
    } else {
2291
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2292
        std::vector<double> ret(num_col);
2293
        for (int i = 0; i < num_col; ++i) {
2294
          ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
2295
2296
2297
2298
        }
        return ret;
      };
    }
Guolin Ke's avatar
Guolin Ke committed
2299
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2300
2301
    const double* data_ptr = reinterpret_cast<const double*>(data);
    if (is_row_major) {
2302
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2303
        std::vector<double> ret(num_col);
2304
        auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
2305
        for (int i = 0; i < num_col; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2306
          ret[i] = static_cast<double>(*(tmp_ptr + i));
2307
2308
2309
2310
        }
        return ret;
      };
    } else {
2311
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2312
        std::vector<double> ret(num_col);
2313
        for (int i = 0; i < num_col; ++i) {
2314
          ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
2315
2316
2317
2318
2319
        }
        return ret;
      };
    }
  }
2320
  Log::Fatal("Unknown data type in RowFunctionFromDenseMatric");
2321
  return nullptr;
2322
2323
2324
2325
}

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
2326
2327
  auto inner_function = RowFunctionFromDenseMatric(data, num_row, num_col, data_type, is_row_major);
  if (inner_function != nullptr) {
2328
    return [inner_function] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2329
2330
      auto raw_values = inner_function(row_idx);
      std::vector<std::pair<int, double>> ret;
Guolin Ke's avatar
Guolin Ke committed
2331
      ret.reserve(raw_values.size());
Guolin Ke's avatar
Guolin Ke committed
2332
      for (int i = 0; i < static_cast<int>(raw_values.size()); ++i) {
Guolin Ke's avatar
Guolin Ke committed
2333
        if (std::fabs(raw_values[i]) > kZeroThreshold || std::isnan(raw_values[i])) {
Guolin Ke's avatar
Guolin Ke committed
2334
          ret.emplace_back(i, raw_values[i]);
2335
        }
Guolin Ke's avatar
Guolin Ke committed
2336
2337
2338
      }
      return ret;
    };
2339
  }
Guolin Ke's avatar
Guolin Ke committed
2340
  return nullptr;
2341
2342
}

2343
2344
2345
2346
2347
2348
2349
// 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
2350
    ret.reserve(raw_values.size());
2351
2352
2353
2354
2355
2356
2357
2358
2359
    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;
  };
}

2360
2361
template<typename T>
std::function<std::vector<std::pair<int, double>>(T idx)>
2362
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
2363
  if (data_type == C_API_DTYPE_FLOAT32) {
2364
    const float* data_ptr = reinterpret_cast<const float*>(data);
Guolin Ke's avatar
Guolin Ke committed
2365
    if (indptr_type == C_API_DTYPE_INT32) {
2366
      const int32_t* ptr_indptr = reinterpret_cast<const int32_t*>(indptr);
2367
      return [=] (T idx) {
2368
2369
2370
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2371
2372
2373
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2374
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2375
          ret.emplace_back(indices[i], data_ptr[i]);
2376
2377
2378
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
2379
    } else if (indptr_type == C_API_DTYPE_INT64) {
2380
      const int64_t* ptr_indptr = reinterpret_cast<const int64_t*>(indptr);
2381
      return [=] (T idx) {
2382
2383
2384
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2385
2386
2387
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2388
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2389
          ret.emplace_back(indices[i], data_ptr[i]);
2390
2391
2392
2393
        }
        return ret;
      };
    }
Guolin Ke's avatar
Guolin Ke committed
2394
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2395
    const double* data_ptr = reinterpret_cast<const double*>(data);
Guolin Ke's avatar
Guolin Ke committed
2396
    if (indptr_type == C_API_DTYPE_INT32) {
2397
      const int32_t* ptr_indptr = reinterpret_cast<const int32_t*>(indptr);
2398
      return [=] (T idx) {
2399
2400
2401
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2402
2403
2404
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2405
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2406
          ret.emplace_back(indices[i], data_ptr[i]);
2407
2408
2409
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
2410
    } else if (indptr_type == C_API_DTYPE_INT64) {
2411
      const int64_t* ptr_indptr = reinterpret_cast<const int64_t*>(indptr);
2412
      return [=] (T idx) {
2413
2414
2415
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2416
2417
2418
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2419
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2420
          ret.emplace_back(indices[i], data_ptr[i]);
2421
2422
2423
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
2424
2425
    }
  }
2426
  Log::Fatal("Unknown data type in RowFunctionFromCSR");
2427
  return nullptr;
2428
2429
}

Guolin Ke's avatar
Guolin Ke committed
2430
std::function<std::pair<int, double>(int idx)>
2431
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
2432
  CHECK(col_idx < ncol_ptr && col_idx >= 0);
Guolin Ke's avatar
Guolin Ke committed
2433
  if (data_type == C_API_DTYPE_FLOAT32) {
2434
    const float* data_ptr = reinterpret_cast<const float*>(data);
Guolin Ke's avatar
Guolin Ke committed
2435
    if (col_ptr_type == C_API_DTYPE_INT32) {
2436
      const int32_t* ptr_col_ptr = reinterpret_cast<const int32_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2437
2438
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2439
2440
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2441
2442
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2443
        }
Guolin Ke's avatar
Guolin Ke committed
2444
2445
2446
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2447
      };
Guolin Ke's avatar
Guolin Ke committed
2448
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2449
      const int64_t* ptr_col_ptr = reinterpret_cast<const int64_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2450
2451
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2452
2453
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2454
2455
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2456
        }
Guolin Ke's avatar
Guolin Ke committed
2457
2458
2459
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2460
      };
Guolin Ke's avatar
Guolin Ke committed
2461
    }
Guolin Ke's avatar
Guolin Ke committed
2462
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2463
    const double* data_ptr = reinterpret_cast<const double*>(data);
Guolin Ke's avatar
Guolin Ke committed
2464
    if (col_ptr_type == C_API_DTYPE_INT32) {
2465
      const int32_t* ptr_col_ptr = reinterpret_cast<const int32_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2466
2467
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2468
2469
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2470
2471
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2472
        }
Guolin Ke's avatar
Guolin Ke committed
2473
2474
2475
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2476
      };
Guolin Ke's avatar
Guolin Ke committed
2477
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2478
      const int64_t* ptr_col_ptr = reinterpret_cast<const int64_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2479
2480
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2481
2482
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2483
2484
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2485
        }
Guolin Ke's avatar
Guolin Ke committed
2486
2487
2488
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2489
      };
Guolin Ke's avatar
Guolin Ke committed
2490
2491
    }
  }
2492
  Log::Fatal("Unknown data type in CSC matrix");
2493
  return nullptr;
2494
2495
}

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