[DataFrame] Implement __array_wrap__ (ray-project#2218)

* Implement __array_wrap__ * Removing unnecessary test
rongou · Jun 11, 2018 · b886cec · b886cec
1 parent 61139e1
commit b886cec
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 10 deletions.
diff --git a/python/ray/dataframe/dataframe.py b/python/ray/dataframe/dataframe.py
@@ -4926,9 +4926,8 @@ def __array__(self, dtype=None):
         return to_pandas(self).__array__(dtype=dtype)
 
     def __array_wrap__(self, result, context=None):
-        raise NotImplementedError(
-            "To contribute to Pandas on Ray, please visit "
-            "github.com/ray-project/ray.")
+        # TODO: This is very inefficient, see also __array__ and as_matrix
+        return to_pandas(self).__array_wrap__(result, context=context)
 
     def __getstate__(self):
         raise NotImplementedError(

diff --git a/python/ray/dataframe/test/test_dataframe.py b/python/ray/dataframe/test/test_dataframe.py
@@ -3191,13 +3191,6 @@ def test___array__(ray_df, pandas_df):
     assert np.array_equal(ray_df.__array__(), pandas_df.__array__())
 
 
-def test___array_wrap__():
-    ray_df = create_test_dataframe()
-
-    with pytest.raises(NotImplementedError):
-        ray_df.__array_wrap__(None)
-
-
 def test___getstate__():
     ray_df = create_test_dataframe()