lgb.Dataset.R 24 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
Dataset <- R6Class(
  "lgb.Dataset",
3
  cloneable = FALSE,
Guolin Ke's avatar
Guolin Ke committed
4
5
6
7
8
9
10
11
  public = list(
    finalize = function() {
      if (!lgb.is.null.handle(private$handle)) {
        lgb.call("LGBM_DatasetFree_R", ret = NULL, private$handle)
        private$handle <- NULL
      }
    },
    initialize = function(data,
12
13
14
                          params              = list(),
                          reference           = NULL,
                          colnames            = NULL,
15
                          categorical_feature = NULL,
16
17
18
19
                          predictor           = NULL,
                          free_raw_data       = TRUE,
                          used_indices        = NULL,
                          info                = list(),
Guolin Ke's avatar
Guolin Ke committed
20
                          ...) {
21
22
23
24
25
      additional_params <- list(...)
      INFO_KEYS <- c('label', 'weight', 'init_score', 'group')
      for (key in names(additional_params)) {
        if (key %in% INFO_KEYS) {
          info[[key]] <- additional_params[[key]]
Guolin Ke's avatar
Guolin Ke committed
26
        } else {
27
          params[[key]] <- additional_params[[key]]
Guolin Ke's avatar
Guolin Ke committed
28
29
30
31
        }
      }
      if (!is.null(reference)) {
        if (!lgb.check.r6.class(reference, "lgb.Dataset")) {
32
          stop("lgb.Dataset: Can only use ", sQuote("lgb.Dataset"), " as reference")
Guolin Ke's avatar
Guolin Ke committed
33
34
35
36
        }
      }
      if (!is.null(predictor)) {
        if (!lgb.check.r6.class(predictor, "lgb.Predictor")) {
37
          stop("lgb.Dataset: Only can use ", sQuote("lgb.Predictor"), " as predictor")
Guolin Ke's avatar
Guolin Ke committed
38
39
        }
      }
40
41
      private$raw_data  <- data
      private$params    <- params
Guolin Ke's avatar
Guolin Ke committed
42
      private$reference <- reference
43
44
      private$colnames  <- colnames

45
      private$categorical_feature <- categorical_feature
46
47
48
49
      private$predictor           <- predictor
      private$free_raw_data       <- free_raw_data
      private$used_indices        <- used_indices
      private$info                <- info
Guolin Ke's avatar
Guolin Ke committed
50
51
52
53
54
55
56
    },
    create_valid = function(data, info = list(),  ...) {
      ret <- Dataset$new(
        data,
        private$params,
        self,
        private$colnames,
57
        private$categorical_feature,
Guolin Ke's avatar
Guolin Ke committed
58
59
60
61
62
63
        private$predictor,
        private$free_raw_data,
        NULL,
        info,
        ...
      )
64
      ret
Guolin Ke's avatar
Guolin Ke committed
65
66
67
68
69
70
71
    },
    construct = function() {
      if (!lgb.is.null.handle(private$handle)) {
        return(self)
      }
      # Get feature names
      cnames <- NULL
72
      if (is.matrix(private$raw_data) || is(private$raw_data, "dgCMatrix")) {
Guolin Ke's avatar
Guolin Ke committed
73
74
75
        cnames <- colnames(private$raw_data)
      }
      # set feature names if not exist
76
      if (is.null(private$colnames) && !is.null(cnames)) {
Guolin Ke's avatar
Guolin Ke committed
77
78
        private$colnames <- as.character(cnames)
      }
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
      # Get categorical feature index
      if (!is.null(private$categorical_feature)) {
        if (typeof(private$categorical_feature) == "character") {
            cate_indices <- as.list(match(private$categorical_feature, private$colnames) - 1)
            if (sum(is.na(cate_indices)) > 0) {
              stop("lgb.self.get.handle: supplied an unknown feature in categorical_feature: ", sQuote(private$categorical_feature[is.na(cate_indices)]))
            }
          } else {
            if (max(private$categorical_feature) > length(private$colnames)) {
              stop("lgb.self.get.handle: supplied a too large value in categorical_feature: ", max(private$categorical_feature), " but only ", length(private$colnames), " features")
            }
            cate_indices <- as.list(private$categorical_feature - 1)
          }
        private$params$categorical_feature <- cate_indices
      }
Guolin Ke's avatar
Guolin Ke committed
94
95
      # Check has header or not
      has_header <- FALSE
96
      if (!is.null(private$params$has_header) ||
Guolin Ke's avatar
Guolin Ke committed
97
98
          !is.null(private$params$header)) {
        if (tolower(as.character(private$params$has_header)) == "true"
99
            ||
Guolin Ke's avatar
Guolin Ke committed
100
101
102
103
104
105
106
107
108
109
110
111
112
113
            tolower(as.character(private$params$header)) == "true") {
          has_header <- TRUE
        }
      }
      # Generate parameter str
      params_str <- lgb.params2str(private$params)
      # get handle of reference dataset
      ref_handle <- NULL
      if (!is.null(private$reference)) {
        ref_handle <- private$reference$.__enclos_env__$private$get_handle()
      }
      handle <- lgb.new.handle()
      # not subset
      if (is.null(private$used_indices)) {
114
115
        if (is.character(private$raw_data)) {
          handle <- lgb.call(
Guolin Ke's avatar
Guolin Ke committed
116
117
118
119
120
121
122
              "LGBM_DatasetCreateFromFile_R",
              ret = handle,
              lgb.c_str(private$raw_data),
              params_str,
              ref_handle
            )
        } else if (is.matrix(private$raw_data)) {
123
          handle <- lgb.call(
Guolin Ke's avatar
Guolin Ke committed
124
125
126
127
128
129
130
131
              "LGBM_DatasetCreateFromMat_R",
              ret = handle,
              private$raw_data,
              nrow(private$raw_data),
              ncol(private$raw_data),
              params_str,
              ref_handle
            )
132
        } else if (is(private$raw_data, "dgCMatrix")) {
Guolin Ke's avatar
Guolin Ke committed
133
134
135
136
137
138
139
140
141
142
143
144
145
          handle <- lgb.call(
            "LGBM_DatasetCreateFromCSC_R",
            ret = handle,
            private$raw_data@p,
            private$raw_data@i,
            private$raw_data@x,
            length(private$raw_data@p),
            length(private$raw_data@x),
            nrow(private$raw_data),
            params_str,
            ref_handle
          )
        } else {
146
147
148
          stop(
            "lgb.Dataset.construct: does not support constructing from ", sQuote(class(private$raw_data))
          )
Guolin Ke's avatar
Guolin Ke committed
149
150
151
152
        }
      } else {
        # construct subset
        if (is.null(private$reference)) {
153
          stop("lgb.Dataset.construct: reference cannot be NULL for constructing data subset")
Guolin Ke's avatar
Guolin Ke committed
154
        }
155
        handle <- lgb.call(
Guolin Ke's avatar
Guolin Ke committed
156
157
158
159
160
161
162
163
164
165
166
            "LGBM_DatasetGetSubset_R",
            ret = handle,
            ref_handle,
            private$used_indices,
            length(private$used_indices),
            params_str
          )
      }
      class(handle) <- "lgb.Dataset.handle"
      private$handle <- handle
      # set feature names
167
168
      if (!is.null(private$colnames)) { self$set_colnames(private$colnames) }

Guolin Ke's avatar
Guolin Ke committed
169
      # load init score
170
      if (!is.null(private$predictor) &&
Guolin Ke's avatar
Guolin Ke committed
171
          is.null(private$used_indices)) {
172
173
        init_score <- private$predictor$predict(private$raw_data, rawscore = TRUE, reshape = TRUE)
        # do not need to transpose, for is col_marjor
Guolin Ke's avatar
Guolin Ke committed
174
175
176
        init_score <- as.vector(init_score)
        private$info$init_score <- init_score
      }
177
      if (isTRUE(private$free_raw_data)) { private$raw_data <- NULL }
Guolin Ke's avatar
Guolin Ke committed
178
179
      if (length(private$info) > 0) {
        # set infos
180
        for (i in seq_along(private$info)) {
Guolin Ke's avatar
Guolin Ke committed
181
182
183
184
185
186
187
          p <- private$info[i]
          self$setinfo(names(p), p[[1]])
        }
      }
      if (is.null(self$getinfo("label"))) {
        stop("lgb.Dataset.construct: label should be set")
      }
188
      self
Guolin Ke's avatar
Guolin Ke committed
189
190
191
    },
    dim = function() {
      if (!lgb.is.null.handle(private$handle)) {
192
193
        num_row <- 0L
        num_col <- 0L
194
195
196

        c(
          lgb.call("LGBM_DatasetGetNumData_R",    ret = num_row, private$handle),
Guolin Ke's avatar
Guolin Ke committed
197
          lgb.call("LGBM_DatasetGetNumFeature_R", ret = num_col, private$handle)
198
199
200
        )
      } else if (is.matrix(private$raw_data) || is(private$raw_data, "dgCMatrix")) {
        dim(private$raw_data)
Guolin Ke's avatar
Guolin Ke committed
201
202
      } else {
        stop(
203
          "dim: cannot get dimensions before dataset has been constructed, please call lgb.Dataset.construct explicitly"
Guolin Ke's avatar
Guolin Ke committed
204
205
206
207
208
        )
      }
    },
    get_colnames = function() {
      if (!lgb.is.null.handle(private$handle)) {
209
210
211
212
213
        cnames <- lgb.call.return.str("LGBM_DatasetGetFeatureNames_R", private$handle)
        private$colnames <- as.character(base::strsplit(cnames, "\t")[[1]])
        private$colnames
      } else if (is.matrix(private$raw_data) || is(private$raw_data, "dgCMatrix")) {
        colnames(private$raw_data)
Guolin Ke's avatar
Guolin Ke committed
214
215
      } else {
        stop(
216
          "dim: cannot get dimensions before dataset has been constructed, please call lgb.Dataset.construct explicitly"
Guolin Ke's avatar
Guolin Ke committed
217
218
219
220
        )
      }
    },
    set_colnames = function(colnames) {
221
      if (is.null(colnames)) { return(self) }
Guolin Ke's avatar
Guolin Ke committed
222
      colnames <- as.character(colnames)
223
      if (length(colnames) == 0) { return(self) }
Guolin Ke's avatar
Guolin Ke committed
224
225
226
227
228
229
230
231
      private$colnames <- colnames
      if (!lgb.is.null.handle(private$handle)) {
        merged_name <- paste0(as.list(private$colnames), collapse = "\t")
        lgb.call("LGBM_DatasetSetFeatureNames_R",
                 ret = NULL,
                 private$handle,
                 lgb.c_str(merged_name))
      }
232
      self
Guolin Ke's avatar
Guolin Ke committed
233
234
    },
    getinfo = function(name) {
235
236
237
238
      INFONAMES <- c("label", "weight", "init_score", "group")
      if (!is.character(name) ||
          length(name) != 1   ||
          !name %in% INFONAMES) {
Guolin Ke's avatar
Guolin Ke committed
239
        stop(
240
          "getinfo: name must one of the following: ", paste0(sQuote(INFONAMES), collapse = ", ")
Guolin Ke's avatar
Guolin Ke committed
241
242
        )
      }
243
      if (is.null(private$info[[name]]) && !lgb.is.null.handle(private$handle)) {
244
        info_len <- 0L
245
246
247
248
        info_len <- lgb.call("LGBM_DatasetGetFieldSize_R",
                             ret = info_len,
                             private$handle,
                             lgb.c_str(name))
Guolin Ke's avatar
Guolin Ke committed
249
250
        if (info_len > 0) {
          ret <- NULL
251
252
253
254
255
          ret <- if (name == "group") { integer(info_len) } else { rep(0.0, info_len) }
          ret <- lgb.call("LGBM_DatasetGetField_R",
                          ret = ret,
                          private$handle,
                          lgb.c_str(name))
Guolin Ke's avatar
Guolin Ke committed
256
257
258
          private$info[[name]] <- ret
        }
      }
259
      private$info[[name]]
Guolin Ke's avatar
Guolin Ke committed
260
261
    },
    setinfo = function(name, info) {
262
263
264
265
      INFONAMES <- c("label", "weight", "init_score", "group")
      if (!is.character(name) ||
          length(name) != 1   ||
          !name %in% INFONAMES) {
Guolin Ke's avatar
Guolin Ke committed
266
        stop(
267
          "setinfo: name must one of the following: ", paste0(sQuote(INFONAMES), collapse = ", ")
Guolin Ke's avatar
Guolin Ke committed
268
269
        )
      }
270
      info <- if (name == "group") { as.integer(info) } else { as.numeric(info) }
Guolin Ke's avatar
Guolin Ke committed
271
      private$info[[name]] <- info
272
      if (!lgb.is.null.handle(private$handle) && !is.null(info)) {
Guolin Ke's avatar
Guolin Ke committed
273
274
275
276
277
278
279
280
281
282
283
        if (length(info) > 0) {
          lgb.call(
            "LGBM_DatasetSetField_R",
            ret = NULL,
            private$handle,
            lgb.c_str(name),
            info,
            length(info)
          )
        }
      }
284
      self
Guolin Ke's avatar
Guolin Ke committed
285
286
    },
    slice = function(idxset, ...) {
287
      Dataset$new(
Guolin Ke's avatar
Guolin Ke committed
288
289
290
291
        NULL,
        private$params,
        self,
        private$colnames,
292
        private$categorical_feature,
Guolin Ke's avatar
Guolin Ke committed
293
294
295
296
297
298
299
        private$predictor,
        private$free_raw_data,
        idxset,
        NULL,
        ...
      )
    },
300
    update_params = function(params) {
Guolin Ke's avatar
Guolin Ke committed
301
      private$params <- modifyList(private$params, params)
302
      self
Guolin Ke's avatar
Guolin Ke committed
303
    },
304
305
306
307
308
309
310
311
312
313
314
315
    set_categorical_feature = function(categorical_feature) {
      if (identical(private$categorical_feature, categorical_feature)) { return(self) }
      if (is.null(private$raw_data)) {
        stop(
          "set_categorical_feature: cannot set categorical feature after freeing raw data,
          please set ", sQuote("free_raw_data = FALSE"), " when you construct lgb.Dataset"
        )
      }
      private$categorical_feature <- categorical_feature
      self$finalize()
      self
    },
Guolin Ke's avatar
Guolin Ke committed
316
    set_reference = function(reference) {
317
      self$set_categorical_feature(reference$.__enclos_env__$private$categorical_feature)
Guolin Ke's avatar
Guolin Ke committed
318
319
      self$set_colnames(reference$get_colnames())
      private$set_predictor(reference$.__enclos_env__$private$predictor)
320
      if (identical(private$reference, reference)) { return(self) }
Guolin Ke's avatar
Guolin Ke committed
321
322
      if (is.null(private$raw_data)) {
        stop(
323
324
          "set_reference: cannot set reference after freeing raw data,
          please set ", sQuote("free_raw_data = FALSE"), " when you construct lgb.Dataset"
Guolin Ke's avatar
Guolin Ke committed
325
326
327
328
        )
      }
      if (!is.null(reference)) {
        if (!lgb.check.r6.class(reference, "lgb.Dataset")) {
329
          stop("set_reference: Can only use lgb.Dataset as a reference")
Guolin Ke's avatar
Guolin Ke committed
330
331
332
333
        }
      }
      private$reference <- reference
      self$finalize()
334
      self
Guolin Ke's avatar
Guolin Ke committed
335
336
337
338
339
340
341
    },
    save_binary = function(fname) {
      self$construct()
      lgb.call("LGBM_DatasetSaveBinary_R",
               ret = NULL,
               private$handle,
               lgb.c_str(fname))
342
      self
Guolin Ke's avatar
Guolin Ke committed
343
344
345
    }
  ),
  private = list(
346
347
348
349
350
    handle              = NULL,
    raw_data            = NULL,
    params              = list(),
    reference           = NULL,
    colnames            = NULL,
351
    categorical_feature = NULL,
352
353
354
355
356
357
358
    predictor           = NULL,
    free_raw_data       = TRUE,
    used_indices        = NULL,
    info                = NULL,
    get_handle          = function() {
      if (lgb.is.null.handle(private$handle)) { self$construct() }
      private$handle
Guolin Ke's avatar
Guolin Ke committed
359
360
    },
    set_predictor = function(predictor) {
361
      if (identical(private$predictor, predictor)) { return(self) }
Guolin Ke's avatar
Guolin Ke committed
362
363
364
      if (is.null(private$raw_data)) {
        stop(
          "set_predictor: cannot set predictor after free raw data,
365
          please set ", sQuote("free_raw_data = FALSE"), " when you construct lgb.Dataset"
Guolin Ke's avatar
Guolin Ke committed
366
367
368
369
        )
      }
      if (!is.null(predictor)) {
        if (!lgb.check.r6.class(predictor, "lgb.Predictor")) {
370
          stop("set_predictor: Can only use lgb.Predictor as predictor")
Guolin Ke's avatar
Guolin Ke committed
371
372
373
374
        }
      }
      private$predictor <- predictor
      self$finalize()
375
      self
Guolin Ke's avatar
Guolin Ke committed
376
377
378
379
    }
  )
)

