c_api.cpp 96.6 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
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
                  out_indices, out_data, data_type, &is_data_float32, num_matrices);
    std::vector<int> row_sizes(num_matrices * nrow);
    std::vector<int64_t> row_matrix_offsets(num_matrices * nrow);
    int64_t row_vector_cnt = 0;
    for (int m = 0; m < num_matrices; ++m) {
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        auto row_vector = agg[i];
        auto row_vector_size = row_vector[m].size();
        // keep track of the row_vector sizes for parallelization
        row_sizes[row_vector_cnt] = static_cast<int>(row_vector_size);
        if (i == 0) {
          row_matrix_offsets[row_vector_cnt] = 0;
        } else {
          row_matrix_offsets[row_vector_cnt] = static_cast<int64_t>(row_sizes[row_vector_cnt - 1] + row_matrix_offsets[row_vector_cnt - 1]);
        }
        row_vector_cnt++;
      }
    }
    // copy vector results to output for each row
    int64_t indptr_index = 0;
    for (int m = 0; m < num_matrices; ++m) {
      if (is_indptr_int32) {
        (reinterpret_cast<int32_t*>(*out_indptr))[indptr_index] = 0;
      } else {
        (reinterpret_cast<int64_t*>(*out_indptr))[indptr_index] = 0;
      }
      indptr_index++;
      int64_t matrix_start_index = m * static_cast<int64_t>(agg.size());
      OMP_INIT_EX();
      #pragma omp parallel for schedule(static)
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        OMP_LOOP_EX_BEGIN();
        auto row_vector = agg[i];
        int64_t row_start_index = matrix_start_index + i;
        int64_t element_index = row_matrix_offsets[row_start_index];
        int64_t indptr_loop_index = indptr_index + i;
        for (auto it = row_vector[m].begin(); it != row_vector[m].end(); ++it) {
          (*out_indices)[element_index] = it->first;
          if (is_data_float32) {
            (reinterpret_cast<float*>(*out_data))[element_index] = static_cast<float>(it->second);
          } else {
            (reinterpret_cast<double*>(*out_data))[element_index] = it->second;
          }
          element_index++;
        }
        int64_t indptr_value = row_matrix_offsets[row_start_index] + row_sizes[row_start_index];
        if (is_indptr_int32) {
          (reinterpret_cast<int32_t*>(*out_indptr))[indptr_loop_index] = static_cast<int32_t>(indptr_value);
        } else {
          (reinterpret_cast<int64_t*>(*out_indptr))[indptr_loop_index] = indptr_value;
        }
        OMP_LOOP_EX_END();
      }
      OMP_THROW_EX();
      indptr_index += static_cast<int64_t>(agg.size());
    }
    out_len[0] = elements_size;
    out_len[1] = indptr_size;
  }

