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

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

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

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

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

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

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

60
61
62
const int PREDICTOR_TYPES = 4;

// Single row predictor to abstract away caching logic
63
class SingleRowPredictorInner {
64
65
66
67
 public:
  PredictFunction predict_function;
  int64_t num_pred_in_one_row;

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

90
  ~SingleRowPredictorInner() {}
91

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

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

109
110
/*!
 * \brief Object to store resources meant for single-row Fast Predict methods.
111
 *
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
 * For legacy reasons this is called `FastConfig` in the public C API.
 *
 * Meant to be used by the *Fast* predict methods only.
 * It stores the configuration and prediction resources for reuse across predictions.
 */
struct SingleRowPredictor {
 public:
  SingleRowPredictor(yamc::alternate::shared_mutex *booster_mutex,
             const char *parameters,
             const int data_type,
             const int32_t num_cols,
             int predict_type,
             Boosting *boosting,
             int start_iter,
             int num_iter) : config(Config::Str2Map(parameters)), data_type(data_type), num_cols(num_cols), single_row_predictor_inner(predict_type, boosting, config, start_iter, num_iter), booster_mutex(booster_mutex) {
    if (!config.predict_disable_shape_check && num_cols != 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.", num_cols, boosting->MaxFeatureIdx() + 1);
    }
  }

  void Predict(std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun,
               double* out_result, int64_t* out_len) const {
    UNIQUE_LOCK(single_row_predictor_mutex)
    yamc::shared_lock<yamc::alternate::shared_mutex> booster_shared_lock(booster_mutex);

    auto one_row = get_row_fun(0);
    single_row_predictor_inner.predict_function(one_row, out_result);

    *out_len = single_row_predictor_inner.num_pred_in_one_row;
  }

 public:
  Config config;
  const int data_type;
  const int32_t num_cols;

 private:
  SingleRowPredictorInner single_row_predictor_inner;

  // Prevent the booster from being modified while we have a predictor relying on it during prediction
  yamc::alternate::shared_mutex *booster_mutex;

  // If several threads try to predict at the same time using the same SingleRowPredictor
  // we want them to still provide correct values, so the mutex is necessary due to the shared
  // resources in the predictor.
  // However the recommended approach is to instantiate one SingleRowPredictor per thread,
  // to avoid contention here.
  mutable yamc::alternate::shared_mutex single_row_predictor_mutex;
};

Guolin Ke's avatar
Guolin Ke committed
163
class Booster {
Nikita Titov's avatar
Nikita Titov committed
164
 public:
Guolin Ke's avatar
Guolin Ke committed
165
  explicit Booster(const char* filename) {
166
    boosting_.reset(Boosting::CreateBoosting("gbdt", filename));
167
168
  }

Guolin Ke's avatar
Guolin Ke committed
169
  Booster(const Dataset* train_data,
170
          const char* parameters) {
Guolin Ke's avatar
Guolin Ke committed
171
    auto param = Config::Str2Map(parameters);
wxchan's avatar
wxchan committed
172
    config_.Set(param);
173
    OMP_SET_NUM_THREADS(config_.num_threads);
Guolin Ke's avatar
Guolin Ke committed
174
    // create boosting
Guolin Ke's avatar
Guolin Ke committed
175
    if (config_.input_model.size() > 0) {
176
177
      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
178
    }
Guolin Ke's avatar
Guolin Ke committed
179

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

182
183
    train_data_ = train_data;
    CreateObjectiveAndMetrics();
Guolin Ke's avatar
Guolin Ke committed
184
    // initialize the boosting
Guolin Ke's avatar
Guolin Ke committed
185
    if (config_.tree_learner == std::string("feature")) {
186
      Log::Fatal("Do not support feature parallel in c api");
187
    }
Guolin Ke's avatar
Guolin Ke committed
188
    if (Network::num_machines() == 1 && config_.tree_learner != std::string("serial")) {
189
      Log::Warning("Only find one worker, will switch to serial tree learner");
Guolin Ke's avatar
Guolin Ke committed
190
      config_.tree_learner = "serial";
191
    }
Guolin Ke's avatar
Guolin Ke committed
192
    boosting_->Init(&config_, train_data_, objective_fun_.get(),
193
                    Common::ConstPtrInVectorWrapper<Metric>(train_metric_));
wxchan's avatar
wxchan committed
194
195
196
  }

  void MergeFrom(const Booster* other) {
197
    UNIQUE_LOCK(mutex_)
wxchan's avatar
wxchan committed
198
    boosting_->MergeFrom(other->boosting_.get());
Guolin Ke's avatar
Guolin Ke committed
199
200
201
202
  }

  ~Booster() {
  }
203

204
  void CreateObjectiveAndMetrics() {
Guolin Ke's avatar
Guolin Ke committed
205
    // create objective function
Guolin Ke's avatar
Guolin Ke committed
206
207
    objective_fun_.reset(ObjectiveFunction::CreateObjectiveFunction(config_.objective,
                                                                    config_));
Guolin Ke's avatar
Guolin Ke committed
208
    if (objective_fun_ == nullptr) {
209
      Log::Info("Using self-defined objective function");
Guolin Ke's avatar
Guolin Ke committed
210
211
212
213
214
215
216
217
    }
    // 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
218
    for (auto metric_type : config_.metric) {
Guolin Ke's avatar
Guolin Ke committed
219
      auto metric = std::unique_ptr<Metric>(
Guolin Ke's avatar
Guolin Ke committed
220
        Metric::CreateMetric(metric_type, config_));
221
222
223
      if (metric == nullptr) {
        continue;
      }
Guolin Ke's avatar
Guolin Ke committed
224
225
226
227
      metric->Init(train_data_->metadata(), train_data_->num_data());
      train_metric_.push_back(std::move(metric));
    }
    train_metric_.shrink_to_fit();
228
229
230
231
  }

