From b886ceca478cdea1b759e247bc7f925e81bfa9b7 Mon Sep 17 00:00:00 2001 From: Devin Petersohn Date: Mon, 11 Jun 2018 08:56:43 -0700 Subject: [PATCH] [DataFrame] Implement __array_wrap__ (#2218) * Implement __array_wrap__ * Removing unnecessary test --- python/ray/dataframe/dataframe.py | 5 ++--- python/ray/dataframe/test/test_dataframe.py | 7 ------- 2 files changed, 2 insertions(+), 10 deletions(-) diff --git a/python/ray/dataframe/dataframe.py b/python/ray/dataframe/dataframe.py index 10260431a..c4ef03599 100644 --- a/python/ray/dataframe/dataframe.py +++ b/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( diff --git a/python/ray/dataframe/test/test_dataframe.py b/python/ray/dataframe/test/test_dataframe.py index 81cf3b420..39e2b0585 100644 --- 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()