ENH: Support MultiIndex columns in parquet pandas-dev#34777

1. Update check to handle MultiIndex columns for parquet format
2. Edit whatsnew entry.
3. Add test for writing MultiIndex columns with string column names

Author: hweecat

doc/source/whatsnew/v1.2.0.rst

@@ -297,7 +297,7 @@ I/O
 - :meth:`to_csv` did not support zip compression for binary file object not having a filename (:issue: `35058`)
 - :meth:`to_csv` and :meth:`read_csv` did not honor `compression` and `encoding` for path-like objects that are internally converted to file-like objects (:issue:`35677`, :issue:`26124`, and :issue:`32392`)
 - :meth:`to_picke` and :meth:`read_pickle` did not support compression for file-objects (:issue:`26237`, :issue:`29054`, and :issue:`29570`)
-- :meth:`to_parquet` did not support MultiIndex for columns in parquet format (:issue:`34777`)
+- :meth:`to_parquet` did not support :class:`MultiIndex` for columns in parquet format (:issue:`34777`)
 
 Plotting
 ^^^^^^^^

pandas/io/parquet.py

@@ -7,7 +7,7 @@
 from pandas.compat._optional import import_optional_dependency
 from pandas.errors import AbstractMethodError
 
-from pandas import DataFrame, get_option
+from pandas import DataFrame, MultiIndex, get_option
 
 from pandas.io.common import get_filepath_or_buffer, is_fsspec_url, stringify_path
 
@@ -54,7 +54,7 @@ def validate_dataframe(df: DataFrame):
             raise ValueError("to_parquet only supports IO with DataFrames")
 
         # must have value column names for all index levels (strings only)
-        if df.columns.nlevels > 1:
+        if isinstance(df.columns, MultiIndex):
             if not all(
                 x.inferred_type in {"string", "empty"} for x in df.columns.levels
             ):

pandas/tests/io/test_parquet.py

@@ -410,11 +410,24 @@ def test_write_multiindex(self, pa):
         check_round_trip(df, engine)
 
     def test_write_column_multiindex(self, engine):
-        # column multi-index
+        # Not able to write column multi-indexes with non-string column names.
         mi_columns = pd.MultiIndex.from_tuples([("a", 1), ("a", 2), ("b", 1)])
         df = pd.DataFrame(np.random.randn(4, 3), columns=mi_columns)
         self.check_error_on_write(df, engine, ValueError)
 
+    def test_write_column_multiindex_string(self, pa):
+        # Not supported in fastparquet as of 0.1.3 or older pyarrow version
+        engine = pa
+
+        # Write column multi-indexes with string column names
+        arrays = [
+            ["bar", "bar", "baz", "baz", "foo", "foo", "qux", "qux"],
+            ["one", "two", "one", "two", "one", "two", "one", "two"],
+        ]
+        df = pd.DataFrame(np.random.randn(8, 8), columns=arrays)
+
+        check_round_trip(df, engine)
+
     def test_multiindex_with_columns(self, pa):
         engine = pa
         dates = pd.date_range("01-Jan-2018", "01-Dec-2018", freq="MS")