  void ResetTrainingData(const Dataset* train_data) {
    if (train_data != train_data_) {
232
      UNIQUE_LOCK(mutex_)
233
234
235
236
237
238
      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
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
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
  static void CheckDatasetResetConfig(
      const Config& old_config,
      const std::unordered_map<std::string, std::string>& new_param) {
    Config new_config;
    new_config.Set(new_param);
    if (new_param.count("data_random_seed") &&
        new_config.data_random_seed != old_config.data_random_seed) {
      Log::Fatal("Cannot change data_random_seed after constructed Dataset handle.");
    }
    if (new_param.count("max_bin") &&
        new_config.max_bin != old_config.max_bin) {
      Log::Fatal("Cannot change max_bin after constructed Dataset handle.");
    }
    if (new_param.count("max_bin_by_feature") &&
        new_config.max_bin_by_feature != old_config.max_bin_by_feature) {
      Log::Fatal(
          "Cannot change max_bin_by_feature after constructed Dataset handle.");
    }
    if (new_param.count("bin_construct_sample_cnt") &&
        new_config.bin_construct_sample_cnt !=
            old_config.bin_construct_sample_cnt) {
      Log::Fatal(
          "Cannot change bin_construct_sample_cnt after constructed Dataset "
          "handle.");
    }
    if (new_param.count("min_data_in_bin") &&
        new_config.min_data_in_bin != old_config.min_data_in_bin) {
      Log::Fatal(
          "Cannot change min_data_in_bin after constructed Dataset handle.");
    }
    if (new_param.count("use_missing") &&
        new_config.use_missing != old_config.use_missing) {
      Log::Fatal("Cannot change use_missing after constructed Dataset handle.");
    }
    if (new_param.count("zero_as_missing") &&
        new_config.zero_as_missing != old_config.zero_as_missing) {
      Log::Fatal(
          "Cannot change zero_as_missing after constructed Dataset handle.");
    }
    if (new_param.count("categorical_feature") &&
        new_config.categorical_feature != old_config.categorical_feature) {
      Log::Fatal(
          "Cannot change categorical_feature after constructed Dataset "
          "handle.");
    }
    if (new_param.count("feature_pre_filter") &&
        new_config.feature_pre_filter != old_config.feature_pre_filter) {
      Log::Fatal(
          "Cannot change feature_pre_filter after constructed Dataset handle.");
    }
    if (new_param.count("is_enable_sparse") &&
        new_config.is_enable_sparse != old_config.is_enable_sparse) {
      Log::Fatal(
          "Cannot change is_enable_sparse after constructed Dataset handle.");
    }
    if (new_param.count("pre_partition") &&
        new_config.pre_partition != old_config.pre_partition) {
      Log::Fatal(
          "Cannot change pre_partition after constructed Dataset handle.");
    }
    if (new_param.count("enable_bundle") &&
        new_config.enable_bundle != old_config.enable_bundle) {
      Log::Fatal(
          "Cannot change enable_bundle after constructed Dataset handle.");
    }
    if (new_param.count("header") && new_config.header != old_config.header) {
      Log::Fatal("Cannot change header after constructed Dataset handle.");
    }
    if (new_param.count("two_round") &&
        new_config.two_round != old_config.two_round) {
      Log::Fatal("Cannot change two_round after constructed Dataset handle.");
    }
    if (new_param.count("label_column") &&
        new_config.label_column != old_config.label_column) {
      Log::Fatal(
          "Cannot change label_column after constructed Dataset handle.");
    }
    if (new_param.count("weight_column") &&
        new_config.weight_column != old_config.weight_column) {
      Log::Fatal(
          "Cannot change weight_column after constructed Dataset handle.");
    }
    if (new_param.count("group_column") &&
        new_config.group_column != old_config.group_column) {
      Log::Fatal(
          "Cannot change group_column after constructed Dataset handle.");
    }
    if (new_param.count("ignore_column") &&
        new_config.ignore_column != old_config.ignore_column) {
      Log::Fatal(
          "Cannot change ignore_column after constructed Dataset handle.");
    }
    if (new_param.count("forcedbins_filename")) {
      Log::Fatal("Cannot change forced bins after constructed Dataset handle.");
    }
    if (new_param.count("min_data_in_leaf") &&
        new_config.min_data_in_leaf < old_config.min_data_in_leaf &&
        old_config.feature_pre_filter) {
      Log::Fatal(
          "Reducing `min_data_in_leaf` with `feature_pre_filter=true` may "
          "cause unexpected behaviour "
          "for features that were pre-filtered by the larger "
          "`min_data_in_leaf`.\n"
          "You need to set `feature_pre_filter=false` to dynamically change "
          "the `min_data_in_leaf`.");
    }
Nikita Titov's avatar
Nikita Titov committed
347
    if (new_param.count("linear_tree") && new_config.linear_tree != old_config.linear_tree) {
348
      Log::Fatal("Cannot change linear_tree after constructed Dataset handle.");
349
    }
Nikita Titov's avatar
Nikita Titov committed
350
351
352
353
    if (new_param.count("precise_float_parser") &&
        new_config.precise_float_parser != old_config.precise_float_parser) {
      Log::Fatal("Cannot change precise_float_parser after constructed Dataset handle.");
    }
354
355
  }

wxchan's avatar
wxchan committed
356
  void ResetConfig(const char* parameters) {
357
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
358
    auto param = Config::Str2Map(parameters);
359
360
361
    Config new_config;
    new_config.Set(param);
    if (param.count("num_class") && new_config.num_class != config_.num_class) {
362
      Log::Fatal("Cannot change num_class during training");
wxchan's avatar
wxchan committed
363
    }
364
    if (param.count("boosting") && new_config.boosting != config_.boosting) {
Guolin Ke's avatar
Guolin Ke committed
365
      Log::Fatal("Cannot change boosting during training");
wxchan's avatar
wxchan committed
366
    }
367
    if (param.count("metric") && new_config.metric != config_.metric) {
368
      Log::Fatal("Cannot change metric during training");
Guolin Ke's avatar
Guolin Ke committed
369
    }
370
371
    CheckDatasetResetConfig(config_, param);

Guolin Ke's avatar
Guolin Ke committed
372
    config_.Set(param);
373

374
    OMP_SET_NUM_THREADS(config_.num_threads);
Guolin Ke's avatar
Guolin Ke committed
375
376
377

    if (param.count("objective")) {
      // create objective function
Guolin Ke's avatar
Guolin Ke committed
378
379
      objective_fun_.reset(ObjectiveFunction::CreateObjectiveFunction(config_.objective,
                                                                      config_));
Guolin Ke's avatar
Guolin Ke committed
380
      if (objective_fun_ == nullptr) {
381
        Log::Info("Using self-defined objective function");
Guolin Ke's avatar
Guolin Ke committed
382
383
384
385
386
      }
      // initialize the objective function
      if (objective_fun_ != nullptr) {
        objective_fun_->Init(train_data_->metadata(), train_data_->num_data());
      }
387
388
      boosting_->ResetTrainingData(train_data_,
                                   objective_fun_.get(), Common::ConstPtrInVectorWrapper<Metric>(train_metric_));
wxchan's avatar
wxchan committed
389
    }
Guolin Ke's avatar
Guolin Ke committed
390

Guolin Ke's avatar
Guolin Ke committed
391
    boosting_->ResetConfig(&config_);
wxchan's avatar
wxchan committed
392
393
394
  }

  void AddValidData(const Dataset* valid_data) {
395
    UNIQUE_LOCK(mutex_)
wxchan's avatar
wxchan committed
396
    valid_metrics_.emplace_back();
Guolin Ke's avatar
Guolin Ke committed
397
398
    for (auto metric_type : config_.metric) {
      auto metric = std::unique_ptr<Metric>(Metric::CreateMetric(metric_type, config_));
399
400
401
      if (metric == nullptr) {
        continue;
      }
wxchan's avatar
wxchan committed
402
403
404
405
406
      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,
407
                               Common::ConstPtrInVectorWrapper<Metric>(valid_metrics_.back()));
wxchan's avatar
wxchan committed
408
  }
Guolin Ke's avatar
Guolin Ke committed
409

410
  bool TrainOneIter() {
411
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
412
    return boosting_->TrainOneIter(nullptr, nullptr);
413
414
  }

Guolin Ke's avatar
Guolin Ke committed
415
  void Refit(const int32_t* leaf_preds, int32_t nrow, int32_t ncol) {
416
    UNIQUE_LOCK(mutex_)
417
    boosting_->RefitTree(leaf_preds, nrow, ncol);
Guolin Ke's avatar
Guolin Ke committed
418
419
  }

420
  bool TrainOneIter(const score_t* gradients, const score_t* hessians) {
421
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
422
    return boosting_->TrainOneIter(gradients, hessians);
423
424
  }

wxchan's avatar
wxchan committed
425
  void RollbackOneIter() {
426
    UNIQUE_LOCK(mutex_)
wxchan's avatar
wxchan committed
427
428
429
    boosting_->RollbackOneIter();
  }

430
  void SetSingleRowPredictorInner(int start_iteration, int num_iteration, int predict_type, const Config& config) {
431
432
433
      UNIQUE_LOCK(mutex_)
      if (single_row_predictor_[predict_type].get() == nullptr ||
          !single_row_predictor_[predict_type]->IsPredictorEqual(config, num_iteration, boosting_.get())) {
434
        single_row_predictor_[predict_type].reset(new SingleRowPredictorInner(predict_type, boosting_.get(),
435
                                                                         config, start_iteration, num_iteration));
436
437
438
      }
  }

439
440
441
442
443
444
445
446
447
448
449
  std::unique_ptr<SingleRowPredictor> InitSingleRowPredictor(int predict_type, int start_iteration, int num_iteration, int data_type, int32_t num_cols, const char *parameters) {
    // Workaround https://github.com/microsoft/LightGBM/issues/6142 by locking here
    // This is only a workaround because if predictors are initialized differently it may still behave incorrectly,
    // and because multiple racing Predictor initializations through LGBM_BoosterPredictForMat suffers from that same issue of Predictor init writing things in the booster.
    // Once #6142 is fixed (predictor doesn't write in the Booster as should have been the case since 1c35c3b9ede9adab8ccc5fd7b4b2b6af188a79f0), this line can be removed.
    UNIQUE_LOCK(mutex_)

    return std::unique_ptr<SingleRowPredictor>(new SingleRowPredictor(
      &mutex_, parameters, data_type, num_cols, predict_type, boosting_.get(), start_iteration, num_iteration));
  }

450
  void PredictSingleRow(int predict_type, int ncol,
451
452
               std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun,
               const Config& config,
453
               double* out_result, int64_t* out_len) const {
454
455
456
    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);
457
    }
458
    UNIQUE_LOCK(mutex_)
459
    const auto& single_row_predictor = single_row_predictor_[predict_type];
460
461
    auto one_row = get_row_fun(0);
    auto pred_wrt_ptr = out_result;
462
    single_row_predictor->predict_function(one_row, pred_wrt_ptr);
463

464
    *out_len = single_row_predictor->num_pred_in_one_row;
465
466
  }

467
  std::shared_ptr<Predictor> CreatePredictor(int start_iteration, int num_iteration, int predict_type, int ncol, const Config& config) const {
468
469
470
    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);
471
    }
Guolin Ke's avatar
Guolin Ke committed
472
473
    bool is_predict_leaf = false;
    bool is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
474
    bool predict_contrib = false;
Guolin Ke's avatar
Guolin Ke committed
475
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
Guolin Ke's avatar
Guolin Ke committed
476
      is_predict_leaf = true;
Guolin Ke's avatar
Guolin Ke committed
477
    } else if (predict_type == C_API_PREDICT_RAW_SCORE) {
Guolin Ke's avatar
Guolin Ke committed
478
      is_raw_score = true;
479
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
Guolin Ke's avatar
Guolin Ke committed
480
      predict_contrib = true;
Guolin Ke's avatar
Guolin Ke committed
481
482
    } else {
      is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
483
    }
Guolin Ke's avatar
Guolin Ke committed
484

485
    return std::make_shared<Predictor>(boosting_.get(), start_iteration, num_iteration, is_raw_score, is_predict_leaf, predict_contrib,
486
                        config.pred_early_stop, config.pred_early_stop_freq, config.pred_early_stop_margin);
487
488
  }

489
  void Predict(int start_iteration, int num_iteration, int predict_type, int nrow, int ncol,
490
491
               std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun,
               const Config& config,
492
493
               double* out_result, int64_t* out_len) const {
    SHARED_LOCK(mutex_);
494
    auto predictor = CreatePredictor(start_iteration, num_iteration, predict_type, ncol, config);
495
496
497
498
499
500
501
    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;
    }
502
    int64_t num_pred_in_one_row = boosting_->NumPredictOneRow(start_iteration, num_iteration, is_predict_leaf, predict_contrib);
503
    auto pred_fun = predictor->GetPredictFunction();
504
    OMP_INIT_EX();
505
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
Guolin Ke's avatar
Guolin Ke committed
506
    for (int i = 0; i < nrow; ++i) {
507
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
508
      auto one_row = get_row_fun(i);
Tony-Y's avatar
Tony-Y committed
509
      auto pred_wrt_ptr = out_result + static_cast<size_t>(num_pred_in_one_row) * i;
Guolin Ke's avatar
Guolin Ke committed
510
      pred_fun(one_row, pred_wrt_ptr);
511
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
512
    }
513
    OMP_THROW_EX();
514
    *out_len = num_pred_in_one_row * nrow;
Guolin Ke's avatar
Guolin Ke committed
515
516
  }

517
  void PredictSparse(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
518
519
520
521
                     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,
522
                     bool* is_data_float32_ptr, int num_matrices) const {
523
    auto predictor = CreatePredictor(start_iteration, num_iteration, predict_type, ncol, config);
524
    auto pred_sparse_fun = predictor->GetPredictSparseFunction();
525
526
    std::vector<std::vector<std::unordered_map<int, double>>>& agg = *agg_ptr;
    OMP_INIT_EX();
527
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
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
    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];
  }

559
  void PredictSparseCSR(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
560
561
562
                        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,
563
564
                        int32_t** out_indices, void** out_data, int data_type) const {
    SHARED_LOCK(mutex_);
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
    // 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;
582
    PredictSparse(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
583
584
585
                  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);
586
    std::vector<int64_t> matrix_offsets(num_matrices);
587
588
589
590
591
592
593
594
595
596
597
598
599
600
    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++;
      }
601
602
603
604
605
606
      if (m == 0) {
        matrix_offsets[m] = 0;
      }
      if (m + 1 < num_matrices) {
        matrix_offsets[m + 1] = static_cast<int64_t>(matrix_offsets[m] + row_matrix_offsets[row_vector_cnt - 1] + row_sizes[row_vector_cnt - 1]);
      }
607
608
609
610
611
612
613
614
615
616
617
618
    }
    // 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();
619
      #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
620
621
622
623
      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;
624
        int64_t element_index = row_matrix_offsets[row_start_index] + matrix_offsets[m];
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
        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;
  }

650
  void PredictSparseCSC(int start_iteration, int num_iteration, int predict_type, int64_t nrow, int ncol,
651
652
653
                        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,
654
655
                        int32_t** out_indices, void** out_data, int data_type) const {
    SHARED_LOCK(mutex_);
656
657
    // Get the number of trees per iteration (for multiclass scenario we output multiple sparse matrices)
    int num_matrices = boosting_->NumModelPerIteration();
658
    auto predictor = CreatePredictor(start_iteration, num_iteration, predict_type, ncol, config);
659
    auto pred_sparse_fun = predictor->GetPredictSparseFunction();
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
    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;
676
    PredictSparse(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
                  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;
      }
724
725
      if (m + 1 < num_matrices) {
        matrix_start_indices[m + 1] = matrix_start_indices[m] + last_column_start_index + last_column_size;
726
      }
727
      col_ptr_index++;
728
    }
729
730
    // Note: we parallelize across matrices instead of rows because of the column_counts[m][col_idx] increment inside the loop
    OMP_INIT_EX();
731
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
732
    for (int m = 0; m < num_matrices; ++m) {
733
      OMP_LOOP_EX_BEGIN();
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
      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;
          }
        }
      }
752
      OMP_LOOP_EX_END();
753
    }
754
    OMP_THROW_EX();
