Repository: dask/dask
Pull request: Support custom_metadata= argument in to_parquet
Adds a `custom_metadata=` argument to `DataFrame.to_parquet`.  If a dictionary is passed in by the user, those key/values will be included in all footer metadata, and in the global _metadata file, along with the usual `b"pandas"` metadata.

~Note that this PR only adds `custom_metadata=` support for the "pyarrow" engine. @martindurant - Do you know if something similar can be accomplished in the "fastparquet" engine?  I have not investigated that end just yet.~ [EDIT: Both pyarrow and fastparquet engines are now supported.]

- [x] Closes #7343
- [x] Tests added / passed
- [x] Passes `black dask` / `flake8 dask`


Problem statement:
Define Column Metadata with to_parquet
When writing out to the `parquet` format, I want to define some column metadata.  Using pyarrow, I can write the metadata to the table using the following:

When writing out to parquet using pandas, I can create the Table using the pyarrow library, and once written, I can update the tables' metdata.  
```python
    metadata = table.schema.metadata
    metadata.update({b"geo": _encode_metadata(geo_metadata)})
```

With dask, can I somehow pass this information in above to the `to_parquet` class?

I see metadata files, but I want to define it on the table itself.

Hints from issue discussion:
cc @rjzamora 
This sounds reasonable to me.  Thanks for the suggestion @achapkowski 

Just to clarify, would an appropriate solution be for dask to support a `custom_metadata=` argument, where the user can specify their own dictionary of metadata that will be added to all footer metadata in the output dataset (including the global _metadata file)?  If the metadata is the same for all output partitions, this shouldn't be too difficult to implement.
A custom_metadata would be great, where you would do something like:

```python
to_parquet(...., custom_meta={'key' : 'values'})
```

Commit message: Merge remote-tracking branch 'upstream/main' into custom-metadata

Patch to review:
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
index d8d5aa4c3a4..9729f587688 100644
--- a/dask/dataframe/io/parquet/arrow.py
+++ b/dask/dataframe/io/parquet/arrow.py
@@ -828,6 +828,7 @@ def write_partition(
         index_cols=None,
         schema=None,
         head=False,
+        custom_metadata=None,
         **kwargs,
     ):
         _meta = None
@@ -839,6 +840,10 @@ def write_partition(
             index_cols = []
 
         t = cls._pandas_to_arrow_table(df, preserve_index=preserve_index, schema=schema)
+        if custom_metadata:
+            _md = t.schema.metadata
+            _md.update(custom_metadata)
+            t = t.replace_schema_metadata(metadata=_md)
 
         if partition_on:
             md_list = _write_partitioned(
diff --git a/dask/dataframe/io/parquet/core.py b/dask/dataframe/io/parquet/core.py
index f98075962c8..9c37edc71d7 100644
--- a/dask/dataframe/io/parquet/core.py
+++ b/dask/dataframe/io/parquet/core.py
@@ -411,6 +411,7 @@ def to_parquet(
     ignore_divisions=False,
     partition_on=None,
     storage_options=None,
+    custom_metadata=None,
     write_metadata_file=True,
     compute=True,
     compute_kwargs=None,
@@ -459,6 +460,10 @@ def to_parquet(
         there will be no global groupby.
     storage_options : dict, optional
         Key/value pairs to be passed on to the file-system backend, if any.
+    custom_metadata : dict, optional
+        Custom key/value metadata to include in all footer metadata (and
+        in the global "_metadata" file, if applicable).  Note that the custom
+        metadata may not contain the reserved b"pandas" key.
     write_metadata_file : bool, optional
         Whether to write the special "_metadata" file.
     compute : bool, optional
@@ -629,6 +634,15 @@ def to_parquet(
     kwargs_pass["compression"] = compression
     kwargs_pass["index_cols"] = index_cols
     kwargs_pass["schema"] = schema
+    if custom_metadata:
+        if b"pandas" in custom_metadata.keys():
+            raise ValueError(
+                "User-defined key/value metadata (custom_metadata) can not "
+                "contain a b'pandas' key.  This key is reserved by Pandas, "
+                "and overwriting the corresponding value can render the "
+                "entire dataset unreadable."
+            )
+        kwargs_pass["custom_metadata"] = custom_metadata
     for d, filename in enumerate(filenames):
         dsk[(name, d)] = (
             apply,
diff --git a/dask/dataframe/io/parquet/fastparquet.py b/dask/dataframe/io/parquet/fastparquet.py
index 594d945ce9f..0a8f271a1dd 100644
--- a/dask/dataframe/io/parquet/fastparquet.py
+++ b/dask/dataframe/io/parquet/fastparquet.py
@@ -898,9 +898,19 @@ def write_partition(
         return_metadata,
         fmd=None,
         compression=None,
+        custom_metadata=None,
         **kwargs,
     ):
+        # Update key/value metadata if necessary
         fmd = copy.copy(fmd)
+        if custom_metadata and fmd is not None:
+            fmd.key_value_metadata.extend(
+                [
+                    fastparquet.parquet_thrift.KeyValue(key=key, value=value)
+                    for key, value in custom_metadata.items()
+                ]
+            )
+
         if not len(df):
             # Write nothing for empty partitions
             rgs = []
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
index b12375600f9..0fe05596430 100644
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -3345,3 +3345,44 @@ def test_roundtrip_rename_columns(tmpdir, engine):
     df1.columns = ["d", "e", "f"]
 
     assert_eq(df1, ddf2.compute())
+
+
+def test_pyarrow_custom_metadata(tmpdir, engine):
+    # Write a parquet dataset with custom metadata
+
+    # Define custom metadata
+    custom_metadata = {b"my_key": b"my_data"}
+
+    # Write parquet dataset
+    path = str(tmpdir)
+    df = pd.DataFrame({"a": range(10), "b": range(10)})
+    dd.from_pandas(df, npartitions=2).to_parquet(
+        path,
+        engine=engine,
+        custom_metadata=custom_metadata,
+    )
+
+    # Check that data is correct
+    assert_eq(df, dd.read_parquet(path, engine=engine))
+
+    # Require pyarrow.parquet to check key/value metadata
+    if pq:
+        # Read footer metadata and _metadata.
+        # Check that it contains keys/values from `custom_metadata`
+        files = glob.glob(os.path.join(path, "*.parquet"))
+        files += [os.path.join(path, "_metadata")]
+        for fn in files:
+            _md = pq.ParquetFile(fn).metadata.metadata
+            for k, v in custom_metadata.items():
+                assert _md[k] == custom_metadata[k]
+
+    # Make sure we raise an error if the custom metadata
+    # includes a b"pandas" key
+    custom_metadata = {b"pandas": b"my_new_pandas_md"}
+    with pytest.raises(ValueError) as e:
+        dd.from_pandas(df, npartitions=2).to_parquet(
+            path,
+            engine=engine,
+            custom_metadata=custom_metadata,
+        )
+    assert "User-defined key/value" in str(e.value)


Code review comments:
diff_hunk: @@ -3345,3 +3345,44 @@ def test_roundtrip_rename_columns(tmpdir, engine):
     df1.columns = ["d", "e", "f"]
 
     assert_eq(df1, ddf2.compute())
+
+
+def test_pyarrow_custom_metadata(tmpdir, engine):
path: dask/dataframe/io/tests/test_parquet.py
text: @user1:
```suggestion
def test_custom_metadata(tmpdir, engine):
```

@author:
Forgot about that - Thanks!