fix bug in set_query

include/LightGBM/dataset.h

@@ -88,7 +88,7 @@ public:
 
   void SetWeights(const float* weights, data_size_t len);
 
-  void SetQueryBoundaries(const data_size_t* query_boundaries, data_size_t len);
+  void SetQuery(const data_size_t* query, data_size_t len);
 
   void SetQueryId(const data_size_t* query_id, data_size_t len);
 

src/io/dataset.cpp

@@ -91,7 +91,7 @@ bool Dataset::SetIntField(const char* field_name, const int* field_data, data_si
   std::string name(field_name);
   name = Common::Trim(name);
   if (name == std::string("query") || name == std::string("group")) {
-    metadata_.SetQueryBoundaries(field_data, num_element);
+    metadata_.SetQuery(field_data, num_element);
   } else if (name == std::string("query_id") || name == std::string("group_id")) {
     metadata_.SetQueryId(field_data, num_element);
   } else {

src/io/metadata.cpp

@@ -312,26 +312,27 @@ void Metadata::SetWeights(const float* weights, data_size_t len) {
   LoadQueryWeights();
 }
 
-void Metadata::SetQueryBoundaries(const data_size_t* query_boundaries, data_size_t len) {
+void Metadata::SetQuery(const data_size_t* query, data_size_t len) {
   std::lock_guard<std::mutex> lock(mutex_);
   // save to nullptr
-  if (query_boundaries == nullptr || len == 0) {
+  if (query == nullptr || len == 0) {
     query_boundaries_.clear();
     num_queries_ = 0;
     return;
   }
   data_size_t sum = 0;
   for (data_size_t i = 0; i < len; ++i) {
-    sum += query_boundaries[i];
+    sum += query[i];
   }
   if (num_data_ != sum) {
     Log::Fatal("sum of query counts is not same with #data");
   }
   if (!query_boundaries_.empty()) { query_boundaries_.clear(); }
   num_queries_ = len;
-  query_boundaries_ = std::vector<data_size_t>(num_queries_);
+  query_boundaries_ = std::vector<data_size_t>(num_queries_ + 1);
+  query_boundaries_[0] = 0;
   for (data_size_t i = 0; i < num_queries_; ++i) {
-    query_boundaries_[i] = query_boundaries[i];
+    query_boundaries_[i + 1] = query_boundaries_[i] + query[i];
   }
   LoadQueryWeights();
 }