755
756
757
758
    out_len[0] = elements_size;
    out_len[1] = col_ptr_size;
  }

759
  void Predict(int start_iteration, int num_iteration, int predict_type, const char* data_filename,
Guolin Ke's avatar
Guolin Ke committed
760
               int data_has_header, const Config& config,
761
762
               const char* result_filename) const {
    SHARED_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
763
764
    bool is_predict_leaf = false;
    bool is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
765
    bool predict_contrib = false;
Guolin Ke's avatar
Guolin Ke committed
766
767
768
769
    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;
770
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
Guolin Ke's avatar
Guolin Ke committed
771
      predict_contrib = true;
Guolin Ke's avatar
Guolin Ke committed
772
773
774
    } else {
      is_raw_score = false;
    }
775
    Predictor predictor(boosting_.get(), start_iteration, num_iteration, is_raw_score, is_predict_leaf, predict_contrib,
776
                        config.pred_early_stop, config.pred_early_stop_freq, config.pred_early_stop_margin);
Guolin Ke's avatar
Guolin Ke committed
777
    bool bool_data_has_header = data_has_header > 0 ? true : false;
Chen Yufei's avatar
Chen Yufei committed
778
779
    predictor.Predict(data_filename, result_filename, bool_data_has_header, config.predict_disable_shape_check,
                      config.precise_float_parser);
Guolin Ke's avatar
Guolin Ke committed
780
781
  }

782
  void GetPredictAt(int data_idx, double* out_result, int64_t* out_len) const {
wxchan's avatar
wxchan committed
783
784
785
    boosting_->GetPredictAt(data_idx, out_result, out_len);
  }

786
  void SaveModelToFile(int start_iteration, int num_iteration, int feature_importance_type, const char* filename) const {
787
    boosting_->SaveModelToFile(start_iteration, num_iteration, feature_importance_type, filename);
Guolin Ke's avatar
Guolin Ke committed
788
  }
789

790
  void LoadModelFromString(const char* model_str) {
791
792
    size_t len = std::strlen(model_str);
    boosting_->LoadModelFromString(model_str, len);
793
794
  }

795
  std::string SaveModelToString(int start_iteration, int num_iteration,
796
                                int feature_importance_type) const {
797
798
    return boosting_->SaveModelToString(start_iteration,
                                        num_iteration, feature_importance_type);
799
800
  }

801
  std::string DumpModel(int start_iteration, int num_iteration,
802
                        int feature_importance_type) const {
803
804
    return boosting_->DumpModel(start_iteration, num_iteration,
                                feature_importance_type);
wxchan's avatar
wxchan committed
805
  }
806

807
  std::vector<double> FeatureImportance(int num_iteration, int importance_type) const {
808
809
810
    return boosting_->FeatureImportance(num_iteration, importance_type);
  }

811
  double UpperBoundValue() const {
812
    SHARED_LOCK(mutex_)
813
814
815
816
    return boosting_->GetUpperBoundValue();
  }

  double LowerBoundValue() const {
817
    SHARED_LOCK(mutex_)
818
819
820
    return boosting_->GetLowerBoundValue();
  }

Guolin Ke's avatar
Guolin Ke committed
821
  double GetLeafValue(int tree_idx, int leaf_idx) const {
822
    SHARED_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
823
    return dynamic_cast<GBDTBase*>(boosting_.get())->GetLeafValue(tree_idx, leaf_idx);
Guolin Ke's avatar
Guolin Ke committed
824
825
826
  }

  void SetLeafValue(int tree_idx, int leaf_idx, double val) {
827
    UNIQUE_LOCK(mutex_)
Guolin Ke's avatar
Guolin Ke committed
828
    dynamic_cast<GBDTBase*>(boosting_.get())->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
829
830
  }

831
  void ShuffleModels(int start_iter, int end_iter) {
832
    UNIQUE_LOCK(mutex_)
833
    boosting_->ShuffleModels(start_iter, end_iter);
834
835
  }

wxchan's avatar
wxchan committed
836
  int GetEvalCounts() const {
837
    SHARED_LOCK(mutex_)
wxchan's avatar
wxchan committed
838
839
840
841
842
843
    int ret = 0;
    for (const auto& metric : train_metric_) {
      ret += static_cast<int>(metric->GetName().size());
    }
    return ret;
  }
844

845
  int GetEvalNames(char** out_strs, const int len, const size_t buffer_len, size_t *out_buffer_len) const {
846
    SHARED_LOCK(mutex_)
847
    *out_buffer_len = 0;
wxchan's avatar
wxchan committed
848
849
850
    int idx = 0;
    for (const auto& metric : train_metric_) {
      for (const auto& name : metric->GetName()) {
851
852
853
854
855
        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
856
857
858
859
860
861
        ++idx;
      }
    }
    return idx;
  }

862
  int GetFeatureNames(char** out_strs, const int len, const size_t buffer_len, size_t *out_buffer_len) const {
863
    SHARED_LOCK(mutex_)
864
    *out_buffer_len = 0;
wxchan's avatar
wxchan committed
865
866
    int idx = 0;
    for (const auto& name : boosting_->FeatureNames()) {
867
868
869
870
871
      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
872
873
874
875
876
      ++idx;
    }
    return idx;
  }

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

Nikita Titov's avatar
Nikita Titov committed
879
 private:
wxchan's avatar
wxchan committed
880
  const Dataset* train_data_;
Guolin Ke's avatar
Guolin Ke committed
881
  std::unique_ptr<Boosting> boosting_;
882
  std::unique_ptr<SingleRowPredictorInner> single_row_predictor_[PREDICTOR_TYPES];
883

Guolin Ke's avatar
Guolin Ke committed
884
  /*! \brief All configs */
Guolin Ke's avatar
Guolin Ke committed
885
  Config config_;
Guolin Ke's avatar
Guolin Ke committed
886
  /*! \brief Metric for training data */
Guolin Ke's avatar
Guolin Ke committed
887
  std::vector<std::unique_ptr<Metric>> train_metric_;
Guolin Ke's avatar
Guolin Ke committed
888
  /*! \brief Metrics for validation data */
Guolin Ke's avatar
Guolin Ke committed
889
  std::vector<std::vector<std::unique_ptr<Metric>>> valid_metrics_;
Guolin Ke's avatar
Guolin Ke committed
890
  /*! \brief Training objective function */
Guolin Ke's avatar
Guolin Ke committed
891
  std::unique_ptr<ObjectiveFunction> objective_fun_;
wxchan's avatar
wxchan committed
892
  /*! \brief mutex for threading safe call */
893
  mutable yamc::alternate::shared_mutex mutex_;
Guolin Ke's avatar
Guolin Ke committed
894
895
};

896
}  // namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
897

898
899
// explicitly declare symbols from LightGBM namespace
using LightGBM::AllgatherFunction;
900
using LightGBM::ArrowChunkedArray;
901
using LightGBM::ArrowTable;
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
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;
917
using LightGBM::SingleRowPredictor;
Guolin Ke's avatar
Guolin Ke committed
918

Guolin Ke's avatar
Guolin Ke committed
919
920
921
922
923
924
925
926
// 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);

927
928
929
std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseRows(const void** data, int num_col, int data_type);

930
931
template<typename T>
std::function<std::vector<std::pair<int, double>>(T idx)>
Guolin Ke's avatar
Guolin Ke committed
932
RowFunctionFromCSR(const void* indptr, int indptr_type, const int32_t* indices,
933
                   const void* data, int data_type, int64_t nindptr, int64_t nelem);
Guolin Ke's avatar
Guolin Ke committed
934
935
936

// Row iterator of on column for CSC matrix
class CSC_RowIterator {
Nikita Titov's avatar
Nikita Titov committed
937
 public:
Guolin Ke's avatar
Guolin Ke committed
938
  CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
939
                  const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx);
Guolin Ke's avatar
Guolin Ke committed
940
941
942
943
944
  ~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
945
946

 private:
Guolin Ke's avatar
Guolin Ke committed
947
948
949
950
951
952
953
954
955
  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
956
const char* LGBM_GetLastError() {
wxchan's avatar
wxchan committed
957
  return LastErrorMsg();
Guolin Ke's avatar
Guolin Ke committed
958
959
}

960
961
962
963
964
965
966
967
968
969
970
971
int LGBM_DumpParamAliases(int64_t buffer_len,
                          int64_t* out_len,
                          char* out_str) {
  API_BEGIN();
  std::string aliases = Config::DumpAliases();
  *out_len = static_cast<int64_t>(aliases.size()) + 1;
  if (*out_len <= buffer_len) {
    std::memcpy(out_str, aliases.c_str(), *out_len);
  }
  API_END();
}

972
973
974
975
976
977
int LGBM_RegisterLogCallback(void (*callback)(const char*)) {
  API_BEGIN();
  Log::ResetCallBack(callback);
  API_END();
}

978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
static inline int SampleCount(int32_t total_nrow, const Config& config) {
  return static_cast<int>(total_nrow < config.bin_construct_sample_cnt ? total_nrow : config.bin_construct_sample_cnt);
}

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

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

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

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

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

1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
int LGBM_ByteBufferGetAt(ByteBufferHandle handle, int32_t index, uint8_t* out_val) {
  API_BEGIN();
  LightGBM::ByteBuffer* byteBuffer = reinterpret_cast<LightGBM::ByteBuffer*>(handle);
  *out_val = byteBuffer->GetAt(index);
  API_END();
}

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

