c_api.cpp 96.5 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();
Guolin Ke's avatar
Guolin Ke committed
1614
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1615
  #ifdef SCORE_T_USE_DOUBLE
1616
  Log::Fatal("Don't support custom loss function when SCORE_T_USE_DOUBLE is enabled");
1617
  #else
1618
1619
1620
1621
1622
  if (ref_booster->TrainOneIter(grad, hess)) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1623
  #endif
1624
  API_END();
1625
1626
}

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

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

1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
int LGBM_BoosterNumModelPerIteration(BoosterHandle handle, int* out_tree_per_iteration) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_tree_per_iteration = ref_booster->GetBoosting()->NumModelPerIteration();
  API_END();
}

int LGBM_BoosterNumberOfTotalModel(BoosterHandle handle, int* out_models) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_models = ref_booster->GetBoosting()->NumberOfTotalModel();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1655
int LGBM_BoosterGetEvalCounts(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1656
1657
1658
1659
1660
1661
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetEvalCounts();
  API_END();
}

1662
1663
1664
1665
1666
1667
int LGBM_BoosterGetEvalNames(BoosterHandle handle,
                             const int len,
                             int* out_len,
                             const size_t buffer_len,
                             size_t* out_buffer_len,
                             char** out_strs) {
wxchan's avatar
wxchan committed
1668
1669
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1670
  *out_len = ref_booster->GetEvalNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1671
1672
1673
  API_END();
}

1674
1675
1676
1677
1678
1679
int LGBM_BoosterGetFeatureNames(BoosterHandle handle,
                                const int len,
                                int* out_len,
                                const size_t buffer_len,
                                size_t* out_buffer_len,
                                char** out_strs) {
wxchan's avatar
wxchan committed
1680
1681
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1682
  *out_len = ref_booster->GetFeatureNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1683
1684
1685
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1686
int LGBM_BoosterGetNumFeature(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1687
1688
1689
1690
1691
1692
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->MaxFeatureIdx() + 1;
  API_END();
}

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

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

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

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

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

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

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

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

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

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

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

1974
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
1975

1976
1977
1978
1979
1980
1981
  *out_fastConfig = fastConfig_ptr.release();
  API_END();
}

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

1997

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

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

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

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

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

2118
  fastConfig_ptr->booster->SetSingleRowPredictor(start_iteration, num_iteration, predict_type, fastConfig_ptr->config);
2119

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

2138

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

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

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

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

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

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

2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
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();
}

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

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

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

Guolin Ke's avatar
Guolin Ke committed
2294
// ---- start of some help functions
2295
2296
2297

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

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

2362
2363
2364
2365
2366
2367
2368
// 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
2369
    ret.reserve(raw_values.size());
2370
2371
2372
2373
2374
2375
2376
2377
2378
    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;
  };
}

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

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

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