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

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

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

29
#include "application/predictor.hpp"
Guolin Ke's avatar
Guolin Ke committed
30

Guolin Ke's avatar
Guolin Ke committed
31
32
namespace LightGBM {

Guolin Ke's avatar
Guolin Ke committed
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
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;

49
50
51
52
53
54
55
56
const int PREDICTOR_TYPES = 4;

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

Guolin Ke's avatar
Guolin Ke committed
57
  SingleRowPredictor(int predict_type, Boosting* boosting, const Config& config, int iter) {
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
    bool is_predict_leaf = false;
    bool is_raw_score = false;
    bool predict_contrib = false;
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
      is_predict_leaf = true;
    } else if (predict_type == C_API_PREDICT_RAW_SCORE) {
      is_raw_score = true;
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
      predict_contrib = true;
    } else {
      is_raw_score = false;
    }
    early_stop_ = config.pred_early_stop;
    early_stop_freq_ = config.pred_early_stop_freq;
    early_stop_margin_ = config.pred_early_stop_margin;
    iter_ = iter;
Guolin Ke's avatar
Guolin Ke committed
74
    predictor_.reset(new Predictor(boosting, iter_, is_raw_score, is_predict_leaf, predict_contrib,
75
                                   early_stop_, early_stop_freq_, early_stop_margin_));
Guolin Ke's avatar
Guolin Ke committed
76
    num_pred_in_one_row = boosting->NumPredictOneRow(iter_, is_predict_leaf, predict_contrib);
77
    predict_function = predictor_->GetPredictFunction();
Guolin Ke's avatar
Guolin Ke committed
78
    num_total_model_ = boosting->NumberOfTotalModel();
79
  }
80

81
  ~SingleRowPredictor() {}
82

Guolin Ke's avatar
Guolin Ke committed
83
  bool IsPredictorEqual(const Config& config, int iter, Boosting* boosting) {
84
85
86
87
88
    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();
89
  }
Guolin Ke's avatar
Guolin Ke committed
90

91
92
93
94
95
96
97
98
99
 private:
  std::unique_ptr<Predictor> predictor_;
  bool early_stop_;
  int early_stop_freq_;
  double early_stop_margin_;
  int iter_;
  int num_total_model_;
};

Guolin Ke's avatar
Guolin Ke committed
100
class Booster {
Nikita Titov's avatar
Nikita Titov committed
101
 public:
Guolin Ke's avatar
Guolin Ke committed
102
  explicit Booster(const char* filename) {
103
    boosting_.reset(Boosting::CreateBoosting("gbdt", filename));
104
105
  }

Guolin Ke's avatar
Guolin Ke committed
106
  Booster(const Dataset* train_data,
107
          const char* parameters) {
Guolin Ke's avatar
Guolin Ke committed
108
    auto param = Config::Str2Map(parameters);
wxchan's avatar
wxchan committed
109
    config_.Set(param);
110
111
112
    if (config_.num_threads > 0) {
      omp_set_num_threads(config_.num_threads);
    }
Guolin Ke's avatar
Guolin Ke committed
113
    // create boosting
Guolin Ke's avatar
Guolin Ke committed
114
    if (config_.input_model.size() > 0) {
115
116
      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
117
    }
Guolin Ke's avatar
Guolin Ke committed
118

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

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

  void MergeFrom(const Booster* other) {
    std::lock_guard<std::mutex> lock(mutex_);
    boosting_->MergeFrom(other->boosting_.get());
Guolin Ke's avatar
Guolin Ke committed
138
139
140
141
  }

  ~Booster() {
  }
142

143
  void CreateObjectiveAndMetrics() {
Guolin Ke's avatar
Guolin Ke committed
144
    // create objective function
Guolin Ke's avatar
Guolin Ke committed
145
146
    objective_fun_.reset(ObjectiveFunction::CreateObjectiveFunction(config_.objective,
                                                                    config_));
Guolin Ke's avatar
Guolin Ke committed
147
148
149
150
151
152
153
154
155
156
    if (objective_fun_ == nullptr) {
      Log::Warning("Using self-defined objective function");
    }
    // initialize the objective function
    if (objective_fun_ != nullptr) {
      objective_fun_->Init(train_data_->metadata(), train_data_->num_data());
    }

    // create training metric
    train_metric_.clear();
Guolin Ke's avatar
Guolin Ke committed
157
    for (auto metric_type : config_.metric) {
Guolin Ke's avatar
Guolin Ke committed
158
      auto metric = std::unique_ptr<Metric>(
Guolin Ke's avatar
Guolin Ke committed
159
        Metric::CreateMetric(metric_type, config_));
Guolin Ke's avatar
Guolin Ke committed
160
161
162
163
164
      if (metric == nullptr) { continue; }
      metric->Init(train_data_->metadata(), train_data_->num_data());
      train_metric_.push_back(std::move(metric));
    }
    train_metric_.shrink_to_fit();
165
166
167
168
169
170
171
172
173
174
175
  }

  void ResetTrainingData(const Dataset* train_data) {
    if (train_data != train_data_) {
      std::lock_guard<std::mutex> lock(mutex_);
      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
176
177
  }

178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
  static void CheckDatasetResetConfig(
      const Config& old_config,
      const std::unordered_map<std::string, std::string>& new_param) {
    Config new_config;
    new_config.Set(new_param);
    if (new_param.count("data_random_seed") &&
        new_config.data_random_seed != old_config.data_random_seed) {
      Log::Fatal("Cannot change data_random_seed after constructed Dataset handle.");
    }
    if (new_param.count("max_bin") &&
        new_config.max_bin != old_config.max_bin) {
      Log::Fatal("Cannot change max_bin after constructed Dataset handle.");
    }
    if (new_param.count("max_bin_by_feature") &&
        new_config.max_bin_by_feature != old_config.max_bin_by_feature) {
      Log::Fatal(
          "Cannot change max_bin_by_feature after constructed Dataset handle.");
    }
    if (new_param.count("bin_construct_sample_cnt") &&
        new_config.bin_construct_sample_cnt !=
            old_config.bin_construct_sample_cnt) {
      Log::Fatal(
          "Cannot change bin_construct_sample_cnt after constructed Dataset "
          "handle.");
    }
    if (new_param.count("min_data_in_bin") &&
        new_config.min_data_in_bin != old_config.min_data_in_bin) {
      Log::Fatal(
          "Cannot change min_data_in_bin after constructed Dataset handle.");
    }
    if (new_param.count("use_missing") &&
        new_config.use_missing != old_config.use_missing) {
      Log::Fatal("Cannot change use_missing after constructed Dataset handle.");
    }
    if (new_param.count("zero_as_missing") &&
        new_config.zero_as_missing != old_config.zero_as_missing) {
      Log::Fatal(
          "Cannot change zero_as_missing after constructed Dataset handle.");
    }
    if (new_param.count("categorical_feature") &&
        new_config.categorical_feature != old_config.categorical_feature) {
      Log::Fatal(
          "Cannot change categorical_feature after constructed Dataset "
          "handle.");
    }
    if (new_param.count("feature_pre_filter") &&
        new_config.feature_pre_filter != old_config.feature_pre_filter) {
      Log::Fatal(
          "Cannot change feature_pre_filter after constructed Dataset handle.");
    }
    if (new_param.count("is_enable_sparse") &&
        new_config.is_enable_sparse != old_config.is_enable_sparse) {
      Log::Fatal(
          "Cannot change is_enable_sparse after constructed Dataset handle.");
    }
    if (new_param.count("pre_partition") &&
        new_config.pre_partition != old_config.pre_partition) {
      Log::Fatal(
          "Cannot change pre_partition after constructed Dataset handle.");
    }
    if (new_param.count("enable_bundle") &&
        new_config.enable_bundle != old_config.enable_bundle) {
      Log::Fatal(
          "Cannot change enable_bundle after constructed Dataset handle.");
    }
    if (new_param.count("header") && new_config.header != old_config.header) {
      Log::Fatal("Cannot change header after constructed Dataset handle.");
    }
    if (new_param.count("two_round") &&
        new_config.two_round != old_config.two_round) {
      Log::Fatal("Cannot change two_round after constructed Dataset handle.");
    }
    if (new_param.count("label_column") &&
        new_config.label_column != old_config.label_column) {
      Log::Fatal(
          "Cannot change label_column after constructed Dataset handle.");
    }
    if (new_param.count("weight_column") &&
        new_config.weight_column != old_config.weight_column) {
      Log::Fatal(
          "Cannot change weight_column after constructed Dataset handle.");
    }
    if (new_param.count("group_column") &&
        new_config.group_column != old_config.group_column) {
      Log::Fatal(
          "Cannot change group_column after constructed Dataset handle.");
    }
    if (new_param.count("ignore_column") &&
        new_config.ignore_column != old_config.ignore_column) {
      Log::Fatal(
          "Cannot change ignore_column after constructed Dataset handle.");
    }
    if (new_param.count("forcedbins_filename")) {
      Log::Fatal("Cannot change forced bins after constructed Dataset handle.");
    }
    if (new_param.count("min_data_in_leaf") &&
        new_config.min_data_in_leaf < old_config.min_data_in_leaf &&
        old_config.feature_pre_filter) {
      Log::Fatal(
          "Reducing `min_data_in_leaf` with `feature_pre_filter=true` may "
          "cause unexpected behaviour "
          "for features that were pre-filtered by the larger "
          "`min_data_in_leaf`.\n"
          "You need to set `feature_pre_filter=false` to dynamically change "
          "the `min_data_in_leaf`.");
    }
  }

wxchan's avatar
wxchan committed
286
  void ResetConfig(const char* parameters) {
Guolin Ke's avatar
Guolin Ke committed
287
    std::lock_guard<std::mutex> lock(mutex_);
Guolin Ke's avatar
Guolin Ke committed
288
    auto param = Config::Str2Map(parameters);
wxchan's avatar
wxchan committed
289
    if (param.count("num_class")) {
290
      Log::Fatal("Cannot change num_class during training");
wxchan's avatar
wxchan committed
291
    }
Guolin Ke's avatar
Guolin Ke committed
292
293
    if (param.count("boosting")) {
      Log::Fatal("Cannot change boosting during training");
wxchan's avatar
wxchan committed
294
    }
Guolin Ke's avatar
Guolin Ke committed
295
    if (param.count("metric")) {
296
      Log::Fatal("Cannot change metric during training");
Guolin Ke's avatar
Guolin Ke committed
297
    }
298
299
    CheckDatasetResetConfig(config_, param);

Guolin Ke's avatar
Guolin Ke committed
300
    config_.Set(param);
301

302
303
304
    if (config_.num_threads > 0) {
      omp_set_num_threads(config_.num_threads);
    }
Guolin Ke's avatar
Guolin Ke committed
305
306
307

    if (param.count("objective")) {
      // create objective function
Guolin Ke's avatar
Guolin Ke committed
308
309
      objective_fun_.reset(ObjectiveFunction::CreateObjectiveFunction(config_.objective,
                                                                      config_));
Guolin Ke's avatar
Guolin Ke committed
310
311
312
313
314
315
316
      if (objective_fun_ == nullptr) {
        Log::Warning("Using self-defined objective function");
      }
      // initialize the objective function
      if (objective_fun_ != nullptr) {
        objective_fun_->Init(train_data_->metadata(), train_data_->num_data());
      }
317
318
      boosting_->ResetTrainingData(train_data_,
                                   objective_fun_.get(), Common::ConstPtrInVectorWrapper<Metric>(train_metric_));
wxchan's avatar
wxchan committed
319
    }
Guolin Ke's avatar
Guolin Ke committed
320

Guolin Ke's avatar
Guolin Ke committed
321
    boosting_->ResetConfig(&config_);
wxchan's avatar
wxchan committed
322
323
324
325
326
  }

  void AddValidData(const Dataset* valid_data) {
    std::lock_guard<std::mutex> lock(mutex_);
    valid_metrics_.emplace_back();
Guolin Ke's avatar
Guolin Ke committed
327
328
    for (auto metric_type : config_.metric) {
      auto metric = std::unique_ptr<Metric>(Metric::CreateMetric(metric_type, config_));
wxchan's avatar
wxchan committed
329
330
331
332
333
334
      if (metric == nullptr) { continue; }
      metric->Init(valid_data->metadata(), valid_data->num_data());
      valid_metrics_.back().push_back(std::move(metric));
    }
    valid_metrics_.back().shrink_to_fit();
    boosting_->AddValidDataset(valid_data,
335
                               Common::ConstPtrInVectorWrapper<Metric>(valid_metrics_.back()));
wxchan's avatar
wxchan committed
336
  }
Guolin Ke's avatar
Guolin Ke committed
337

338
  bool TrainOneIter() {
wxchan's avatar
wxchan committed
339
    std::lock_guard<std::mutex> lock(mutex_);
Guolin Ke's avatar
Guolin Ke committed
340
    return boosting_->TrainOneIter(nullptr, nullptr);
341
342
  }

Guolin Ke's avatar
Guolin Ke committed
343
344
345
346
347
348
349
350
351
352
353
  void Refit(const int32_t* leaf_preds, int32_t nrow, int32_t ncol) {
    std::lock_guard<std::mutex> lock(mutex_);
    std::vector<std::vector<int32_t>> v_leaf_preds(nrow, std::vector<int32_t>(ncol, 0));
    for (int i = 0; i < nrow; ++i) {
      for (int j = 0; j < ncol; ++j) {
        v_leaf_preds[i][j] = leaf_preds[i * ncol + j];
      }
    }
    boosting_->RefitTree(v_leaf_preds);
  }

354
  bool TrainOneIter(const score_t* gradients, const score_t* hessians) {
wxchan's avatar
wxchan committed
355
    std::lock_guard<std::mutex> lock(mutex_);
Guolin Ke's avatar
Guolin Ke committed
356
    return boosting_->TrainOneIter(gradients, hessians);
357
358
  }

wxchan's avatar
wxchan committed
359
360
361
362
363
  void RollbackOneIter() {
    std::lock_guard<std::mutex> lock(mutex_);
    boosting_->RollbackOneIter();
  }

364
  void PredictSingleRow(int num_iteration, int predict_type, int ncol,
365
366
367
               std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun,
               const Config& config,
               double* out_result, int64_t* out_len) {
368
369
370
    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);
371
    }
372
    std::lock_guard<std::mutex> lock(mutex_);
373
    if (single_row_predictor_[predict_type].get() == nullptr ||
Guolin Ke's avatar
Guolin Ke committed
374
375
        !single_row_predictor_[predict_type]->IsPredictorEqual(config, num_iteration, boosting_.get())) {
      single_row_predictor_[predict_type].reset(new SingleRowPredictor(predict_type, boosting_.get(),
376
                                                                       config, num_iteration));
377
378
379
    }
    auto one_row = get_row_fun(0);
    auto pred_wrt_ptr = out_result;
380
    single_row_predictor_[predict_type]->predict_function(one_row, pred_wrt_ptr);
381

382
    *out_len = single_row_predictor_[predict_type]->num_pred_in_one_row;
383
384
  }

385
  Predictor CreatePredictor(int num_iteration, int predict_type, int ncol, const Config& config) {
386
387
388
    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);
389
    }