Guolin Ke's avatar
Guolin Ke committed
1036
int LGBM_DatasetCreateFromFile(const char* filename,
1037
1038
1039
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
1040
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1041
1042
  auto param = Config::Str2Map(parameters);
  Config config;
1043
  config.Set(param);
1044
  OMP_SET_NUM_THREADS(config.num_threads);
1045
  DatasetLoader loader(config, nullptr, 1, filename);
Guolin Ke's avatar
Guolin Ke committed
1046
  if (reference == nullptr) {
1047
    if (Network::num_machines() == 1) {
1048
      *out = loader.LoadFromFile(filename);
1049
    } else {
1050
      *out = loader.LoadFromFile(filename, Network::rank(), Network::num_machines());
1051
    }
Guolin Ke's avatar
Guolin Ke committed
1052
  } else {
1053
    *out = loader.LoadFromFileAlignWithOtherDataset(filename,
1054
                                                    reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
1055
  }
1056
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1057
1058
}

Guolin Ke's avatar
Guolin Ke committed
1059
int LGBM_DatasetCreateFromSampledColumn(double** sample_data,
1060
1061
1062
1063
                                        int** sample_indices,
                                        int32_t ncol,
                                        const int* num_per_col,
                                        int32_t num_sample_row,
1064
1065
                                        int32_t num_local_row,
                                        int64_t num_dist_row,
1066
1067
                                        const char* parameters,
                                        DatasetHandle* out) {
1068
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1069
1070
  auto param = Config::Str2Map(parameters);
  Config config;
1071
  config.Set(param);
1072
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1073
  DatasetLoader loader(config, nullptr, 1, nullptr);
1074
1075
1076
1077
  *out = loader.ConstructFromSampleData(sample_data,
                                        sample_indices,
                                        ncol,
                                        num_per_col,
1078
                                        num_sample_row,
1079
1080
                                        static_cast<data_size_t>(num_local_row),
                                        num_dist_row);
1081
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1082
1083
}

Guolin Ke's avatar
Guolin Ke committed
1084
int LGBM_DatasetCreateByReference(const DatasetHandle reference,
1085
1086
                                  int64_t num_total_row,
                                  DatasetHandle* out) {
Guolin Ke's avatar
Guolin Ke committed
1087
1088
  API_BEGIN();
  std::unique_ptr<Dataset> ret;
1089
1090
1091
1092
1093
  data_size_t nrows = static_cast<data_size_t>(num_total_row);
  ret.reset(new Dataset(nrows));
  const Dataset* reference_dataset = reinterpret_cast<const Dataset*>(reference);
  ret->CreateValid(reference_dataset);
  ret->InitByReference(nrows, reference_dataset);
Guolin Ke's avatar
Guolin Ke committed
1094
1095
1096
1097
  *out = ret.release();
  API_END();
}

1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
int LGBM_DatasetCreateFromSerializedReference(const void* ref_buffer,
                                              int32_t ref_buffer_size,
                                              int64_t num_row,
                                              int32_t num_classes,
                                              const char* parameters,
                                              DatasetHandle* out) {
  API_BEGIN();
  auto param = Config::Str2Map(parameters);
  Config config;
  config.Set(param);
  OMP_SET_NUM_THREADS(config.num_threads);
  DatasetLoader loader(config, nullptr, 1, nullptr);
  *out = loader.LoadFromSerializedReference(static_cast<const char*>(ref_buffer),
    static_cast<size_t>(ref_buffer_size),
    static_cast<data_size_t>(num_row),
    num_classes);
  API_END();
}

1117
1118
1119
1120
1121
int LGBM_DatasetInitStreaming(DatasetHandle dataset,
                              int32_t has_weights,
                              int32_t has_init_scores,
                              int32_t has_queries,
                              int32_t nclasses,
1122
1123
                              int32_t nthreads,
                              int32_t omp_max_threads) {
1124
1125
1126
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto num_data = p_dataset->num_data();
1127
  p_dataset->InitStreaming(num_data, has_weights, has_init_scores, has_queries, nclasses, nthreads, omp_max_threads);
1128
1129
1130
1131
  p_dataset->set_wait_for_manual_finish(true);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1132
int LGBM_DatasetPushRows(DatasetHandle dataset,
1133
1134
1135
1136
1137
                         const void* data,
                         int data_type,
                         int32_t nrow,
                         int32_t ncol,
                         int32_t start_row) {
Guolin Ke's avatar
Guolin Ke committed
1138
1139
1140
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto get_row_fun = RowFunctionFromDenseMatric(data, nrow, ncol, data_type, 1);
1141
1142
1143
  if (p_dataset->has_raw()) {
    p_dataset->ResizeRaw(p_dataset->num_numeric_features() + nrow);
  }
1144
  OMP_INIT_EX();
1145
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1146
  for (int i = 0; i < nrow; ++i) {
1147
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1148
1149
1150
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(tid, start_row + i, one_row);
1151
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1152
  }
1153
  OMP_THROW_EX();
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
  if (!p_dataset->wait_for_manual_finish() && (start_row + nrow == p_dataset->num_data())) {
    p_dataset->FinishLoad();
  }
  API_END();
}

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

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

1186
  OMP_INIT_EX();
1187
#pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
1188
1189
1190
  for (int i = 0; i < nrow; ++i) {
    OMP_LOOP_EX_BEGIN();
    // convert internal thread id to be unique based on external thread id
1191
    const int internal_tid = omp_get_thread_num() + (max_omp_threads * tid);
1192
1193
1194
1195
1196
1197
1198
1199
1200
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(internal_tid, start_row + i, one_row);
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();

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

  if (!p_dataset->wait_for_manual_finish() && (start_row + nrow == p_dataset->num_data())) {
Guolin Ke's avatar
Guolin Ke committed
1201
1202
1203
1204
1205
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1206
int LGBM_DatasetPushRowsByCSR(DatasetHandle dataset,
1207
1208
1209
1210
1211
1212
1213
1214
1215
                              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
1216
1217
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
1218
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
1219
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
1220
1221
1222
  if (p_dataset->has_raw()) {
    p_dataset->ResizeRaw(p_dataset->num_numeric_features() + nrow);
  }
1223
  OMP_INIT_EX();
1224
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1225
  for (int i = 0; i < nrow; ++i) {
1226
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1227
1228
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
1229
    p_dataset->PushOneRow(tid, static_cast<data_size_t>(start_row + i), one_row);
1230
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1231
  }
1232
  OMP_THROW_EX();
1233
  if (!p_dataset->wait_for_manual_finish() && (start_row + nrow == static_cast<int64_t>(p_dataset->num_data()))) {
Guolin Ke's avatar
Guolin Ke committed
1234
1235
    p_dataset->FinishLoad();
  }
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
  API_END();
}

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

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

1269
  OMP_INIT_EX();
1270
#pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
1271
1272
1273
  for (int i = 0; i < nrow; ++i) {
    OMP_LOOP_EX_BEGIN();
    // convert internal thread id to be unique based on external thread id
1274
    const int internal_tid = omp_get_thread_num() + (max_omp_threads * tid);
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(internal_tid, static_cast<data_size_t>(start_row + i), one_row);
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();

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

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

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

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

Guolin Ke's avatar
Guolin Ke committed
1303
int LGBM_DatasetCreateFromMat(const void* data,
1304
1305
1306
1307
1308
1309
1310
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
1311
1312
1313
1314
1315
  return LGBM_DatasetCreateFromMats(1,
                                    &data,
                                    data_type,
                                    &nrow,
                                    ncol,
1316
                                    &is_row_major,
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
                                    parameters,
                                    reference,
                                    out);
}

int LGBM_DatasetCreateFromMats(int32_t nmat,
                               const void** data,
                               int data_type,
                               int32_t* nrow,
                               int32_t ncol,
1327
                               int* is_row_major,
1328
1329
1330
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
1331
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1332
1333
  auto param = Config::Str2Map(parameters);
  Config config;
1334
  config.Set(param);
1335
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1336
  std::unique_ptr<Dataset> ret;
1337
1338
1339
1340
1341
1342
1343
  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) {
1344
    get_row_fun.push_back(RowFunctionFromDenseMatric(data[j], nrow[j], ncol, data_type, is_row_major[j]));
1345
  }
1346

Guolin Ke's avatar
Guolin Ke committed
1347
1348
  if (reference == nullptr) {
    // sample data first
1349
1350
    auto sample_indices = CreateSampleIndices(total_nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
1351
    std::vector<std::vector<double>> sample_values(ncol);
Guolin Ke's avatar
Guolin Ke committed
1352
    std::vector<std::vector<int>> sample_idx(ncol);
1353
1354
1355

    int offset = 0;
    int j = 0;
Guolin Ke's avatar
Guolin Ke committed
1356
    for (size_t i = 0; i < sample_indices.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1357
      auto idx = sample_indices[i];
1358
1359
1360
1361
      while ((idx - offset) >= nrow[j]) {
        offset += nrow[j];
        ++j;
      }
1362

1363
1364
1365
1366
1367
      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
1368
        }
Guolin Ke's avatar
Guolin Ke committed
1369
1370
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1371
    DatasetLoader loader(config, nullptr, 1, nullptr);
1372
1373
1374
1375
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             ncol,
                                             VectorSize<double>(sample_values).data(),
1376
1377
1378
                                             sample_cnt,
                                             total_nrow,
                                             total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1379
  } else {
1380
    ret.reset(new Dataset(total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1381
    ret->CreateValid(
1382
      reinterpret_cast<const Dataset*>(reference));
1383
1384
1385
    if (ret->has_raw()) {
      ret->ResizeRaw(total_nrow);
    }
Guolin Ke's avatar
Guolin Ke committed
1386
  }
1387
1388
1389
  int32_t start_row = 0;
  for (int j = 0; j < nmat; ++j) {
    OMP_INIT_EX();
1390
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
    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
1401
1402
  }
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1403
  *out = ret.release();
1404
  API_END();
1405
1406
}

Guolin Ke's avatar
Guolin Ke committed
1407
int LGBM_DatasetCreateFromCSR(const void* indptr,
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
                              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) {
1418
  API_BEGIN();
1419
1420
1421
1422
1423
  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
1424
1425
  auto param = Config::Str2Map(parameters);
  Config config;
1426
  config.Set(param);
1427
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1428
  std::unique_ptr<Dataset> ret;
1429
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1430
1431
1432
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
  if (reference == nullptr) {
    // sample data first
1433
1434
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1435
1436
    std::vector<std::vector<double>> sample_values(num_col);
    std::vector<std::vector<int>> sample_idx(num_col);
1437
1438
1439
1440
    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
1441
        CHECK_LT(inner_data.first, num_col);
Guolin Ke's avatar
Guolin Ke committed
1442
        if (std::fabs(inner_data.second) > kZeroThreshold || std::isnan(inner_data.second)) {
Guolin Ke's avatar
Guolin Ke committed
1443
1444
          sample_values[inner_data.first].emplace_back(inner_data.second);
          sample_idx[inner_data.first].emplace_back(static_cast<int>(i));
1445
1446
1447
        }
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1448
    DatasetLoader loader(config, nullptr, 1, nullptr);
1449
1450
1451
1452
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             static_cast<int>(num_col),
                                             VectorSize<double>(sample_values).data(),
1453
1454
1455
                                             sample_cnt,
                                             nrow,
                                             nrow));
1456
  } else {
1457
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1458
    ret->CreateValid(
1459
      reinterpret_cast<const Dataset*>(reference));
1460
1461
1462
    if (ret->has_raw()) {
      ret->ResizeRaw(nrow);
    }
1463
  }
1464
  OMP_INIT_EX();
1465
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
1466
  for (int i = 0; i < static_cast<int>(nindptr - 1); ++i) {
1467
    OMP_LOOP_EX_BEGIN();
1468
1469
1470
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    ret->PushOneRow(tid, i, one_row);
1471
    OMP_LOOP_EX_END();
1472
  }
1473
  OMP_THROW_EX();
1474
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1475
  *out = ret.release();
1476
  API_END();
1477
1478
}

1479
int LGBM_DatasetCreateFromCSRFunc(void* get_row_funptr,
1480
1481
1482
1483
1484
                                  int num_rows,
                                  int64_t num_col,
                                  const char* parameters,
                                  const DatasetHandle reference,
                                  DatasetHandle* out) {
1485
  API_BEGIN();
1486
1487
1488
1489
1490
  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.");
  }
1491
1492
1493
1494
  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);
1495
  OMP_SET_NUM_THREADS(config.num_threads);
1496
1497
1498
1499
  std::unique_ptr<Dataset> ret;
  int32_t nrow = num_rows;
  if (reference == nullptr) {
    // sample data first
1500
1501
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
1502
1503
1504
1505
1506
1507
1508
1509
    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
1510
        CHECK_LT(inner_data.first, num_col);
1511
1512
1513
1514
1515
1516
1517
        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);
1518
1519
1520
1521
    ret.reset(loader.ConstructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                             Vector2Ptr<int>(&sample_idx).data(),
                                             static_cast<int>(num_col),
                                             VectorSize<double>(sample_values).data(),
1522
1523
1524
                                             sample_cnt,
                                             nrow,
                                             nrow));
1525
1526
1527
1528
  } else {
    ret.reset(new Dataset(nrow));
    ret->CreateValid(
      reinterpret_cast<const Dataset*>(reference));
1529
1530
1531
    if (ret->has_raw()) {
      ret->ResizeRaw(nrow);
    }
1532
  }
1533

1534
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1535
  std::vector<std::pair<int, double>> thread_buffer;
1536
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static) private(thread_buffer)
1537
1538
1539
  for (int i = 0; i < num_rows; ++i) {
    OMP_LOOP_EX_BEGIN();
    {
1540
      const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1541
1542
      get_row_fun(i, thread_buffer);
      ret->PushOneRow(tid, i, thread_buffer);
1543
1544
1545
1546
1547
1548
1549
1550
1551
    }
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();
  ret->FinishLoad();
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1552
int LGBM_DatasetCreateFromCSC(const void* col_ptr,
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
                              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) {
1563
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1564
1565
  auto param = Config::Str2Map(parameters);
  Config config;
1566
  config.Set(param);
1567
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
1568
  std::unique_ptr<Dataset> ret;
Guolin Ke's avatar
Guolin Ke committed
1569
1570
1571
  int32_t nrow = static_cast<int32_t>(num_row);
  if (reference == nullptr) {
    // sample data first
1572
1573
    auto sample_indices = CreateSampleIndices(nrow, config);
    int sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
1574
    std::vector<std::vector<double>> sample_values(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
1575
    std::vector<std::vector<int>> sample_idx(ncol_ptr - 1);
1576
    OMP_INIT_EX();
1577
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1578
    for (int i = 0; i < static_cast<int>(sample_values.size()); ++i) {
1579
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1580
1581
1582
      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
1583
        if (std::fabs(val) > kZeroThreshold || std::isnan(val)) {
Guolin Ke's avatar
Guolin Ke committed
1584
1585
          sample_values[i].emplace_back(val);
          sample_idx[i].emplace_back(j);
Guolin Ke's avatar
Guolin Ke committed
1586
1587
        }
      }
1588
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1589
    }
1590
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1591
    DatasetLoader loader(config, nullptr, 1, nullptr);
1592
1593
1594
1595
    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(),
1596
1597
1598
                                             sample_cnt,
                                             nrow,
                                             nrow));
Guolin Ke's avatar
Guolin Ke committed
1599
  } else {
1600
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1601
    ret->CreateValid(
1602
      reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
1603
  }
1604
  OMP_INIT_EX();
1605
  #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
1606
  for (int i = 0; i < static_cast<int>(ncol_ptr - 1); ++i) {
1607
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1608
    const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1609
    int feature_idx = ret->InnerFeatureIndex(i);
1610
1611
1612
    if (feature_idx < 0) {
      continue;
    }
Guolin Ke's avatar
Guolin Ke committed
1613
1614
    int group = ret->Feature2Group(feature_idx);
    int sub_feature = ret->Feture2SubFeature(feature_idx);
Guolin Ke's avatar
Guolin Ke committed
1615
    CSC_RowIterator col_it(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, i);
Guolin Ke's avatar
Guolin Ke committed
1616
1617
1618
1619
1620
1621
1622
    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
1623
1624
1625
        if (row_idx < 0) {
          break;
        }
1626
        ret->PushOneData(tid, row_idx, group, feature_idx, sub_feature, pair.second);
Guolin Ke's avatar
Guolin Ke committed
1627
1628
1629
1630
      }
    } else {
      for (int row_idx = 0; row_idx < nrow; ++row_idx) {
        auto val = col_it.Get(row_idx);
1631
        ret->PushOneData(tid, row_idx, group, feature_idx, sub_feature, val);
Guolin Ke's avatar
Guolin Ke committed
1632
      }
Guolin Ke's avatar
Guolin Ke committed
1633
    }
1634
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1635
  }
