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

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

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

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

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

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

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

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

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

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

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

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

89
  ~SingleRowPredictor() {}
90

Guolin Ke's avatar
Guolin Ke committed
91
  bool IsPredictorEqual(const Config& config, int iter, Boosting* boosting) {
92
93
94
95
96
    return early_stop_ == config.pred_early_stop &&
      early_stop_freq_ == config.pred_early_stop_freq &&
      early_stop_margin_ == config.pred_early_stop_margin &&
      iter_ == iter &&
      num_total_model_ == boosting->NumberOfTotalModel();
97
  }
Guolin Ke's avatar
Guolin Ke committed
98

99
100
101
102
103
104
105
106
107
 private:
  std::unique_ptr<Predictor> predictor_;
  bool early_stop_;
  int early_stop_freq_;
  double early_stop_margin_;
  int iter_;
  int num_total_model_;
};

Guolin Ke's avatar
Guolin Ke committed
108
class Booster {
Nikita Titov's avatar
Nikita Titov committed
109
 public:
Guolin Ke's avatar
Guolin Ke committed
110
  explicit Booster(const char* filename) {
111
    boosting_.reset(Boosting::CreateBoosting("gbdt", filename));
112
113
  }

Guolin Ke's avatar
Guolin Ke committed
114
  Booster(const Dataset* train_data,
115
          const char* parameters) {
Guolin Ke's avatar
Guolin Ke committed
116
    auto param = Config::Str2Map(parameters);
wxchan's avatar
wxchan committed
117
    config_.Set(param);
118
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
  void SetSingleRowPredictor(int start_iteration, int num_iteration, int predict_type, const Config& config) {
373
374
375
376
      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(),
377
                                                                         config, start_iteration, num_iteration));
378
379
380
381
      }
  }

  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 start_iteration, 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

416
    Predictor predictor(boosting_.get(), start_iteration, 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
    return predictor;
  }

421
  void Predict(int start_iteration, int num_iteration, int predict_type, int nrow, int ncol,
422
423
               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
    auto predictor = CreatePredictor(start_iteration, num_iteration, predict_type, ncol, config);
427
428
429
430
431
432
433
    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;
    }
434
    int64_t num_pred_in_one_row = boosting_->NumPredictOneRow(start_iteration, num_iteration, is_predict_leaf, predict_contrib);
Guolin Ke's avatar
Guolin Ke committed
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
  void PredictSparse(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
450
451
452
453
                     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
    auto predictor = CreatePredictor(start_iteration, num_iteration, predict_type, ncol, config);
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
    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];
  }

491
  void PredictSparseCSR(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
492
493
494
                        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
    // 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;
514
    PredictSparse(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
515
516
517
                  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);
518
    std::vector<int64_t> matrix_offsets(num_matrices);
519
520
521
522
523
524
525
526
527
528
529
530
531
532
    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++;
      }
533
534
535
536
537
538
      if (m == 0) {
        matrix_offsets[m] = 0;
      }
      if (m + 1 < num_matrices) {
        matrix_offsets[m + 1] = static_cast<int64_t>(matrix_offsets[m] + row_matrix_offsets[row_vector_cnt - 1] + row_sizes[row_vector_cnt - 1]);
      }
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
    }
    // 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;
556
        int64_t element_index = row_matrix_offsets[row_start_index] + matrix_offsets[m];
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
        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;
  }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

827
}  // namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
828

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

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

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

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

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

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

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

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

919

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

942

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

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

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

Guolin Ke's avatar
Guolin Ke committed
1057
1058
  if (reference == nullptr) {
    // sample data first
Guolin Ke's avatar
Guolin Ke committed
1059
    Random rand(config.data_random_seed);
1060
1061
    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);
1062
    sample_cnt = static_cast<int>(sample_indices.size());
1063
    std::vector<std::vector<double>> sample_values(ncol);
Guolin Ke's avatar
Guolin Ke committed
1064
    std::vector<std::vector<int>> sample_idx(ncol);
1065
1066
1067

    int offset = 0;
    int j = 0;
Guolin Ke's avatar
Guolin Ke committed
1068
    for (size_t i = 0; i < sample_indices.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1069
      auto idx = sample_indices[i];
1070
1071
1072
1073
      while ((idx - offset) >= nrow[j]) {
        offset += nrow[j];
        ++j;
      }
1074

1075
1076
1077
1078
1079
      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
1080
        }
Guolin Ke's avatar
Guolin Ke committed
1081
1082
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1083
    DatasetLoader loader(config, nullptr, 1, nullptr);
