Added feature types for Rover data frames

ankane · ankane · commit f3bb6411cad6 · 2025-10-18T13:54:11.000-07:00
diff --git a/lib/xgboost/dmatrix.rb b/lib/xgboost/dmatrix.rb
@@ -34,6 +34,17 @@ def initialize(data, label: nil, weight: nil, missing: Float::NAN)
       elsif rover?(data)
         nrow, ncol = data.shape
         feature_names = data.keys
+        feature_types =
+          data.types.map do |_, v|
+            v = v.to_s
+            if v.start_with?("int") || v.start_with?("uint")
+              "int"
+            elsif v.start_with?("float")
+              "float"
+            else
+              raise Error, "Unknown feature type: #{v}"
+            end
+          end
         data = data.to_numo
       else
         nrow = data.count
diff --git a/test/dmatrix_test.rb b/test/dmatrix_test.rb
@@ -84,9 +84,8 @@ def test_rover
     dataset = XGBoost::DMatrix.new(data, label: label)
     names = ["x0", "x1", "x2", "x3"]
     assert_equal names, dataset.feature_names
-    # TODO add types
-    # types = ["float", "float", "float", "int"]
-    # assert_equal types, dataset.feature_types
+    types = ["float", "float", "float", "int"]
+    assert_equal types, dataset.feature_types
   end
 
   def test_invalid