wxchan's avatar
wxchan committed
380
#' Construct lgb.Dataset object
Guolin Ke's avatar
Guolin Ke committed
381
#'
wxchan's avatar
wxchan committed
382
#' Construct lgb.Dataset object from dense matrix, sparse matrix
Guolin Ke's avatar
Guolin Ke committed
383
384
385
386
387
388
#' or local file (that was created previously by saving an \code{lgb.Dataset}).
#'
#' @param data a \code{matrix} object, a \code{dgCMatrix} object or a character representing a filename
#' @param params a list of parameters
#' @param reference reference dataset
#' @param colnames names of columns
389
#' @param categorical_feature categorical features
Guolin Ke's avatar
Guolin Ke committed
390
391
392
393
394
#' @param free_raw_data TRUE for need to free raw data after construct
#' @param info a list of information of the lgb.Dataset object
#' @param ... other information to pass to \code{info} or parameters pass to \code{params}
#' @return constructed dataset
#' @examples
395
396
397
398
399
400
401
402
#' \dontrun{
#'   data(agaricus.train, package='lightgbm')
#'   train <- agaricus.train
#'   dtrain <- lgb.Dataset(train$data, label=train$label)
#'   lgb.Dataset.save(dtrain, 'lgb.Dataset.data')
#'   dtrain <- lgb.Dataset('lgb.Dataset.data')
#'   lgb.Dataset.construct(dtrain)
#' }
Guolin Ke's avatar
Guolin Ke committed
403
404
#' @export
lgb.Dataset <- function(data,
405
406
407
                        params              = list(),
                        reference           = NULL,
                        colnames            = NULL,
408
                        categorical_feature = NULL,
409
410
                        free_raw_data       = TRUE,
                        info                = list(),
Guolin Ke's avatar
Guolin Ke committed
411
412
413
414
415
416
                        ...) {
  Dataset$new(
    data,
    params,
    reference,
    colnames,
417
    categorical_feature,
Guolin Ke's avatar
Guolin Ke committed
418
419
420
421
422
423
424
425
426
    NULL,
    free_raw_data,
    NULL,
    info,
    ...
  )
}


