Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: enable manual column optimization with dask-awkward #1394

Merged
merged 2 commits into from
Feb 28, 2025

Conversation

pfackeldey
Copy link
Collaborator

@pfackeldey pfackeldey commented Feb 27, 2025

Enables dask-contrib/dask-awkward#565 for uproot.dask.

Example:

import uproot
import dask_awkward as dak


# foo.root has 3 branches: "A", "B", "C"
array = uproot.dask({"foo.root": "tree_name"})

# get dask key in graph:
key, _ = next(iter(array.dask.keys()))

# manual column optimization for "A" and "C"
needs = {key: ["A", "C"]}
optimized_array = dak.manual.optimize_columns(array, needs)

# this only loads "A" & "C" from disk in a single read call (`tree.arrays([...])`), 
# it doesn't load "B", and avoids typetracing entirely (it's already optimized...)
concrete_array = optimized_array.compute()

key corresponds to the key of the from-uproot IO layer. needs can have multiple keys that point to various projectable IO-layers.

@pfackeldey pfackeldey requested a review from ianna February 27, 2025 19:53
Copy link
Collaborator

@ianna ianna left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@pfackeldey - Thanks! Looks good to me. Please, merge it. Thanks.

@pfackeldey pfackeldey merged commit 472cf39 into main Feb 28, 2025
26 checks passed
@pfackeldey pfackeldey deleted the pfackeldey/manual_column_optimization branch February 28, 2025 18:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants