lib/xgboost/dmatrix.rb in xgb-0.3.0 vs lib/xgboost/dmatrix.rb in xgb-0.3.1

@@ -1,8 +1,8 @@
 module XGBoost
   class DMatrix
-    attr_reader :data
+    attr_reader :data, :feature_names, :feature_types
 
     def initialize(data, label: nil, weight: nil, missing: Float::NAN)
       @data = data
 
       @handle = ::FFI::MemoryPointer.new(:pointer)
@@ -13,10 +13,22 @@
           ncol = data.column_count
           flat_data = data.to_a.flatten
         elsif daru?(data)
           nrow, ncol = data.shape
           flat_data = data.map_rows(&:to_a).flatten
+          @feature_names = data.each_vector.map(&:name)
+          @feature_types =
+            data.each_vector.map(&:db_type).map do |v|
+              case v
+              when "INTEGER"
+                "int"
+              when "DOUBLE"
+                "float"
+              else
+                raise Error, "Unknown feature type: #{v}"
+              end
+            end
         elsif narray?(data)
           nrow, ncol = data.shape
           flat_data = data.flatten.to_a
         else
           nrow = data.count
@@ -28,9 +40,11 @@
         c_data = ::FFI::MemoryPointer.new(:float, nrow * ncol)
         c_data.write_array_of_float(flat_data)
         check_result FFI.XGDMatrixCreateFromMat(c_data, nrow, ncol, missing, @handle)
 
         ObjectSpace.define_finalizer(self, self.class.finalize(handle_pointer))
+
+        @feature_names ||= ncol.times.map { |i| "f#{i}" }
       end
 
       self.label = label if label
       self.weight = weight if weight
     end