wxchan's avatar
wxchan committed
427
#' Construct validation data
Guolin Ke's avatar
Guolin Ke committed
428
#'
wxchan's avatar
wxchan committed
429
#' Construct validation data according to training data
Guolin Ke's avatar
Guolin Ke committed
430
431
432
433
434
435
436
#'
#' @param dataset \code{lgb.Dataset} object, training data
#' @param data a \code{matrix} object, a \code{dgCMatrix} object or a character representing a filename
#' @param info a list of information of the lgb.Dataset object
#' @param ... other information to pass to \code{info}.
#' @return constructed dataset
#' @examples
437
438
439
440
441
442
443
444
#' \dontrun{
#'   data(agaricus.train, package='lightgbm')
#'   train <- agaricus.train
#'   dtrain <- lgb.Dataset(train$data, label=train$label)
#'   data(agaricus.test, package='lightgbm')
#'   test <- agaricus.test
#'   dtest <- lgb.Dataset.create.valid(dtrain, test$data, label=test$label)
#' }
Guolin Ke's avatar
Guolin Ke committed
445
#' @export
446
447
448
lgb.Dataset.create.valid <- function(dataset, data, info = list(),  ...) {
  if (!lgb.is.Dataset(dataset)) {
    stop("lgb.Dataset.create.valid: input data should be an lgb.Dataset object")
Guolin Ke's avatar
Guolin Ke committed
449
  }
450
451
  dataset$create_valid(data, info, ...)
}
Guolin Ke's avatar
Guolin Ke committed
452

