[DataFrame] Implement __array_wrap__ (#2218)

* Implement __array_wrap__

* Removing unnecessary test

python/ray/dataframe/dataframe.py

@@ -4926,9 +4926,8 @@ class DataFrame(object):
         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(

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()