feat: support `shape_touched` from Dask #900

agoose77 · 2023-10-03T16:42:56Z

This PR at present represents a hack of the existing uproot remapping loader to support dask-awkward's new buffer projection interface.

@lgray the Coffea form_mapping object that is passed to uproot.dask needs to satisfy the ImplementsFormMapping interface:

class ImplementsFormMapping(Protocol):
    def __call__(self, form: Form) -> tuple[Form, ImplementsFormMappingInfo]:
        ...

The idea is that the form mapping generates both the remapped form and any auxiliary state required for projection. The return tuple's second item, ImplementsFormMappingInfo has the following protocol:

class ImplementsFormMappingInfo(Protocol):
    @property
    def behavior(self) -> dict | None:
        ...

    buffer_key: str | Callable

    def parse_buffer_key(self, buffer_key: str) -> tuple[str, str]:
        ...

    def keys_for_buffer_keys(self, buffer_keys: frozenset[str]) -> frozenset[str]:
        ...

    def load_buffers(
        self,
        tree: HasBranches,
        keys: frozenset[str],
        start: int,
        stop: int,
        options: Any,
    ) -> Mapping[str, AwkArray]:
        ...

This PR hacks the _map_schema_uproot object to provide this interface, but it could be cleaner; one could remove any unused code here.

for more information, see https://pre-commit.ci

…uched' into agoose77/feat-support-shape-touched

src/coffea/processor/executor.py

agoose77 · 2023-10-04T10:49:47Z

@lgray this PR implements the protocols that uproot expects for form remapping in factory.py for the Uproot source. I'm not yet au-fait with how the parquet equivalent is handled here, but you'd probably need to replicate the logic in uproot._dask.UprootReadMixin for the parquet IO function.

Can I hand this over to you?

…uched' into agoose77/feat-support-shape-touched

lgray · 2023-10-04T12:12:04Z

Yes I can handle the parquet reader stuff - no big deal there.

I'll give this a try today.

for more information, see https://pre-commit.ci

…uched' into agoose77/feat-support-shape-touched

lgray · 2023-10-07T00:30:14Z

@agoose77 we're just missing a release of dask-histogram and then I can mess with pins here. I'll let you handle the pre-release vs. full release state of things.

Once we get everything to a full release we can merge this one.

pyproject.toml

Co-authored-by: Angus Hollands <[email protected]>

lgray · 2023-10-07T18:12:58Z

For the two new protocols defined we should update dask-awkward since that's where the present protocol ImplementsFormTransformation is kept and I have no problem with changing that.

agoose77 · 2023-10-07T20:30:18Z

For the two new protocols defined we should update dask-awkward since that's where the present protocol ImplementsFormTransformation is kept and I have no problem with changing that.

Could you clarify what you mean here @lgray ? :)

lgray · 2023-10-08T01:34:30Z

@agoose77

This:
https://github.com/dask-contrib/dask-awkward/blob/main/src/dask_awkward/lib/io/io.py#L45-L61

agoose77 · 2023-10-08T11:02:39Z

I've made a PR to remove it, as it's a better fit for uproot! :)

agoose77 · 2023-10-11T07:27:44Z

src/coffea/nanoevents/factory.py

+                    buffer_key(form_key=form_key, attribute=attribute, form=None)
+                ]
+
+        return TranslateBufferKeys()

    def create_column_mapping_and_key(self, tree, start, stop, interp_options):


@lgray can we remove this from the base definition, or does someone out there likely use it?

interface is completely private and rather obscure to boot, so I am fine to remove it.

tests/test_jetmet_tools.py

agoose77 · 2023-10-11T07:36:56Z

tests/test_jetmet_tools.py

@@ -837,9 +837,9 @@ def test_corrected_jets_factory(optimization_enabled):
            **{name: evaluator[name] for name in jec_stack_names[5:6]}
        )

-        print(dak.necessary_columns(jets.eta))
+        print(dak.report_necessary_columns(jets.eta))


@lgray note that report_necessary_columns reports whatever the IO layer thinks a "column" is. It doesn't necessarily mean "dotted field" right now, rather I envisage it talking about the concept of a "key" or "column" that can be read from the IO source.

@douglasdavis this is something I didn't fully think about when we made necessary_columns shim report_necessary_columns.

The difference will be that for remapped forms, the report_necessary_columns (and by proxy, necessary_columns) will report the underlying ROOT file keys, not the remapped keys.

It therefore might mean that we should remove that alias for necessary_columns.

I actually prefer the new behavior since users will probably expect to get a list of columns in the file they're reading! But seeing the remapped keys is nice for debugging.

Can we make it so that we get a version or two of dask-awkward before dak.necessary_columns is gone and give migration instructions? I know a few end-users that use that function.

@lgray I'm not 100% clear if we're on the same page.

report_necessary_columns is a new function, so I'm comfortable with its semantics deviating from ak.necessary_columns.

It would be possible to restore the old behavior, but it will require a reasonable amount of work on the uproot side (I am guessing).

ak.necessary_columns is currently an alias for this new function, but for uproot this wouldn't be a non-breaking change; if previously it would report ROOT_KEY.record.foo.bar, now it would report ROOT_KEY. Is this OK? I don't see it being that useful if it breaks existing users.