1084
1085
1086
1087
1088
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             ncol,
                                             VectorSize<double>(sample_values).data(),
                                             sample_cnt, total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1089
  } else {
1090
    ret.reset(new Dataset(total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1091
    ret->CreateValid(
1092
      reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
1093
  }
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
  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
1108
1109
  }
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1110
  *out = ret.release();
1111
  API_END();
1112
1113
}

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

1185
int LGBM_DatasetCreateFromCSRFunc(void* get_row_funptr,
1186
1187
1188
1189
1190
                                  int num_rows,
                                  int64_t num_col,
                                  const char* parameters,
                                  const DatasetHandle reference,
                                  DatasetHandle* out) {
1191
  API_BEGIN();
1192
1193
1194
1195
1196
  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.");
  }
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
  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
1220
        CHECK_LT(inner_data.first, num_col);
1221
1222
1223
1224
1225
1226
1227
        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);
1228
1229
1230
1231
1232
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             static_cast<int>(num_col),
                                             VectorSize<double>(sample_values).data(),
                                             sample_cnt, nrow));
1233
1234
1235
1236
1237
  } else {
    ret.reset(new Dataset(nrow));
    ret->CreateValid(
      reinterpret_cast<const Dataset*>(reference));
  }
1238

1239
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1240
1241
  std::vector<std::pair<int, double>> thread_buffer;
  #pragma omp parallel for schedule(static) private(thread_buffer)
1242
1243
1244
  for (int i = 0; i < num_rows; ++i) {
    OMP_LOOP_EX_BEGIN();
    {
1245
      const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1246
1247
      get_row_fun(i, thread_buffer);
      ret->PushOneRow(tid, i, thread_buffer);
1248
1249
1250
1251
1252
1253
1254
1255
1256
    }
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();
  ret->FinishLoad();
  *out = ret.release();
  API_END();
}

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

Guolin Ke's avatar
Guolin Ke committed
1345
int LGBM_DatasetGetSubset(
1346
  const DatasetHandle handle,
wxchan's avatar
wxchan committed
1347
1348
1349
  const int32_t* used_row_indices,
  int32_t num_used_row_indices,
  const char* parameters,
Guolin Ke's avatar
typo  
Guolin Ke committed
1350
  DatasetHandle* out) {
wxchan's avatar
wxchan committed
1351
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1352
1353
  auto param = Config::Str2Map(parameters);
  Config config;
1354
1355
1356
1357
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
1358
  auto full_dataset = reinterpret_cast<const Dataset*>(handle);
1359
  CHECK_GT(num_used_row_indices, 0);
1360
1361
  const int32_t lower = 0;
  const int32_t upper = full_dataset->num_data() - 1;
1362
  CheckElementsIntervalClosed(used_row_indices, lower, upper, num_used_row_indices, "Used indices of subset");
1363
1364
1365
  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
1366
  auto ret = std::unique_ptr<Dataset>(new Dataset(num_used_row_indices));
1367
  ret->CopyFeatureMapperFrom(full_dataset);
1368
  ret->CopySubrow(full_dataset, used_row_indices, num_used_row_indices, true);
wxchan's avatar
wxchan committed
1369
1370
1371
1372
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1373
int LGBM_DatasetSetFeatureNames(
Guolin Ke's avatar
typo  
Guolin Ke committed
1374
  DatasetHandle handle,
Guolin Ke's avatar
Guolin Ke committed
1375
  const char** feature_names,
Guolin Ke's avatar
Guolin Ke committed
1376
  int num_feature_names) {
Guolin Ke's avatar
Guolin Ke committed
1377
1378
1379
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  std::vector<std::string> feature_names_str;
Guolin Ke's avatar
Guolin Ke committed
1380
  for (int i = 0; i < num_feature_names; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1381
1382
1383
1384
1385
1386
    feature_names_str.emplace_back(feature_names[i]);
  }
  dataset->set_feature_names(feature_names_str);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1387
int LGBM_DatasetGetFeatureNames(
1388
1389
1390
1391
1392
1393
    DatasetHandle handle,
    const int len,
    int* num_feature_names,
    const size_t buffer_len,
    size_t* out_buffer_len,
    char** feature_names) {
1394
  API_BEGIN();
1395
  *out_buffer_len = 0;
1396
1397
  auto dataset = reinterpret_cast<Dataset*>(handle);
  auto inside_feature_name = dataset->feature_names();
Guolin Ke's avatar
Guolin Ke committed
1398
1399
  *num_feature_names = static_cast<int>(inside_feature_name.size());
  for (int i = 0; i < *num_feature_names; ++i) {
1400
1401
1402
1403
1404
    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);
1405
1406
1407
1408
  }
  API_END();
}

1409
1410
1411
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1412
int LGBM_DatasetFree(DatasetHandle handle) {
1413
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1414
  delete reinterpret_cast<Dataset*>(handle);
1415
  API_END();
1416
1417
}

Guolin Ke's avatar
Guolin Ke committed
1418
int LGBM_DatasetSaveBinary(DatasetHandle handle,
1419
                           const char* filename) {
1420
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1421
1422
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->SaveBinaryFile(filename);
1423
  API_END();
1424
1425
}

1426
1427
1428
1429
1430
1431
1432
1433
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
1434
int LGBM_DatasetSetField(DatasetHandle handle,
1435
1436
1437
1438
                         const char* field_name,
                         const void* field_data,
                         int num_element,
                         int type) {
1439
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1440
  auto dataset = reinterpret_cast<Dataset*>(handle);
1441
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1442
  if (type == C_API_DTYPE_FLOAT32) {
Guolin Ke's avatar
Guolin Ke committed
1443
    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
1444
  } else if (type == C_API_DTYPE_INT32) {
Guolin Ke's avatar
Guolin Ke committed
1445
    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
1446
1447
  } 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));