575
  void PredictSparseCSC(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
576
577
578
                        std::function<std::vector<std::pair<int, double>>(int64_t row_idx)> get_row_fun,
                        const Config& config,
                        int64_t* out_len, void** out_col_ptr, int col_ptr_type,
579
580
                        int32_t** out_indices, void** out_data, int data_type) const {
    SHARED_LOCK(mutex_);
581
582
    // Get the number of trees per iteration (for multiclass scenario we output multiple sparse matrices)
    int num_matrices = boosting_->NumModelPerIteration();
583
    auto predictor = CreatePredictor(start_iteration, num_iteration, predict_type, ncol, config);
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
    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;
601
    PredictSparse(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
                  out_indices, out_data, data_type, &is_data_float32, num_matrices);
    // calculate number of elements per column to construct
    // the CSC matrix with random access
    std::vector<std::vector<int64_t>> column_sizes(num_matrices);
    for (int m = 0; m < num_matrices; ++m) {
      column_sizes[m] = std::vector<int64_t>(num_output_cols, 0);
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        auto row_vector = agg[i];
        for (auto it = row_vector[m].begin(); it != row_vector[m].end(); ++it) {
          column_sizes[m][it->first] += 1;
        }
      }
    }
    // keep track of column counts
    std::vector<std::vector<int64_t>> column_counts(num_matrices);
    // keep track of beginning index for each column
    std::vector<std::vector<int64_t>> column_start_indices(num_matrices);
    // keep track of beginning index for each matrix
    std::vector<int64_t> matrix_start_indices(num_matrices, 0);
    int col_ptr_index = 0;
    for (int m = 0; m < num_matrices; ++m) {
      int64_t col_ptr_value = 0;
      column_start_indices[m] = std::vector<int64_t>(num_output_cols, 0);
      column_counts[m] = std::vector<int64_t>(num_output_cols, 0);
      if (is_col_ptr_int32) {
        (reinterpret_cast<int32_t*>(*out_col_ptr))[col_ptr_index] = static_cast<int32_t>(col_ptr_value);
      } else {
        (reinterpret_cast<int64_t*>(*out_col_ptr))[col_ptr_index] = col_ptr_value;
      }
      col_ptr_index++;
      for (int64_t i = 1; i < static_cast<int64_t>(column_sizes[m].size()); ++i) {
        column_start_indices[m][i] = column_sizes[m][i - 1] + column_start_indices[m][i - 1];
        if (is_col_ptr_int32) {
          (reinterpret_cast<int32_t*>(*out_col_ptr))[col_ptr_index] = static_cast<int32_t>(column_start_indices[m][i]);
        } else {
          (reinterpret_cast<int64_t*>(*out_col_ptr))[col_ptr_index] = column_start_indices[m][i];
        }
        col_ptr_index++;
      }
      int64_t last_elem_index = static_cast<int64_t>(column_sizes[m].size()) - 1;
      int64_t last_column_start_index = column_start_indices[m][last_elem_index];
      int64_t last_column_size = column_sizes[m][last_elem_index];
      if (is_col_ptr_int32) {
        (reinterpret_cast<int32_t*>(*out_col_ptr))[col_ptr_index] = static_cast<int32_t>(last_column_start_index + last_column_size);
      } else {
        (reinterpret_cast<int64_t*>(*out_col_ptr))[col_ptr_index] = last_column_start_index + last_column_size;
      }
      if (m != 0) {
        matrix_start_indices[m] = matrix_start_indices[m - 1] +
          last_column_start_index +
          last_column_size;
      }
    }
    for (int m = 0; m < num_matrices; ++m) {
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        auto row_vector = agg[i];
        for (auto it = row_vector[m].begin(); it != row_vector[m].end(); ++it) {
          int64_t col_idx = it->first;
          int64_t element_index = column_start_indices[m][col_idx] +
            matrix_start_indices[m] +
            column_counts[m][col_idx];
          // store the row index
          (*out_indices)[element_index] = static_cast<int32_t>(i);
          // update column count
          column_counts[m][col_idx]++;
          if (is_data_float32) {
            (reinterpret_cast<float*>(*out_data))[element_index] = static_cast<float>(it->second);
          } else {
            (reinterpret_cast<double*>(*out_data))[element_index] = it->second;
          }
        }
      }
    }
    out_len[0] = elements_size;
    out_len[1] = col_ptr_size;
  }

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

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

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

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

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

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

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

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

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

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

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

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

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

764
  int GetEvalNames(char** out_strs, const int len, const size_t buffer_len, size_t *out_buffer_len) const {
765
    SHARED_LOCK(mutex_)
766
    *out_buffer_len = 0;
wxchan's avatar
wxchan committed
767
768
769
    int idx = 0;
    for (const auto& metric : train_metric_) {
      for (const auto& name : metric->GetName()) {
770
771
772
773
774
        if (idx < len) {
          std::memcpy(out_strs[idx], name.c_str(), std::min(name.size() + 1, buffer_len));
          out_strs[idx][buffer_len - 1] = '\0';
        }
        *out_buffer_len = std::max(name.size() + 1, *out_buffer_len);
wxchan's avatar
wxchan committed
775
776
777
778
779
780
        ++idx;
      }
    }
    return idx;
  }

781
  int GetFeatureNames(char** out_strs, const int len, const size_t buffer_len, size_t *out_buffer_len) const {
782
    SHARED_LOCK(mutex_)
783
    *out_buffer_len = 0;
wxchan's avatar
wxchan committed
784
785
    int idx = 0;
    for (const auto& name : boosting_->FeatureNames()) {
786
787
788
789
790
      if (idx < len) {
        std::memcpy(out_strs[idx], name.c_str(), std::min(name.size() + 1, buffer_len));
        out_strs[idx][buffer_len - 1] = '\0';
      }
      *out_buffer_len = std::max(name.size() + 1, *out_buffer_len);
wxchan's avatar
wxchan committed
791
792
793
794
795
      ++idx;
    }
    return idx;
  }

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

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

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

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

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