453
#' Construct Dataset explicitly
Guolin Ke's avatar
Guolin Ke committed
454
455
456
#'
#' @param dataset Object of class \code{lgb.Dataset}
#' @examples
457
458
459
460
461
462
#' \dontrun{
#'   data(agaricus.train, package='lightgbm')
#'   train <- agaricus.train
#'   dtrain <- lgb.Dataset(train$data, label=train$label)
#'   lgb.Dataset.construct(dtrain)
#' }
Guolin Ke's avatar
Guolin Ke committed
463
464
#' @export
lgb.Dataset.construct <- function(dataset) {
465
466
  if (!lgb.is.Dataset(dataset)) {
    stop("lgb.Dataset.construct: input data should be an lgb.Dataset object")
Guolin Ke's avatar
Guolin Ke committed
467
  }
468
  dataset$construct()
Guolin Ke's avatar
Guolin Ke committed
469
470
}

471
#' Dimensions of an lgb.Dataset
Guolin Ke's avatar
Guolin Ke committed
472
473
474
475
476
477
478
479
480
481
482
#'
#' Returns a vector of numbers of rows and of columns in an \code{lgb.Dataset}.
#' @param x Object of class \code{lgb.Dataset}
#' @param ... other parameters
#' @return a vector of numbers of rows and of columns
#'
#' @details
#' Note: since \code{nrow} and \code{ncol} internally use \code{dim}, they can also
#' be directly used with an \code{lgb.Dataset} object.
#'
#' @examples
483
484
485
486
487
488
489
490
491
#' dontrun{
#'   data(agaricus.train, package='lightgbm')
#'   train <- agaricus.train
#'   dtrain <- lgb.Dataset(train$data, label=train$label)
#'
#'   stopifnot(nrow(dtrain) == nrow(train$data))
#'   stopifnot(ncol(dtrain) == ncol(train$data))
#'   stopifnot(all(dim(dtrain) == dim(train$data)))
#' }
Guolin Ke's avatar
Guolin Ke committed
492
493
494
#' @rdname dim
#' @export
dim.lgb.Dataset <- function(x, ...) {
495
496
  if (!lgb.is.Dataset(x)) {
    stop("dim.lgb.Dataset: input data should be an lgb.Dataset object")
Guolin Ke's avatar
Guolin Ke committed
497
  }
498
  x$dim()
Guolin Ke's avatar
Guolin Ke committed
499
500
501
502
503
}