OK I suppose for my most typical use case I never see keys with ROOT_KEY.record.foo.bar formatting, it's really only ever ROOT_KEY and any structure is up in the form being remapped to. So they appear the same to me.

lgray · 2023-10-11T08:50:37Z

@nsmith- I think this one is ready for review now. If you could give it a bit of time before the end of the week it would be appreciated.

agoose77 · 2023-10-11T12:51:07Z

src/coffea/nanoevents/factory.py

+    def keys_for_buffer_keys(self, buffer_keys):
+        base_columns = set()
+        for buffer_key in buffer_keys:
+            form_key, attribute = self.parse_buffer_key(buffer_key)
+            operands = urllib.parse.unquote(form_key).split(",")
+
+            it_operands = iter(operands)
+            next(it_operands)


This function behaves like extract_form_key_base_columns, but supports operands e.g. attribute lookup.

agoose77 · 2023-10-11T12:51:48Z

src/coffea/nanoevents/factory.py

@@ -147,8 +175,17 @@ def create_column_mapping_and_key(self, tree, start, stop, interp_options):
            use_ak_forth=True,
        )
        mapping.preload_column_source(partition_key[0], partition_key[1], tree)


I've deliberately tried not to touch this logic at all. Hopefully, by presenting the correct interface to uproot, internal refactoring can happen here in future without any side effects.

lgray · 2023-10-12T07:08:13Z

suspicious new failure...

ikrommyd · 2023-10-14T01:48:31Z

I'd like to report that I've been using this all day on LPC running stuff with the newest awkward, uproot, dask-awkward and dask-histogram and have not ran into any errors.

agoose77 · 2023-10-14T08:30:48Z

Todo: bump the uproot and dask-awkward versions

nb: safer to repin in coffea for users because of numba's sliding window, very easy to get a mismatch

lgray · 2023-10-16T12:24:34Z

@nsmith- a gentle reminder for a review, please!

nsmith-

We need to do some cleanup, but it can happen later I think.
@agoose77 does the protocol you defined in the initial comment have some explanatory docstrings available somewhere?

nsmith- · 2023-10-17T13:59:34Z

src/coffea/processor/executor.py

