Reverting lgb.train(), using helper in lgb.Dataset

Author: walkerjameschris

R-package/R/lgb.Dataset.R

@@ -21,6 +21,17 @@ NULL
   return(c("label", "weight", "init_score", "group"))
 }
 
+# [description] A more robust version of identical() for checking lgb.Dataset equivalence
+# [return] A single boolean.
+.datasets_are_equal <- function(ds1, ds2) {
+  handler1 <- ds1$.__enclos_env__$private$get_handle
+  handler2 <- ds2$.__enclos_env__$private$get_handle
+  if (is.null(handler1) || is.null(handler2)) {
+    return(FALSE)
+  }
+  return(identical(handler1(), handler2()))
+}
+
 #' @importFrom methods is
 #' @importFrom R6 R6Class
 #' @importFrom utils modifyList
@@ -647,7 +658,12 @@ Dataset <- R6::R6Class(
     set_reference = function(reference) {
 
       # setting reference to this same Dataset object doesn't require any changes
-      if (identical(private$reference, reference)) {
+      if (.datasets_are_equal(private$reference, reference)) {
+        return(invisible(self))
+      }
+
+      # The new reference and the existing object may be slices of a parent dataset
+      if (.datasets_are_equal(private$reference, reference$.__enclos_env__$private$reference)) {
         return(invisible(self))
       }
 

R-package/R/lgb.train.R

@@ -184,14 +184,9 @@ lgb.train <- function(params = list(),
         next
       }
 
-      # Update parameters
+      # Update parameters, data
       valid_data$update_params(params)
-
-      # No need to set reference if one exists
-      if (is.null(valid_data$.__enclos_env__$private$reference)) {
-        valid_data$set_reference(data)
-      }
-
+      valid_data$set_reference(data)
       reduced_valid_sets[[key]] <- valid_data
 
     }