lib/xgb/dmatrix.rb in xgb-0.1.1 vs lib/xgb/dmatrix.rb in xgb-0.1.2

@@ -25,34 +25,47 @@
         end
 
         c_data = ::FFI::MemoryPointer.new(:float, nrow * ncol)
         c_data.put_array_of_float(0, flat_data)
         check_result FFI.XGDMatrixCreateFromMat(c_data, nrow, ncol, missing, @handle)
+
+        ObjectSpace.define_finalizer(self, self.class.finalize(handle_pointer))
       end
 
       set_float_info("label", label) if label
       set_float_info("weight", weight) if weight
     end
 
+    def self.finalize(pointer)
+      # must use proc instead of stabby lambda
+      proc { FFI.XGDMatrixFree(pointer) }
+    end
+
     def label
       float_info("label")
     end
 
     def weight
       float_info("weight")
     end
 
+    def group=(group)
+      c_data = ::FFI::MemoryPointer.new(:int, group.size)
+      c_data.put_array_of_int(0, group)
+      check_result FFI.XGDMatrixSetGroup(handle_pointer, c_data, group.size)
+    end
+
     def num_row
-      out = ::FFI::MemoryPointer.new(:ulong)
+      out = ::FFI::MemoryPointer.new(:uint64)
       check_result FFI.XGDMatrixNumRow(handle_pointer, out)
-      out.read_ulong
+      out.read_uint64
     end
 
     def num_col
-      out = ::FFI::MemoryPointer.new(:ulong)
+      out = ::FFI::MemoryPointer.new(:uint64)
       check_result FFI.XGDMatrixNumCol(handle_pointer, out)
-      out.read_ulong
+      out.read_uint64
     end
 
     def slice(rindex)
       res = DMatrix.new(nil)
       idxset = ::FFI::MemoryPointer.new(:int, rindex.count)
@@ -74,19 +87,11 @@
     end
 
     private
 
     def set_float_info(field, data)
-      data =
-        if matrix?(data)
-          data.to_a[0]
-        elsif daru_vector?(data) || narray?(data)
-          data.to_a
-        else
-          data
-        end
-
+      data = data.to_a unless data.is_a?(Array)
       c_data = ::FFI::MemoryPointer.new(:float, data.size)
       c_data.put_array_of_float(0, data)
       check_result FFI.XGDMatrixSetFloatInfo(handle_pointer, field.to_s, c_data, data.size)
     end
 
@@ -102,13 +107,9 @@
       defined?(Matrix) && data.is_a?(Matrix)
     end
 
     def daru?(data)
       defined?(Daru::DataFrame) && data.is_a?(Daru::DataFrame)
-    end
-
-    def daru_vector?(data)
-      defined?(Daru::Vector) && data.is_a?(Daru::Vector)
     end
 
     def narray?(data)
       defined?(Numo::NArray) && data.is_a?(Numo::NArray)
     end