Guolin Ke's avatar
Guolin Ke committed
835
836
837
838
839
840
841
842
// some help functions used to convert data

std::function<std::vector<double>(int row_idx)>
RowFunctionFromDenseMatric(const void* data, int num_row, int num_col, int data_type, int is_row_major);

std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseMatric(const void* data, int num_row, int num_col, int data_type, int is_row_major);

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

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

// Row iterator of on column for CSC matrix
class CSC_RowIterator {
Nikita Titov's avatar
Nikita Titov committed
853
 public:
Guolin Ke's avatar
Guolin Ke committed
854
  CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
855
                  const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx);
Guolin Ke's avatar
Guolin Ke committed
856
857
858
859
860
  ~CSC_RowIterator() {}
  // return value at idx, only can access by ascent order
  double Get(int idx);
  // return next non-zero pair, if index < 0, means no more data
  std::pair<int, double> NextNonZero();
Nikita Titov's avatar
Nikita Titov committed
861
862

 private:
Guolin Ke's avatar
Guolin Ke committed
863
864
865
866
867
868
869
870
871
  int nonzero_idx_ = 0;
  int cur_idx_ = -1;
  double cur_val_ = 0.0f;
  bool is_end_ = false;
  std::function<std::pair<int, double>(int idx)> iter_fun_;
};

// start of c_api functions

Guolin Ke's avatar
Guolin Ke committed
872
const char* LGBM_GetLastError() {
wxchan's avatar
wxchan committed
873
  return LastErrorMsg();
Guolin Ke's avatar
Guolin Ke committed
874
875
}

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

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

907

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

930

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

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

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

Guolin Ke's avatar
Guolin Ke committed
998
int LGBM_DatasetCreateFromMat(const void* data,
999
1000
1001
1002
1003
1004
1005
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
  return LGBM_DatasetCreateFromMats(1,
                                    &data,
                                    data_type,
                                    &nrow,
                                    ncol,
                                    is_row_major,
                                    parameters,
                                    reference,
                                    out);
}


