Added to_numo method to DataFrame

CHANGELOG.md

@@ -1,6 +1,6 @@
 ## 0.3.1 (unreleased)
 
-- Added `to_numo` method to `Series`
+- Added `to_numo` method to `Series` and `DataFrame`
 
 ## 0.3.0 (2023-02-15)
 

ext/polars/src/dataframe.rs

@@ -6,6 +6,7 @@ use polars::io::mmap::ReaderBytes;
 use polars::io::RowCount;
 use polars::prelude::pivot::{pivot, pivot_stable};
 use polars::prelude::*;
+use polars_core::utils::try_get_supertype;
 use std::cell::RefCell;
 use std::io::{BufWriter, Cursor};
 use std::ops::Deref;
@@ -493,6 +494,25 @@ impl RbDataFrame {
         .into()
     }
 
+    pub fn to_numo(&self) -> Option<Value> {
+        let mut st = None;
+        for s in self.df.borrow().iter() {
+            let dt_i = s.dtype();
+            match st {
+                None => st = Some(dt_i.clone()),
+                Some(ref mut st) => {
+                    *st = try_get_supertype(st, dt_i).ok()?;
+                }
+            }
+        }
+        let st = st?;
+
+        match st {
+            // TODO
+            _ => None,
+        }
+    }
+
     pub fn write_parquet(
         &self,
         rb_f: Value,

ext/polars/src/lib.rs

@@ -88,6 +88,7 @@ fn init() -> RbResult<()> {
     class.define_method("write_ipc", method!(RbDataFrame::write_ipc, 2))?;
     class.define_method("row_tuple", method!(RbDataFrame::row_tuple, 1))?;
     class.define_method("row_tuples", method!(RbDataFrame::row_tuples, 0))?;
+    class.define_method("to_numo", method!(RbDataFrame::to_numo, 0))?;
     class.define_method("write_parquet", method!(RbDataFrame::write_parquet, 5))?;
     class.define_method("add", method!(RbDataFrame::add, 1))?;
     class.define_method("sub", method!(RbDataFrame::sub, 1))?;

lib/polars/data_frame.rb

@@ -699,8 +699,26 @@ def to_hashes
       end
     end
 
-    # def to_numo
-    # end
+    # Convert DataFrame to a 2D Numo array.
+    #
+    # This operation clones data.
+    #
+    # @return [Numo::NArray]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"foo" => [1, 2, 3], "bar" => [6, 7, 8], "ham" => ["a", "b", "c"]}
+    #   )
+    #   df.to_numo.class
+    #   # => Numo::RObject
+    def to_numo
+      out = _df.to_numo
+      if out.nil?
+        Numo::NArray.vstack(width.times.map { |i| to_series(i).to_numo }).transpose
+      else
+        out
+      end
+    end
 
     # no to_pandas
 

test/numo_test.rb

@@ -30,4 +30,10 @@ def test_series_str
     assert_kind_of Numo::RObject, s.to_numo
     assert_equal s.to_a, s.to_numo.to_a
   end
+
+  def test_data_frame
+    df = Polars::DataFrame.new({"a" => [1, 2, 3], "b" => ["one", "two", "three"]})
+    assert_kind_of Numo::RObject, df.to_numo
+    assert_equal [[1, "one"], [2, "two"], [3, "three"]], df.to_numo.to_a
+  end
 end