Enable dask-expr #383

dougbrn · 2024-03-11T21:59:11Z

Change Description

Resolves #382.

My PR includes a link to the issue that I am addressing

Solution Description

This PR enables the use of the recently introduced dask-expr backend to dask with TAPE. For the most part things just needed small tweaks, but there's a couple things that stood out:

_meta and divisions properties are now protected, so I had to update cases where we were directly setting those. In the case of divisions, I opted to just let Dask set it's divisions without our intervention as in some cases we were trying to use the source divisions for batch results, but the source divisions should be what batch finds regardless (as it's a groupby on the source table).
divisions were a bit finicky to be set, I believe I tracked this down to a kwarg change in set_index but we should watch out for any cases where divisions are not being set. It seemed more finicky for single partition tables as well, so I updated a few unit tests to produce tables with >1 division.
This PR switches TAPE to support only the dask-expr backend. I am not sure how much would break of the user tried the old backend, but I've included some checks in initialization to force the use of dask-expr. We should release a minor version of TAPE (v0.4.0) once this is included to easily denote the break point between the two backends.

Code Quality

My code builds (or compiles) cleanly without any errors or warnings
My code contains relevant comments and necessary documentation

Project-Specific Pull Request Checklists

I have added a function that requires a sync_tables command, and have added the neccesary sync_tables call

Bug Fix Checklist

My fix includes a new test that breaks as a result of the bug (if possible)
My change includes a breaking change
- My change includes backwards compatibility and deprecation warnings (if possible)

New Feature Checklist

I have added or updated the docstrings associated with my feature using the NumPy docstring format
I have updated the tutorial to highlight my new feature (if appropriate)
I have added unit/End-to-End (E2E) test cases to cover my new feature
My change includes a breaking change
- My change includes backwards compatibility and deprecation warnings (if possible)

Documentation Change Checklist

Any updated docstrings use the NumPy docstring format

Build/CI Change Checklist

If required or optional dependencies have changed (including version numbers), I have updated the README to reflect this
If this is a new CI setup, I have added the associated badge to the README

Other Change Checklist

Any new or updated docstrings use the NumPy docstring format.
I have updated the tutorial to highlight my new feature (if appropriate)
I have added unit/End-to-End (E2E) test cases to cover any changes
My change includes a breaking change
- My change includes backwards compatibility and deprecation warnings (if possible)

github-actions · 2024-03-11T22:01:09Z

Before [`dae414b`]	After [`caea38f`]	Ratio	Benchmark (Parameter)
44.6±0.5ms	32.6±0.2ms	0.73	benchmarks.time_batch
47.3±0.1ms	33.5±0.1ms	0.71	benchmarks.time_prune_sync_workflow

Click here to view all benchmarks.

codecov · 2024-03-20T20:13:56Z

Codecov Report

Attention: Patch coverage is 97.11538% with 3 lines in your changes are missing coverage. Please review.

Project coverage is 95.77%. Comparing base (dae414b) to head (4ee4856).

Files	Patch %	Lines
src/tape/ensemble_frame.py	95.94%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #383      +/-   ##
==========================================
+ Coverage   95.55%   95.77%   +0.21%     
==========================================
  Files          25       25              
  Lines        1710     1751      +41     
==========================================
+ Hits         1634     1677      +43     
+ Misses         76       74       -2

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

review-notebook-app · 2024-03-21T20:16:21Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

src/tape/ensemble_frame.py

wilsonbb

Some minor questions and comments but overall this looks really good!

docs/tutorials/binning_slowly_changing_sources.ipynb

wilsonbb · 2024-03-26T21:53:58Z

src/tape/ensemble.py

+        self.update_frame(
+            self.source._propagate_metadata(result)
+        )  # propagate source metadata and update frame


Is there some method we can override in the EnsembleFrame or some method to register to the dispatcher that can do this for us whenever a user calls concat? If it's non-obvious, we could file a small issue to look into it

I'm not sure, maybe? I opted to just do this as this was the only time I've seen concat used, so I wasn't keen on wrapping it just for this. If you know of a better way happy to change, or an issue is also good. This is one of a few cases where I was wondering if we could have a better way to wrap the API where the only needed change is this metadata propagation...

"if we could have a better way to wrap the API where the only needed change is this metadata propagation..."

Yeah an issue for this sounds good to me. I'm also a bit frustrated by this

wilsonbb · 2024-03-26T21:55:41Z

src/tape/ensemble.py