int LGBM_DatasetCreateFromMats(int32_t nmat,
                               const void** data,
                               int data_type,
                               int32_t* nrow,
                               int32_t ncol,
                               int is_row_major,
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
1027
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1028
1029
  auto param = Config::Str2Map(parameters);
  Config config;
1030
1031
1032
1033
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1034
  std::unique_ptr<Dataset> ret;
1035
1036
1037
1038
1039
1040
1041
1042
1043
  int32_t total_nrow = 0;
  for (int j = 0; j < nmat; ++j) {
    total_nrow += nrow[j];
  }

  std::vector<std::function<std::vector<double>(int row_idx)>> get_row_fun;
  for (int j = 0; j < nmat; ++j) {
    get_row_fun.push_back(RowFunctionFromDenseMatric(data[j], nrow[j], ncol, data_type, is_row_major));
  }
1044

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

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

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

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

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

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

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

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

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

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

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

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

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

1414
1415
1416
1417
1418
1419
1420
1421
int LGBM_DatasetDumpText(DatasetHandle handle,
                         const char* filename) {
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->DumpTextFile(filename);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1422
int LGBM_DatasetSetField(DatasetHandle handle,
1423
1424
1425
1426
                         const char* field_name,
                         const void* field_data,
                         int num_element,
                         int type) {
1427
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1428
  auto dataset = reinterpret_cast<Dataset*>(handle);
1429
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1430
  if (type == C_API_DTYPE_FLOAT32) {
Guolin Ke's avatar
Guolin Ke committed
1431
    is_success = dataset->SetFloatField(field_name, reinterpret_cast<const float*>(field_data), static_cast<int32_t>(num_element));
Guolin Ke's avatar
Guolin Ke committed
1432
  } else if (type == C_API_DTYPE_INT32) {
Guolin Ke's avatar
Guolin Ke committed
1433
    is_success = dataset->SetIntField(field_name, reinterpret_cast<const int*>(field_data), static_cast<int32_t>(num_element));
Guolin Ke's avatar
Guolin Ke committed
1434
1435
  } else if (type == C_API_DTYPE_FLOAT64) {
    is_success = dataset->SetDoubleField(field_name, reinterpret_cast<const double*>(field_data), static_cast<int32_t>(num_element));
1436
  }
1437
  if (!is_success) { Log::Fatal("Input data type error or field not found"); }
1438
  API_END();
1439
1440
}

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

1464
int LGBM_DatasetUpdateParamChecking(const char* old_parameters, const char* new_parameters) {
1465
  API_BEGIN();
1466
1467
1468
1469
1470
  auto old_param = Config::Str2Map(old_parameters);
  Config old_config;
  old_config.Set(old_param);
  auto new_param = Config::Str2Map(new_parameters);
  Booster::CheckDatasetResetConfig(old_config, new_param);
1471
1472
1473
  API_END();
}

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

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

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

1499
1500
// ---- start of booster

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

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

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

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

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

Guolin Ke's avatar
Guolin Ke committed
1550
int LGBM_BoosterMerge(BoosterHandle handle,
1551
                      BoosterHandle other_handle) {
wxchan's avatar
wxchan committed
1552
1553
1554
1555
1556
1557
1558
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  Booster* ref_other_booster = reinterpret_cast<Booster*>(other_handle);
  ref_booster->MergeFrom(ref_other_booster);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1559
int LGBM_BoosterAddValidData(BoosterHandle handle,
1560
                             const DatasetHandle valid_data) {
wxchan's avatar
wxchan committed
1561
1562
1563
1564
1565
1566
1567
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  const Dataset* p_dataset = reinterpret_cast<const Dataset*>(valid_data);
  ref_booster->AddValidData(p_dataset);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1568
int LGBM_BoosterResetTrainingData(BoosterHandle handle,
1569
                                  const DatasetHandle train_data) {
wxchan's avatar
wxchan committed
1570
1571
1572
1573
1574
1575
1576
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  const Dataset* p_dataset = reinterpret_cast<const Dataset*>(train_data);
  ref_booster->ResetTrainingData(p_dataset);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1577
int LGBM_BoosterResetParameter(BoosterHandle handle, const char* parameters) {
wxchan's avatar
wxchan committed
1578
1579
1580
1581
1582
1583
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->ResetConfig(parameters);
  API_END();
}

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

Guolin Ke's avatar
Guolin Ke committed
1591
1592
1593
1594
1595
1596
1597
int LGBM_BoosterRefit(BoosterHandle handle, const int32_t* leaf_preds, int32_t nrow, int32_t ncol) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->Refit(leaf_preds, nrow, ncol);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1598
int LGBM_BoosterUpdateOneIter(BoosterHandle handle, int* is_finished) {
1599
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1600
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1601
1602
1603
1604
1605
  if (ref_booster->TrainOneIter()) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1606
  API_END();
1607
1608
}

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

Guolin Ke's avatar
Guolin Ke committed
1631
int LGBM_BoosterRollbackOneIter(BoosterHandle handle) {
wxchan's avatar
wxchan committed
1632
1633
1634
1635
1636
1637
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->RollbackOneIter();
  API_END();
}

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

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

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

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

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

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

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

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

Guolin Ke's avatar
Guolin Ke committed
1752
int LGBM_BoosterCalcNumPredict(BoosterHandle handle,
1753
1754
                               int num_row,
                               int predict_type,
1755
                               int start_iteration,
1756
1757
                               int num_iteration,
                               int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1758
1759
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1760
  *out_len = static_cast<int64_t>(num_row) * ref_booster->GetBoosting()->NumPredictOneRow(start_iteration,
1761
    num_iteration, predict_type == C_API_PREDICT_LEAF_INDEX, predict_type == C_API_PREDICT_CONTRIB);
Guolin Ke's avatar
Guolin Ke committed
1762
1763
1764
  API_END();
}

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

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

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

1918
int LGBM_BoosterPredictForCSRSingleRow(BoosterHandle handle,
1919
1920
1921
1922
1923
1924
1925
                                       const void* indptr,
                                       int indptr_type,
                                       const int32_t* indices,
                                       const void* data,
                                       int data_type,
                                       int64_t nindptr,
                                       int64_t nelem,
1926
                                       int64_t num_col,
1927
                                       int predict_type,
1928
                                       int start_iteration,
1929
1930
1931
1932
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
1933
  API_BEGIN();
1934
1935
1936
1937
1938
  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.");
  }
1939
1940
1941
1942
1943
1944
1945
  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);
1946
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1947
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
1948
  ref_booster->PredictSingleRow(predict_type, static_cast<int32_t>(num_col), get_row_fun, config, out_result, out_len);
1949
1950
1951
  API_END();
}

1952
int LGBM_BoosterPredictForCSRSingleRowFastInit(BoosterHandle handle,
1953
                                               const int predict_type,
1954
                                               const int start_iteration,
1955
                                               const int num_iteration,
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
                                               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,
1970
    predict_type,
1971
1972
1973
1974
1975
1976
1977
    data_type,
    static_cast<int32_t>(num_col)));

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

1978
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
1979

1980
1981
1982
1983
1984
1985
  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

int LGBM_BoosterPredictForCSRSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* indptr,
1986
                                           const int indptr_type,
1987
1988
                                           const int32_t* indices,
                                           const void* data,
1989
1990
                                           const int64_t nindptr,
                                           const int64_t nelem,
1991
1992
1993
1994
1995
                                           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);
1996
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
1997
1998
1999
2000
                                        get_row_fun, fastConfig->config, out_result, out_len);
  API_END();
}

2001

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

Guolin Ke's avatar
Guolin Ke committed
2051
int LGBM_BoosterPredictForMat(BoosterHandle handle,
2052
2053
2054
2055
2056
2057
                              const void* data,
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              int predict_type,
2058
                              int start_iteration,
2059
                              int num_iteration,
2060
                              const char* parameter,
2061
2062
                              int64_t* out_len,
                              double* out_result) {
2063
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2064
2065
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2066
2067
2068
2069
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
2070
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2071
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, nrow, ncol, data_type, is_row_major);
2072
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
2073
                       config, out_result, out_len);
2074
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2075
}
2076

2077
int LGBM_BoosterPredictForMatSingleRow(BoosterHandle handle,
2078
2079
2080
2081
2082
                                       const void* data,
                                       int data_type,
                                       int32_t ncol,
                                       int is_row_major,
                                       int predict_type,
2083
                                       int start_iteration,
2084
2085
2086
2087
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2088
2089
2090
2091
2092
2093
2094
2095
2096
  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);
2097
  ref_booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, config);
2098
  ref_booster->PredictSingleRow(predict_type, ncol, get_row_fun, config, out_result, out_len);
2099
2100
2101
  API_END();
}

2102
int LGBM_BoosterPredictForMatSingleRowFastInit(BoosterHandle handle,
2103
                                               const int predict_type,
2104
                                               const int start_iteration,
2105
                                               const int num_iteration,
2106
2107
2108
2109
2110
2111
2112
2113
                                               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,
2114
    predict_type,
2115
2116
2117
2118
2119
2120
2121
    data_type,
    ncol));

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

2122
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2123

2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
  *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);
2136
  fastConfig->booster->PredictSingleRow(fastConfig->predict_type, fastConfig->ncol,
2137
2138
2139
2140
2141
                                        get_row_fun, fastConfig->config,
                                        out_result, out_len);
  API_END();
}

2142

2143
2144
2145
2146
2147
2148
int LGBM_BoosterPredictForMats(BoosterHandle handle,
                               const void** data,
                               int data_type,
                               int32_t nrow,
                               int32_t ncol,
                               int predict_type,
2149
                               int start_iteration,
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
                               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);
