From 24a9d11ef3f53e2a89837588603e3c229834659b Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Mon, 12 Feb 2024 09:07:15 -0600 Subject: [PATCH 1/5] fix: projected AwkwardInputLayers have __dask_tokenize__ --- src/dask_awkward/layers/layers.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/dask_awkward/layers/layers.py b/src/dask_awkward/layers/layers.py index 89c49c41..3b493118 100644 --- a/src/dask_awkward/layers/layers.py +++ b/src/dask_awkward/layers/layers.py @@ -117,6 +117,19 @@ def io_func_implements_report(func: ImplementsIOFunction) -> bool: return hasattr(func, "return_report") +class AwkwardTokenizable: + + def __init__(self, ret_func, parent_name): + self.parent_name = parent_name + self.ret_func = ret_func + + def __dask_tokenize__(self): + return (AwkwardTokenizable, self.parent_name) + + def __call__(self, *_, **__): + return self.ret_func + + class AwkwardInputLayer(AwkwardBlockwiseLayer): """A layer known to perform IO and produce Awkward arrays @@ -230,7 +243,7 @@ def prepare_for_projection(self) -> tuple[AwkwardInputLayer, TypeTracerReport, T new_input_layer = AwkwardInputLayer( name=self.name, inputs=[None][: int(list(self.numblocks.values())[0][0])], - io_func=lambda *_, **__: new_return, + io_func=AwkwardTokenizable(new_return, self.name), label=self.label, produces_tasks=self.produces_tasks, creation_info=self.creation_info, From c75b65c1275c8cdb2ef0f878300f3538331c37e4 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Mon, 12 Feb 2024 09:27:25 -0600 Subject: [PATCH 2/5] use a literal Co-authored-by: Martin Durant --- src/dask_awkward/layers/layers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/dask_awkward/layers/layers.py b/src/dask_awkward/layers/layers.py index 3b493118..ab50791e 100644 --- a/src/dask_awkward/layers/layers.py +++ b/src/dask_awkward/layers/layers.py @@ -124,7 +124,7 @@ def __init__(self, ret_func, parent_name): self.ret_func = ret_func def __dask_tokenize__(self): - return (AwkwardTokenizable, self.parent_name) + return ("AwkwardTokenizable", self.parent_name) def __call__(self, *_, **__): return self.ret_func From 7df3cdba232a1bc1d0ea31a1245595be6e62bd85 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Mon, 12 Feb 2024 09:46:25 -0600 Subject: [PATCH 3/5] ret_val not ret_func --- src/dask_awkward/layers/layers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/dask_awkward/layers/layers.py b/src/dask_awkward/layers/layers.py index ab50791e..92441443 100644 --- a/src/dask_awkward/layers/layers.py +++ b/src/dask_awkward/layers/layers.py @@ -119,15 +119,15 @@ def io_func_implements_report(func: ImplementsIOFunction) -> bool: class AwkwardTokenizable: - def __init__(self, ret_func, parent_name): + def __init__(self, ret_val, parent_name): self.parent_name = parent_name - self.ret_func = ret_func + self.ret_val = ret_val def __dask_tokenize__(self): return ("AwkwardTokenizable", self.parent_name) def __call__(self, *_, **__): - return self.ret_func + return self.ret_val class AwkwardInputLayer(AwkwardBlockwiseLayer): From ba147c36b4fb4d51e18563d8796ee2edb68fc408 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Mon, 12 Feb 2024 09:58:07 -0600 Subject: [PATCH 4/5] harden test again shifts in basic column touching --- tests/test_inspect.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_inspect.py b/tests/test_inspect.py index 6ffc7ed9..fdcfa7bc 100644 --- a/tests/test_inspect.py +++ b/tests/test_inspect.py @@ -97,7 +97,10 @@ def test_basic_root_works(): .fill(events.MET_pt) ) + columns = list(dak.necessary_columns(q1_hist).values())[0] + assert columns == frozenset({"MET_pt"}) dask.compute(q1_hist) + def test_sample(daa): From 1aadcfcc3e0dfabf9487b0b914e0b1cb7352fdcf Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 12 Feb 2024 16:00:21 +0000 Subject: [PATCH 5/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_inspect.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_inspect.py b/tests/test_inspect.py index fdcfa7bc..b4e2651b 100644 --- a/tests/test_inspect.py +++ b/tests/test_inspect.py @@ -100,7 +100,6 @@ def test_basic_root_works(): columns = list(dak.necessary_columns(q1_hist).values())[0] assert columns == frozenset({"MET_pt"}) dask.compute(q1_hist) - def test_sample(daa):