lib/lightgbm/booster.rb in lightgbm-0.1.4 vs lib/lightgbm/booster.rb in lightgbm-0.1.5

@@ -114,19 +114,26 @@
       out.read_int
     end
 
     # TODO support different prediction types
     def predict(input, num_iteration: nil, **params)
-      raise TypeError unless input.is_a?(Array)
+      input =
+        if daru?(input)
+          input.map_rows(&:to_a)
+        else
+          input.to_a
+        end
 
       singular = !input.first.is_a?(Array)
       input = [input] if singular
 
       num_iteration ||= best_iteration
       num_class ||= num_class()
 
+      flat_input = input.flatten
+      handle_missing(flat_input)
       data = ::FFI::MemoryPointer.new(:float, input.count * input.first.count)
-      data.put_array_of_float(0, input.flatten)
+      data.put_array_of_float(0, flat_input)
 
       out_len = ::FFI::MemoryPointer.new(:int64)
       out_result = ::FFI::MemoryPointer.new(:double, num_class * input.count)
       check_result FFI.LGBM_BoosterPredictForMat(handle_pointer, data, 0, input.count, input.first.count, 1, 0, num_iteration, params_str(params), out_len, out_result)
       out = out_result.read_array_of_double(out_len.read_int64)