2163
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, out_result, out_len);
2164
2165
2166
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2167
int LGBM_BoosterSaveModel(BoosterHandle handle,
2168
                          int start_iteration,
2169
                          int num_iteration,
2170
                          int feature_importance_type,
2171
                          const char* filename) {
2172
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2173
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2174
2175
  ref_booster->SaveModelToFile(start_iteration, num_iteration,
                               feature_importance_type, filename);
wxchan's avatar
wxchan committed
2176
2177
2178
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2179
int LGBM_BoosterSaveModelToString(BoosterHandle handle,
2180
                                  int start_iteration,
2181
                                  int num_iteration,
2182
                                  int feature_importance_type,
2183
                                  int64_t buffer_len,
2184
                                  int64_t* out_len,
2185
                                  char* out_str) {
2186
2187
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2188
2189
  std::string model = ref_booster->SaveModelToString(
      start_iteration, num_iteration, feature_importance_type);
2190
  *out_len = static_cast<int64_t>(model.size()) + 1;
2191
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2192
    std::memcpy(out_str, model.c_str(), *out_len);
2193
2194
2195
2196
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2197
int LGBM_BoosterDumpModel(BoosterHandle handle,
2198
                          int start_iteration,
2199
                          int num_iteration,
2200
                          int feature_importance_type,
2201
2202
                          int64_t buffer_len,
                          int64_t* out_len,
2203
                          char* out_str) {
wxchan's avatar
wxchan committed
2204
2205
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2206
2207
  std::string model = ref_booster->DumpModel(start_iteration, num_iteration,
                                             feature_importance_type);
2208
  *out_len = static_cast<int64_t>(model.size()) + 1;
wxchan's avatar
wxchan committed
2209
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2210
    std::memcpy(out_str, model.c_str(), *out_len);
wxchan's avatar
wxchan committed
2211
  }
2212
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2213
}
2214

Guolin Ke's avatar
Guolin Ke committed
2215
int LGBM_BoosterGetLeafValue(BoosterHandle handle,
2216
2217
2218
                             int tree_idx,
                             int leaf_idx,
                             double* out_val) {
Guolin Ke's avatar
Guolin Ke committed
2219
2220
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2221
  *out_val = static_cast<double>(ref_booster->GetLeafValue(tree_idx, leaf_idx));
Guolin Ke's avatar
Guolin Ke committed
2222
2223
2224
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2225
int LGBM_BoosterSetLeafValue(BoosterHandle handle,
2226
2227
2228
                             int tree_idx,
                             int leaf_idx,
                             double val) {
Guolin Ke's avatar
Guolin Ke committed
2229
2230
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2231
  ref_booster->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
2232
2233
2234
  API_END();
}

2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
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();
}

2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
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();
}

2266
2267
2268
2269
2270
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
2271
  Config config;
2272
  config.machines = RemoveQuotationSymbol(std::string(machines));
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
  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();
}

2288
2289
2290
int LGBM_NetworkInitWithFunctions(int num_machines, int rank,
                                  void* reduce_scatter_ext_fun,
                                  void* allgather_ext_fun) {
ww's avatar
ww committed
2291
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2292
  if (num_machines > 1) {
2293
    Network::Init(num_machines, rank, (ReduceScatterFunction)reduce_scatter_ext_fun, (AllgatherFunction)allgather_ext_fun);
ww's avatar
ww committed
2294
2295
2296
  }
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
2297

Guolin Ke's avatar
Guolin Ke committed
2298
// ---- start of some help functions
2299
2300
2301

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
2302
  if (data_type == C_API_DTYPE_FLOAT32) {
2303
2304
    const float* data_ptr = reinterpret_cast<const float*>(data);
    if (is_row_major) {
2305
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2306
        std::vector<double> ret(num_col);
2307
        auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
2308
        for (int i = 0; i < num_col; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2309
          ret[i] = static_cast<double>(*(tmp_ptr + i));
2310
2311
2312
2313
        }
        return ret;
      };
    } else {
2314
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2315
        std::vector<double> ret(num_col);
2316
        for (int i = 0; i < num_col; ++i) {
2317
          ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
2318
2319
2320
2321
        }
        return ret;
      };
    }
Guolin Ke's avatar
Guolin Ke committed
2322
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2323
2324
    const double* data_ptr = reinterpret_cast<const double*>(data);
    if (is_row_major) {
2325
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2326
        std::vector<double> ret(num_col);
2327
        auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
2328
        for (int i = 0; i < num_col; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2329
          ret[i] = static_cast<double>(*(tmp_ptr + i));
2330
2331
2332
2333
        }
        return ret;
      };
    } else {
2334
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2335
        std::vector<double> ret(num_col);
2336
        for (int i = 0; i < num_col; ++i) {
2337
          ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
2338
2339
2340
2341
2342
        }
        return ret;
      };
    }
  }