Guolin Ke's avatar
Guolin Ke committed
390
391
    bool is_predict_leaf = false;
    bool is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
392
    bool predict_contrib = false;
Guolin Ke's avatar
Guolin Ke committed
393
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
Guolin Ke's avatar
Guolin Ke committed
394
      is_predict_leaf = true;
Guolin Ke's avatar
Guolin Ke committed
395
    } else if (predict_type == C_API_PREDICT_RAW_SCORE) {
Guolin Ke's avatar
Guolin Ke committed
396
      is_raw_score = true;
397
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
Guolin Ke's avatar
Guolin Ke committed
398
      predict_contrib = true;
Guolin Ke's avatar
Guolin Ke committed
399
400
    } else {
      is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
401
    }
Guolin Ke's avatar
Guolin Ke committed
402

Guolin Ke's avatar
Guolin Ke committed
403
    Predictor predictor(boosting_.get(), num_iteration, is_raw_score, is_predict_leaf, predict_contrib,
404
                        config.pred_early_stop, config.pred_early_stop_freq, config.pred_early_stop_margin);
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
    return predictor;
  }

  void Predict(int num_iteration, int predict_type, int nrow, int ncol,
               std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun,
               const Config& config,
               double* out_result, int64_t* out_len) {
    std::lock_guard<std::mutex> lock(mutex_);
    auto predictor = CreatePredictor(num_iteration, predict_type, ncol, config);
    bool is_predict_leaf = false;
    bool predict_contrib = false;
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
      is_predict_leaf = true;
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
      predict_contrib = true;
    }
Guolin Ke's avatar
Guolin Ke committed
421
    int64_t num_pred_in_one_row = boosting_->NumPredictOneRow(num_iteration, is_predict_leaf, predict_contrib);
Guolin Ke's avatar
Guolin Ke committed
422
    auto pred_fun = predictor.GetPredictFunction();
423
424
    OMP_INIT_EX();
    #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
425
    for (int i = 0; i < nrow; ++i) {
426
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
427
      auto one_row = get_row_fun(i);
Tony-Y's avatar
Tony-Y committed
428
      auto pred_wrt_ptr = out_result + static_cast<size_t>(num_pred_in_one_row) * i;
Guolin Ke's avatar
Guolin Ke committed
429
      pred_fun(one_row, pred_wrt_ptr);
430
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
431
    }
432
    OMP_THROW_EX();
433
    *out_len = num_pred_in_one_row * nrow;
Guolin Ke's avatar
Guolin Ke committed
434
435
  }

