From 3792f191b25c4d3e0eb8fcd1a650b1809430de2d Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Tue, 12 Nov 2024 12:04:06 -0800
Subject: [PATCH] Changed Dataset to use column names for feature names with
 Rover and Daru

---
 CHANGELOG.md            | 2 +-
 lib/lightgbm/dataset.rb | 4 ++--
 test/dataset_test.rb    | 8 +-------
 3 files changed, 4 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6c99c23..d1c448c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,7 +1,7 @@
 ## 0.4.0 (unreleased)
 
 - Added support for hashes to `predict` method
-- Added support for `feature_names: "auto"` to `Dataset`
+- Changed `Dataset` to use column names for feature names with Rover and Daru
 - Dropped support for Ruby < 3.1
 
 ## 0.3.4 (2024-07-28)
diff --git a/lib/lightgbm/dataset.rb b/lib/lightgbm/dataset.rb
index 9426e19..e6ad269 100644
--- a/lib/lightgbm/dataset.rb
+++ b/lib/lightgbm/dataset.rb
@@ -11,7 +11,7 @@ def initialize(data, label: nil, weight: nil, group: nil, params: nil, reference
       @reference = reference
       @used_indices = used_indices
       @categorical_feature = categorical_feature
-      @feature_name = feature_name || feature_names
+      @feature_name = feature_name || feature_names || "auto"
 
       construct
     end
@@ -178,7 +178,7 @@ def construct
       self.label = @label if @label
       self.weight = @weight if @weight
       self.group = @group if @group
-      self.feature_name = @feature_name if @feature_name
+      self.feature_name = @feature_name if @feature_name && @feature_name != "auto"
     end
 
     def dump_text(filename)
diff --git a/test/dataset_test.rb b/test/dataset_test.rb
index 28d4be6..70a3c9f 100644
--- a/test/dataset_test.rb
+++ b/test/dataset_test.rb
@@ -61,10 +61,7 @@ def test_daru
     label = data["y"]
     data = data.delete_vector("y")
     dataset = LightGBM::Dataset.new(data, label: label)
-    assert_equal ["Column_0", "Column_1", "Column_2", "Column_3"], dataset.feature_name
-
-    dataset = LightGBM::Dataset.new(data, label: label, feature_name: "auto")
-    assert_equal ["x0", "x1", "x2", "x3"], dataset.feature_names
+    assert_equal ["x0", "x1", "x2", "x3"], dataset.feature_name
   end
 
   def test_numo
@@ -84,9 +81,6 @@ def test_rover
     data = Rover.read_csv(data_path)
     label = data.delete("y")
     dataset = LightGBM::Dataset.new(data, label: label)
-    assert_equal ["Column_0", "Column_1", "Column_2", "Column_3"], dataset.feature_name
-
-    dataset = LightGBM::Dataset.new(data, label: label, feature_name: "auto")
     assert_equal ["x0", "x1", "x2", "x3"], dataset.feature_name
   end