2343
  Log::Fatal("Unknown data type in RowFunctionFromDenseMatric");
2344
  return nullptr;
2345
2346
2347
2348
}

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
2349
2350
  auto inner_function = RowFunctionFromDenseMatric(data, num_row, num_col, data_type, is_row_major);
  if (inner_function != nullptr) {
2351
    return [inner_function] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2352
2353
      auto raw_values = inner_function(row_idx);
      std::vector<std::pair<int, double>> ret;
Guolin Ke's avatar
Guolin Ke committed
2354
      ret.reserve(raw_values.size());
Guolin Ke's avatar
Guolin Ke committed
2355
      for (int i = 0; i < static_cast<int>(raw_values.size()); ++i) {
Guolin Ke's avatar
Guolin Ke committed
2356
        if (std::fabs(raw_values[i]) > kZeroThreshold || std::isnan(raw_values[i])) {
Guolin Ke's avatar
Guolin Ke committed
2357
          ret.emplace_back(i, raw_values[i]);
2358
        }
Guolin Ke's avatar
Guolin Ke committed
2359
2360
2361
      }
      return ret;
    };
2362
  }
Guolin Ke's avatar
Guolin Ke committed
2363
  return nullptr;
2364
2365
}

2366
2367
2368
2369
2370
2371
2372
// 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
2373
    ret.reserve(raw_values.size());
2374
2375
2376
2377
2378
2379
2380
2381
2382
    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;
  };
}

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

Guolin Ke's avatar
Guolin Ke committed
2453
std::function<std::pair<int, double>(int idx)>
2454
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
2455
  CHECK(col_idx < ncol_ptr && col_idx >= 0);
Guolin Ke's avatar
Guolin Ke committed
2456
  if (data_type == C_API_DTYPE_FLOAT32) {
2457
    const float* data_ptr = reinterpret_cast<const float*>(data);
Guolin Ke's avatar
Guolin Ke committed
2458
    if (col_ptr_type == C_API_DTYPE_INT32) {
2459
      const int32_t* ptr_col_ptr = reinterpret_cast<const int32_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2460
2461
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2462
2463
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2464
2465
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2466
        }
Guolin Ke's avatar
Guolin Ke committed
2467
2468
2469
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2470
      };
Guolin Ke's avatar
Guolin Ke committed
2471
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2472
      const int64_t* ptr_col_ptr = reinterpret_cast<const int64_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2473
2474
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2475
2476
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2477
2478
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2479
        }
Guolin Ke's avatar
Guolin Ke committed
2480
2481
2482
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2483
      };
Guolin Ke's avatar
Guolin Ke committed
2484
    }
Guolin Ke's avatar
Guolin Ke committed
2485
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2486
    const double* data_ptr = reinterpret_cast<const double*>(data);
Guolin Ke's avatar
Guolin Ke committed
2487
    if (col_ptr_type == C_API_DTYPE_INT32) {
2488
      const int32_t* ptr_col_ptr = reinterpret_cast<const int32_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2489
2490
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2491
2492
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2493
2494
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2495
        }
Guolin Ke's avatar
Guolin Ke committed
2496
2497
2498
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2499
      };
Guolin Ke's avatar
Guolin Ke committed
2500
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2501
      const int64_t* ptr_col_ptr = reinterpret_cast<const int64_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2502
2503
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2504
2505
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2506
2507
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2508
        }
Guolin Ke's avatar
Guolin Ke committed
2509
2510
2511
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2512
      };
Guolin Ke's avatar
Guolin Ke committed
2513
2514
    }
  }
2515
  Log::Fatal("Unknown data type in CSC matrix");
2516
  return nullptr;
2517
2518
}

Guolin Ke's avatar
Guolin Ke committed
2519
CSC_RowIterator::CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
2520
                                 const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx) {
Guolin Ke's avatar
Guolin Ke committed
2521
2522
2523
2524
2525
2526
2527
2528
2529
  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;
2530
    }
Guolin Ke's avatar
Guolin Ke committed
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
    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;
2548
    }
Guolin Ke's avatar
Guolin Ke committed
2549
2550
2551
    return ret;
  } else {
    return std::make_pair(-1, 0.0);
2552
  }
Guolin Ke's avatar
Guolin Ke committed
2553
}