Repository: benchopt/benchopt
Pull request #637 (2023-08-03 20:46:47)

ENH add helper to store metadata in parquet files
### Checks before merging PR
- [x] ~added documentation for any new feature~
- [x] added unit test
- [x] edited the [what's new](../../whatsnew.rst) (if applicable)


codecov-commenter: ## [Codecov](https://app.codecov.io/gh/benchopt/benchopt/pull/637?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=benchopt) Report
> Merging [#637](https://app.codecov.io/gh/benchopt/benchopt/pull/637?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=benchopt) (e886911) into [main](https://app.codecov.io/gh/benchopt/benchopt/commit/4368a19b8e0720d5fc26bc897f378d56105ebfbf?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=benchopt) (4368a19) will **decrease** coverage by `0.15%`.
> The diff coverage is `35.00%`.

:exclamation: Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please [install the Github App Integration](https://github.com/apps/codecov) for your organization. [Read more](https://about.codecov.io/blog/codecov-is-updating-its-github-integration/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=benchopt).

```diff
@@            Coverage Diff             @@
##             main     #637      +/-   ##
==========================================
- Coverage   55.73%   55.59%   -0.15%     
==========================================
  Files          44       45       +1     
  Lines        2869     2889      +20     
  Branches      528      528              
==========================================
+ Hits         1599     1606       +7     
- Misses       1153     1166      +13     
  Partials      117      117              
```




Patch:
diff --git a/.gitignore b/.gitignore
index e4d5835..a017afa 100644
--- a/.gitignore
+++ b/.gitignore
@@ -24,7 +24,7 @@ benchopt_front_page_convergence.pdf
 /benchmarks/
 /results/
 /ci_template.github/
-
+/communication_materials/
 
 # IDE specific folders
 .vscode
diff --git a/benchopt/base.py b/benchopt/base.py
index 4cea888..172ed80 100644
--- a/benchopt/base.py
+++ b/benchopt/base.py
@@ -233,12 +233,15 @@ class BaseSolver(ParametrizedNameMixin, DependenciesMixin, ABC):
             self._output.progress('caching warmup times.')
 
         if self._solver_strategy == "callback":
+            stopping_criterion = (
+                SingleRunCriterion(stop_val=stop_val)
+                .get_runner_instance(solver=self)
+            )
             run_once_cb = _Callback(
                 lambda x: {'objective_value': 1},
-                {},
-                SingleRunCriterion(stop_val=stop_val).get_runner_instance(
-                    solver=self
-                )
+                solver=self,
+                meta={},
+                stopping_criterion=stopping_criterion
             )
             run_once_cb.start()
             self.run(run_once_cb)
diff --git a/benchopt/callback.py b/benchopt/callback.py
index 8e7f5f2..a4d365e 100644
--- a/benchopt/callback.py
+++ b/benchopt/callback.py
@@ -11,6 +11,9 @@ class _Callback:
     ----------
     objective : instance of BaseObjective
         The objective to minimize.
+    solver : instance of BaseSolver
+        The solver that is currently run. This allows to retrieve the results
+        with :method:`benchopt.base.Solver.get_result` in the callback.
     meta : dict
         Metadata passed to store in Cost results.
         Contains objective and data names, problem dimension, etc.
diff --git a/benchopt/utils/parquet.py b/benchopt/utils/parquet.py
new file mode 100644
index 0000000..f0f43fd
--- /dev/null
+++ b/benchopt/utils/parquet.py
@@ -0,0 +1,61 @@
+import json
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+JSON_KEY = b"_meta_json"
+
+
+def to_parquet(df, path, metadata=None):
+    """Write a ``pandas.DataFrame`` in a parquet file, with optional metadata.
+
+    Parameters
+    ----------
+    df: pd.DataFrame
+        DataFrame to write in the parquet file.
+    path: str | Path
+        Path to write the parquet file.
+    metadata: dict or None
+        Metadata to store in the parquet file. This metadata should be
+        serializable with json.
+    """
+    table = pa.Table.from_pandas(df)
+    new_metadata = {
+        JSON_KEY: json.dumps(metadata).encode("utf-8"),
+        **table.schema.metadata
+    }
+    table = table.replace_schema_metadata(new_metadata)
+    pq.write_table(table, path)
+
+
+def update_metadata(path, metadata):
+    """Update metadata in an existing parquet file.
+
+    Parameters
+    ----------
+    path: str | Path
+        Path of the parquet file to update.
+    metadata: dict or None
+        Metadata to store in the parquet file. This metadata should be
+        serializable with json.
+    """
+    table = pq.read_table(path)
+    new_metadata = {
+        **table.schema.metadata
+    }
+    metadata_ = json.loads(new_metadata.get(JSON_KEY, b"{}").decode("utf-8"))
+    metadata_.update(metadata)
+    new_metadata[JSON_KEY] = json.dumps(metadata_).encode("utf-8")
+    table = table.replace_schema_metadata(new_metadata)
+    pq.write_table(table, path)
+
+
+def get_metadata(path):
+    """Retrieve metadata embedded using the ``to_parquet`` function.
+
+    Parameters
+    ----------
+    path: str | Path
+        Path of the parquet file to read from.
+    """
+    meta = pq.read_metadata(path)
+    return json.loads(meta.metadata[JSON_KEY].decode("utf-8"))
diff --git a/doc/conf.py b/doc/conf.py
index 729ea70..3c22434 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -87,7 +87,7 @@ language = None
 exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
 
 # The name of the Pygments (syntax highlighting) style to use.
-pygments_style = 'sphinx'
+# pygments_style = 'sphinx'
 
 # If true, `todo` and `todoList` produce output, else they produce nothing.
 todo_include_todos = False
@@ -121,6 +121,9 @@ html_theme_options = {
     "use_fullscreen_button": False,
     "repository_url": "https://github.com/benchopt/benchopt",
     "home_page_in_toc": True,
+    # for a complete list of themes refer to https://pygments.org/styles/
+    "pygment_light_style": "colorful",
+    "pygment_dark_style": "gruvbox-dark",
 }
 
 html_sidebars = {
diff --git a/doc/whats_new.rst b/doc/whats_new.rst
index 491a69a..0823bc9 100644
--- a/doc/whats_new.rst
+++ b/doc/whats_new.rst
@@ -44,6 +44,13 @@ DOC
 - Reformatting and enriching the documentation for easy onboarding.
   By `Badr Moufad`_ and `Mathurin Massias`_ (:gh:`619`, :gh:`629`).
 
+Internals
+~~~~~~~~~
+
+- Add helper to store and retrieve metadata in parquet files. This will
+  allow storing per-run plotting information.
+  By `Thomas Moreau`_ (:gh`637`).
+
 
 .. _changes_1_4:
 


Test patch:
diff --git a/benchopt/tests/test_benchmark_features.py b/benchopt/tests/test_benchmark_features.py
index a695b55..a26f023 100644
--- a/benchopt/tests/test_benchmark_features.py
+++ b/benchopt/tests/test_benchmark_features.py
@@ -135,6 +135,71 @@ def test_ignore_hidden_files():
         ], 'benchopt', standalone_mode=False)
 
 
+@pytest.mark.parametrize("n_iter", [1, 2, 5])
+def test_run_once_iteration(n_iter):
+
+    solver1 = f"""from benchopt import BaseSolver
+    import numpy as np
+
+    class Solver(BaseSolver):
+        name = 'solver1'
+        sampling_strategy = 'iteration'
+
+        def set_objective(self, X, y, lmbd):
+            self.n_features = X.shape[1]
+            self.run_once({n_iter})
+
+        def run(self, n_iter): print(f"RUNONCE({{n_iter}})")
+
+        def get_result(self, **data):
+            return {{'beta': np.zeros(self.n_features)}}
+    """
+
+    with temp_benchmark(solvers=[solver1]) as benchmark:
+        with CaptureRunOutput() as out:
+            run([
+                str(benchmark.benchmark_dir),
+                *'-s solver1 -d test-dataset -n 0 -r 1 --no-plot'.split(),
+                *'-o dummy*[reg=0.5]'.split()
+            ], standalone_mode=False)
+        out.check_output(rf"RUNONCE\({n_iter}\)", repetition=1)
+
+
+@pytest.mark.parametrize("n_iter", [1, 2, 5])
+def test_run_once_callback(n_iter):
+
+    solver1 = f"""from benchopt import BaseSolver
+    import numpy as np
+
+    class Solver(BaseSolver):
+        name = 'solver1'
+        sampling_strategy = 'callback'
+
+        def set_objective(self, X, y, lmbd):
+            self.n_features = X.shape[1]
+            self.run_once({n_iter})
+
+        def run(self, cb):
+            i = 0
+            while cb():
+                i += 1
+            print(f"RUNONCE({{i}})")
+
+        def get_result(self, **data):
+            return {{'beta': np.zeros(self.n_features)}}
+    """
+
+    with temp_benchmark(solvers=[solver1]) as benchmark:
+        with CaptureRunOutput() as out:
+            run([
+                str(benchmark.benchmark_dir),
+                *'-s solver1 -d test-dataset -n 0 -r 1 --no-plot'.split(),
+                *'-o dummy*[reg=0.5]'.split()
+            ], standalone_mode=False)
+
+        out.check_output(rf"RUNONCE\({n_iter}\)", repetition=1)
+
+
 ##############################################################################
 # Test for deprecated features in 1.5
 ##############################################################################
diff --git a/benchopt/utils/tests/test_parquet.py b/benchopt/utils/tests/test_parquet.py
new file mode 100644
index 0000000..15b6428
--- /dev/null
+++ b/benchopt/utils/tests/test_parquet.py
@@ -0,0 +1,39 @@
+import json
+import tempfile
+import pandas as pd
+
+from benchopt.utils.parquet import to_parquet
+from benchopt.utils.parquet import get_metadata
+from benchopt.utils.parquet import update_metadata
+
+
+def test_parquet_metadata():
+    df = pd.DataFrame({
+        'a': range(4),
+        'b': list('abcd'),
+        'c': [.1, .2, .3, .4],
+    })
+
+    with tempfile.NamedTemporaryFile("rb+", suffix=".pq") as f:
+
+        path = f.name
+        metadata = {'test': 'info', 'plot_configs': [{}, {}, {}]}
+        to_parquet(df, path, metadata)
+
+        # Make sure the data is readable
+        df_new = pd.read_parquet(path)
+        assert (df_new == df).all().all()
+
+        # Check that the metadata can be retrieved
+        assert json.dumps(get_metadata(path)) == json.dumps(metadata)
+
+        meta_update = {'test': 'ofni', 'a': 'b'}
+        metadata.update(meta_update)
+        update_metadata(path, meta_update)
+
+        # Make sure the data did not change
+        df_new = pd.read_parquet(path)
+        assert df.equals(df_new)
+
+        # Check that the metadata has been changed correctly
+        assert json.dumps(get_metadata(path)) == json.dumps(metadata)