Skip to content

Commit

Permalink
Add a __dataframe__ method to _VaexDataFrame
Browse files Browse the repository at this point in the history
This will align the implementation with those in other libraries,
xref data-apis/dataframe-api#80
  • Loading branch information
rgommers authored and maartenbreddels committed Sep 23, 2022
1 parent 573d1e1 commit 8966c31
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 1 deletion.
5 changes: 4 additions & 1 deletion packages/vaex-core/vaex/dataframe_protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
For a background and spec, see:
* https://data-apis.org/blog/dataframe_protocol_rfc/
* https://data-apis.org/dataframe-protocol/latest/index.html
Notes
-----
- Interpreting a raw pointer (as in ``Buffer.ptr``) is annoying and unsafe to
Expand Down Expand Up @@ -701,6 +701,9 @@ def __init__(self, df: vaex.dataframe.DataFrame, nan_as_null: bool = False, allo
self._nan_as_null = nan_as_null
self._allow_copy = allow_copy

def __dataframe__(self, nan_as_null : bool = False, allow_copy : bool = True) -> "_VaexDataFrame":
return _VaexDataFrame(self._df, nan_as_null=nan_as_null, allow_copy=allow_copy)

@property
def metadata(self) -> Dict[str, Any]:
return {}
Expand Down
2 changes: 2 additions & 0 deletions tests/dataframe_protocol_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -294,6 +294,8 @@ def test_VaexDataFrame():
assert df2.select_columns((0, 2))._df[:, 0].tolist() == df2.select_columns_by_name(("x", "z"))._df[:, 0].tolist()
assert df2.select_columns((0, 2))._df[:, 1].tolist() == df2.select_columns_by_name(("x", "z"))._df[:, 1].tolist()

assert_dataframe_equal(df2.__dataframe__(), df)


def test_chunks(df_factory):
x = np.arange(10)
Expand Down

0 comments on commit 8966c31

Please sign in to comment.