1636
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1637
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1638
  *out = ret.release();
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
  API_END();
}

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

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

  std::unique_ptr<Dataset> ret;

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

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

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

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

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

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

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

  ret->FinishLoad();
  *out = ret.release();
1731
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1732
1733
}

Guolin Ke's avatar
Guolin Ke committed
1734
int LGBM_DatasetGetSubset(
1735
  const DatasetHandle handle,
wxchan's avatar
wxchan committed
1736
1737
1738
  const int32_t* used_row_indices,
  int32_t num_used_row_indices,
  const char* parameters,
Guolin Ke's avatar
typo  
Guolin Ke committed
1739
  DatasetHandle* out) {
wxchan's avatar
wxchan committed
1740
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1741
1742
  auto param = Config::Str2Map(parameters);
  Config config;
1743
  config.Set(param);
1744
  OMP_SET_NUM_THREADS(config.num_threads);
1745
  auto full_dataset = reinterpret_cast<const Dataset*>(handle);
1746
  CHECK_GT(num_used_row_indices, 0);
1747
1748
  const int32_t lower = 0;
  const int32_t upper = full_dataset->num_data() - 1;
1749
  CheckElementsIntervalClosed(used_row_indices, lower, upper, num_used_row_indices, "Used indices of subset");
1750
1751
1752
  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
1753
  auto ret = std::unique_ptr<Dataset>(new Dataset(num_used_row_indices));
1754
  ret->CopyFeatureMapperFrom(full_dataset);
1755
  ret->CopySubrow(full_dataset, used_row_indices, num_used_row_indices, true);
wxchan's avatar
wxchan committed
1756
1757
1758
1759
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1760
int LGBM_DatasetSetFeatureNames(
Guolin Ke's avatar
typo  
Guolin Ke committed
1761
  DatasetHandle handle,
Guolin Ke's avatar
Guolin Ke committed
1762
  const char** feature_names,
Guolin Ke's avatar
Guolin Ke committed
1763
  int num_feature_names) {
Guolin Ke's avatar
Guolin Ke committed
1764
1765
1766
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  std::vector<std::string> feature_names_str;
Guolin Ke's avatar
Guolin Ke committed
1767
  for (int i = 0; i < num_feature_names; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1768
1769
1770
1771
1772
1773
    feature_names_str.emplace_back(feature_names[i]);
  }
  dataset->set_feature_names(feature_names_str);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1774
int LGBM_DatasetGetFeatureNames(
1775
1776
1777
1778
1779
1780
    DatasetHandle handle,
    const int len,
    int* num_feature_names,
    const size_t buffer_len,
    size_t* out_buffer_len,
    char** feature_names) {
1781
  API_BEGIN();
1782
  *out_buffer_len = 0;
1783
1784
  auto dataset = reinterpret_cast<Dataset*>(handle);
  auto inside_feature_name = dataset->feature_names();
Guolin Ke's avatar
Guolin Ke committed
1785
1786
  *num_feature_names = static_cast<int>(inside_feature_name.size());
  for (int i = 0; i < *num_feature_names; ++i) {
1787
1788
1789
1790
1791
    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);
1792
1793
1794
1795
  }
  API_END();
}

1796
1797
1798
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1799
int LGBM_DatasetFree(DatasetHandle handle) {
1800
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1801
  delete reinterpret_cast<Dataset*>(handle);
1802
  API_END();
1803
1804
}

Guolin Ke's avatar
Guolin Ke committed
1805
int LGBM_DatasetSaveBinary(DatasetHandle handle,
1806
                           const char* filename) {
1807
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1808
1809
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->SaveBinaryFile(filename);
1810
  API_END();
1811
1812
}

1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
int LGBM_DatasetSerializeReferenceToBinary(DatasetHandle handle,
                                           ByteBufferHandle* out,
                                           int32_t* out_len) {
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  std::unique_ptr<LightGBM::ByteBuffer> ret;
  ret.reset(new LightGBM::ByteBuffer());
  dataset->SerializeReference(ret.get());
  *out_len = static_cast<int32_t>(ret->GetSize());
  *out = ret.release();
  API_END();
}

1826
1827
1828
1829
1830
1831
1832
1833
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
1834
int LGBM_DatasetSetField(DatasetHandle handle,
1835
1836
1837
1838
                         const char* field_name,
                         const void* field_data,
                         int num_element,
                         int type) {
1839
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1840
  auto dataset = reinterpret_cast<Dataset*>(handle);
1841
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1842
  if (type == C_API_DTYPE_FLOAT32) {
Guolin Ke's avatar
Guolin Ke committed
1843
    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
1844
  } else if (type == C_API_DTYPE_INT32) {
Guolin Ke's avatar
Guolin Ke committed
1845
    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
1846
1847
  } 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));
1848
  }
1849
1850
1851
  if (!is_success) {
    Log::Fatal("Input data type error or field not found");
  }
1852
  API_END();
1853
1854
}

1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
int LGBM_DatasetSetFieldFromArrow(DatasetHandle handle,
                                  const char* field_name,
                                  int64_t n_chunks,
                                  const ArrowArray* chunks,
                                  const ArrowSchema* schema) {
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  ArrowChunkedArray ca(n_chunks, chunks, schema);
  auto is_success = dataset->SetFieldFromArrow(field_name, ca);
  if (!is_success) {
    Log::Fatal("Input field is not supported");
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1870
int LGBM_DatasetGetField(DatasetHandle handle,
1871
1872
1873
1874
                         const char* field_name,
                         int* out_len,
                         const void** out_ptr,
                         int* out_type) {
1875
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1876
  auto dataset = reinterpret_cast<Dataset*>(handle);
1877
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1878
  if (dataset->GetFloatField(field_name, out_len, reinterpret_cast<const float**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1879
    *out_type = C_API_DTYPE_FLOAT32;
1880
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1881
  } else if (dataset->GetIntField(field_name, out_len, reinterpret_cast<const int**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1882
    *out_type = C_API_DTYPE_INT32;
1883
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1884
1885
1886
  } 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
1887
  }
1888
1889
1890
1891
1892
1893
  if (!is_success) {
    Log::Fatal("Field not found");
  }
  if (*out_ptr == nullptr) {
    *out_len = 0;
  }
1894
  API_END();
1895
1896
}

1897
int LGBM_DatasetUpdateParamChecking(const char* old_parameters, const char* new_parameters) {
1898
  API_BEGIN();
1899
1900
1901
1902
1903
  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);
1904
1905
1906
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1907
int LGBM_DatasetGetNumData(DatasetHandle handle,
1908
                           int* out) {
1909
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1910
1911
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_data();
1912
  API_END();
1913
1914
}

Guolin Ke's avatar
Guolin Ke committed
1915
int LGBM_DatasetGetNumFeature(DatasetHandle handle,
1916
                              int* out) {
1917
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1918
1919
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_total_features();
1920
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1921
}
1922

1923
1924
1925
1926
1927
int LGBM_DatasetGetFeatureNumBin(DatasetHandle handle,
                                 int feature,
                                 int* out) {
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
1928
1929
1930
1931
1932
  int num_features = dataset->num_total_features();
  if (feature < 0 || feature >= num_features) {
    Log::Fatal("Tried to retrieve number of bins for feature index %d, "
               "but the valid feature indices are [0, %d].", feature, num_features - 1);
  }
1933
1934
1935
1936
1937
1938
1939
1940
1941
  int inner_idx = dataset->InnerFeatureIndex(feature);
  if (inner_idx >= 0) {
    *out = dataset->FeatureNumBin(inner_idx);
  } else {
    *out = 0;
  }
  API_END();
}

1942
1943
1944
1945
1946
int LGBM_DatasetAddFeaturesFrom(DatasetHandle target,
                                DatasetHandle source) {
  API_BEGIN();
  auto target_d = reinterpret_cast<Dataset*>(target);
  auto source_d = reinterpret_cast<Dataset*>(source);
1947
  target_d->AddFeaturesFrom(source_d);
1948
1949
1950
  API_END();
}

1951
1952
// ---- start of booster

Guolin Ke's avatar
Guolin Ke committed
1953
int LGBM_BoosterCreate(const DatasetHandle train_data,
1954
1955
                       const char* parameters,
                       BoosterHandle* out) {
1956
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1957
  const Dataset* p_train_data = reinterpret_cast<const Dataset*>(train_data);
wxchan's avatar
wxchan committed
1958
1959
  auto ret = std::unique_ptr<Booster>(new Booster(p_train_data, parameters));
  *out = ret.release();
1960
  API_END();
1961
1962
}

Guolin Ke's avatar
Guolin Ke committed
1963
int LGBM_BoosterCreateFromModelfile(
1964
  const char* filename,
Guolin Ke's avatar
Guolin Ke committed
1965
  int* out_num_iterations,
1966
  BoosterHandle* out) {
1967
  API_BEGIN();
wxchan's avatar
wxchan committed
1968
  auto ret = std::unique_ptr<Booster>(new Booster(filename));
Guolin Ke's avatar
Guolin Ke committed
1969
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
wxchan's avatar
wxchan committed
1970
  *out = ret.release();
1971
  API_END();
1972
1973
}

Guolin Ke's avatar
Guolin Ke committed
1974
int LGBM_BoosterLoadModelFromString(
1975
1976
1977
1978
  const char* model_str,
  int* out_num_iterations,
  BoosterHandle* out) {
  API_BEGIN();
wxchan's avatar
wxchan committed
1979
  auto ret = std::unique_ptr<Booster>(new Booster(nullptr));
1980
1981
1982
1983
1984
1985
  ret->LoadModelFromString(model_str);
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
  *out = ret.release();
  API_END();
}

1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
int LGBM_BoosterGetLoadedParam(
  BoosterHandle handle,
  int64_t buffer_len,
  int64_t* out_len,
  char* out_str) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  std::string params = ref_booster->GetBoosting()->GetLoadedParam();
  *out_len = static_cast<int64_t>(params.size()) + 1;
  if (*out_len <= buffer_len) {
    std::memcpy(out_str, params.c_str(), *out_len);
  }
  API_END();
}

2001
2002
2003
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
2004
int LGBM_BoosterFree(BoosterHandle handle) {
2005
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2006
  delete reinterpret_cast<Booster*>(handle);
2007
  API_END();
2008
2009
}

2010
int LGBM_BoosterShuffleModels(BoosterHandle handle, int start_iter, int end_iter) {
2011
2012
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2013
  ref_booster->ShuffleModels(start_iter, end_iter);
2014
2015
2016
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2017
int LGBM_BoosterMerge(BoosterHandle handle,
2018
                      BoosterHandle other_handle) {
wxchan's avatar
wxchan committed
2019
2020
2021
2022
2023
2024
2025
  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
2026
int LGBM_BoosterAddValidData(BoosterHandle handle,
2027
                             const DatasetHandle valid_data) {
wxchan's avatar
wxchan committed
2028
2029
2030
2031
2032
2033
2034
  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
2035
int LGBM_BoosterResetTrainingData(BoosterHandle handle,
2036
                                  const DatasetHandle train_data) {
wxchan's avatar
wxchan committed
2037
2038
2039
2040
2041
2042
2043
  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
2044
int LGBM_BoosterResetParameter(BoosterHandle handle, const char* parameters) {
wxchan's avatar
wxchan committed
2045
2046
2047
2048
2049
2050
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->ResetConfig(parameters);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2051
int LGBM_BoosterGetNumClasses(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
2052
2053
2054
2055
2056
2057
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->NumberOfClasses();
  API_END();
}

2058
int LGBM_BoosterGetLinear(BoosterHandle handle, int* out) {
2059
2060
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2061
2062
2063
2064
2065
  if (ref_booster->GetBoosting()->IsLinear()) {
    *out = 1;
  } else {
    *out = 0;
  }
2066
2067
2068
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2069
2070
2071
2072
2073
2074
2075
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();
}

2076
int LGBM_BoosterUpdateOneIter(BoosterHandle handle, int* produced_empty_tree) {
2077
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2078
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2079
  if (ref_booster->TrainOneIter()) {
2080
    *produced_empty_tree = 1;
2081
  } else {
2082
    *produced_empty_tree = 0;
2083
  }
2084
  API_END();
2085
2086
}

Guolin Ke's avatar
Guolin Ke committed
2087
int LGBM_BoosterUpdateOneIterCustom(BoosterHandle handle,
2088
2089
                                    const float* grad,
                                    const float* hess,
2090
                                    int* produced_empty_tree) {
2091
  API_BEGIN();
2092
  #ifdef SCORE_T_USE_DOUBLE
2093
2094
2095
  (void) handle;       // UNUSED VARIABLE
  (void) grad;         // UNUSED VARIABLE
  (void) hess;         // UNUSED VARIABLE
2096
  (void) produced_empty_tree;  // UNUSED VARIABLE
2097
  Log::Fatal("Don't support custom loss function when SCORE_T_USE_DOUBLE is enabled");
2098
  #else
2099
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2100
  if (ref_booster->TrainOneIter(grad, hess)) {
2101
    *produced_empty_tree = 1;
2102
  } else {
2103
    *produced_empty_tree = 0;
2104
  }
2105
  #endif
2106
  API_END();
2107
2108
}

Guolin Ke's avatar
Guolin Ke committed
2109
int LGBM_BoosterRollbackOneIter(BoosterHandle handle) {
wxchan's avatar
wxchan committed
2110
2111
2112
2113
2114
2115
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->RollbackOneIter();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2116
int LGBM_BoosterGetCurrentIteration(BoosterHandle handle, int* out_iteration) {
wxchan's avatar
wxchan committed
2117
2118
2119
2120
2121
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_iteration = ref_booster->GetBoosting()->GetCurrentIteration();
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
2122

2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
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
2137
int LGBM_BoosterGetEvalCounts(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
2138
2139
2140
2141
2142
2143
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetEvalCounts();
  API_END();
}

2144
2145
2146
2147
2148
2149
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
2150
2151
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2152
  *out_len = ref_booster->GetEvalNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
2153
2154
2155
  API_END();
}

2156
2157
2158
2159
2160
2161
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
2162
2163
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2164
  *out_len = ref_booster->GetFeatureNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
2165
2166
2167
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2168
int LGBM_BoosterGetNumFeature(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
2169
2170
2171
2172
2173
2174
  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
2175
int LGBM_BoosterGetEval(BoosterHandle handle,
2176
2177
2178
                        int data_idx,
                        int* out_len,
                        double* out_results) {
2179
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2180
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2181
  auto boosting = ref_booster->GetBoosting();
wxchan's avatar
wxchan committed
2182
  auto result_buf = boosting->GetEvalAt(data_idx);
Guolin Ke's avatar
Guolin Ke committed
2183
  *out_len = static_cast<int>(result_buf.size());
2184
  for (size_t i = 0; i < result_buf.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
2185
    (out_results)[i] = static_cast<double>(result_buf[i]);
2186
  }
2187
  API_END();
2188
2189
}

Guolin Ke's avatar
Guolin Ke committed
2190
int LGBM_BoosterGetNumPredict(BoosterHandle handle,
2191
2192
                              int data_idx,
                              int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
2193
2194
2195
2196
2197
2198
  API_BEGIN();
  auto boosting = reinterpret_cast<Booster*>(handle)->GetBoosting();
  *out_len = boosting->GetNumPredictAt(data_idx);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2199
int LGBM_BoosterGetPredict(BoosterHandle handle,
2200
2201
2202
                           int data_idx,
                           int64_t* out_len,
                           double* out_result) {
2203
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2204
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2205
  ref_booster->GetPredictAt(data_idx, out_result, out_len);
2206
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2207
2208
}

Guolin Ke's avatar
Guolin Ke committed
2209
int LGBM_BoosterPredictForFile(BoosterHandle handle,
2210
2211
2212
                               const char* data_filename,
                               int data_has_header,
                               int predict_type,
2213
                               int start_iteration,
2214
                               int num_iteration,
2215
                               const char* parameter,
2216
                               const char* result_filename) {
2217
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2218
2219
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2220
  config.Set(param);
2221
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
2222
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2223
  ref_booster->Predict(start_iteration, num_iteration, predict_type, data_filename, data_has_header,
Guolin Ke's avatar
Guolin Ke committed
2224
                       config, result_filename);
2225
  API_END();
2226
2227
}

Guolin Ke's avatar
Guolin Ke committed
2228
int LGBM_BoosterCalcNumPredict(BoosterHandle handle,
2229
2230
                               int num_row,
                               int predict_type,
2231
                               int start_iteration,
2232
2233
                               int num_iteration,
                               int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
2234
2235
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2236
  *out_len = static_cast<int64_t>(num_row) * ref_booster->GetBoosting()->NumPredictOneRow(start_iteration,
2237
    num_iteration, predict_type == C_API_PREDICT_LEAF_INDEX, predict_type == C_API_PREDICT_CONTRIB);
Guolin Ke's avatar
Guolin Ke committed
2238
2239
2240
  API_END();
}

2241
2242
2243
2244
2245
2246
// Naming: In future versions of LightGBM, public API named around `FastConfig` should be made named around
// `SingleRowPredictor`, because it is specific to single row prediction, and doesn't actually hold only config.
// For now this is kept as `FastConfig` for backwards compatibility.
// At the same time, one should consider removing the old non-fast single row public API that stores its Predictor
// in the Booster, because that will enable removing these Predictors from the Booster, and associated initialization
// code.
2247
2248
int LGBM_FastConfigFree(FastConfigHandle fastConfig) {
  API_BEGIN();
2249
  delete reinterpret_cast<SingleRowPredictor*>(fastConfig);
2250
2251
2252
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2253
int LGBM_BoosterPredictForCSR(BoosterHandle handle,
2254
2255
2256
2257
2258
2259
2260
                              const void* indptr,
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
2261
                              int64_t num_col,
2262
                              int predict_type,
2263
                              int start_iteration,
2264
                              int num_iteration,
2265
                              const char* parameter,
2266
2267
                              int64_t* out_len,
                              double* out_result) {
2268
  API_BEGIN();
2269
2270
2271
2272
2273
  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
2274
2275
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2276
  config.Set(param);
2277
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
2278
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2279
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
2280
  int nrow = static_cast<int>(nindptr - 1);
2281
  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
2282
                       config, out_result, out_len);
2283
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2284
}
2285

2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
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,
2296
                                    int start_iteration,
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
                                    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);
2309
  OMP_SET_NUM_THREADS(config.num_threads);
2310
2311
2312
2313
2314
2315
2316
2317
  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;
2318
    ref_booster->PredictSparseCSR(start_iteration, num_iteration, predict_type, nrow, static_cast<int>(num_col_or_row), get_row_fun,
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
                                  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;
    };
2342
    ref_booster->PredictSparseCSC(start_iteration, num_iteration, predict_type, num_col_or_row, ncol, get_row_fun, config,
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
                                  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) {
2353
    delete[] reinterpret_cast<int32_t*>(indptr);
2354
  } else if (indptr_type == C_API_DTYPE_INT64) {
2355
    delete[] reinterpret_cast<int64_t*>(indptr);
2356
2357
2358
  } else {
    Log::Fatal("Unknown indptr type in LGBM_BoosterFreePredictSparse");
  }
2359
  delete[] indices;
2360
  if (data_type == C_API_DTYPE_FLOAT32) {
2361
    delete[] reinterpret_cast<float*>(data);
2362
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2363
    delete[] reinterpret_cast<double*>(data);
2364
2365
2366
2367
2368
2369
  } else {
    Log::Fatal("Unknown data type in LGBM_BoosterFreePredictSparse");
  }
  API_END();
}

2370
int LGBM_BoosterPredictForCSRSingleRow(BoosterHandle handle,
2371
2372
2373
2374
2375
2376
2377
                                       const void* indptr,
                                       int indptr_type,
                                       const int32_t* indices,
                                       const void* data,
                                       int data_type,
                                       int64_t nindptr,
                                       int64_t nelem,
2378
                                       int64_t num_col,
2379
                                       int predict_type,
2380
                                       int start_iteration,
2381
2382
2383
2384
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2385
  API_BEGIN();
2386
2387
2388
2389
2390
  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.");
  }
2391
2392
2393
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
2394
  OMP_SET_NUM_THREADS(config.num_threads);
2395
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2396
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
2397
  ref_booster->SetSingleRowPredictorInner(start_iteration, num_iteration, predict_type, config);
2398
  ref_booster->PredictSingleRow(predict_type, static_cast<int32_t>(num_col), get_row_fun, config, out_result, out_len);
2399
2400
2401
  API_END();
}

2402
int LGBM_BoosterPredictForCSRSingleRowFastInit(BoosterHandle handle,
2403
                                               const int predict_type,
2404
                                               const int start_iteration,
2405
                                               const int num_iteration,
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
                                               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.");
  }

2417
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2418

2419
2420
  std::unique_ptr<SingleRowPredictor> single_row_predictor =
    ref_booster->InitSingleRowPredictor(start_iteration, num_iteration, predict_type, data_type, static_cast<int32_t>(num_col), parameter);
2421

2422
  OMP_SET_NUM_THREADS(single_row_predictor->config.num_threads);
2423

2424
  *out_fastConfig = single_row_predictor.release();
2425
2426
2427
2428
2429
  API_END();
}

int LGBM_BoosterPredictForCSRSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* indptr,
2430
                                           const int indptr_type,
2431
2432
                                           const int32_t* indices,
                                           const void* data,
2433
2434
                                           const int64_t nindptr,
                                           const int64_t nelem,
2435
2436
2437
                                           int64_t* out_len,
                                           double* out_result) {
  API_BEGIN();
2438
2439
2440
  SingleRowPredictor *single_row_predictor = reinterpret_cast<SingleRowPredictor*>(fastConfig_handle);
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, single_row_predictor->data_type, nindptr, nelem);
  single_row_predictor->Predict(get_row_fun, out_result, out_len);
2441
2442
2443
  API_END();
}

2444

Guolin Ke's avatar
Guolin Ke committed
2445
int LGBM_BoosterPredictForCSC(BoosterHandle handle,
2446
2447
2448
2449
2450
2451
2452
2453
2454
                              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,
2455
                              int start_iteration,
2456
                              int num_iteration,
2457
                              const char* parameter,
2458
2459
                              int64_t* out_len,
                              double* out_result) {
Guolin Ke's avatar
Guolin Ke committed
2460
2461
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2462
2463
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2464
  config.Set(param);
2465
  OMP_SET_NUM_THREADS(config.num_threads);
2466
  int num_threads = OMP_NUM_THREADS();
Guolin Ke's avatar
Guolin Ke committed
2467
  int ncol = static_cast<int>(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
2468
2469
2470
2471
2472
  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
2473
2474
  }
  std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun =
Guolin Ke's avatar
Guolin Ke committed
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
      [&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;
      };
2487
  ref_booster->Predict(start_iteration, num_iteration, predict_type, static_cast<int>(num_row), ncol, get_row_fun, config,
cbecker's avatar
cbecker committed
2488
                       out_result, out_len);
Guolin Ke's avatar
Guolin Ke committed
2489
2490
2491
  API_END();
}

2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
int LGBM_BoosterValidateFeatureNames(BoosterHandle handle,
                                     const char** data_names,
                                     int data_num_features) {
  API_BEGIN();
  int booster_num_features;
  size_t out_buffer_len;
  LGBM_BoosterGetFeatureNames(handle, 0, &booster_num_features, 0, &out_buffer_len, nullptr);
  if (booster_num_features != data_num_features) {
    Log::Fatal("Model was trained on %d features, but got %d input features to predict.", booster_num_features, data_num_features);
  }
  std::vector<std::vector<char>> tmp_names(booster_num_features, std::vector<char>(out_buffer_len));
  std::vector<char*> booster_names = Vector2Ptr(&tmp_names);
  LGBM_BoosterGetFeatureNames(handle, data_num_features, &booster_num_features, out_buffer_len, &out_buffer_len, booster_names.data());
  for (int i = 0; i < booster_num_features; ++i) {
    if (strcmp(data_names[i], booster_names[i]) != 0) {
      Log::Fatal("Expected '%s' at position %d but found '%s'", booster_names[i], i, data_names[i]);
    }
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2513
int LGBM_BoosterPredictForMat(BoosterHandle handle,
2514
2515
2516
2517
2518
2519
                              const void* data,
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              int predict_type,
2520
                              int start_iteration,
2521
                              int num_iteration,
2522
                              const char* parameter,
2523
2524
                              int64_t* out_len,
                              double* out_result) {
2525
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2526
2527
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2528
  config.Set(param);
2529
  OMP_SET_NUM_THREADS(config.num_threads);
Guolin Ke's avatar
Guolin Ke committed
2530
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2531
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, nrow, ncol, data_type, is_row_major);
2532
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
2533
                       config, out_result, out_len);
2534
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2535
}
2536

2537
int LGBM_BoosterPredictForMatSingleRow(BoosterHandle handle,
2538
2539
2540
2541
2542
                                       const void* data,
                                       int data_type,
                                       int32_t ncol,
                                       int is_row_major,
                                       int predict_type,
2543
                                       int start_iteration,
2544
2545
2546
2547
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2548
2549
2550
2551
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
2552
  OMP_SET_NUM_THREADS(config.num_threads);
2553
2554
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, ncol, data_type, is_row_major);
2555
  ref_booster->SetSingleRowPredictorInner(start_iteration, num_iteration, predict_type, config);
2556
  ref_booster->PredictSingleRow(predict_type, ncol, get_row_fun, config, out_result, out_len);
2557
2558
2559
  API_END();
}

2560
int LGBM_BoosterPredictForMatSingleRowFastInit(BoosterHandle handle,
2561
                                               const int predict_type,
2562
                                               const int start_iteration,
2563
                                               const int num_iteration,
2564
2565
2566
2567
2568
                                               const int data_type,
                                               const int32_t ncol,
                                               const char* parameter,
                                               FastConfigHandle *out_fastConfig) {
  API_BEGIN();
2569
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2570

2571
2572
  std::unique_ptr<SingleRowPredictor> single_row_predictor =
    ref_booster->InitSingleRowPredictor(predict_type, start_iteration, num_iteration, data_type, ncol, parameter);
2573

2574
  OMP_SET_NUM_THREADS(single_row_predictor->config.num_threads);
2575

2576
  *out_fastConfig = single_row_predictor.release();
2577
2578
2579
2580
2581
2582
2583
2584
  API_END();
}

int LGBM_BoosterPredictForMatSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* data,
                                           int64_t* out_len,
                                           double* out_result) {
  API_BEGIN();
2585
  SingleRowPredictor *single_row_predictor = reinterpret_cast<SingleRowPredictor*>(fastConfig_handle);
2586
  // Single row in row-major format:
2587
2588
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, single_row_predictor->num_cols, single_row_predictor->data_type, 1);
  single_row_predictor->Predict(get_row_fun, out_result, out_len);
2589
2590
2591
  API_END();
}

2592

2593
2594
2595
2596
2597
2598
int LGBM_BoosterPredictForMats(BoosterHandle handle,
                               const void** data,
                               int data_type,
                               int32_t nrow,
                               int32_t ncol,
                               int predict_type,
2599
                               int start_iteration,
2600
2601
2602
2603
2604
2605
2606
2607
                               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);
2608
  OMP_SET_NUM_THREADS(config.num_threads);
2609
2610
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseRows(data, ncol, data_type);
2611
  ref_booster->Predict(start_iteration, num_iteration, predict_type, nrow, ncol, get_row_fun, config, out_result, out_len);
2612
2613
2614
  API_END();
}

2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
int LGBM_BoosterPredictForArrow(BoosterHandle handle,
                                int64_t n_chunks,
                                const ArrowArray* chunks,
                                const ArrowSchema* schema,
                                int predict_type,
                                int start_iteration,
                                int num_iteration,
                                const char* parameter,
                                int64_t* out_len,
                                double* out_result) {
  API_BEGIN();

  // Apply the configuration
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  OMP_SET_NUM_THREADS(config.num_threads);

  // Set up chunked array and iterators for all columns
  ArrowTable table(n_chunks, chunks, schema);
  std::vector<ArrowChunkedArray::Iterator<double>> its;
  its.reserve(table.get_num_columns());
  for (int64_t j = 0; j < table.get_num_columns(); ++j) {
    its.emplace_back(table.get_column(j).begin<double>());
  }

  // Build row function
  auto num_columns = table.get_num_columns();
  auto row_fn = [num_columns, &its] (int row_idx) {
    std::vector<std::pair<int, double>> result;
    result.reserve(num_columns);
    for (int64_t j = 0; j < num_columns; ++j) {
      result.emplace_back(static_cast<int>(j), its[j][row_idx]);
    }
    return result;
  };

  // Run prediction
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->Predict(start_iteration,
                       num_iteration,
                       predict_type,
                       static_cast<int>(table.get_num_rows()),
                       static_cast<int>(table.get_num_columns()),
                       row_fn,
                       config,
                       out_result,
                       out_len);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2666
int LGBM_BoosterSaveModel(BoosterHandle handle,
2667
                          int start_iteration,
2668
                          int num_iteration,
2669
                          int feature_importance_type,
2670
                          const char* filename) {
2671
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2672
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2673
2674
  ref_booster->SaveModelToFile(start_iteration, num_iteration,
                               feature_importance_type, filename);
wxchan's avatar
wxchan committed
2675
2676
2677
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2678
int LGBM_BoosterSaveModelToString(BoosterHandle handle,
2679
                                  int start_iteration,
2680
                                  int num_iteration,
2681
                                  int feature_importance_type,
2682
                                  int64_t buffer_len,
2683
                                  int64_t* out_len,
2684
                                  char* out_str) {
2685
2686
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2687
2688
  std::string model = ref_booster->SaveModelToString(
      start_iteration, num_iteration, feature_importance_type);
2689
  *out_len = static_cast<int64_t>(model.size()) + 1;
2690
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2691
    std::memcpy(out_str, model.c_str(), *out_len);
2692
2693
2694
2695
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2696
int LGBM_BoosterDumpModel(BoosterHandle handle,
2697
                          int start_iteration,
2698
                          int num_iteration,
2699
                          int feature_importance_type,
2700
2701
                          int64_t buffer_len,
                          int64_t* out_len,
2702
                          char* out_str) {
wxchan's avatar
wxchan committed
2703
2704
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2705
2706
  std::string model = ref_booster->DumpModel(start_iteration, num_iteration,
                                             feature_importance_type);
2707
  *out_len = static_cast<int64_t>(model.size()) + 1;
wxchan's avatar
wxchan committed
2708
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2709
    std::memcpy(out_str, model.c_str(), *out_len);
wxchan's avatar
wxchan committed
2710
  }
2711
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2712
}
2713

Guolin Ke's avatar
Guolin Ke committed
2714
int LGBM_BoosterGetLeafValue(BoosterHandle handle,
2715
2716
2717
                             int tree_idx,
                             int leaf_idx,
                             double* out_val) {
Guolin Ke's avatar
Guolin Ke committed
2718
2719
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2720
  *out_val = static_cast<double>(ref_booster->GetLeafValue(tree_idx, leaf_idx));
Guolin Ke's avatar
Guolin Ke committed
2721
2722
2723
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2724
int LGBM_BoosterSetLeafValue(BoosterHandle handle,
2725
2726
2727
                             int tree_idx,
                             int leaf_idx,
                             double val) {
Guolin Ke's avatar
Guolin Ke committed
2728
2729
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2730
  ref_booster->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
2731
2732
2733
  API_END();
}

2734
2735
2736
2737
2738
2739
2740
2741
2742
2743
2744
2745
2746
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();
}

2747
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
2758
2759
2760
2761
2762
2763
2764
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();
}

2765
2766
2767
2768
2769
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
2770
  Config config;
2771
  config.machines = RemoveQuotationSymbol(std::string(machines));
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
  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();
}