+        result = self.source._propagate_metadata(
+            result.reset_index().set_index(self._id_col).drop(columns=[tmp_time_col])
+        )


Do we know where in this chain we lose the metadata? (groupby, reset_index, set_index, drop, etc)

I think it's the drop call, but it could also possibly be the reset_index call

Ah, took a look again and reminded myself. It's actually before any of these calls, right above this theres a groupby->aggregate. The meta is lost when doing this, I think technically at the groupby stage as it returns a dask groupby object.

tests/tape_tests/test_ensemble.py

tests/tape_tests/test_ensemble_frame.py

src/tape/ensemble_frame.py

wilsonbb

Thanks for the changes, looks good to me!

dougbrn force-pushed the use_dask_expr branch from 6daaa52 to da39b86 Compare March 18, 2024 17:45

dougbrn commented Mar 22, 2024

View reviewed changes

src/tape/ensemble_frame.py Outdated Show resolved Hide resolved

dougbrn added 16 commits March 26, 2024 10:05

some initial test changes

f049b07

WIP: add expr implementation

bc97f2a

WIP: continue adding features

7bd5aff

WIP: more expr

72c524f

batch_by_band fix

04230bc

small tweaks

5fdb42a

fixing more expressions issues

3114576

full local unit test passes

5af7d99

fix formatting

630e717

limit dropna test to just divisions-enabled

64e4498

limit to just divisions

ea6d143

fix _meta setting

53e61f9

fix binning nb

2494898

replace ensemble_frame

969b761

tidy

c76fbc8

fix codebase todos

51e8562

dougbrn force-pushed the use_dask_expr branch from 6b720b3 to 51e8562 Compare March 26, 2024 17:05

dougbrn added 3 commits March 26, 2024 10:17

remove manual divisions setting

3795269

ensure dask expressions backend

d9f11f0

tidy

dbd5b4c

dougbrn changed the title ~~WIP: Enable dask-expr~~ Enable dask-expr Mar 26, 2024

dougbrn marked this pull request as ready for review March 26, 2024 20:32

dougbrn requested a review from wilsonbb March 26, 2024 20:37

wilsonbb reviewed Mar 26, 2024

View reviewed changes

address review comments

4ee4856

dougbrn requested a review from wilsonbb March 27, 2024 20:20

wilsonbb approved these changes Mar 27, 2024

View reviewed changes

dougbrn merged commit 750fe4b into main Mar 27, 2024
10 checks passed

dougbrn mentioned this pull request Mar 28, 2024

Clear divisions warnings in docs #410

Closed

dougbrn deleted the use_dask_expr branch April 4, 2024 17:16

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enable dask-expr #383

Enable dask-expr #383

dougbrn commented Mar 11, 2024 •

edited

Loading

github-actions bot commented Mar 11, 2024 •

edited

Loading

codecov bot commented Mar 20, 2024 •

edited

Loading

review-notebook-app bot commented Mar 21, 2024

wilsonbb left a comment

wilsonbb Mar 26, 2024

dougbrn Mar 26, 2024

wilsonbb Mar 27, 2024

wilsonbb Mar 26, 2024

dougbrn Mar 26, 2024

dougbrn Mar 27, 2024

wilsonbb left a comment

Enable dask-expr #383

Enable dask-expr #383

Conversation

dougbrn commented Mar 11, 2024 • edited Loading

Change Description

Solution Description

Code Quality

Project-Specific Pull Request Checklists

Bug Fix Checklist

New Feature Checklist

Documentation Change Checklist

Build/CI Change Checklist

Other Change Checklist

github-actions bot commented Mar 11, 2024 • edited Loading

codecov bot commented Mar 20, 2024 • edited Loading

Codecov Report

review-notebook-app bot commented Mar 21, 2024

wilsonbb left a comment

Choose a reason for hiding this comment

wilsonbb Mar 26, 2024

Choose a reason for hiding this comment

dougbrn Mar 26, 2024

Choose a reason for hiding this comment

wilsonbb Mar 27, 2024

Choose a reason for hiding this comment

wilsonbb Mar 26, 2024

Choose a reason for hiding this comment

dougbrn Mar 26, 2024

Choose a reason for hiding this comment

dougbrn Mar 27, 2024

Choose a reason for hiding this comment

wilsonbb left a comment

Choose a reason for hiding this comment

dougbrn commented Mar 11, 2024 •

edited

Loading

github-actions bot commented Mar 11, 2024 •

edited

Loading

codecov bot commented Mar 20, 2024 •

edited

Loading