@@ -1734,11 +1734,11 @@ def _work_function(
                    metrics = {}
                    if isinstance(file, uproot.ReadOnlyDirectory):
                        metrics["bytesread"] = file.file.source.num_requested_bytes
+                    # metrics["data_and_shape_buffers"] = set(materialized)
+                    # metrics["shape_only_buffers"] = set(materialized)


Are these placeholders for something in the future?

This is just to get tests to run - executor is going to be significantly changed in a subsequent PR.

nsmith- · 2023-10-17T14:02:18Z

src/coffea/nanoevents/factory.py

+    def parse_buffer_key(self, buffer_key):
+        prefix, attribute, form_key = buffer_key.rsplit("/", maxsplit=2)
+        if attribute == "offsets":
+            return (form_key[: -len("%2C%21offsets")], attribute)


This seems to also be in the previous code as well, but some of these string splitting and urllib name mangling things were put into https://github.com/CoffeaTeam/coffea/blob/master/src/coffea/nanoevents/util.py and could be used here

nsmith- · 2023-10-17T14:04:43Z

src/coffea/nanoevents/factory.py

@@ -68,7 +67,7 @@ def _key_formatter(prefix, form_key, form, attribute):
    return prefix + f"/{attribute}/{form_key}"


-class _map_schema_base(ImplementsFormTransformation):
+class _map_schema_base:  # ImplementsFormMapping, ImplementsFormMappingInfo


Any reason not to inherit the protocol(s)?

They weren't exported yet! I think a new release includes them. Let me check.

for more information, see https://pre-commit.ci

agoose77 and others added 2 commits October 3, 2023 17:41

wip: initial commit

b8fc7fe

[pre-commit.ci] auto fixes from pre-commit.com hooks

1b4bd50

for more information, see https://pre-commit.ci

This was referenced Oct 3, 2023

feat: implement shape_touched optimisation dask-contrib/dask-awkward#379

Merged

feat: add support for shape touching in Dask scikit-hep/uproot5#966

Merged

agoose77 added 4 commits October 4, 2023 10:18

fix: rename function

daa8529

Merge remote-tracking branch 'agoose77/agoose77/feat-support-shape-to…

df97041

…uched' into agoose77/feat-support-shape-touched

fix: use report_necessary_buffers

66c8710

Merge branch 'master' into agoose77/feat-support-shape-touched

f31815b

agoose77 commented Oct 4, 2023

View reviewed changes

src/coffea/processor/executor.py Show resolved Hide resolved

agoose77 added 2 commits October 4, 2023 12:19

fix: properly parse form keys

2353a23

Merge remote-tracking branch 'agoose77/agoose77/feat-support-shape-to…

17a6f2f

…uched' into agoose77/feat-support-shape-touched

agoose77 and others added 4 commits October 5, 2023 14:05

hack: convert Content to array

bd07d03

[pre-commit.ci] auto fixes from pre-commit.com hooks

a6848a0

for more information, see https://pre-commit.ci

fix: ensure layout nodes converted to arrays

9c90205

Merge remote-tracking branch 'agoose77/agoose77/feat-support-shape-to…

379a508

…uched' into agoose77/feat-support-shape-touched

adjust coffea pins to latest releases and pre-releases

04b5a1a

agoose77 commented Oct 7, 2023

View reviewed changes

pyproject.toml Outdated Show resolved Hide resolved

lgray and others added 2 commits October 7, 2023 09:36

use pytorch-only triton image

f19c11b

streamline version requirements

7051d2e

Co-authored-by: Angus Hollands <[email protected]>

lgray marked this pull request as ready for review October 7, 2023 15:41

Merge branch 'master' into agoose77/feat-support-shape-touched

3fe957b

lgray mentioned this pull request Oct 7, 2023

build: pin dask-awkward >= 2023.9.3, awkward >= 2.4.3 #891

Closed

Merge branch 'master' into agoose77/feat-support-shape-touched

d3fb1c9

remove deprecated interface definition

0d9c913

agoose77 commented Oct 11, 2023

View reviewed changes

tests/test_jetmet_tools.py Outdated Show resolved Hide resolved

agoose77 commented Oct 11, 2023

View reviewed changes

tests/test_jetmet_tools.py Outdated Show resolved Hide resolved

agoose77 added 2 commits October 11, 2023 08:28

Update tests/test_jetmet_tools.py

bb4df59

Update tests/test_jetmet_tools.py

e0694ad

agoose77 commented Oct 11, 2023

View reviewed changes

remove further remnants of old remapping interface

e7384f9

refactor: make key translation obvious

92efdb2

agoose77 commented Oct 11, 2023

View reviewed changes

fix typo from refactor

8ae3cd5

lgray added 4 commits October 14, 2023 10:22

update pins (note uncapped numpy and numba skooch)

45a0060

try to convince pip to upgrade numpy upon installing coffea

14d2cc2

be more insistent

ab3599e

numba 0.58 pins numpy from above < 1.26

faff41e

nb: safer to repin in coffea for users because of numba's sliding window, very easy to get a mismatch

lgray mentioned this pull request Oct 14, 2023

fix: remove binder setup problems in coffea 2023 #908

Merged

lgray mentioned this pull request Oct 17, 2023

fix: added hasattr check to NanoAODEvents __repr__ #909

Merged

nsmith- approved these changes Oct 17, 2023

View reviewed changes

lgray mentioned this pull request Oct 17, 2023

Implement protocols defined in #900 #910

Closed

lgray and others added 3 commits October 17, 2023 10:22

clean up usage of quoted ",!offsets"

864f709

[pre-commit.ci] auto fixes from pre-commit.com hooks

9b96f7b

for more information, see https://pre-commit.ci

flake8 lint

0a525d0

lgray enabled auto-merge October 17, 2023 15:29

lgray merged commit 8662c9f into scikit-hep:master Oct 17, 2023
13 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: support `shape_touched` from Dask #900

feat: support `shape_touched` from Dask #900

agoose77 commented Oct 3, 2023 •

edited

Loading

agoose77 commented Oct 4, 2023

lgray commented Oct 4, 2023

lgray commented Oct 7, 2023

lgray commented Oct 7, 2023

agoose77 commented Oct 7, 2023

lgray commented Oct 8, 2023

agoose77 commented Oct 8, 2023

agoose77 Oct 11, 2023

lgray Oct 11, 2023

agoose77 Oct 11, 2023

agoose77 Oct 11, 2023

lgray Oct 11, 2023

lgray Oct 11, 2023

agoose77 Oct 11, 2023 •

edited

Loading

lgray Oct 12, 2023

lgray commented Oct 11, 2023

agoose77 Oct 11, 2023

agoose77 Oct 11, 2023

lgray commented Oct 12, 2023

ikrommyd commented Oct 14, 2023

agoose77 commented Oct 14, 2023

lgray commented Oct 16, 2023

nsmith- left a comment

nsmith- Oct 17, 2023

lgray Oct 17, 2023

nsmith- Oct 17, 2023

nsmith- Oct 17, 2023

agoose77 Oct 17, 2023

feat: support shape_touched from Dask #900

feat: support shape_touched from Dask #900

Conversation

agoose77 commented Oct 3, 2023 • edited Loading

agoose77 commented Oct 4, 2023

lgray commented Oct 4, 2023

lgray commented Oct 7, 2023

lgray commented Oct 7, 2023

agoose77 commented Oct 7, 2023

lgray commented Oct 8, 2023

agoose77 commented Oct 8, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

agoose77 Oct 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lgray commented Oct 11, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lgray commented Oct 12, 2023

ikrommyd commented Oct 14, 2023

agoose77 commented Oct 14, 2023

lgray commented Oct 16, 2023

nsmith- left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

feat: support `shape_touched` from Dask #900

feat: support `shape_touched` from Dask #900

agoose77 commented Oct 3, 2023 •

edited

Loading

agoose77 Oct 11, 2023 •

edited

Loading