1448
  }
1449
  if (!is_success) { Log::Fatal("Input data type error or field not found"); }
1450
  API_END();
1451
1452
}

Guolin Ke's avatar
Guolin Ke committed
1453
int LGBM_DatasetGetField(DatasetHandle handle,
1454
1455
1456
1457
                         const char* field_name,
                         int* out_len,
                         const void** out_ptr,
                         int* out_type) {
1458
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1459
  auto dataset = reinterpret_cast<Dataset*>(handle);
1460
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1461
  if (dataset->GetFloatField(field_name, out_len, reinterpret_cast<const float**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1462
    *out_type = C_API_DTYPE_FLOAT32;
1463
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1464
  } else if (dataset->GetIntField(field_name, out_len, reinterpret_cast<const int**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1465
    *out_type = C_API_DTYPE_INT32;
1466
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1467
1468
1469
  } 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
1470
  }
1471
  if (!is_success) { Log::Fatal("Field not found"); }
wxchan's avatar
wxchan committed
1472
  if (*out_ptr == nullptr) { *out_len = 0; }
1473
  API_END();
1474
1475
}

1476
int LGBM_DatasetUpdateParamChecking(const char* old_parameters, const char* new_parameters) {
1477
  API_BEGIN();
1478
1479
1480
1481
1482
  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);
1483
1484
1485
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1486
int LGBM_DatasetGetNumData(DatasetHandle handle,
1487
                           int* out) {
1488
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1489
1490
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_data();
1491
  API_END();
1492
1493
}

Guolin Ke's avatar
Guolin Ke committed
1494
int LGBM_DatasetGetNumFeature(DatasetHandle handle,
1495
                              int* out) {
1496
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1497
1498
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_total_features();
1499
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1500
}
1501

1502
1503
1504
1505
1506
int LGBM_DatasetAddFeaturesFrom(DatasetHandle target,
                                DatasetHandle source) {
  API_BEGIN();
  auto target_d = reinterpret_cast<Dataset*>(target);
  auto source_d = reinterpret_cast<Dataset*>(source);
1507
  target_d->AddFeaturesFrom(source_d);
1508
1509
1510
  API_END();
}

1511
1512
// ---- start of booster

Guolin Ke's avatar
Guolin Ke committed
1513
int LGBM_BoosterCreate(const DatasetHandle train_data,
1514
1515
                       const char* parameters,
                       BoosterHandle* out) {
1516
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1517
  const Dataset* p_train_data = reinterpret_cast<const Dataset*>(train_data);
wxchan's avatar
wxchan committed
1518
1519
  auto ret = std::unique_ptr<Booster>(new Booster(p_train_data, parameters));
  *out = ret.release();
1520
  API_END();
1521
1522
}

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

Guolin Ke's avatar
Guolin Ke committed
1534
int LGBM_BoosterLoadModelFromString(
1535
1536
1537
1538
  const char* model_str,
  int* out_num_iterations,
  BoosterHandle* out) {
  API_BEGIN();
wxchan's avatar
wxchan committed
1539
  auto ret = std::unique_ptr<Booster>(new Booster(nullptr));
1540
1541
1542
1543
1544
1545
  ret->LoadModelFromString(model_str);
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
  *out = ret.release();
  API_END();
}

1546
1547
1548
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1549
int LGBM_BoosterFree(BoosterHandle handle) {
1550
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1551
  delete reinterpret_cast<Booster*>(handle);
1552
  API_END();
1553
1554
}

1555
int LGBM_BoosterShuffleModels(BoosterHandle handle, int start_iter, int end_iter) {
1556
1557
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1558
  ref_booster->ShuffleModels(start_iter, end_iter);
1559
1560
1561
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1562
int LGBM_BoosterMerge(BoosterHandle handle,
1563
                      BoosterHandle other_handle) {
wxchan's avatar
wxchan committed
1564
1565
1566
1567
1568
1569
1570
  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
1571
int LGBM_BoosterAddValidData(BoosterHandle handle,
1572
                             const DatasetHandle valid_data) {
wxchan's avatar
wxchan committed
1573
1574
1575
1576
1577
1578
1579
  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
1580
int LGBM_BoosterResetTrainingData(BoosterHandle handle,
1581
                                  const DatasetHandle train_data) {
wxchan's avatar
wxchan committed
1582
1583
1584
1585
1586
1587
1588
  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
1589
int LGBM_BoosterResetParameter(BoosterHandle handle, const char* parameters) {
wxchan's avatar
wxchan committed
1590
1591
1592
1593
1594
1595
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->ResetConfig(parameters);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1596
int LGBM_BoosterGetNumClasses(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1597
1598
1599
1600
1601
1602
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->NumberOfClasses();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1603
1604
1605
1606
1607
1608
1609
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
1610
int LGBM_BoosterUpdateOneIter(BoosterHandle handle, int* is_finished) {
1611
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1612
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1613
1614
1615
1616
1617
  if (ref_booster->TrainOneIter()) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1618
  API_END();
1619
1620
}

Guolin Ke's avatar
Guolin Ke committed
1621
int LGBM_BoosterUpdateOneIterCustom(BoosterHandle handle,
1622
1623
1624
                                    const float* grad,
                                    const float* hess,
                                    int* is_finished) {
1625
  API_BEGIN();
1626
  #ifdef SCORE_T_USE_DOUBLE
1627
1628
1629
1630
  (void) handle;       // UNUSED VARIABLE
  (void) grad;         // UNUSED VARIABLE
  (void) hess;         // UNUSED VARIABLE
  (void) is_finished;  // UNUSED VARIABLE
1631
  Log::Fatal("Don't support custom loss function when SCORE_T_USE_DOUBLE is enabled");
1632
  #else
1633
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1634
1635
1636
1637
1638
  if (ref_booster->TrainOneIter(grad, hess)) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1639
  #endif
1640
  API_END();
1641
1642
}

Guolin Ke's avatar
Guolin Ke committed
1643
int LGBM_BoosterRollbackOneIter(BoosterHandle handle) {
wxchan's avatar
wxchan committed
1644
1645
1646
1647
1648
1649
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->RollbackOneIter();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1650
int LGBM_BoosterGetCurrentIteration(BoosterHandle handle, int* out_iteration) {
wxchan's avatar
wxchan committed
1651
1652
1653
1654
1655
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_iteration = ref_booster->GetBoosting()->GetCurrentIteration();
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
1656

1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
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
1671
int LGBM_BoosterGetEvalCounts(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1672
1673
1674
1675
1676
1677
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetEvalCounts();
  API_END();
}

1678
1679
1680
1681
1682
1683
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
1684
1685
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1686
  *out_len = ref_booster->GetEvalNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1687
1688
1689
  API_END();
}

1690
1691
1692
1693
1694
1695
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
1696
1697
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1698
  *out_len = ref_booster->GetFeatureNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1699
1700
1701
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1702
int LGBM_BoosterGetNumFeature(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1703
1704
1705
1706
1707
1708
  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
1709
int LGBM_BoosterGetEval(BoosterHandle handle,
1710
1711
1712
                        int data_idx,
                        int* out_len,
                        double* out_results) {
1713
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1714
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1715
  auto boosting = ref_booster->GetBoosting();
wxchan's avatar
wxchan committed
1716
  auto result_buf = boosting->GetEvalAt(data_idx);
Guolin Ke's avatar
Guolin Ke committed
1717
  *out_len = static_cast<int>(result_buf.size());
1718
  for (size_t i = 0; i < result_buf.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1719
    (out_results)[i] = static_cast<double>(result_buf[i]);
1720
  }
1721
  API_END();
1722
1723
}

Guolin Ke's avatar
Guolin Ke committed
1724
int LGBM_BoosterGetNumPredict(BoosterHandle handle,
1725
1726
                              int data_idx,
                              int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1727
1728
1729
1730
1731
1732
  API_BEGIN();
  auto boosting = reinterpret_cast<Booster*>(handle)->GetBoosting();
  *out_len = boosting->GetNumPredictAt(data_idx);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1733
int LGBM_BoosterGetPredict(BoosterHandle handle,
1734
1735
1736
                           int data_idx,
                           int64_t* out_len,
                           double* out_result) {
1737
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1738
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1739
  ref_booster->GetPredictAt(data_idx, out_result, out_len);
1740
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1741
1742
}

Guolin Ke's avatar
Guolin Ke committed
1743
int LGBM_BoosterPredictForFile(BoosterHandle handle,
1744
1745
1746
                               const char* data_filename,
                               int data_has_header,
                               int predict_type,
1747
                               int start_iteration,
1748
                               int num_iteration,
1749
                               const char* parameter,
1750
                               const char* result_filename) {
1751
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1752
1753
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1754
1755
1756
1757
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1758
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1759
  ref_booster->Predict(start_iteration, num_iteration, predict_type, data_filename, data_has_header,
Guolin Ke's avatar
Guolin Ke committed
1760
                       config, result_filename);
1761
  API_END();
1762
1763
}

Guolin Ke's avatar
Guolin Ke committed
1764
int LGBM_BoosterCalcNumPredict(BoosterHandle handle,
1765
1766
                               int num_row,
                               int predict_type,
1767
                               int start_iteration,
1768
1769
                               int num_iteration,
                               int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1770
1771
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1772
  *out_len = static_cast<int64_t>(num_row) * ref_booster->GetBoosting()->NumPredictOneRow(start_iteration,
1773
    num_iteration, predict_type == C_API_PREDICT_LEAF_INDEX, predict_type == C_API_PREDICT_CONTRIB);
Guolin Ke's avatar
Guolin Ke committed
1774
1775
1776
  API_END();
}

1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
/*!
 * \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,
1790
             const int predict_type_,
1791
             const int data_type_,
1792
             const int32_t num_cols) : booster(booster_ptr), predict_type(predict_type_), data_type(data_type_), ncol(num_cols) {
1793
1794
1795
1796
1797
    config.Set(Config::Str2Map(parameter));
  }

  Booster* const booster;
  Config config;
1798
  const int predict_type;
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
  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
1809
int LGBM_BoosterPredictForCSR(BoosterHandle handle,
1810
1811
1812
1813
1814
1815
1816
                              const void* indptr,
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
1817
                              int64_t num_col,
1818
                              int predict_type,
1819
                              int start_iteration,
1820
                              int num_iteration,
1821
                              const char* parameter,
1822
1823
                              int64_t* out_len,
                              double* out_result) {
1824
  API_BEGIN();
1825
1826
1827
1828
1829
  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
1830
1831
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1832
1833
1834
1835
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1836
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1837
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
1838
  int nrow = static_cast<int>(nindptr - 1);
1839
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, static_cast<int>(num_col), get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
1840
                       config, out_result, out_len);
1841
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1842
}
1843

1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
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,
1854
                                    int start_iteration,
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
                                    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;
1878
    ref_booster->PredictSparseCSR(start_iteration, num_iteration, predict_type, nrow, static_cast<int>(num_col_or_row), get_row_fun,
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
                                  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;
    };
1902
    ref_booster->PredictSparseCSC(start_iteration, num_iteration, predict_type, num_col_or_row, ncol, get_row_fun, config,
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
                                  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();
}

1930
int LGBM_BoosterPredictForCSRSingleRow(BoosterHandle handle,
1931
1932
1933
1934
1935
1936
1937
                                       const void* indptr,
                                       int indptr_type,
                                       const int32_t* indices,
                                       const void* data,
                                       int data_type,
                                       int64_t nindptr,
                                       int64_t nelem,
1938
                                       int64_t num_col,
1939
                                       int predict_type,
1940
                                       int start_iteration,
1941
1942
1943
1944
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
1945
  API_BEGIN();
1946
1947
1948
1949
1950
  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.");
  }
1951
1952
1953
1954
1955
1956
1957
  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);
1958
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1959
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
1960
  ref_booster->PredictSingleRow(predict_type, static_cast<int32_t>(num_col), get_row_fun, config, out_result, out_len);
1961
1962
1963
  API_END();
}

1964
int LGBM_BoosterPredictForCSRSingleRowFastInit(BoosterHandle handle,
1965
                                               const int predict_type,
1966
                                               const int start_iteration,
1967
                                               const int num_iteration,
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
                                               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,
1982
    predict_type,
1983
1984
1985
1986
1987
1988
1989
    data_type,
    static_cast<int32_t>(num_col)));

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

1990
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
1991

1992
1993
1994
1995
1996
1997
  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

int LGBM_BoosterPredictForCSRSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* indptr,
1998
                                           const int indptr_type,
1999
2000
                                           const int32_t* indices,
                                           const void* data,
2001
2002
                                           const int64_t nindptr,
                                           const int64_t nelem,
2003
2004
2005
2006
2007
                                           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);
2008
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
2009
2010
2011
2012
                                        get_row_fun, fastConfig->config, out_result, out_len);
  API_END();
}

2013

Guolin Ke's avatar
Guolin Ke committed
2014
int LGBM_BoosterPredictForCSC(BoosterHandle handle,
2015
2016
2017
2018
2019
2020
2021
2022
2023
                              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,
2024
                              int start_iteration,
2025
                              int num_iteration,
2026
                              const char* parameter,
2027
2028
                              int64_t* out_len,
                              double* out_result) {
Guolin Ke's avatar
Guolin Ke committed
2029
2030
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2031
2032
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2033
2034
2035
2036
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
2037
  int num_threads = OMP_NUM_THREADS();
Guolin Ke's avatar
Guolin Ke committed
2038
  int ncol = static_cast<int>(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
2039
2040
2041
2042
2043
  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
2044
2045
  }
  std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun =
Guolin Ke's avatar
Guolin Ke committed
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
      [&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;
      };
2058
  ref_booster->Predict(start_iteration, num_iteration, predict_type, static_cast<int>(num_row), ncol, get_row_fun, config,
cbecker's avatar
cbecker committed
2059
                       out_result, out_len);
Guolin Ke's avatar
Guolin Ke committed
2060
2061
2062
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2063
int LGBM_BoosterPredictForMat(BoosterHandle handle,
2064
2065
2066
2067
2068
2069
                              const void* data,
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              int predict_type,
2070
                              int start_iteration,
2071
                              int num_iteration,
2072
                              const char* parameter,
2073
2074
                              int64_t* out_len,
                              double* out_result) {
2075
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2076
2077
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2078
2079
2080
2081
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
2082
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2083
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, nrow, ncol, data_type, is_row_major);
2084
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
2085
                       config, out_result, out_len);
2086
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2087
}
2088

2089
int LGBM_BoosterPredictForMatSingleRow(BoosterHandle handle,
2090
2091
2092
2093
2094
                                       const void* data,
                                       int data_type,
                                       int32_t ncol,
                                       int is_row_major,
                                       int predict_type,
2095
                                       int start_iteration,
2096
2097
2098
2099
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2100
2101
2102
2103
2104
2105
2106
2107
2108
  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);
2109
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
2110
  ref_booster->PredictSingleRow(predict_type, ncol, get_row_fun, config, out_result, out_len);
2111
2112
2113
  API_END();
}

2114
int LGBM_BoosterPredictForMatSingleRowFastInit(BoosterHandle handle,
2115
                                               const int predict_type,
2116
                                               const int start_iteration,
2117
                                               const int num_iteration,
2118
2119
2120
2121
2122
2123
2124
2125
                                               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,
2126
    predict_type,
2127
2128
2129
2130
2131
2132
2133
    data_type,
    ncol));

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

2134
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2135

2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

int LGBM_BoosterPredictForMatSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* data,
                                           int64_t* out_len,
                                           double* out_result) {
  API_BEGIN();
  FastConfig *fastConfig = reinterpret_cast<FastConfig*>(fastConfig_handle);
  // Single row in row-major format:
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, fastConfig->ncol, fastConfig->data_type, 1);
2148
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
2149
2150
2151
2152
2153
                                        get_row_fun, fastConfig->config,
                                        out_result, out_len);
  API_END();
}

2154

2155
2156
2157
2158
2159
2160
int LGBM_BoosterPredictForMats(BoosterHandle handle,
                               const void** data,
                               int data_type,
                               int32_t nrow,
                               int32_t ncol,
                               int predict_type,
2161
                               int start_iteration,
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
                               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);
2175
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, out_result, out_len);
2176
2177
2178
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2179
int LGBM_BoosterSaveModel(BoosterHandle handle,
2180
                          int start_iteration,
2181
                          int num_iteration,
2182
                          int feature_importance_type,
2183
                          const char* filename) {
2184
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2185
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2186
2187
  ref_booster->SaveModelToFile(start_iteration, num_iteration,
                               feature_importance_type, filename);
wxchan's avatar
wxchan committed
2188
2189
2190
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2191
int LGBM_BoosterSaveModelToString(BoosterHandle handle,
2192
                                  int start_iteration,
2193
                                  int num_iteration,
2194
                                  int feature_importance_type,
2195
                                  int64_t buffer_len,
2196
                                  int64_t* out_len,
2197
                                  char* out_str) {
2198
2199
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2200
2201
  std::string model = ref_booster->SaveModelToString(
      start_iteration, num_iteration, feature_importance_type);
2202
  *out_len = static_cast<int64_t>(model.size()) + 1;
2203
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2204
    std::memcpy(out_str, model.c_str(), *out_len);
2205
2206
2207
2208
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2209
int LGBM_BoosterDumpModel(BoosterHandle handle,
2210
                          int start_iteration,
2211
                          int num_iteration,
2212
                          int feature_importance_type,
2213
2214
                          int64_t buffer_len,
                          int64_t* out_len,
2215
                          char* out_str) {
wxchan's avatar
wxchan committed
2216
2217
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2218
2219
  std::string model = ref_booster->DumpModel(start_iteration, num_iteration,
                                             feature_importance_type);
2220
  *out_len = static_cast<int64_t>(model.size()) + 1;
wxchan's avatar
wxchan committed
2221
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2222
    std::memcpy(out_str, model.c_str(), *out_len);
wxchan's avatar
wxchan committed
2223
  }
2224
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2225
}
2226

Guolin Ke's avatar
Guolin Ke committed
2227
int LGBM_BoosterGetLeafValue(BoosterHandle handle,
2228
2229
2230
                             int tree_idx,
                             int leaf_idx,
                             double* out_val) {
Guolin Ke's avatar
Guolin Ke committed
2231
2232
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2233
  *out_val = static_cast<double>(ref_booster->GetLeafValue(tree_idx, leaf_idx));
Guolin Ke's avatar
Guolin Ke committed
2234
2235
2236
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2237
int LGBM_BoosterSetLeafValue(BoosterHandle handle,
2238
2239
2240
                             int tree_idx,
                             int leaf_idx,
                             double val) {
Guolin Ke's avatar
Guolin Ke committed
2241
2242
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2243
  ref_booster->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
2244
2245
2246
  API_END();
}

2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
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();
}

2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
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();
}

2278
2279
2280
2281
2282
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
2283
  Config config;
2284
  config.machines = RemoveQuotationSymbol(std::string(machines));
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
  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();
}

2300
2301
2302
int LGBM_NetworkInitWithFunctions(int num_machines, int rank,
                                  void* reduce_scatter_ext_fun,
                                  void* allgather_ext_fun) {
ww's avatar
ww committed
2303
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2304
  if (num_machines > 1) {
2305
    Network::Init(num_machines, rank, (ReduceScatterFunction)reduce_scatter_ext_fun, (AllgatherFunction)allgather_ext_fun);
ww's avatar
ww committed
2306
2307
2308
  }
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
2309

Guolin Ke's avatar
Guolin Ke committed
2310
// ---- start of some help functions
2311
2312
2313

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
2314
  if (data_type == C_API_DTYPE_FLOAT32) {
2315
2316
    const float* data_ptr = reinterpret_cast<const float*>(data);
    if (is_row_major) {
2317
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2318
        std::vector<double> ret(num_col);
2319
        auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
2320
        for (int i = 0; i < num_col; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2321
          ret[i] = static_cast<double>(*(tmp_ptr + i));
2322
2323
2324
2325
        }
        return ret;
      };
    } else {
2326
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2327
        std::vector<double> ret(num_col);
2328
        for (int i = 0; i < num_col; ++i) {
2329
          ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
2330
2331
2332
2333
        }
        return ret;
      };
    }
Guolin Ke's avatar
Guolin Ke committed
2334
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2335
2336
    const double* data_ptr = reinterpret_cast<const double*>(data);
    if (is_row_major) {
2337
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2338
        std::vector<double> ret(num_col);
2339
        auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
2340
        for (int i = 0; i < num_col; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2341
          ret[i] = static_cast<double>(*(tmp_ptr + i));
2342
2343
2344
2345
        }
        return ret;
      };
    } else {
2346
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2347
        std::vector<double> ret(num_col);
2348
        for (int i = 0; i < num_col; ++i) {
2349
          ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
2350
2351
2352
2353
2354
        }
        return ret;
      };
    }
  }
2355
  Log::Fatal("Unknown data type in RowFunctionFromDenseMatric");
2356
  return nullptr;
2357
2358
2359
2360
}

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
2361
2362
  auto inner_function = RowFunctionFromDenseMatric(data, num_row, num_col, data_type, is_row_major);
  if (inner_function != nullptr) {
2363
    return [inner_function] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2364
2365
      auto raw_values = inner_function(row_idx);
      std::vector<std::pair<int, double>> ret;
Guolin Ke's avatar
Guolin Ke committed
2366
      ret.reserve(raw_values.size());
Guolin Ke's avatar
Guolin Ke committed
2367
      for (int i = 0; i < static_cast<int>(raw_values.size()); ++i) {
Guolin Ke's avatar
Guolin Ke committed
2368
        if (std::fabs(raw_values[i]) > kZeroThreshold || std::isnan(raw_values[i])) {
Guolin Ke's avatar
Guolin Ke committed
2369
          ret.emplace_back(i, raw_values[i]);
2370
        }
Guolin Ke's avatar
Guolin Ke committed
2371
2372
2373
      }
      return ret;
    };
2374
  }
Guolin Ke's avatar
Guolin Ke committed
2375
  return nullptr;
2376
2377
}

2378
2379
2380
2381
2382
2383
2384
// 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
2385
    ret.reserve(raw_values.size());
2386
2387
2388
2389
2390
2391
2392
2393
2394
    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;
  };
}

2395
2396
template<typename T>
std::function<std::vector<std::pair<int, double>>(T idx)>
2397
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
2398
  if (data_type == C_API_DTYPE_FLOAT32) {
2399
    const float* data_ptr = reinterpret_cast<const float*>(data);
Guolin Ke's avatar
Guolin Ke committed
2400
    if (indptr_type == C_API_DTYPE_INT32) {
2401
      const int32_t* ptr_indptr = reinterpret_cast<const int32_t*>(indptr);
2402
      return [=] (T idx) {
2403
2404
2405
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2406
2407
2408
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2409
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2410
          ret.emplace_back(indices[i], data_ptr[i]);
2411
2412
2413
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
2414
    } else if (indptr_type == C_API_DTYPE_INT64) {
2415
      const int64_t* ptr_indptr = reinterpret_cast<const int64_t*>(indptr);
2416
      return [=] (T idx) {
2417
2418
2419
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2420
2421
2422
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2423
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2424
          ret.emplace_back(indices[i], data_ptr[i]);
2425
2426
2427
2428
        }
        return ret;
      };
    }
Guolin Ke's avatar
Guolin Ke committed
2429
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2430
    const double* data_ptr = reinterpret_cast<const double*>(data);
Guolin Ke's avatar
Guolin Ke committed
2431
    if (indptr_type == C_API_DTYPE_INT32) {
2432
      const int32_t* ptr_indptr = reinterpret_cast<const int32_t*>(indptr);
2433
      return [=] (T idx) {
2434
2435
2436
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2437
2438
2439
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2440
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2441
          ret.emplace_back(indices[i], data_ptr[i]);
2442
2443
2444
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
2445
    } else if (indptr_type == C_API_DTYPE_INT64) {
2446
      const int64_t* ptr_indptr = reinterpret_cast<const int64_t*>(indptr);
2447
      return [=] (T idx) {
2448
2449
2450
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2451
2452
2453
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2454
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2455
          ret.emplace_back(indices[i], data_ptr[i]);
2456
2457
2458
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
2459
2460
    }
  }
2461
  Log::Fatal("Unknown data type in RowFunctionFromCSR");
2462
  return nullptr;
2463
2464
}

Guolin Ke's avatar
Guolin Ke committed
2465
std::function<std::pair<int, double>(int idx)>
2466
IterateFunctionFromCSC(const void* col_ptr, int col_ptr_type, const int32_t* indices, const void* data, int data_type, int64_t ncol_ptr, int64_t , int col_idx) {
Guolin Ke's avatar
Guolin Ke committed
2467
  CHECK(col_idx < ncol_ptr && col_idx >= 0);
Guolin Ke's avatar
Guolin Ke committed
2468
  if (data_type == C_API_DTYPE_FLOAT32) {
2469
    const float* data_ptr = reinterpret_cast<const float*>(data);
Guolin Ke's avatar
Guolin Ke committed
2470
    if (col_ptr_type == C_API_DTYPE_INT32) {
2471
      const int32_t* ptr_col_ptr = reinterpret_cast<const int32_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2472
2473
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2474
2475
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2476
2477
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2478
        }
Guolin Ke's avatar
Guolin Ke committed
2479
2480
2481
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2482
      };
Guolin Ke's avatar
Guolin Ke committed
2483
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2484
      const int64_t* ptr_col_ptr = reinterpret_cast<const int64_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2485
2486
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2487
2488
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2489
2490
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2491
        }
Guolin Ke's avatar
Guolin Ke committed
2492
2493
2494
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2495
      };
Guolin Ke's avatar
Guolin Ke committed
2496
    }
Guolin Ke's avatar
Guolin Ke committed
2497
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2498
    const double* data_ptr = reinterpret_cast<const double*>(data);
Guolin Ke's avatar
Guolin Ke committed
2499
    if (col_ptr_type == C_API_DTYPE_INT32) {
2500
      const int32_t* ptr_col_ptr = reinterpret_cast<const int32_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2501
2502
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2503
2504
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2505
2506
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2507
        }
Guolin Ke's avatar
Guolin Ke committed
2508
2509
2510
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2511
      };
Guolin Ke's avatar
Guolin Ke committed
2512
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2513
      const int64_t* ptr_col_ptr = reinterpret_cast<const int64_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2514
2515
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2516
2517
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2518
2519
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2520
        }
Guolin Ke's avatar
Guolin Ke committed
2521
2522
2523
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2524
      };
Guolin Ke's avatar
Guolin Ke committed
2525
2526
    }
  }
2527
  Log::Fatal("Unknown data type in CSC matrix");
2528
  return nullptr;
2529
2530
}

Guolin Ke's avatar
Guolin Ke committed
2531
CSC_RowIterator::CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
2532
                                 const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx) {
Guolin Ke's avatar
Guolin Ke committed
2533
2534
2535
2536
2537
2538
2539
2540
2541
  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;
2542
    }
Guolin Ke's avatar
Guolin Ke committed
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
    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;
2560
    }
Guolin Ke's avatar
Guolin Ke committed
2561
2562
2563
    return ret;
  } else {
    return std::make_pair(-1, 0.0);
2564
  }
Guolin Ke's avatar
Guolin Ke committed
2565
}