Skip to content

Commit

Permalink
make which optimizations to run more fine tunable
Browse files Browse the repository at this point in the history
  • Loading branch information
douglasdavis committed Aug 16, 2023
1 parent 93d38d1 commit 3029d6a
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 7 deletions.
17 changes: 14 additions & 3 deletions src/dask_awkward/awkward.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,22 @@ 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 blockwise 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
Expand Down
9 changes: 5 additions & 4 deletions src/dask_awkward/lib/optimize.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,10 +76,11 @@ def optimize(
"""
if dask.config.get("awkward.optimization.enabled", default=False):
dsk = optimize_columns(dsk) # type: ignore

# blockwise layer chaining optimization.
dsk = rewrite_layer_chains(dsk)
which = dask.config.get("awkward.optimization.which", default=[])
if "columns" in which:
dsk = optimize_columns(dsk) # type: ignore
if "layer-chains" in which:
dsk = rewrite_layer_chains(dsk)

return dsk

Expand Down

0 comments on commit 3029d6a

Please sign in to comment.