ENH: Read and write pandas attrs to parquet with pyarrow engine

Author: Alan D. Snow

doc/source/whatsnew/v1.3.0.rst

@@ -228,7 +228,7 @@ Other enhancements
 - Constructing a :class:`DataFrame` or :class:`Series` with the ``data`` argument being a Python iterable that is *not* a NumPy ``ndarray`` consisting of NumPy scalars will now result in a dtype with a precision the maximum of the NumPy scalars; this was already the case when ``data`` is a NumPy ``ndarray`` (:issue:`40908`)
 - Add keyword ``sort`` to :func:`pivot_table` to allow non-sorting of the result (:issue:`39143`)
 - Add keyword ``dropna`` to :meth:`DataFrame.value_counts` to allow counting rows that include ``NA`` values (:issue:`41325`)
--
+- Read and write :class:`DataFrame` and :class:`Series` attrs to parquet with pyarrow engine (:issue:`20521`)
 
 .. ---------------------------------------------------------------------------
 

pandas/io/parquet.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import io
+import json
 import os
 from typing import (
     Any,
@@ -154,6 +155,31 @@ def __init__(self):
 
         self.api = pyarrow
 
+    @staticmethod
+    def _write_attrs(table, df: DataFrame):
+        schema_metadata = dict(table.schema.metadata or {})
+        pandas_metadata = json.loads(schema_metadata.get(b"pandas", "{}"))
+        column_attrs = {}
+        for col in df.columns:
+            attrs = df[col].attrs
+            if not attrs or not isinstance(col, str):
+                continue
+            column_attrs[col] = attrs
+        pandas_metadata.update(
+            attrs=df.attrs,
+            column_attrs=column_attrs,
+        )
+        schema_metadata[b"pandas"] = json.dumps(pandas_metadata)
+        return table.replace_schema_metadata(schema_metadata)
+
+    @staticmethod
+    def _read_attrs(table, df: DataFrame):
+        pandas_metadata = json.loads(table.schema.metadata[b"pandas"])
+        df.attrs = pandas_metadata.get("attrs", {})
+        col_attrs = pandas_metadata.get("column_attrs", {})
+        for col, attrs in col_attrs.items():
+            df[col].attrs = attrs
+
     def write(
         self,
         df: DataFrame,
@@ -171,6 +197,7 @@ def write(
             from_pandas_kwargs["preserve_index"] = index
 
         table = self.api.Table.from_pandas(df, **from_pandas_kwargs)
+        table = self._write_attrs(table, df)
 
         path_or_handle, handles, kwargs["filesystem"] = _get_path_or_handle(
             path,
@@ -242,9 +269,11 @@ def read(
             mode="rb",
         )
         try:
-            result = self.api.parquet.read_table(
+            table = self.api.parquet.read_table(
                 path_or_handle, columns=columns, **kwargs
-            ).to_pandas(**to_pandas_kwargs)
+            )
+            result = table.to_pandas(**to_pandas_kwargs)
+            self._read_attrs(table, result)
             if manager == "array":
                 result = result._as_manager("array", copy=False)
             return result

pandas/tests/io/test_parquet.py

@@ -907,6 +907,26 @@ def test_read_parquet_manager(self, pa, using_array_manager):
         else:
             assert isinstance(result._mgr, pd.core.internals.BlockManager)
 
+    @td.skip_if_no("pyarrow")
+    def test_read_write_attrs(self, pa):
+        pdf = pd.DataFrame({"a": [1]})
+        pdf.attrs = {"name": "my custom dataset"}
+        pdf.a.attrs = {
+            "long_name": "Description about data",
+            "nodata": -1,
+            "units": "metre",
+        }
+        with tm.ensure_clean() as path:
+            pdf.to_parquet(path)
+            result = read_parquet(path)
+
+        assert result.attrs == {"name": "my custom dataset"}
+        assert result["a"].attrs == {
+            "long_name": "Description about data",
+            "nodata": -1,
+            "units": "metre",
+        }
+
 
 class TestParquetFastParquet(Base):
     def test_basic(self, fp, df_full):