2787
2788
2789
int LGBM_NetworkInitWithFunctions(int num_machines, int rank,
                                  void* reduce_scatter_ext_fun,
                                  void* allgather_ext_fun) {
ww's avatar
ww committed
2790
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2791
  if (num_machines > 1) {
2792
    Network::Init(num_machines, rank, (ReduceScatterFunction)reduce_scatter_ext_fun, (AllgatherFunction)allgather_ext_fun);
ww's avatar
ww committed
2793
2794
2795
  }
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
2796

2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812
2813
int LGBM_SetMaxThreads(int num_threads) {
  API_BEGIN();
  if (num_threads <= 0) {
    LGBM_MAX_NUM_THREADS = -1;
  } else {
    LGBM_MAX_NUM_THREADS = num_threads;
  }
  API_END();
}

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


Guolin Ke's avatar
Guolin Ke committed
2814
// ---- start of some help functions
2815

2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840

template<typename T>
std::function<std::vector<double>(int row_idx)>
RowFunctionFromDenseMatric_helper(const void* data, int num_row, int num_col, int is_row_major) {
  const T* data_ptr = reinterpret_cast<const T*>(data);
  if (is_row_major) {
    return [=] (int row_idx) {
      std::vector<double> ret(num_col);
      auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
      for (int i = 0; i < num_col; ++i) {
        ret[i] = static_cast<double>(*(tmp_ptr + i));
      }
      return ret;
    };
  } else {
    return [=] (int row_idx) {
      std::vector<double> ret(num_col);
      for (int i = 0; i < num_col; ++i) {
        ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
      }
      return ret;
    };
  }
}

2841
2842
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
2843
  if (data_type == C_API_DTYPE_FLOAT32) {
2844
    return RowFunctionFromDenseMatric_helper<float>(data, num_row, num_col, is_row_major);
Guolin Ke's avatar
Guolin Ke committed
2845
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2846
    return RowFunctionFromDenseMatric_helper<double>(data, num_row, num_col, is_row_major);
2847
  }
2848
  Log::Fatal("Unknown data type in RowFunctionFromDenseMatric");
2849
  return nullptr;
2850
2851
2852
2853
}

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
2854
2855
  auto inner_function = RowFunctionFromDenseMatric(data, num_row, num_col, data_type, is_row_major);
  if (inner_function != nullptr) {
2856
    return [inner_function] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2857
2858
      auto raw_values = inner_function(row_idx);
      std::vector<std::pair<int, double>> ret;
Guolin Ke's avatar
Guolin Ke committed
2859
      ret.reserve(raw_values.size());
Guolin Ke's avatar
Guolin Ke committed
2860
      for (int i = 0; i < static_cast<int>(raw_values.size()); ++i) {
Guolin Ke's avatar
Guolin Ke committed
2861
        if (std::fabs(raw_values[i]) > kZeroThreshold || std::isnan(raw_values[i])) {
Guolin Ke's avatar
Guolin Ke committed
2862
          ret.emplace_back(i, raw_values[i]);
2863
        }
Guolin Ke's avatar
Guolin Ke committed
2864
2865
2866
      }
      return ret;
    };
2867
  }
Guolin Ke's avatar
Guolin Ke committed
2868
  return nullptr;
2869
2870
}

2871
2872
2873
2874
2875
2876
2877
// 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
2878
    ret.reserve(raw_values.size());
2879
2880
2881
2882
2883
2884
2885
2886
2887
    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;
  };
}

2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900
2901
2902
2903
2904
2905
2906
template<typename T, typename T1, typename T2>
std::function<std::vector<std::pair<int, double>>(T idx)>
RowFunctionFromCSR_helper(const void* indptr, const int32_t* indices, const void* data) {
  const T1* data_ptr = reinterpret_cast<const T1*>(data);
  const T2* ptr_indptr = reinterpret_cast<const T2*>(indptr);
  return [=] (T idx) {
    std::vector<std::pair<int, double>> ret;
    int64_t start = ptr_indptr[idx];
    int64_t end = ptr_indptr[idx + 1];
    if (end - start > 0)  {
      ret.reserve(end - start);
    }
    for (int64_t i = start; i < end; ++i) {
      ret.emplace_back(indices[i], data_ptr[i]);
    }
    return ret;
  };
}

2907
2908
template<typename T>
std::function<std::vector<std::pair<int, double>>(T idx)>
2909
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
2910
2911
  if (data_type == C_API_DTYPE_FLOAT32) {
    if (indptr_type == C_API_DTYPE_INT32) {
2912
     return RowFunctionFromCSR_helper<T, float, int32_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2913
    } else if (indptr_type == C_API_DTYPE_INT64) {
2914
     return RowFunctionFromCSR_helper<T, float, int64_t>(indptr, indices, data);
2915
    }
Guolin Ke's avatar
Guolin Ke committed
2916
2917
  } else if (data_type == C_API_DTYPE_FLOAT64) {
    if (indptr_type == C_API_DTYPE_INT32) {
2918
     return RowFunctionFromCSR_helper<T, double, int32_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2919
    } else if (indptr_type == C_API_DTYPE_INT64) {
2920
     return RowFunctionFromCSR_helper<T, double, int64_t>(indptr, indices, data);
Guolin Ke's avatar
Guolin Ke committed
2921
2922
    }
  }
2923
  Log::Fatal("Unknown data type in RowFunctionFromCSR");
2924
  return nullptr;
2925
2926
}

2927
2928
2929
2930
2931
2932
2933
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945


template <typename T1, typename T2>
std::function<std::pair<int, double>(int idx)> IterateFunctionFromCSC_helper(const void* col_ptr, const int32_t* indices, const void* data, int col_idx) {
  const T1* data_ptr = reinterpret_cast<const T1*>(data);
  const T2* ptr_col_ptr = reinterpret_cast<const T2*>(col_ptr);
  int64_t start = ptr_col_ptr[col_idx];
  int64_t end = ptr_col_ptr[col_idx + 1];
  return [=] (int offset) {
    int64_t i = static_cast<int64_t>(start + offset);
    if (i >= end) {
      return std::make_pair(-1, 0.0);
    }
    int idx = static_cast<int>(indices[i]);
    double val = static_cast<double>(data_ptr[i]);
    return std::make_pair(idx, val);
  };
}

Guolin Ke's avatar
Guolin Ke committed
2946
std::function<std::pair<int, double>(int idx)>
2947
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
2948
  CHECK(col_idx < ncol_ptr && col_idx >= 0);
Guolin Ke's avatar
Guolin Ke committed
2949
2950
  if (data_type == C_API_DTYPE_FLOAT32) {
    if (col_ptr_type == C_API_DTYPE_INT32) {
2951
      return IterateFunctionFromCSC_helper<float, int32_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2952
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2953
      return IterateFunctionFromCSC_helper<float, int64_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2954
    }
Guolin Ke's avatar
Guolin Ke committed
2955
2956
  } else if (data_type == C_API_DTYPE_FLOAT64) {
    if (col_ptr_type == C_API_DTYPE_INT32) {
2957
      return IterateFunctionFromCSC_helper<double, int32_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2958
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2959
      return IterateFunctionFromCSC_helper<double, int64_t>(col_ptr, indices, data, col_idx);
Guolin Ke's avatar
Guolin Ke committed
2960
2961
    }
  }
2962
  Log::Fatal("Unknown data type in CSC matrix");
2963
  return nullptr;
2964
2965
}

Guolin Ke's avatar
Guolin Ke committed
2966
CSC_RowIterator::CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
2967
                                 const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx) {
Guolin Ke's avatar
Guolin Ke committed
2968
2969
2970
2971
2972
2973
2974
2975
2976
  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;
2977
    }
Guolin Ke's avatar
Guolin Ke committed
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990
2991
2992
2993
2994
    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;
2995
    }
Guolin Ke's avatar
Guolin Ke committed
2996
2997
2998
    return ret;
  } else {
    return std::make_pair(-1, 0.0);
2999
  }
Guolin Ke's avatar
Guolin Ke committed
3000
}