feat(connect): add repartition support #3382
Open
Conversation
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
This was referenced Nov 21, 2024
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
2 times, most recently
from
e4dadeb to
9bcb827
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
2 times, most recently
from
e1aa22e to
c27e74e
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
9bcb827 to
886d014
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
c27e74e to
cf1ba9d
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
886d014 to
14d1d50
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
cf1ba9d to
581eb36
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
2 times, most recently
from
7302bfb to
62ae066
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
581eb36 to
59a7046
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
62ae066 to
6e58760
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
59a7046 to
54df4fd
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
6e58760 to
65d47db
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
54df4fd to
3d8335e
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
65d47db to
670c495
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
3d8335e to
7819f2a
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
670c495 to
c57e270
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
7819f2a to
a97b69b
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
c57e270 to
7bdd623
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
a97b69b to
5957c06
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
7bdd623 to
8fcea97
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
5957c06 to
353a6e1
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
8fcea97 to
3038c0f
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
353a6e1 to
eb6a937
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
3038c0f to
31e8bc8
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
f18c35d to
5e7d068
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
b7848ce to
0de7b03
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
5e7d068 to
003929c
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
0de7b03 to
3cb7252
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
003929c to
a40e010
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
3cb7252 to
21fc2ec
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
a40e010 to
9f2e53a
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
21fc2ec to
f2a1799
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
9f2e53a to
538bb93
Compare
andrewgazelka
force-pushed
the
andrew/connect-intersect-union
branch
from
f2a1799 to
11fefd5
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
538bb93 to
a269fea
Compare
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
a269fea to
456be41
Compare
andrewgazelka
changed the base branch from
andrew/connect-intersect-union
to
main
andrewgazelka
changed the title
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
456be41 to
5196ab7
Compare
CodSpeed Performance ReportMerging #3382 will degrade performances by 32.09%Comparing Summary
Benchmarks breakdown
|
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #3382 +/- ##
=======================================
Coverage 77.86% 77.86%
=======================================
Files 719 720 +1
Lines 88459 88488 +29
=======================================
+ Hits 68877 68902 +25
- Misses 19582 19586 +4
|
Comment on lines
+28
to
+37
| let shuffle = shuffle.unwrap_or(true); | ||
|
|
||
| if !shuffle { | ||
| bail!("Repartitioning without shuffling is not yet supported"); | ||
| } | ||
|
|
||
| plan.builder = plan | ||
| .builder | ||
| .random_shuffle(Some(num_partitions)) | ||
| .wrap_err("Failed to apply random_shuffle to logical plan")?; |
There was a problem hiding this comment.
we should support non shuffle based repartitions too.
in daft, we do this via hash_repartition
Comment on lines
+3
to
+13
| def test_repartition(spark_session): | ||
| # Create a simple DataFrame | ||
| df = spark_session.range(10) | ||
|
|
||
| # Test repartitioning to 2 partitions | ||
| repartitioned = df.repartition(2) | ||
|
|
||
| # Verify data is preserved after repartitioning | ||
| original_data = sorted(df.collect()) | ||
| repartitioned_data = sorted(repartitioned.collect()) | ||
| assert repartitioned_data == original_data, "Data should be preserved after repartitioning" |
There was a problem hiding this comment.
i'd instead just look at the plan and assert that it has a partition node.
ex: (note: this is using native spark, so our explain will be a little different)
d = [{'idx': 1}]
df = spark.createDataFrame(d).repartition(2)
df.explain()
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- Exchange RoundRobinPartitioning(2), REPARTITION_BY_NUM, [plan_id=6]
+- Scan ExistingRDD[idx#0L]There was a problem hiding this comment.
i guess this is fine for now considering we don't have explain yet
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
6 times, most recently
from
7d2d9d6 to
655aec2
Compare
- [ ] the test is not great but idk how to do it better since rdd does not work with spark connect (I think) - [ ] do we want to support non-shuffle repartitioning?
andrewgazelka
force-pushed
the
andrew/connect-partition
branch
from
655aec2 to
36eba74
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
not work with spark connect (I think)