436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
  void PredictSparse(int num_iteration, int predict_type, int64_t nrow, int ncol,
                     std::function<std::vector<std::pair<int, double>>(int64_t row_idx)> get_row_fun,
                     const Config& config, int64_t* out_elements_size,
                     std::vector<std::vector<std::unordered_map<int, double>>>* agg_ptr,
                     int32_t** out_indices, void** out_data, int data_type,
                     bool* is_data_float32_ptr, int num_matrices) {
    auto predictor = CreatePredictor(num_iteration, predict_type, ncol, config);
    auto pred_sparse_fun = predictor.GetPredictSparseFunction();
    std::vector<std::vector<std::unordered_map<int, double>>>& agg = *agg_ptr;
    OMP_INIT_EX();
    #pragma omp parallel for schedule(static)
    for (int64_t i = 0; i < nrow; ++i) {
      OMP_LOOP_EX_BEGIN();
      auto one_row = get_row_fun(i);
      agg[i] = std::vector<std::unordered_map<int, double>>(num_matrices);
      pred_sparse_fun(one_row, &agg[i]);
      OMP_LOOP_EX_END();
    }
    OMP_THROW_EX();
    // calculate the nonzero data and indices size
    int64_t elements_size = 0;
    for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
      auto row_vector = agg[i];
      for (int j = 0; j < static_cast<int>(row_vector.size()); ++j) {
        elements_size += static_cast<int64_t>(row_vector[j].size());
      }
    }
    *out_elements_size = elements_size;
    *is_data_float32_ptr = false;
    // allocate data and indices arrays
    if (data_type == C_API_DTYPE_FLOAT32) {
      *out_data = new float[elements_size];
      *is_data_float32_ptr = true;
    } else if (data_type == C_API_DTYPE_FLOAT64) {
      *out_data = new double[elements_size];
    } else {
      Log::Fatal("Unknown data type in PredictSparse");
      return;
    }
    *out_indices = new int32_t[elements_size];
  }

  void PredictSparseCSR(int num_iteration, int predict_type, int64_t nrow, int ncol,
                        std::function<std::vector<std::pair<int, double>>(int64_t row_idx)> get_row_fun,
                        const Config& config,
                        int64_t* out_len, void** out_indptr, int indptr_type,
                        int32_t** out_indices, void** out_data, int data_type) {
    std::lock_guard<std::mutex> lock(mutex_);
    // Get the number of trees per iteration (for multiclass scenario we output multiple sparse matrices)
    int num_matrices = boosting_->NumModelPerIteration();
    bool is_indptr_int32 = false;
    bool is_data_float32 = false;
    int64_t indptr_size = (nrow + 1) * num_matrices;
    if (indptr_type == C_API_DTYPE_INT32) {
      *out_indptr = new int32_t[indptr_size];
      is_indptr_int32 = true;
    } else if (indptr_type == C_API_DTYPE_INT64) {
      *out_indptr = new int64_t[indptr_size];
    } else {
      Log::Fatal("Unknown indptr type in PredictSparseCSR");
      return;
    }
    // aggregated per row feature contribution results
    std::vector<std::vector<std::unordered_map<int, double>>> agg(nrow);
    int64_t elements_size = 0;
    PredictSparse(num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
                  out_indices, out_data, data_type, &is_data_float32, num_matrices);
    std::vector<int> row_sizes(num_matrices * nrow);
    std::vector<int64_t> row_matrix_offsets(num_matrices * nrow);
    int64_t row_vector_cnt = 0;
    for (int m = 0; m < num_matrices; ++m) {
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        auto row_vector = agg[i];
        auto row_vector_size = row_vector[m].size();
        // keep track of the row_vector sizes for parallelization
        row_sizes[row_vector_cnt] = static_cast<int>(row_vector_size);
        if (i == 0) {
          row_matrix_offsets[row_vector_cnt] = 0;
        } else {
          row_matrix_offsets[row_vector_cnt] = static_cast<int64_t>(row_sizes[row_vector_cnt - 1] + row_matrix_offsets[row_vector_cnt - 1]);
        }
        row_vector_cnt++;
      }
    }
    // copy vector results to output for each row
    int64_t indptr_index = 0;
    for (int m = 0; m < num_matrices; ++m) {
      if (is_indptr_int32) {
        (reinterpret_cast<int32_t*>(*out_indptr))[indptr_index] = 0;
      } else {
        (reinterpret_cast<int64_t*>(*out_indptr))[indptr_index] = 0;
      }
      indptr_index++;
      int64_t matrix_start_index = m * static_cast<int64_t>(agg.size());
      OMP_INIT_EX();
      #pragma omp parallel for schedule(static)
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        OMP_LOOP_EX_BEGIN();
        auto row_vector = agg[i];
        int64_t row_start_index = matrix_start_index + i;
        int64_t element_index = row_matrix_offsets[row_start_index];
        int64_t indptr_loop_index = indptr_index + i;
        for (auto it = row_vector[m].begin(); it != row_vector[m].end(); ++it) {
          (*out_indices)[element_index] = it->first;
          if (is_data_float32) {
            (reinterpret_cast<float*>(*out_data))[element_index] = static_cast<float>(it->second);
          } else {
            (reinterpret_cast<double*>(*out_data))[element_index] = it->second;
          }
          element_index++;
        }
        int64_t indptr_value = row_matrix_offsets[row_start_index] + row_sizes[row_start_index];
        if (is_indptr_int32) {
          (reinterpret_cast<int32_t*>(*out_indptr))[indptr_loop_index] = static_cast<int32_t>(indptr_value);
        } else {
          (reinterpret_cast<int64_t*>(*out_indptr))[indptr_loop_index] = indptr_value;
        }
        OMP_LOOP_EX_END();
      }
      OMP_THROW_EX();
      indptr_index += static_cast<int64_t>(agg.size());
    }
    out_len[0] = elements_size;
    out_len[1] = indptr_size;
  }

  void PredictSparseCSC(int num_iteration, int predict_type, int64_t nrow, int ncol,
                        std::function<std::vector<std::pair<int, double>>(int64_t row_idx)> get_row_fun,
                        const Config& config,
                        int64_t* out_len, void** out_col_ptr, int col_ptr_type,
                        int32_t** out_indices, void** out_data, int data_type) {
    std::lock_guard<std::mutex> lock(mutex_);
    // Get the number of trees per iteration (for multiclass scenario we output multiple sparse matrices)
    int num_matrices = boosting_->NumModelPerIteration();
    auto predictor = CreatePredictor(num_iteration, predict_type, ncol, config);
    auto pred_sparse_fun = predictor.GetPredictSparseFunction();
    bool is_col_ptr_int32 = false;
    bool is_data_float32 = false;
    int num_output_cols = ncol + 1;
    int col_ptr_size = (num_output_cols + 1) * num_matrices;
    if (col_ptr_type == C_API_DTYPE_INT32) {
      *out_col_ptr = new int32_t[col_ptr_size];
      is_col_ptr_int32 = true;
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
      *out_col_ptr = new int64_t[col_ptr_size];
    } else {
      Log::Fatal("Unknown col_ptr type in PredictSparseCSC");
      return;
    }
    // aggregated per row feature contribution results
    std::vector<std::vector<std::unordered_map<int, double>>> agg(nrow);
    int64_t elements_size = 0;
    PredictSparse(num_iteration, predict_type, nrow, ncol, get_row_fun, config, &elements_size, &agg,
                  out_indices, out_data, data_type, &is_data_float32, num_matrices);
    // calculate number of elements per column to construct
    // the CSC matrix with random access
    std::vector<std::vector<int64_t>> column_sizes(num_matrices);
    for (int m = 0; m < num_matrices; ++m) {
      column_sizes[m] = std::vector<int64_t>(num_output_cols, 0);
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        auto row_vector = agg[i];
        for (auto it = row_vector[m].begin(); it != row_vector[m].end(); ++it) {
          column_sizes[m][it->first] += 1;
        }
      }
    }
    // keep track of column counts
    std::vector<std::vector<int64_t>> column_counts(num_matrices);
    // keep track of beginning index for each column
    std::vector<std::vector<int64_t>> column_start_indices(num_matrices);
    // keep track of beginning index for each matrix
    std::vector<int64_t> matrix_start_indices(num_matrices, 0);
    int col_ptr_index = 0;
    for (int m = 0; m < num_matrices; ++m) {
      int64_t col_ptr_value = 0;
      column_start_indices[m] = std::vector<int64_t>(num_output_cols, 0);
      column_counts[m] = std::vector<int64_t>(num_output_cols, 0);
      if (is_col_ptr_int32) {
        (reinterpret_cast<int32_t*>(*out_col_ptr))[col_ptr_index] = static_cast<int32_t>(col_ptr_value);
      } else {
        (reinterpret_cast<int64_t*>(*out_col_ptr))[col_ptr_index] = col_ptr_value;
      }
      col_ptr_index++;
      for (int64_t i = 1; i < static_cast<int64_t>(column_sizes[m].size()); ++i) {
        column_start_indices[m][i] = column_sizes[m][i - 1] + column_start_indices[m][i - 1];
        if (is_col_ptr_int32) {
          (reinterpret_cast<int32_t*>(*out_col_ptr))[col_ptr_index] = static_cast<int32_t>(column_start_indices[m][i]);
        } else {
          (reinterpret_cast<int64_t*>(*out_col_ptr))[col_ptr_index] = column_start_indices[m][i];
        }
        col_ptr_index++;
      }
      int64_t last_elem_index = static_cast<int64_t>(column_sizes[m].size()) - 1;
      int64_t last_column_start_index = column_start_indices[m][last_elem_index];
      int64_t last_column_size = column_sizes[m][last_elem_index];
      if (is_col_ptr_int32) {
        (reinterpret_cast<int32_t*>(*out_col_ptr))[col_ptr_index] = static_cast<int32_t>(last_column_start_index + last_column_size);
      } else {
        (reinterpret_cast<int64_t*>(*out_col_ptr))[col_ptr_index] = last_column_start_index + last_column_size;
      }
      if (m != 0) {
        matrix_start_indices[m] = matrix_start_indices[m - 1] +
          last_column_start_index +
          last_column_size;
      }
    }
    for (int m = 0; m < num_matrices; ++m) {
      for (int64_t i = 0; i < static_cast<int64_t>(agg.size()); ++i) {
        auto row_vector = agg[i];
        for (auto it = row_vector[m].begin(); it != row_vector[m].end(); ++it) {
          int64_t col_idx = it->first;
          int64_t element_index = column_start_indices[m][col_idx] +
            matrix_start_indices[m] +
            column_counts[m][col_idx];
          // store the row index
          (*out_indices)[element_index] = static_cast<int32_t>(i);
          // update column count
          column_counts[m][col_idx]++;
          if (is_data_float32) {
            (reinterpret_cast<float*>(*out_data))[element_index] = static_cast<float>(it->second);
          } else {
            (reinterpret_cast<double*>(*out_data))[element_index] = it->second;
          }
        }
      }
    }
    out_len[0] = elements_size;
    out_len[1] = col_ptr_size;
  }

Guolin Ke's avatar
Guolin Ke committed
666
  void Predict(int num_iteration, int predict_type, const char* data_filename,
Guolin Ke's avatar
Guolin Ke committed
667
               int data_has_header, const Config& config,
cbecker's avatar
cbecker committed
668
               const char* result_filename) {
Guolin Ke's avatar
Guolin Ke committed
669
670
671
    std::lock_guard<std::mutex> lock(mutex_);
    bool is_predict_leaf = false;
    bool is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
672
    bool predict_contrib = false;
Guolin Ke's avatar
Guolin Ke committed
673
674
675
676
    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;
677
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
Guolin Ke's avatar
Guolin Ke committed
678
      predict_contrib = true;
Guolin Ke's avatar
Guolin Ke committed
679
680
681
    } else {
      is_raw_score = false;
    }
