fix #1153

microsoft · Dec 29, 2017 · 3b13f95 · 3b13f95
1 parent 12f55f7
commit 3b13f95
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/src/c_api.cpp b/src/c_api.cpp
@@ -555,23 +555,19 @@ int LGBM_DatasetCreateFromCSR(const void* indptr,
     int sample_cnt = static_cast<int>(nrow < config.io_config.bin_construct_sample_cnt ? nrow : config.io_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;
-    std::vector<std::vector<int>> sample_idx;
+    std::vector<std::vector<double>> sample_values(num_col);
+    std::vector<std::vector<int>> sample_idx(num_col);
     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) {
-        if (static_cast<size_t>(inner_data.first) >= sample_values.size()) {
-          sample_values.resize(inner_data.first + 1);
-          sample_idx.resize(inner_data.first + 1);
-        }
+        CHECK(static_cast<size_t>(inner_data.first) < num_col);
         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));
         }
       }
     }
-    CHECK(num_col >= static_cast<int>(sample_values.size()));
     DatasetLoader loader(config.io_config, nullptr, 1, nullptr);
     ret.reset(loader.CostructFromSampleData(Common::Vector2Ptr<double>(sample_values).data(),
                                             Common::Vector2Ptr<int>(sample_idx).data(),