#' Handling of column names of \code{lgb.Dataset}
#'
#' Only column names are supported for \code{lgb.Dataset}, thus setting of
504
#' row names would have no effect and returned row names would be NULL.
Guolin Ke's avatar
Guolin Ke committed
505
506
507
508
509
510
511
512
513
514
#'
#' @param x object of class \code{lgb.Dataset}
#' @param value a list of two elements: the first one is ignored
#'        and the second one is column names
#'
#' @details
#' Generic \code{dimnames} methods are used by \code{colnames}.
#' Since row names are irrelevant, it is recommended to use \code{colnames} directly.
#'
#' @examples
515
516
517
518
519
520
521
522
523
524
#' dontrun{
#'   data(agaricus.train, package='lightgbm')
#'   train <- agaricus.train
#'   dtrain <- lgb.Dataset(train$data, label=train$label)
#'   lgb.Dataset.construct(dtrain)
#'   dimnames(dtrain)
#'   colnames(dtrain)
#'   colnames(dtrain) <- make.names(1:ncol(train$data))
#'   print(dtrain, verbose=TRUE)
#' }
Guolin Ke's avatar
Guolin Ke committed
525
526
527
#' @rdname dimnames.lgb.Dataset
#' @export
dimnames.lgb.Dataset <- function(x) {
528
529
  if (!lgb.is.Dataset(x)) {
    stop("dimnames.lgb.Dataset: input data should be an lgb.Dataset object")
Guolin Ke's avatar
Guolin Ke committed
530
  }
531
  list(NULL, x$get_colnames())
Guolin Ke's avatar
Guolin Ke committed
532
533
534
535
536
537
}

