-
-
Notifications
You must be signed in to change notification settings - Fork 27
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
Adjust class hierarchy for overlapping stuff #429
Open
phofl
wants to merge
2
commits into
dask:main
Choose a base branch
from
phofl:blockwise_overlapping
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
2 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
|
@@ -1093,7 +1093,19 @@ def _task(self, index: int): | |||||
return self.value | ||||||
|
||||||
|
||||||
class Blockwise(Expr): | ||||||
class BlockwiseOverlapping(Expr): | ||||||
"""Super class that operates like a blockwise op | ||||||
|
||||||
We require information from neighboring partitions, so we can't prune partitions | ||||||
before lowering, but the spec is the same as for Blockwise ops, we don't reoder | ||||||
things, alignment stays consistent, ... | ||||||
|
||||||
""" | ||||||
|
||||||
pass | ||||||
|
||||||
|
||||||
class Blockwise(BlockwiseOverlapping): | ||||||
"""Super-class for block-wise operations | ||||||
|
||||||
This is fairly generic, and includes definitions for `_meta`, `divisions`, | ||||||
|
@@ -1276,7 +1288,7 @@ def _task(self, index: int): | |||||
) | ||||||
|
||||||
|
||||||
class MapOverlap(MapPartitions): | ||||||
class MapOverlap(BlockwiseOverlapping): | ||||||
_parameters = [ | ||||||
"frame", | ||||||
"func", | ||||||
|
@@ -1296,6 +1308,29 @@ class MapOverlap(MapPartitions): | |||||
"clear_divisions": False, | ||||||
} | ||||||
|
||||||
def _broadcast_dep(self, dep: Expr): | ||||||
return dep.npartitions == 1 | ||||||
|
||||||
@property | ||||||
def args(self): | ||||||
return [self.frame] + self.operands[len(self._parameters) :] | ||||||
|
||||||
def _divisions(self): | ||||||
# Unknown divisions | ||||||
if self.clear_divisions: | ||||||
return (None,) * (self.frame.npartitions + 1) | ||||||
|
||||||
# (Possibly) known divisions | ||||||
dfs = [arg for arg in self.args if isinstance(arg, Expr)] | ||||||
return _get_divisions_map_partitions( | ||||||
True, # Partitions must already be "aligned" | ||||||
self.transform_divisions, | ||||||
dfs, | ||||||
self.func, | ||||||
self.args, | ||||||
self.kwargs, | ||||||
) | ||||||
|
||||||
@functools.cached_property | ||||||
def _kwargs(self) -> dict: | ||||||
kwargs = self.kwargs | ||||||
|
@@ -2454,7 +2489,7 @@ def non_blockwise_ancestors(expr): | |||||
e = stack.pop() | ||||||
if isinstance(e, IO): | ||||||
yield e | ||||||
elif isinstance(e, Blockwise): | ||||||
elif isinstance(e, BlockwiseOverlapping): | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. If this is the only application, then this seems lighter weight to me
Suggested change
|
||||||
dependencies = e.dependencies() | ||||||
stack.extend([expr for expr in dependencies if not is_broadcastable(expr)]) | ||||||
else: | ||||||
|
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why does this exist? Is there some shared structure between these two that we need to handle?
My intuition here is just to have
At least from this PR it doesn't seem like we're using the intermediate classes.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not not yet,
are_co_aligned
can leverage this for exampleEvery check that only cares about the partitioning layout should fall back to this base class
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
My recommendation would be to wait until we have a demonstrable need, and even in that case to consider using
(Blockwise, MapOverlap)
instead if possible.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Updated the PR with the use-case
It's a bit similar to how Blockwise and Elemwise relate to each other