From 4cbd7549a30d7390abf529810d99aa5cb49836cb Mon Sep 17 00:00:00 2001
From: Doug Davis <ddavis@ddavis.io>
Date: Wed, 16 Aug 2023 16:01:22 -0500
Subject: [PATCH] feat: configuration can handle list of optimizations to run
 (#341)

---
 docs/ht-configuration.rst        |  8 +++++++-
 docs/me-optimization.rst         | 21 +++++++++++++++++++--
 src/dask_awkward/awkward.yaml    | 15 ++++++++++++---
 src/dask_awkward/lib/optimize.py | 11 ++++++-----
 4 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/docs/ht-configuration.rst b/docs/ht-configuration.rst
index c7d4d9a4..c39fac86 100644
--- a/docs/ht-configuration.rst
+++ b/docs/ht-configuration.rst
@@ -45,7 +45,13 @@ For example, they can be set with the form:
    with dask.config.set({"awkward.optimization.<option>": False}):
        ...
 
-- ``enabled`` (default: ``True``): Enable dask-awkward specific optimizations.
+- ``enabled`` (default: ``True``): Enable dask-awkward specific
+  optimizations. More fine tuning can be handled with the ``which``
+  option.
+- ``which`` (default: ``[columns, layer-chains]``): Which of the
+  optimizations to run. The default setting is to run all available
+  optimizations. (if ``enabled`` is set to ``False`` this option is
+  ignored).
 - ``on-fail`` (default: ``warn``): When set to ``warn`` throw a
   warning of the optimization fails and continue without performing
   the optimization. If set to ``raise``, raise an exception at
diff --git a/docs/me-optimization.rst b/docs/me-optimization.rst
index 0d2117d8..cef3fb61 100644
--- a/docs/me-optimization.rst
+++ b/docs/me-optimization.rst
@@ -11,13 +11,30 @@ we benefit from downstream in dask-awkward. You can read more about
 Dask optimization in general :doc:`at this section of the Dask docs
 <dask:optimize>`.
 
+dask-awkward Optimizations
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+There are two optimizations implemented in the dask-awkward code. One
+is the ``layer-chains`` optimization that fuses adjacent task graph
+layers together (if they are compatible with each other). This is a
+relatively simple optimization that just simplifies the task graph.
+The other optimization is the ``columns`` (or "necessary columns")
+optimization; which is a bit more technical and described in a
+follow-up section.
+
+One can configure which optimizations to run at compute-time; read
+more optimization. More information can be found in the
+:ref:`configuration section
+<ht-configuration:Optimization specific table>` of the docs.
+
+
 Necessary Columns
 ^^^^^^^^^^^^^^^^^
 
 We have one dask-awkward specific optimization that targets efficient
 data access from disk. We call it the "necessary columns"
 optimization. This optimization will execute the task graph *without
-operating on real data*. The data-less executation of the graph helps
+operating on real data*. The data-less execution of the graph helps
 determine which parts of a dataset sitting on disk are actually
 required to read in order to successfully complete the compute.
 
@@ -103,7 +120,7 @@ parameter:
   at compute time.
 - ``"warn"`` (the default): fail with a warning but let the compute
   continue without the necessary columns optimization (can reduce
-  performance by reading unncessary data from disk).
+  performance by reading unnecessary data from disk).
 
 One can also use the ``columns=`` argument (with
 :func:`~dask_awkward.from_parquet`, for example) to manually define
diff --git a/src/dask_awkward/awkward.yaml b/src/dask_awkward/awkward.yaml
index 27ee3123..31f4edca 100644
--- a/src/dask_awkward/awkward.yaml
+++ b/src/dask_awkward/awkward.yaml
@@ -15,11 +15,20 @@ awkward:
   # Optimization specific configuration
   optimization:
 
-    # If true dask-awkward specific optimizations will be run. This is
-    # currently limited to determining necessary columns and applying
-    # column projection.
+    # If true dask-awkward specific optimizations will be run.
     enabled: True
 
+    # Which of the optimizations do we want to run; options include:
+    # - "columns":
+    #   Run the optimization step that determines which columns are
+    #   necessary for a computation and only read those columns
+    #   (either via Parquet, ROOT, or JSONSchema) from the data that
+    #   is on disk.
+    # - "layer-chains":
+    #   Fuse adjacent blockwise layers ("layer chains") into a single
+    #   layer in the task graph.
+    which: [columns, layer-chains]
+
     # This option controls whether or not a warning is thrown, an
     # exception is raised, or if nothing is done if a dask-awkward
     # specific optimization fails (right now this is only the column
diff --git a/src/dask_awkward/lib/optimize.py b/src/dask_awkward/lib/optimize.py
index f2dbe201..0842b63d 100644
--- a/src/dask_awkward/lib/optimize.py
+++ b/src/dask_awkward/lib/optimize.py
@@ -75,11 +75,12 @@ def optimize(
     input layers.
 
     """
-    if dask.config.get("awkward.optimization.enabled", default=False):
-        dsk = optimize_columns(dsk)  # type: ignore
-
-        # blockwise layer chaining optimization.
-        dsk = rewrite_layer_chains(dsk)
+    if dask.config.get("awkward.optimization.enabled"):
+        which = dask.config.get("awkward.optimization.which")
+        if "columns" in which:
+            dsk = optimize_columns(dsk)  # type: ignore
+        if "layer-chains" in which:
+            dsk = rewrite_layer_chains(dsk)
 
     return dsk