#' @rdname dimnames.lgb.Dataset
#' @export
`dimnames<-.lgb.Dataset` <- function(x, value) {
  if (!is.list(value) || length(value) != 2L)
538
539
    stop("invalid ", sQuote("value"), " given: must be a list of two elements")
  if (!is.null(value[[1L]])) { stop("lgb.Dataset does not have rownames") }
Guolin Ke's avatar
Guolin Ke committed
540
541
542
543
544
545
  if (is.null(value[[2]])) {
    x$set_colnames(NULL)
    return(x)
  }
  if (ncol(x) != length(value[[2]]))
    stop("can't assign ",
546
547
548
         sQuote(length(value[[2]])),
         " colnames to an lgb.Dataset with ",
         sQuote(ncol(x)), " columns")
Guolin Ke's avatar
Guolin Ke committed
549
  x$set_colnames(value[[2]])
550
  x
Guolin Ke's avatar
Guolin Ke committed
551
552
}

553
#' Slice a dataset
Guolin Ke's avatar
Guolin Ke committed
554
#'
555
#' Get a new \code{lgb.Dataset} containing the specified rows of
Guolin Ke's avatar
Guolin Ke committed
556
557
558
559
560
561
562
563
#' orginal lgb.Dataset object
#'
#' @param dataset Object of class "lgb.Dataset"
#' @param idxset a integer vector of indices of rows needed
#' @param ... other parameters (currently not used)
#' @return constructed sub dataset
#'
#' @examples
564
565
566
567
#' \dontrun{
#'   data(agaricus.train, package='lightgbm')
#'   train <- agaricus.train
#'   dtrain <- lgb.Dataset(train$data, label=train$label)
Guolin Ke's avatar
Guolin Ke committed
568
#'
569
570
571
#'   dsub <- slice(dtrain, 1:42)
#'   labels1 <- getinfo(dsub, 'label')
#' }
Guolin Ke's avatar
Guolin Ke committed
572
#' @export
573
slice <- function(dataset, ...) { UseMethod("slice") }
Guolin Ke's avatar
Guolin Ke committed
574
575
576
577

