Add metadata attribute to DataFrame and Column

steff456 · rgommers · commit d9419b17c3fb · 2021-08-23T22:41:33.000+02:00
This PR adds a metadata attribute that can be used to store
library-specific things. For example, libraries like Vaex will be able
to store expressions for its virtual columns there.

- Add metadata attribute to the DataFrame class
- Add metadata attribute to the Column class
- Add a test
diff --git a/protocol/dataframe_protocol.py b/protocol/dataframe_protocol.py
@@ -288,6 +288,13 @@ def null_count(self) -> Optional[int]:
         """
         pass
 
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        """
+        The metadata for the column. See `DataFrame.metadata` for more details.
+        """
+        pass
+
     def num_chunks(self) -> int:
         """
         Return the number of chunks the column consists of.
@@ -362,6 +369,19 @@ def __dataframe__(self, nan_as_null : bool = False) -> dict:
             "version": 0        # Version number of the protocol
         }
 
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        """
+        The metadata for the data frame, as a dictionary with string keys. The
+        contents of `metadata` may be anything, they are meant for a library
+        to store information that it needs to, e.g., roundtrip losslessly or
+        for two implementations to share data that is not (yet) part of the
+        interchange protocol specification. For avoiding collisions with other
+        entries, please add name the keys with the name of the library
+        followed by a period and the desired name, e.g, ``pandas.indexcol``.
+        """
+        pass
+
     def num_columns(self) -> int:
         """
         Return the number of columns in the DataFrame.
@@ -429,4 +449,3 @@ def get_chunks(self, n_chunks : Optional[int] = None) -> Iterable[DataFrame]:
         before yielding it.
         """
         pass
-
diff --git a/protocol/pandas_implementation.py b/protocol/pandas_implementation.py
@@ -196,7 +196,7 @@ def convert_string_column(col : ColumnObject) -> np.ndarray:
             v = mbuf[i/8]
             if null_value == 1:
                 v = ~v
-            
+
             if v & (1<<(i%8)):
                 str_list.append(np.nan)
                 continue
@@ -498,6 +498,13 @@ def null_count(self) -> int:
         """
         return self._col.isna().sum()
 
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        """
+        Store specific metadata of the column.
+        """
+        return {}
+
     def num_chunks(self) -> int:
         """
         Return the number of chunks the column consists of.
@@ -682,6 +689,12 @@ def __init__(self, df : pd.DataFrame, nan_as_null : bool = False) -> None:
         # dtypes is added, this value should be propagated to columns.
         self._nan_as_null = nan_as_null
 
+    @property
+    def metadata(self):
+        # `index` isn't a regular column, and the protocol doesn't support row
+        # labels - so we export it as Pandas-specific metadata here.
+        return {"pandas.index": self._df.index}
+
     def num_columns(self) -> int:
         return len(self._df.columns)
 
@@ -777,6 +790,21 @@ def test_string_dtype():
     assert col.describe_null == (4, 0)
     assert col.num_chunks() == 1
 
+def test_metadata():
+    df = pd.DataFrame({'A': [1, 2, 3, 4],'B': [1, 2, 3, 4]})
+
+    # Check the metadata from the dataframe
+    df_metadata = df.__dataframe__().metadata
+    expected = {"pandas.index": df.index}
+    for key in df_metadata:
+        assert all(df_metadata[key] == expected[key])
+
+    # Check the metadata from the column
+    col_metadata = df.__dataframe__().get_column(0).metadata
+    expected = {}
+    for key in col_metadata:
+        assert col_metadata[key] == expected[key]
+
     df2 = from_dataframe(df)
     tm.assert_frame_equal(df, df2)
 
@@ -787,4 +815,4 @@ def test_string_dtype():
     test_mixed_intfloat()
     test_noncontiguous_columns()
     test_string_dtype()
-
+    test_metadata()
