[R-package] allow use of categorical_features in Dataset when raw data does not have column names (fixes #4374)

R-package/R/lgb.Dataset.R

@@ -156,31 +156,32 @@ Dataset <- R6::R6Class(
         # Check for character name
         if (is.character(private$categorical_feature)) {
 
-            cate_indices <- as.list(match(private$categorical_feature, private$colnames) - 1L)
+          cate_indices <- as.list(match(private$categorical_feature, private$colnames) - 1L)
+
+          # Provided indices, but some indices are missing?
+          if (sum(is.na(cate_indices)) > 0L) {
+            stop(
+              "lgb.self.get.handle: supplied an unknown feature in categorical_feature: "
+              , sQuote(private$categorical_feature[is.na(cate_indices)])
+            )
+          }
 
-            # Provided indices, but some indices are missing?
-            if (sum(is.na(cate_indices)) > 0L) {
-              stop(
-                "lgb.self.get.handle: supplied an unknown feature in categorical_feature: "
-                , sQuote(private$categorical_feature[is.na(cate_indices)])
-              )
-            }
+        } else {
 
-          } else {
+          # Check if more categorical features were output over the feature space
+          data_is_matrix <- is.matrix(private$raw_data) || methods::is(private$raw_data, "dgCMatrix")
+          if (data_is_matrix && max(private$categorical_feature) > ncol(private$raw_data)) {
+            stop(
+              "lgb.self.get.handle: supplied a too large value in categorical_feature: "
+              , max(private$categorical_feature)
+              , " but only "
+              , ncol(private$raw_data)
+              , " features"
+            )
+          }
 
-            # Check if more categorical features were output over the feature space
-            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"
-              )
-            }
-
-            # Store indices as [0, n-1] indexed instead of [1, n] indexed
-            cate_indices <- as.list(private$categorical_feature - 1L)
+          # Store indices as [0, n-1] indexed instead of [1, n] indexed
+          cate_indices <- as.list(private$categorical_feature - 1L)
 
           }
 

R-package/tests/testthat/test_dataset.R

@@ -548,3 +548,19 @@ test_that("lgb.Dataset$get_feature_num_bin() works", {
   actual_num_bins <- sapply(1L:5L, ds$get_feature_num_bin)
   expect_identical(actual_num_bins, expected_num_bins)
 })
+
+test_that("lgb.Dataset can be constructed with categorical features and without colnames", {
+  # check that dataset can be constructed
+  raw_mat <- data.matrix(rep(c(0L, 1L), 50L))
+  ds <- lgb.Dataset(raw_mat, categorical_feature = 1L)$construct()
+  sparse_mat <- Matrix::Matrix(raw_mat, sparse = TRUE)
+  expect_true(methods::is(sparse_mat, "dgCMatrix"))
+  ds2 <- lgb.Dataset(sparse_mat, categorical_feature = 1L)$construct()
+  # check that the column names are NULL
+  expect_null(ds$.__enclos_env__$private$colnames)
+  expect_null(ds2$.__enclos_env__$private$colnames)
+  # check for error when index is greater than the number of columns
+  expect_error({
+    lgb.Dataset(raw_mat, categorical_feature = 2L)$construct()
+  }, regexp = "supplied a too large value in categorical_feature: 2 but only 1 features")
+})