#' @rdname slice
#' @export
slice.lgb.Dataset <- function(dataset, idxset, ...) {
578
579
  if (!lgb.is.Dataset(dataset)) {
    stop("slice.lgb.Dataset: input dataset should be an lgb.Dataset object")
Guolin Ke's avatar
Guolin Ke committed
580
  }
581
  dataset$slice(idxset, ...)
Guolin Ke's avatar
Guolin Ke committed
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
}


#' Get information of an lgb.Dataset object
#'
#' @param dataset Object of class \code{lgb.Dataset}
#' @param name the name of the information field to get (see details)
#' @param ... other parameters
#' @return info data
#'
#' @details
#' The \code{name} field can be one of the following:
#'
#' \itemize{
#'     \item \code{label}: label lightgbm learn from ;
#'     \item \code{weight}: to do a weight rescale ;
#'     \item \code{group}: group size
#'     \item \code{init_score}: initial score is the base prediction lightgbm will boost from ;
#' }
#'
#' @examples
603
604
605
606
607
608
609
610
611
612
613
#' \dontrun{
#'   data(agaricus.train, package='lightgbm')
#'   train <- agaricus.train
#'   dtrain <- lgb.Dataset(train$data, label=train$label)
#'   lgb.Dataset.construct(dtrain)
#'   labels <- getinfo(dtrain, 'label')
#'   setinfo(dtrain, 'label', 1-labels)
#'
#'   labels2 <- getinfo(dtrain, 'label')
#'   stopifnot(all(labels2 == 1-labels))
#' }
Guolin Ke's avatar
Guolin Ke committed
614
#' @export
615
getinfo <- function(dataset, ...) { UseMethod("getinfo") }
Guolin Ke's avatar
Guolin Ke committed
616
617
618
619

#' @rdname getinfo
#' @export
getinfo.lgb.Dataset <- function(dataset, name, ...) {
620
621
  if (!lgb.is.Dataset(dataset)) {
    stop("getinfo.lgb.Dataset: input dataset should be an lgb.Dataset object")
Guolin Ke's avatar
Guolin Ke committed
622
  }
623
  dataset$getinfo(name)
Guolin Ke's avatar
Guolin Ke committed
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
}

#' Set information of an lgb.Dataset object
#'
#' @param dataset Object of class "lgb.Dataset"
#' @param name the name of the field to get
#' @param info the specific field of information to set
#' @param ... other parameters
#' @return passed object
#'
#' @details
#' The \code{name} field can be one of the following:
#'
#' \itemize{
#'     \item \code{label}: label lightgbm learn from ;
#'     \item \code{weight}: to do a weight rescale ;
#'     \item \code{init_score}: initial score is the base prediction lightgbm will boost from ;
#'     \item \code{group}.
#' }
#'
#' @examples
645
646
647
648
649
650
651
652
653
654
#' \dontrun{
#'   data(agaricus.train, package='lightgbm')
#'   train <- agaricus.train
#'   dtrain <- lgb.Dataset(train$data, label=train$label)
#'   lgb.Dataset.construct(dtrain)
#'   labels <- getinfo(dtrain, 'label')
#'   setinfo(dtrain, 'label', 1-labels)
#'   labels2 <- getinfo(dtrain, 'label')
#'   stopifnot(all.equal(labels2, 1-labels))
#' }
Guolin Ke's avatar
Guolin Ke committed
655
#' @export
656
setinfo <- function(dataset, ...) { UseMethod("setinfo") }
Guolin Ke's avatar
Guolin Ke committed
657
658
659
660

