From 405b05d58a1347989dcc8a3380884b829ea7718a Mon Sep 17 00:00:00 2001 From: Peter Schafhalter Date: Tue, 10 Apr 2018 10:19:33 -0700 Subject: [PATCH] [DataFrame] Implemented __getattr__ (#1753) * __getattr__ accesses columns * Added test --- python/ray/dataframe/dataframe.py | 16 ++++++++++++++++ python/ray/dataframe/test/test_dataframe.py | 17 +++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/python/ray/dataframe/dataframe.py b/python/ray/dataframe/dataframe.py index cfb852f7bdeb..317d71de78a0 100644 --- a/python/ray/dataframe/dataframe.py +++ b/python/ray/dataframe/dataframe.py @@ -2992,6 +2992,22 @@ def _getitem_indiv_col(self, key, part): lambda df: df.__getitem__(index), self._col_partitions[part]) + def __getattr__(self, key): + """After regular attribute access, looks up the name in the columns + + Args: + key (str): Attribute name. + + Returns: + The value of the attribute. + """ + try: + return object.__getattribute__(self, key) + except AttributeError as e: + if key in self.columns: + return self[key] + raise e + def __setitem__(self, key, value): raise NotImplementedError( "To contribute to Pandas on Ray, please visit " diff --git a/python/ray/dataframe/test/test_dataframe.py b/python/ray/dataframe/test/test_dataframe.py index 2e8701f30701..46b42d28626b 100644 --- a/python/ray/dataframe/test/test_dataframe.py +++ b/python/ray/dataframe/test/test_dataframe.py @@ -2763,6 +2763,23 @@ def test___getitem__(ray_df, pd_df): assert pd_col.equals(ray_col) +def test___getattr__(): + df = create_test_dataframe() + + col = df.__getattr__("col1") + assert isinstance(col, pd.Series) + + col = getattr(df, "col1") + assert isinstance(col, pd.Series) + + col = df.col1 + assert isinstance(col, pd.Series) + + # Check that lookup in column doesn't override other attributes + df2 = df.rename(index=str, columns={"col5": "columns"}) + assert isinstance(df2.columns, pd.Index) + + def test___setitem__(): ray_df = create_test_dataframe()