From 3dd7de92dc468c3b35293f750c73b560aa6067fb Mon Sep 17 00:00:00 2001 From: Angus Hollands Date: Fri, 25 Aug 2023 15:48:16 +0100 Subject: [PATCH] refactor: use `ak.typetracer.touch_data` --- src/dask_awkward/lib/optimize.py | 5 ++--- src/dask_awkward/lib/structure.py | 7 ++----- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/src/dask_awkward/lib/optimize.py b/src/dask_awkward/lib/optimize.py index 0842b63d..70435c0b 100644 --- a/src/dask_awkward/lib/optimize.py +++ b/src/dask_awkward/lib/optimize.py @@ -192,8 +192,7 @@ def _touch_all_data(*args, **kwargs): import awkward as ak for arg in args + tuple(kwargs.values()): - if isinstance(arg, ak.Array): - arg.layout._touch_data(recursive=True) + ak.typetracer.touch_data(arg) def _mock_output(layer): @@ -393,7 +392,7 @@ def _get_column_reports(dsk: HighLevelGraph) -> dict[str, Any]: results = get_sync(hlg, leaf_layers_keys) for out in results: if isinstance(out, (ak.Array, ak.Record)): - out.layout._touch_data(recursive=True) + ak.typetracer.touch_data(out) except Exception as err: on_fail = dask.config.get("awkward.optimization.on-fail") # this is the default, throw a warning but skip the optimization. diff --git a/src/dask_awkward/lib/structure.py b/src/dask_awkward/lib/structure.py index 780753ba..5e601686 100644 --- a/src/dask_awkward/lib/structure.py +++ b/src/dask_awkward/lib/structure.py @@ -82,11 +82,8 @@ def __call__(self, *arrays): # FIXME: with proper typetracer/form rehydration support we # should not need to manually touch this when it's a # typetracer - for a in arrays: - if ak.backend(a) == "typetracer": - a.layout._touch_data(recursive=True) - - return ak.argcartesian(list(arrays), **self.kwargs) + arrays = [ak.typetracer.touch_data(a) for a in arrays] + return ak.argcartesian(arrays, **self.kwargs) @borrow_docstring(ak.argcartesian)