#' @rdname setinfo
#' @export
setinfo.lgb.Dataset <- function(dataset, name, info, ...) {
661
662
  if (!lgb.is.Dataset(dataset)) {
    stop("setinfo.lgb.Dataset: input dataset should be an lgb.Dataset object")
Guolin Ke's avatar
Guolin Ke committed
663
  }
664
  dataset$setinfo(name, info)
Guolin Ke's avatar
Guolin Ke committed
665
666
}

667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
#' Set categorical feature of \code{lgb.Dataset}
#'
#' @param dataset object of class \code{lgb.Dataset}
#' @param categorical_feature categorical features
#' @return passed dataset
#' @examples
#' \dontrun{
#'   data(agaricus.train, package='lightgbm')
#'   train <- agaricus.train
#'   dtrain <- lgb.Dataset(train$data, label=train$label)
#'   lgb.Dataset.save(dtrain, 'lgb.Dataset.data')
#'   dtrain <- lgb.Dataset('lgb.Dataset.data')
#'   lgb.Dataset.set.categorical(dtrain, 1:2)
#' }
#' @rdname lgb.Dataset.set.categorical
#' @export
lgb.Dataset.set.categorical <- function(dataset, categorical_feature) {
  if (!lgb.is.Dataset(dataset)) {
    stop("lgb.Dataset.set.categorical: input dataset should be an lgb.Dataset object")
  }
  dataset$set_categorical_feature(categorical_feature)
}

690
#' Set reference of \code{lgb.Dataset}
Guolin Ke's avatar
Guolin Ke committed
691
#'
692
#' If you want to use validation data, you should set reference to training data
Guolin Ke's avatar
Guolin Ke committed
693
694
695
696
697
#'
#' @param dataset object of class \code{lgb.Dataset}
#' @param reference object of class \code{lgb.Dataset}
#' @return passed dataset
#' @examples
698
699
700
701
702
703
704
705
706
#' \dontrun{
#'   data(agaricus.train, package='lightgbm')
#'   train <- agaricus.train
#'   dtrain <- lgb.Dataset(train$data, label=train$label)
#'   data(agaricus.test, package='lightgbm')
#'   test <- agaricus.test
#'   dtest <- lgb.Dataset(test$data, test=train$label)
#'   lgb.Dataset.set.reference(dtest, dtrain)
#' }
Guolin Ke's avatar
Guolin Ke committed
707
708
709
#' @rdname lgb.Dataset.set.reference
#' @export
lgb.Dataset.set.reference <- function(dataset, reference) {
710
711
  if (!lgb.is.Dataset(dataset)) {
    stop("lgb.Dataset.set.reference: input dataset should be an lgb.Dataset object")
Guolin Ke's avatar
Guolin Ke committed
712
  }
713
  dataset$set_reference(reference)
Guolin Ke's avatar
Guolin Ke committed
714
715
}

716
717
#' Save \code{lgb.Dataset} to a binary file
#'
Guolin Ke's avatar
Guolin Ke committed
718
719
720
721
#' @param dataset object of class \code{lgb.Dataset}
#' @param fname object filename of output file
#' @return passed dataset
#' @examples
722
723
724
725
726
727
#' \dontrun{
#'   data(agaricus.train, package='lightgbm')
#'   train <- agaricus.train
#'   dtrain <- lgb.Dataset(train$data, label=train$label)
#'   lgb.Dataset.save(dtrain, "data.bin")
#' }
Guolin Ke's avatar
Guolin Ke committed
728
729
730
#' @rdname lgb.Dataset.save
#' @export
lgb.Dataset.save <- function(dataset, fname) {
731
732
  if (!lgb.is.Dataset(dataset)) {
    stop("lgb.Dataset.set: input dataset should be an lgb.Dataset object")
Guolin Ke's avatar
Guolin Ke committed
733
  }
734
735
  if (!is.character(fname)) {
    stop("lgb.Dataset.set: fname should be a character or a file connection")
Guolin Ke's avatar
Guolin Ke committed
736
  }
737
  dataset$save_binary(fname)
Guolin Ke's avatar
Guolin Ke committed
738
}