Guolin Ke's avatar
Guolin Ke committed
682
    Predictor predictor(boosting_.get(), num_iteration, is_raw_score, is_predict_leaf, predict_contrib,
683
                        config.pred_early_stop, config.pred_early_stop_freq, config.pred_early_stop_margin);
Guolin Ke's avatar
Guolin Ke committed
684
    bool bool_data_has_header = data_has_header > 0 ? true : false;
685
    predictor.Predict(data_filename, result_filename, bool_data_has_header, config.predict_disable_shape_check);
Guolin Ke's avatar
Guolin Ke committed
686
687
  }

Guolin Ke's avatar
Guolin Ke committed
688
  void GetPredictAt(int data_idx, double* out_result, int64_t* out_len) {
wxchan's avatar
wxchan committed
689
690
691
    boosting_->GetPredictAt(data_idx, out_result, out_len);
  }

692
693
  void SaveModelToFile(int start_iteration, int num_iteration, int feature_importance_type, const char* filename) {
    boosting_->SaveModelToFile(start_iteration, num_iteration, feature_importance_type, filename);
Guolin Ke's avatar
Guolin Ke committed
694
  }
695

696
  void LoadModelFromString(const char* model_str) {
697
698
    size_t len = std::strlen(model_str);
    boosting_->LoadModelFromString(model_str, len);
699
700
  }

701
702
703
704
  std::string SaveModelToString(int start_iteration, int num_iteration,
                                int feature_importance_type) {
    return boosting_->SaveModelToString(start_iteration,
                                        num_iteration, feature_importance_type);
705
706
  }

707
708
709
710
  std::string DumpModel(int start_iteration, int num_iteration,
                        int feature_importance_type) {
    return boosting_->DumpModel(start_iteration, num_iteration,
                                feature_importance_type);
wxchan's avatar
wxchan committed
711
  }
712

713
714
715
716
  std::vector<double> FeatureImportance(int num_iteration, int importance_type) {
    return boosting_->FeatureImportance(num_iteration, importance_type);
  }

717
718
719
720
721
722
723
724
725
726
  double UpperBoundValue() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return boosting_->GetUpperBoundValue();
  }

  double LowerBoundValue() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return boosting_->GetLowerBoundValue();
  }

Guolin Ke's avatar
Guolin Ke committed
727
  double GetLeafValue(int tree_idx, int leaf_idx) const {
Guolin Ke's avatar
Guolin Ke committed
728
    return dynamic_cast<GBDTBase*>(boosting_.get())->GetLeafValue(tree_idx, leaf_idx);
Guolin Ke's avatar
Guolin Ke committed
729
730
731
732
  }

  void SetLeafValue(int tree_idx, int leaf_idx, double val) {
    std::lock_guard<std::mutex> lock(mutex_);
Guolin Ke's avatar
Guolin Ke committed
733
    dynamic_cast<GBDTBase*>(boosting_.get())->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
734
735
  }

736
  void ShuffleModels(int start_iter, int end_iter) {
737
    std::lock_guard<std::mutex> lock(mutex_);
738
    boosting_->ShuffleModels(start_iter, end_iter);
739
740
  }

wxchan's avatar
wxchan committed
741
742
743
744
745
746
747
  int GetEvalCounts() const {
    int ret = 0;
    for (const auto& metric : train_metric_) {
      ret += static_cast<int>(metric->GetName().size());
    }
    return ret;
  }
748

749
750
  int GetEvalNames(char** out_strs, const int len, const size_t buffer_len, size_t *out_buffer_len) const {
    *out_buffer_len = 0;
wxchan's avatar
wxchan committed
751
752
753
    int idx = 0;
    for (const auto& metric : train_metric_) {
      for (const auto& name : metric->GetName()) {
754
755
756
757
758
        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
759
760
761
762
763
764
        ++idx;
      }
    }
    return idx;
  }

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

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

Nikita Titov's avatar
Nikita Titov committed
781
 private:
wxchan's avatar
wxchan committed
782
  const Dataset* train_data_;
Guolin Ke's avatar
Guolin Ke committed
783
  std::unique_ptr<Boosting> boosting_;
784
  std::unique_ptr<SingleRowPredictor> single_row_predictor_[PREDICTOR_TYPES];
785

Guolin Ke's avatar
Guolin Ke committed
786
  /*! \brief All configs */
Guolin Ke's avatar
Guolin Ke committed
787
  Config config_;
Guolin Ke's avatar
Guolin Ke committed
788
  /*! \brief Metric for training data */
Guolin Ke's avatar
Guolin Ke committed
789
  std::vector<std::unique_ptr<Metric>> train_metric_;
Guolin Ke's avatar
Guolin Ke committed
790
  /*! \brief Metrics for validation data */
Guolin Ke's avatar
Guolin Ke committed
791
  std::vector<std::vector<std::unique_ptr<Metric>>> valid_metrics_;
Guolin Ke's avatar
Guolin Ke committed
792
  /*! \brief Training objective function */
Guolin Ke's avatar
Guolin Ke committed
793
  std::unique_ptr<ObjectiveFunction> objective_fun_;
wxchan's avatar
wxchan committed
794
  /*! \brief mutex for threading safe call */
795
  mutable std::mutex mutex_;
Guolin Ke's avatar
Guolin Ke committed
796
797
};

798
}  // namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
799

800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
// explicitly declare symbols from LightGBM namespace
using LightGBM::AllgatherFunction;
using LightGBM::Booster;
using LightGBM::Common::CheckElementsIntervalClosed;
using LightGBM::Common::RemoveQuotationSymbol;
using LightGBM::Common::Vector2Ptr;
using LightGBM::Common::VectorSize;
using LightGBM::Config;
using LightGBM::data_size_t;
using LightGBM::Dataset;
using LightGBM::DatasetLoader;
using LightGBM::kZeroThreshold;
using LightGBM::LGBM_APIHandleException;
using LightGBM::Log;
using LightGBM::Network;
using LightGBM::Random;
using LightGBM::ReduceScatterFunction;
Guolin Ke's avatar
Guolin Ke committed
817

Guolin Ke's avatar
Guolin Ke committed
818
819
820
821
822
823
824
825
// 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);

826
827
828
std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseRows(const void** data, int num_col, int data_type);

829
830
template<typename T>
std::function<std::vector<std::pair<int, double>>(T idx)>
Guolin Ke's avatar
Guolin Ke committed
831
RowFunctionFromCSR(const void* indptr, int indptr_type, const int32_t* indices,
832
                   const void* data, int data_type, int64_t nindptr, int64_t nelem);
Guolin Ke's avatar
Guolin Ke committed
833
834
835

// Row iterator of on column for CSC matrix
class CSC_RowIterator {
Nikita Titov's avatar
Nikita Titov committed
836
 public:
Guolin Ke's avatar
Guolin Ke committed
837
  CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
838
                  const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx);
Guolin Ke's avatar
Guolin Ke committed
839
840
841
842
843
  ~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
844
845

 private:
Guolin Ke's avatar
Guolin Ke committed
846
847
848
849
850
851
852
853
854
  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
855
const char* LGBM_GetLastError() {
wxchan's avatar
wxchan committed
856
  return LastErrorMsg();
Guolin Ke's avatar
Guolin Ke committed
857
858
}

859
860
861
862
863
864
int LGBM_RegisterLogCallback(void (*callback)(const char*)) {
  API_BEGIN();
  Log::ResetCallBack(callback);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
865
int LGBM_DatasetCreateFromFile(const char* filename,
866
867
868
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
869
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
870
871
  auto param = Config::Str2Map(parameters);
  Config config;
872
873
874
875
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
876
  DatasetLoader loader(config, nullptr, 1, filename);
Guolin Ke's avatar
Guolin Ke committed
877
  if (reference == nullptr) {
878
    if (Network::num_machines() == 1) {
879
      *out = loader.LoadFromFile(filename);
880
    } else {
881
      *out = loader.LoadFromFile(filename, Network::rank(), Network::num_machines());
882
    }
Guolin Ke's avatar
Guolin Ke committed
883
  } else {
884
    *out = loader.LoadFromFileAlignWithOtherDataset(filename,
885
                                                    reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
886
  }
887
  API_END();
Guolin Ke's avatar
Guolin Ke committed
888
889
}

890

Guolin Ke's avatar
Guolin Ke committed
891
int LGBM_DatasetCreateFromSampledColumn(double** sample_data,
892
893
894
895
896
897
898
                                        int** sample_indices,
                                        int32_t ncol,
                                        const int* num_per_col,
                                        int32_t num_sample_row,
                                        int32_t num_total_row,
                                        const char* parameters,
                                        DatasetHandle* out) {
899
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
900
901
  auto param = Config::Str2Map(parameters);
  Config config;
902
903
904
905
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
906
  DatasetLoader loader(config, nullptr, 1, nullptr);
907
908
909
910
  *out = loader.CostructFromSampleData(sample_data, sample_indices, ncol, num_per_col,
                                       num_sample_row,
                                       static_cast<data_size_t>(num_total_row));
  API_END();
Guolin Ke's avatar
Guolin Ke committed
911
912
}

913

Guolin Ke's avatar
Guolin Ke committed
914
int LGBM_DatasetCreateByReference(const DatasetHandle reference,
915
916
                                  int64_t num_total_row,
                                  DatasetHandle* out) {
Guolin Ke's avatar
Guolin Ke committed
917
918
919
920
921
922
923
924
  API_BEGIN();
  std::unique_ptr<Dataset> ret;
  ret.reset(new Dataset(static_cast<data_size_t>(num_total_row)));
  ret->CreateValid(reinterpret_cast<const Dataset*>(reference));
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
925
int LGBM_DatasetPushRows(DatasetHandle dataset,
926
927
928
929
930
                         const void* data,
                         int data_type,
                         int32_t nrow,
                         int32_t ncol,
                         int32_t start_row) {
Guolin Ke's avatar
Guolin Ke committed
931
932
933
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto get_row_fun = RowFunctionFromDenseMatric(data, nrow, ncol, data_type, 1);
934
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
935
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
936
  for (int i = 0; i < nrow; ++i) {
937
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
938
939
940
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(tid, start_row + i, one_row);
941
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
942
  }
943
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
944
945
946
947
948
949
  if (start_row + nrow == p_dataset->num_data()) {
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
950
int LGBM_DatasetPushRowsByCSR(DatasetHandle dataset,
951
952
953
954
955
956
957
958
959
                              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
960
961
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
962
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
963
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
964
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
965
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
966
  for (int i = 0; i < nrow; ++i) {
967
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
968
969
970
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(tid,
971
                          static_cast<data_size_t>(start_row + i), one_row);
972
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
973
  }
974
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
975
976
977
978
979
980
  if (start_row + nrow == static_cast<int64_t>(p_dataset->num_data())) {
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
981
int LGBM_DatasetCreateFromMat(const void* data,
982
983
984
985
986
987
988
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
  return LGBM_DatasetCreateFromMats(1,
                                    &data,
                                    data_type,
                                    &nrow,
                                    ncol,
                                    is_row_major,
                                    parameters,
                                    reference,
                                    out);
}


int LGBM_DatasetCreateFromMats(int32_t nmat,
                               const void** data,
                               int data_type,
                               int32_t* nrow,
                               int32_t ncol,
                               int is_row_major,
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
1010
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1011
1012
  auto param = Config::Str2Map(parameters);
  Config config;
1013
1014
1015
1016
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1017
  std::unique_ptr<Dataset> ret;
1018
1019
1020
1021
1022
1023
1024
1025
1026
  int32_t total_nrow = 0;
  for (int j = 0; j < nmat; ++j) {
    total_nrow += nrow[j];
  }

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

Guolin Ke's avatar
Guolin Ke committed
1028
1029
  if (reference == nullptr) {
    // sample data first
Guolin Ke's avatar
Guolin Ke committed
1030
    Random rand(config.data_random_seed);
1031
1032
    int sample_cnt = static_cast<int>(total_nrow < config.bin_construct_sample_cnt ? total_nrow : config.bin_construct_sample_cnt);
    auto sample_indices = rand.Sample(total_nrow, sample_cnt);
1033
    sample_cnt = static_cast<int>(sample_indices.size());
1034
    std::vector<std::vector<double>> sample_values(ncol);
Guolin Ke's avatar
Guolin Ke committed
1035
    std::vector<std::vector<int>> sample_idx(ncol);
1036
1037
1038

    int offset = 0;
    int j = 0;
Guolin Ke's avatar
Guolin Ke committed
1039
    for (size_t i = 0; i < sample_indices.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1040
      auto idx = sample_indices[i];
1041
1042
1043
1044
      while ((idx - offset) >= nrow[j]) {
        offset += nrow[j];
        ++j;
      }
1045

1046
1047
1048
1049
1050
      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
1051
        }
Guolin Ke's avatar
Guolin Ke committed
1052
1053
      }
    }
Guolin Ke's avatar
Guolin Ke committed
1054
    DatasetLoader loader(config, nullptr, 1, nullptr);
1055
1056
    ret.reset(loader.CostructFromSampleData(Vector2Ptr<double>(&sample_values).data(),
                                            Vector2Ptr<int>(&sample_idx).data(),
1057
                                            ncol,
1058
                                            VectorSize<double>(sample_values).data(),
1059
                                            sample_cnt, total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1060
  } else {
1061
    ret.reset(new Dataset(total_nrow));
Guolin Ke's avatar
Guolin Ke committed
1062
    ret->CreateValid(
1063
      reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
1064
  }
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
  int32_t start_row = 0;
  for (int j = 0; j < nmat; ++j) {
    OMP_INIT_EX();
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < nrow[j]; ++i) {
      OMP_LOOP_EX_BEGIN();
      const int tid = omp_get_thread_num();
      auto one_row = get_row_fun[j](i);
      ret->PushOneRow(tid, start_row + i, one_row);
      OMP_LOOP_EX_END();
    }
    OMP_THROW_EX();

    start_row += nrow[j];
Guolin Ke's avatar
Guolin Ke committed
1079
1080
  }
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1081
  *out = ret.release();
1082
  API_END();
1083
1084
}

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

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

1210
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1211
1212
  std::vector<std::pair<int, double>> thread_buffer;
  #pragma omp parallel for schedule(static) private(thread_buffer)
1213
1214
1215
  for (int i = 0; i < num_rows; ++i) {
    OMP_LOOP_EX_BEGIN();
    {
1216
      const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1217
1218
      get_row_fun(i, thread_buffer);
      ret->PushOneRow(tid, i, thread_buffer);
1219
1220
1221
1222
1223
1224
1225
1226
1227
    }
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();
  ret->FinishLoad();
  *out = ret.release();
  API_END();
}

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

Guolin Ke's avatar
Guolin Ke committed
1316
int LGBM_DatasetGetSubset(
1317
  const DatasetHandle handle,
wxchan's avatar
wxchan committed
1318
1319
1320
  const int32_t* used_row_indices,
  int32_t num_used_row_indices,
  const char* parameters,
Guolin Ke's avatar
typo  
Guolin Ke committed
1321
  DatasetHandle* out) {
wxchan's avatar
wxchan committed
1322
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1323
1324
  auto param = Config::Str2Map(parameters);
  Config config;
1325
1326
1327
1328
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
1329
  auto full_dataset = reinterpret_cast<const Dataset*>(handle);
1330
  CHECK_GT(num_used_row_indices, 0);
1331
1332
  const int32_t lower = 0;
  const int32_t upper = full_dataset->num_data() - 1;
1333
  CheckElementsIntervalClosed(used_row_indices, lower, upper, num_used_row_indices, "Used indices of subset");
1334
1335
1336
  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
1337
  auto ret = std::unique_ptr<Dataset>(new Dataset(num_used_row_indices));
1338
  ret->CopyFeatureMapperFrom(full_dataset);
1339
  ret->CopySubrow(full_dataset, used_row_indices, num_used_row_indices, true);
wxchan's avatar
wxchan committed
1340
1341
1342
1343
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1344
int LGBM_DatasetSetFeatureNames(
Guolin Ke's avatar
typo  
Guolin Ke committed
1345
  DatasetHandle handle,
Guolin Ke's avatar
Guolin Ke committed
1346
  const char** feature_names,
Guolin Ke's avatar
Guolin Ke committed
1347
  int num_feature_names) {
Guolin Ke's avatar
Guolin Ke committed
1348
1349
1350
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  std::vector<std::string> feature_names_str;
Guolin Ke's avatar
Guolin Ke committed
1351
  for (int i = 0; i < num_feature_names; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1352
1353
1354
1355
1356
1357
    feature_names_str.emplace_back(feature_names[i]);
  }
  dataset->set_feature_names(feature_names_str);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1358
int LGBM_DatasetGetFeatureNames(
1359
1360
1361
1362
1363
1364
    DatasetHandle handle,
    const int len,
    int* num_feature_names,
    const size_t buffer_len,
    size_t* out_buffer_len,
    char** feature_names) {
1365
  API_BEGIN();
1366
  *out_buffer_len = 0;
1367
1368
  auto dataset = reinterpret_cast<Dataset*>(handle);
  auto inside_feature_name = dataset->feature_names();
Guolin Ke's avatar
Guolin Ke committed
1369
1370
  *num_feature_names = static_cast<int>(inside_feature_name.size());
  for (int i = 0; i < *num_feature_names; ++i) {
1371
1372
1373
1374
1375
    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);
1376
1377
1378
1379
  }
  API_END();
}

1380
1381
1382
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1383
int LGBM_DatasetFree(DatasetHandle handle) {
1384
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1385
  delete reinterpret_cast<Dataset*>(handle);
1386
  API_END();
1387
1388
}

Guolin Ke's avatar
Guolin Ke committed
1389
int LGBM_DatasetSaveBinary(DatasetHandle handle,
1390
                           const char* filename) {
1391
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1392
1393
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->SaveBinaryFile(filename);
1394
  API_END();
1395
1396
}

1397
1398
1399
1400
1401
1402
1403
1404
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
1405
int LGBM_DatasetSetField(DatasetHandle handle,
1406
1407
1408
1409
                         const char* field_name,
                         const void* field_data,
                         int num_element,
                         int type) {
1410
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1411
  auto dataset = reinterpret_cast<Dataset*>(handle);
1412
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1413
  if (type == C_API_DTYPE_FLOAT32) {
Guolin Ke's avatar
Guolin Ke committed
1414
    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
1415
  } else if (type == C_API_DTYPE_INT32) {
Guolin Ke's avatar
Guolin Ke committed
1416
    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
1417
1418
  } 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));
1419
  }
1420
  if (!is_success) { Log::Fatal("Input data type error or field not found"); }
1421
  API_END();
1422
1423
}

Guolin Ke's avatar
Guolin Ke committed
1424
int LGBM_DatasetGetField(DatasetHandle handle,
1425
1426
1427
1428
                         const char* field_name,
                         int* out_len,
                         const void** out_ptr,
                         int* out_type) {
1429
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1430
  auto dataset = reinterpret_cast<Dataset*>(handle);
1431
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1432
  if (dataset->GetFloatField(field_name, out_len, reinterpret_cast<const float**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1433
    *out_type = C_API_DTYPE_FLOAT32;
1434
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1435
  } else if (dataset->GetIntField(field_name, out_len, reinterpret_cast<const int**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1436
    *out_type = C_API_DTYPE_INT32;
1437
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1438
1439
1440
  } 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
1441
  }
1442
  if (!is_success) { Log::Fatal("Field not found"); }
wxchan's avatar
wxchan committed
1443
  if (*out_ptr == nullptr) { *out_len = 0; }
1444
  API_END();
1445
1446
}

1447
int LGBM_DatasetUpdateParamChecking(const char* old_parameters, const char* new_parameters) {
1448
  API_BEGIN();
1449
1450
1451
1452
1453
  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);
1454
1455
1456
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1457
int LGBM_DatasetGetNumData(DatasetHandle handle,
1458
                           int* out) {
1459
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1460
1461
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_data();
1462
  API_END();
1463
1464
}

Guolin Ke's avatar
Guolin Ke committed
1465
int LGBM_DatasetGetNumFeature(DatasetHandle handle,
1466
                              int* out) {
1467
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1468
1469
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_total_features();
1470
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1471
}
1472

1473
1474
1475
1476
1477
int LGBM_DatasetAddFeaturesFrom(DatasetHandle target,
                                DatasetHandle source) {
  API_BEGIN();
  auto target_d = reinterpret_cast<Dataset*>(target);
  auto source_d = reinterpret_cast<Dataset*>(source);
1478
  target_d->AddFeaturesFrom(source_d);
1479
1480
1481
  API_END();
}

1482
1483
// ---- start of booster

Guolin Ke's avatar
Guolin Ke committed
1484
int LGBM_BoosterCreate(const DatasetHandle train_data,
1485
1486
                       const char* parameters,
                       BoosterHandle* out) {
1487
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1488
  const Dataset* p_train_data = reinterpret_cast<const Dataset*>(train_data);
wxchan's avatar
wxchan committed
1489
1490
  auto ret = std::unique_ptr<Booster>(new Booster(p_train_data, parameters));
  *out = ret.release();
1491
  API_END();
1492
1493
}

Guolin Ke's avatar
Guolin Ke committed
1494
int LGBM_BoosterCreateFromModelfile(
1495
  const char* filename,
Guolin Ke's avatar
Guolin Ke committed
1496
  int* out_num_iterations,
1497
  BoosterHandle* out) {
1498
  API_BEGIN();
wxchan's avatar
wxchan committed
1499
  auto ret = std::unique_ptr<Booster>(new Booster(filename));
Guolin Ke's avatar
Guolin Ke committed
1500
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
wxchan's avatar
wxchan committed
1501
  *out = ret.release();
1502
  API_END();
1503
1504
}

Guolin Ke's avatar
Guolin Ke committed
1505
int LGBM_BoosterLoadModelFromString(
1506
1507
1508
1509
  const char* model_str,
  int* out_num_iterations,
  BoosterHandle* out) {
  API_BEGIN();
wxchan's avatar
wxchan committed
1510
  auto ret = std::unique_ptr<Booster>(new Booster(nullptr));
1511
1512
1513
1514
1515
1516
  ret->LoadModelFromString(model_str);
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
  *out = ret.release();
  API_END();
}

1517
1518
1519
#ifdef _MSC_VER
  #pragma warning(disable : 4702)
#endif
Guolin Ke's avatar
Guolin Ke committed
1520
int LGBM_BoosterFree(BoosterHandle handle) {
1521
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1522
  delete reinterpret_cast<Booster*>(handle);
1523
  API_END();
1524
1525
}

1526
int LGBM_BoosterShuffleModels(BoosterHandle handle, int start_iter, int end_iter) {
1527
1528
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1529
  ref_booster->ShuffleModels(start_iter, end_iter);
1530
1531
1532
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1533
int LGBM_BoosterMerge(BoosterHandle handle,
1534
                      BoosterHandle other_handle) {
wxchan's avatar
wxchan committed
1535
1536
1537
1538
1539
1540
1541
  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
1542
int LGBM_BoosterAddValidData(BoosterHandle handle,
1543
                             const DatasetHandle valid_data) {
wxchan's avatar
wxchan committed
1544
1545
1546
1547
1548
1549
1550
  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
1551
int LGBM_BoosterResetTrainingData(BoosterHandle handle,
1552
                                  const DatasetHandle train_data) {
wxchan's avatar
wxchan committed
1553
1554
1555
1556
1557
1558
1559
  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
1560
int LGBM_BoosterResetParameter(BoosterHandle handle, const char* parameters) {
wxchan's avatar
wxchan committed
1561
1562
1563
1564
1565
1566
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->ResetConfig(parameters);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1567
int LGBM_BoosterGetNumClasses(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1568
1569
1570
1571
1572
1573
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->NumberOfClasses();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1574
1575
1576
1577
1578
1579
1580
int LGBM_BoosterRefit(BoosterHandle handle, const int32_t* leaf_preds, int32_t nrow, int32_t ncol) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->Refit(leaf_preds, nrow, ncol);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1581
int LGBM_BoosterUpdateOneIter(BoosterHandle handle, int* is_finished) {
1582
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1583
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1584
1585
1586
1587
1588
  if (ref_booster->TrainOneIter()) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1589
  API_END();
1590
1591
}

Guolin Ke's avatar
Guolin Ke committed
1592
int LGBM_BoosterUpdateOneIterCustom(BoosterHandle handle,
1593
1594
1595
                                    const float* grad,
                                    const float* hess,
                                    int* is_finished) {
1596
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1597
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1598
  #ifdef SCORE_T_USE_DOUBLE
1599
  Log::Fatal("Don't support custom loss function when SCORE_T_USE_DOUBLE is enabled");
1600
  #else
1601
1602
1603
1604
1605
  if (ref_booster->TrainOneIter(grad, hess)) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1606
  #endif
1607
  API_END();
1608
1609
}

Guolin Ke's avatar
Guolin Ke committed
1610
int LGBM_BoosterRollbackOneIter(BoosterHandle handle) {
wxchan's avatar
wxchan committed
1611
1612
1613
1614
1615
1616
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->RollbackOneIter();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1617
int LGBM_BoosterGetCurrentIteration(BoosterHandle handle, int* out_iteration) {
wxchan's avatar
wxchan committed
1618
1619
1620
1621
1622
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_iteration = ref_booster->GetBoosting()->GetCurrentIteration();
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
1623

1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
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
1638
int LGBM_BoosterGetEvalCounts(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1639
1640
1641
1642
1643
1644
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetEvalCounts();
  API_END();
}

1645
1646
1647
1648
1649
1650
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
1651
1652
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1653
  *out_len = ref_booster->GetEvalNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1654
1655
1656
  API_END();
}

1657
1658
1659
1660
1661
1662
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
1663
1664
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1665
  *out_len = ref_booster->GetFeatureNames(out_strs, len, buffer_len, out_buffer_len);
wxchan's avatar
wxchan committed
1666
1667
1668
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1669
int LGBM_BoosterGetNumFeature(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1670
1671
1672
1673
1674
1675
  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
1676
int LGBM_BoosterGetEval(BoosterHandle handle,
1677
1678
1679
                        int data_idx,
                        int* out_len,
                        double* out_results) {
1680
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1681
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1682
  auto boosting = ref_booster->GetBoosting();
wxchan's avatar
wxchan committed
1683
  auto result_buf = boosting->GetEvalAt(data_idx);
Guolin Ke's avatar
Guolin Ke committed
1684
  *out_len = static_cast<int>(result_buf.size());
1685
  for (size_t i = 0; i < result_buf.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1686
    (out_results)[i] = static_cast<double>(result_buf[i]);
1687
  }
1688
  API_END();
1689
1690
}

Guolin Ke's avatar
Guolin Ke committed
1691
int LGBM_BoosterGetNumPredict(BoosterHandle handle,
1692
1693
                              int data_idx,
                              int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1694
1695
1696
1697
1698
1699
  API_BEGIN();
  auto boosting = reinterpret_cast<Booster*>(handle)->GetBoosting();
  *out_len = boosting->GetNumPredictAt(data_idx);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1700
int LGBM_BoosterGetPredict(BoosterHandle handle,
1701
1702
1703
                           int data_idx,
                           int64_t* out_len,
                           double* out_result) {
1704
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1705
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1706
  ref_booster->GetPredictAt(data_idx, out_result, out_len);
1707
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1708
1709
}

Guolin Ke's avatar
Guolin Ke committed
1710
int LGBM_BoosterPredictForFile(BoosterHandle handle,
1711
1712
1713
1714
                               const char* data_filename,
                               int data_has_header,
                               int predict_type,
                               int num_iteration,
1715
                               const char* parameter,
1716
                               const char* result_filename) {
1717
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1718
1719
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1720
1721
1722
1723
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1724
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
cbecker's avatar
cbecker committed
1725
  ref_booster->Predict(num_iteration, predict_type, data_filename, data_has_header,
Guolin Ke's avatar
Guolin Ke committed
1726
                       config, result_filename);
1727
  API_END();
1728
1729
}

Guolin Ke's avatar
Guolin Ke committed
1730
int LGBM_BoosterCalcNumPredict(BoosterHandle handle,
1731
1732
1733
1734
                               int num_row,
                               int predict_type,
                               int num_iteration,
                               int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1735
1736
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1737
1738
  *out_len = static_cast<int64_t>(num_row) * ref_booster->GetBoosting()->NumPredictOneRow(
    num_iteration, predict_type == C_API_PREDICT_LEAF_INDEX, predict_type == C_API_PREDICT_CONTRIB);
Guolin Ke's avatar
Guolin Ke committed
1739
1740
1741
  API_END();
}

1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
/*!
 * \brief Object to store resources meant for single-row Fast Predict methods.
 *
 * Meant to be used as a basic struct by the *Fast* predict methods only.
 * It stores the configuration resources for reuse during prediction.
 *
 * Even the row function is stored. We score the instance at the same memory
 * address all the time. One just replaces the feature values at that address
 * and scores again with the *Fast* methods.
 */
struct FastConfig {
  FastConfig(Booster *const booster_ptr,
             const char *parameter,
             const int data_type_,
             const int32_t num_cols) : booster(booster_ptr), data_type(data_type_), ncol(num_cols) {
    config.Set(Config::Str2Map(parameter));
  }

  Booster* const booster;
  Config config;
  const int data_type;
  const int32_t ncol;
};

int LGBM_FastConfigFree(FastConfigHandle fastConfig) {
  API_BEGIN();
  delete reinterpret_cast<FastConfig*>(fastConfig);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1772
int LGBM_BoosterPredictForCSR(BoosterHandle handle,
1773
1774
1775
1776
1777
1778
1779
                              const void* indptr,
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
1780
                              int64_t num_col,
1781
1782
                              int predict_type,
                              int num_iteration,
1783
                              const char* parameter,
1784
1785
                              int64_t* out_len,
                              double* out_result) {
1786
  API_BEGIN();
1787
1788
1789
1790
1791
  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
1792
1793
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1794
1795
1796
1797
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1798
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1799
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
1800
  int nrow = static_cast<int>(nindptr - 1);
1801
  ref_booster->Predict(num_iteration, predict_type, nrow, static_cast<int>(num_col), get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
1802
                       config, out_result, out_len);
1803
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1804
}
1805

1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
int LGBM_BoosterPredictSparseOutput(BoosterHandle handle,
                                    const void* indptr,
                                    int indptr_type,
                                    const int32_t* indices,
                                    const void* data,
                                    int data_type,
                                    int64_t nindptr,
                                    int64_t nelem,
                                    int64_t num_col_or_row,
                                    int predict_type,
                                    int num_iteration,
                                    const char* parameter,
                                    int matrix_type,
                                    int64_t* out_len,
                                    void** out_indptr,
                                    int32_t** out_indices,
                                    void** out_data) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  if (matrix_type == C_API_MATRIX_TYPE_CSR) {
    if (num_col_or_row <= 0) {
      Log::Fatal("The number of columns should be greater than zero.");
    } else if (num_col_or_row >= INT32_MAX) {
      Log::Fatal("The number of columns should be smaller than INT32_MAX.");
    }
    auto get_row_fun = RowFunctionFromCSR<int64_t>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
    int64_t nrow = nindptr - 1;
    ref_booster->PredictSparseCSR(num_iteration, predict_type, nrow, static_cast<int>(num_col_or_row), get_row_fun,
                                  config, out_len, out_indptr, indptr_type, out_indices, out_data, data_type);
  } else if (matrix_type == C_API_MATRIX_TYPE_CSC) {
    int num_threads = OMP_NUM_THREADS();
    int ncol = static_cast<int>(nindptr - 1);
    std::vector<std::vector<CSC_RowIterator>> iterators(num_threads, std::vector<CSC_RowIterator>());
    for (int i = 0; i < num_threads; ++i) {
      for (int j = 0; j < ncol; ++j) {
        iterators[i].emplace_back(indptr, indptr_type, indices, data, data_type, nindptr, nelem, j);
      }
    }
    std::function<std::vector<std::pair<int, double>>(int64_t row_idx)> get_row_fun =
      [&iterators, ncol](int64_t i) {
      std::vector<std::pair<int, double>> one_row;
      one_row.reserve(ncol);
      const int tid = omp_get_thread_num();
      for (int j = 0; j < ncol; ++j) {
        auto val = iterators[tid][j].Get(static_cast<int>(i));
        if (std::fabs(val) > kZeroThreshold || std::isnan(val)) {
          one_row.emplace_back(j, val);
        }
      }
      return one_row;
    };
    ref_booster->PredictSparseCSC(num_iteration, predict_type, num_col_or_row, ncol, get_row_fun, config,
                                  out_len, out_indptr, indptr_type, out_indices, out_data, data_type);
  } else {
    Log::Fatal("Unknown matrix type in LGBM_BoosterPredictSparseOutput");
  }
  API_END();
}

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

1891
int LGBM_BoosterPredictForCSRSingleRow(BoosterHandle handle,
1892
1893
1894
1895
1896
1897
1898
                                       const void* indptr,
                                       int indptr_type,
                                       const int32_t* indices,
                                       const void* data,
                                       int data_type,
                                       int64_t nindptr,
                                       int64_t nelem,
1899
                                       int64_t num_col,
1900
1901
1902
1903
1904
                                       int predict_type,
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
1905
  API_BEGIN();
1906
1907
1908
1909
1910
  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.");
  }
1911
1912
1913
1914
1915
1916
1917
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1918
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1919
  ref_booster->PredictSingleRow(num_iteration, predict_type, static_cast<int32_t>(num_col), get_row_fun, config, out_result, out_len);
1920
1921
1922
  API_END();
}

1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
int LGBM_BoosterPredictForCSRSingleRowFastInit(BoosterHandle handle,
                                               const int data_type,
                                               const int64_t num_col,
                                               const char* parameter,
                                               FastConfigHandle *out_fastConfig) {
  API_BEGIN();
  if (num_col <= 0) {
    Log::Fatal("The number of columns should be greater than zero.");
  } else if (num_col >= INT32_MAX) {
    Log::Fatal("The number of columns should be smaller than INT32_MAX.");
  }

  auto fastConfig_ptr = std::unique_ptr<FastConfig>(new FastConfig(
    reinterpret_cast<Booster*>(handle),
    parameter,
    data_type,
    static_cast<int32_t>(num_col)));

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

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

int LGBM_BoosterPredictForCSRSingleRowFast(FastConfigHandle fastConfig_handle,
                                           const void* indptr,
                                           int indptr_type,
                                           const int32_t* indices,
                                           const void* data,
                                           int64_t nindptr,
                                           int64_t nelem,
                                           int predict_type,
                                           int num_iteration,
                                           int64_t* out_len,
                                           double* out_result) {
  API_BEGIN();
  FastConfig *fastConfig = reinterpret_cast<FastConfig*>(fastConfig_handle);
  auto get_row_fun = RowFunctionFromCSR<int>(indptr, indptr_type, indices, data, fastConfig->data_type, nindptr, nelem);
  fastConfig->booster->PredictSingleRow(num_iteration, predict_type, fastConfig->ncol,
                                        get_row_fun, fastConfig->config, out_result, out_len);
  API_END();
}

1968

Guolin Ke's avatar
Guolin Ke committed
1969
int LGBM_BoosterPredictForCSC(BoosterHandle handle,
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
                              const void* col_ptr,
                              int col_ptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t ncol_ptr,
                              int64_t nelem,
                              int64_t num_row,
                              int predict_type,
                              int num_iteration,
1980
                              const char* parameter,
1981
1982
                              int64_t* out_len,
                              double* out_result) {
Guolin Ke's avatar
Guolin Ke committed
1983
1984
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
1985
1986
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1987
1988
1989
1990
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
1991
  int num_threads = OMP_NUM_THREADS();
Guolin Ke's avatar
Guolin Ke committed
1992
  int ncol = static_cast<int>(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
1993
1994
1995
1996
1997
  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
1998
1999
  }
  std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun =
Guolin Ke's avatar
Guolin Ke committed
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
      [&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;
      };
2012
  ref_booster->Predict(num_iteration, predict_type, static_cast<int>(num_row), ncol, get_row_fun, config,
cbecker's avatar
cbecker committed
2013
                       out_result, out_len);
Guolin Ke's avatar
Guolin Ke committed
2014
2015
2016
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2017
int LGBM_BoosterPredictForMat(BoosterHandle handle,
2018
2019
2020
2021
2022
2023
2024
                              const void* data,
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              int predict_type,
                              int num_iteration,
2025
                              const char* parameter,
2026
2027
                              int64_t* out_len,
                              double* out_result) {
2028
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2029
2030
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
2031
2032
2033
2034
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
2035
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2036
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, nrow, ncol, data_type, is_row_major);
2037
  ref_booster->Predict(num_iteration, predict_type, nrow, ncol, get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
2038
                       config, out_result, out_len);
2039
  API_END();
Guolin Ke's avatar
Guolin Ke committed
2040
}
2041

2042
int LGBM_BoosterPredictForMatSingleRow(BoosterHandle handle,
2043
2044
2045
2046
2047
2048
2049
2050
2051
                                       const void* data,
                                       int data_type,
                                       int32_t ncol,
                                       int is_row_major,
                                       int predict_type,
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
2052
2053
2054
2055
2056
2057
2058
2059
2060
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, ncol, data_type, is_row_major);
2061
  ref_booster->PredictSingleRow(num_iteration, predict_type, ncol, get_row_fun, config, out_result, out_len);
2062
2063
2064
  API_END();
}

2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
int LGBM_BoosterPredictForMatSingleRowFastInit(BoosterHandle handle,
                                               const int data_type,
                                               const int32_t ncol,
                                               const char* parameter,
                                               FastConfigHandle *out_fastConfig) {
  API_BEGIN();
  auto fastConfig_ptr = std::unique_ptr<FastConfig>(new FastConfig(
    reinterpret_cast<Booster*>(handle),
    parameter,
    data_type,
    ncol));

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

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

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

2101

2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
int LGBM_BoosterPredictForMats(BoosterHandle handle,
                               const void** data,
                               int data_type,
                               int32_t nrow,
                               int32_t ncol,
                               int predict_type,
                               int num_iteration,
                               const char* parameter,
                               int64_t* out_len,
                               double* out_result) {
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseRows(data, ncol, data_type);
2121
  ref_booster->Predict(num_iteration, predict_type, nrow, ncol, get_row_fun, config, out_result, out_len);
2122
2123
2124
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2125
int LGBM_BoosterSaveModel(BoosterHandle handle,
2126
                          int start_iteration,
2127
                          int num_iteration,
2128
                          int feature_importance_type,
2129
                          const char* filename) {
2130
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2131
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2132
2133
  ref_booster->SaveModelToFile(start_iteration, num_iteration,
                               feature_importance_type, filename);
wxchan's avatar
wxchan committed
2134
2135
2136
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2137
int LGBM_BoosterSaveModelToString(BoosterHandle handle,
2138
                                  int start_iteration,
2139
                                  int num_iteration,
2140
                                  int feature_importance_type,
2141
                                  int64_t buffer_len,
2142
                                  int64_t* out_len,
2143
                                  char* out_str) {
2144
2145
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
2146
2147
  std::string model = ref_booster->SaveModelToString(
      start_iteration, num_iteration, feature_importance_type);
2148
  *out_len = static_cast<int64_t>(model.size()) + 1;
2149
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
2150
    std::memcpy(out_str, model.c_str(), *out_len);
2151
2152
2153
2154
  }
  API_END();
}

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

Guolin Ke's avatar
Guolin Ke committed
2173
int LGBM_BoosterGetLeafValue(BoosterHandle handle,
2174
2175
2176
                             int tree_idx,
                             int leaf_idx,
                             double* out_val) {
Guolin Ke's avatar
Guolin Ke committed
2177
2178
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2179
  *out_val = static_cast<double>(ref_booster->GetLeafValue(tree_idx, leaf_idx));
Guolin Ke's avatar
Guolin Ke committed
2180
2181
2182
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
2183
int LGBM_BoosterSetLeafValue(BoosterHandle handle,
2184
2185
2186
                             int tree_idx,
                             int leaf_idx,
                             double val) {
Guolin Ke's avatar
Guolin Ke committed
2187
2188
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
2189
  ref_booster->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
2190
2191
2192
  API_END();
}

2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
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();
}

2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
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();
}

2224
2225
2226
2227
2228
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
2229
  Config config;
2230
  config.machines = RemoveQuotationSymbol(std::string(machines));
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
  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();
}

2246
2247
2248
int LGBM_NetworkInitWithFunctions(int num_machines, int rank,
                                  void* reduce_scatter_ext_fun,
                                  void* allgather_ext_fun) {
ww's avatar
ww committed
2249
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
2250
  if (num_machines > 1) {
2251
    Network::Init(num_machines, rank, (ReduceScatterFunction)reduce_scatter_ext_fun, (AllgatherFunction)allgather_ext_fun);
ww's avatar
ww committed
2252
2253
2254
  }
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
2255

Guolin Ke's avatar
Guolin Ke committed
2256
// ---- start of some help functions
2257
2258
2259

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
2260
  if (data_type == C_API_DTYPE_FLOAT32) {
2261
2262
    const float* data_ptr = reinterpret_cast<const float*>(data);
    if (is_row_major) {
2263
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2264
        std::vector<double> ret(num_col);
2265
        auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
2266
        for (int i = 0; i < num_col; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2267
          ret[i] = static_cast<double>(*(tmp_ptr + i));
2268
2269
2270
2271
        }
        return ret;
      };
    } else {
2272
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2273
        std::vector<double> ret(num_col);
2274
        for (int i = 0; i < num_col; ++i) {
2275
          ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
2276
2277
2278
2279
        }
        return ret;
      };
    }
Guolin Ke's avatar
Guolin Ke committed
2280
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2281
2282
    const double* data_ptr = reinterpret_cast<const double*>(data);
    if (is_row_major) {
2283
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2284
        std::vector<double> ret(num_col);
2285
        auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
2286
        for (int i = 0; i < num_col; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2287
          ret[i] = static_cast<double>(*(tmp_ptr + i));
2288
2289
2290
2291
        }
        return ret;
      };
    } else {
2292
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2293
        std::vector<double> ret(num_col);
2294
        for (int i = 0; i < num_col; ++i) {
2295
          ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
2296
2297
2298
2299
2300
        }
        return ret;
      };
    }
  }
2301
  Log::Fatal("Unknown data type in RowFunctionFromDenseMatric");
2302
  return nullptr;
2303
2304
2305
2306
}

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
2307
2308
  auto inner_function = RowFunctionFromDenseMatric(data, num_row, num_col, data_type, is_row_major);
  if (inner_function != nullptr) {
2309
    return [inner_function] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
2310
2311
      auto raw_values = inner_function(row_idx);
      std::vector<std::pair<int, double>> ret;
Guolin Ke's avatar
Guolin Ke committed
2312
      ret.reserve(raw_values.size());
Guolin Ke's avatar
Guolin Ke committed
2313
      for (int i = 0; i < static_cast<int>(raw_values.size()); ++i) {
Guolin Ke's avatar
Guolin Ke committed
2314
        if (std::fabs(raw_values[i]) > kZeroThreshold || std::isnan(raw_values[i])) {
Guolin Ke's avatar
Guolin Ke committed
2315
          ret.emplace_back(i, raw_values[i]);
2316
        }
Guolin Ke's avatar
Guolin Ke committed
2317
2318
2319
      }
      return ret;
    };
2320
  }
Guolin Ke's avatar
Guolin Ke committed
2321
  return nullptr;
2322
2323
}

2324
2325
2326
2327
2328
2329
2330
// 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
2331
    ret.reserve(raw_values.size());
2332
2333
2334
2335
2336
2337
2338
2339
2340
    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;
  };
}

2341
2342
template<typename T>
std::function<std::vector<std::pair<int, double>>(T idx)>
2343
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
2344
  if (data_type == C_API_DTYPE_FLOAT32) {
2345
    const float* data_ptr = reinterpret_cast<const float*>(data);
Guolin Ke's avatar
Guolin Ke committed
2346
    if (indptr_type == C_API_DTYPE_INT32) {
2347
      const int32_t* ptr_indptr = reinterpret_cast<const int32_t*>(indptr);
2348
      return [=] (T idx) {
2349
2350
2351
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2352
2353
2354
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2355
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2356
          ret.emplace_back(indices[i], data_ptr[i]);
2357
2358
2359
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
2360
    } else if (indptr_type == C_API_DTYPE_INT64) {
2361
      const int64_t* ptr_indptr = reinterpret_cast<const int64_t*>(indptr);
2362
      return [=] (T idx) {
2363
2364
2365
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2366
2367
2368
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2369
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2370
          ret.emplace_back(indices[i], data_ptr[i]);
2371
2372
2373
2374
        }
        return ret;
      };
    }
Guolin Ke's avatar
Guolin Ke committed
2375
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2376
    const double* data_ptr = reinterpret_cast<const double*>(data);
Guolin Ke's avatar
Guolin Ke committed
2377
    if (indptr_type == C_API_DTYPE_INT32) {
2378
      const int32_t* ptr_indptr = reinterpret_cast<const int32_t*>(indptr);
2379
      return [=] (T idx) {
2380
2381
2382
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2383
2384
2385
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2386
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2387
          ret.emplace_back(indices[i], data_ptr[i]);
2388
2389
2390
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
2391
    } else if (indptr_type == C_API_DTYPE_INT64) {
2392
      const int64_t* ptr_indptr = reinterpret_cast<const int64_t*>(indptr);
2393
      return [=] (T idx) {
2394
2395
2396
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
2397
2398
2399
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
2400
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
2401
          ret.emplace_back(indices[i], data_ptr[i]);
2402
2403
2404
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
2405
2406
    }
  }
2407
  Log::Fatal("Unknown data type in RowFunctionFromCSR");
2408
  return nullptr;
2409
2410
}

Guolin Ke's avatar
Guolin Ke committed
2411
std::function<std::pair<int, double>(int idx)>
2412
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
2413
  CHECK(col_idx < ncol_ptr && col_idx >= 0);
Guolin Ke's avatar
Guolin Ke committed
2414
  if (data_type == C_API_DTYPE_FLOAT32) {
2415
    const float* data_ptr = reinterpret_cast<const float*>(data);
Guolin Ke's avatar
Guolin Ke committed
2416
    if (col_ptr_type == C_API_DTYPE_INT32) {
2417
      const int32_t* ptr_col_ptr = reinterpret_cast<const int32_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2418
2419
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2420
2421
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2422
2423
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2424
        }
Guolin Ke's avatar
Guolin Ke committed
2425
2426
2427
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2428
      };
Guolin Ke's avatar
Guolin Ke committed
2429
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2430
      const int64_t* ptr_col_ptr = reinterpret_cast<const int64_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2431
2432
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2433
2434
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2435
2436
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2437
        }
Guolin Ke's avatar
Guolin Ke committed
2438
2439
2440
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2441
      };
Guolin Ke's avatar
Guolin Ke committed
2442
    }
Guolin Ke's avatar
Guolin Ke committed
2443
  } else if (data_type == C_API_DTYPE_FLOAT64) {
2444
    const double* data_ptr = reinterpret_cast<const double*>(data);
Guolin Ke's avatar
Guolin Ke committed
2445
    if (col_ptr_type == C_API_DTYPE_INT32) {
2446
      const int32_t* ptr_col_ptr = reinterpret_cast<const int32_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2447
2448
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2449
2450
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2451
2452
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2453
        }
Guolin Ke's avatar
Guolin Ke committed
2454
2455
2456
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2457
      };
Guolin Ke's avatar
Guolin Ke committed
2458
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
2459
      const int64_t* ptr_col_ptr = reinterpret_cast<const int64_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
2460
2461
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
2462
2463
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
2464
2465
        if (i >= end) {
          return std::make_pair(-1, 0.0);
2466
        }
Guolin Ke's avatar
Guolin Ke committed
2467
2468
2469
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
2470
      };
Guolin Ke's avatar
Guolin Ke committed
2471
2472
    }
  }
2473
  Log::Fatal("Unknown data type in CSC matrix");
2474
  return nullptr;
2475
2476
}

Guolin Ke's avatar
Guolin Ke committed
2477
CSC_RowIterator::CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
2478
                                 const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx) {
Guolin Ke's avatar
Guolin Ke committed
2479
2480
2481
2482
2483
2484
2485
2486
2487
  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;
2488
    }
Guolin Ke's avatar
Guolin Ke committed
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
    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;
2506
    }
Guolin Ke's avatar
Guolin Ke committed
2507
2508
2509
    return ret;
  } else {
    return std::make_pair(-1, 0.0);
2510
  }
Guolin Ke's avatar
Guolin Ke committed
2511
}