Skip to content

Improve shuffling performance #385

New issue
New issue
Closed
#717
Closed
Improve shuffling performance#385
#717

Description

@gabotechs
gabotechs
opened on Mar 23, 2026

The following PR added support for micro benchmarking shuffle operations:

The benchmarks shipped in that PR have several cases, and all of them shuffle the same amount of data, but the output tasks and partitioning is different. I'd expect all cases to take the same time, as they shuffle the same amount of data, but there is some variations:

Benchmarking shuffle/stream/1producer_1consumer_8partitions_1000000rows_1024batch_Nonecompression: 
                        time:   [90.704 ms 91.729 ms 94.261 ms]
shuffle/stream/1producer_1consumer_8partitions_1000000rows_1024batch_Some(LZ4_FRAME)compression: 
                        time:   [115.04 ms 125.10 ms 139.33 ms]
Benchmarking shuffle/stream/8producer_1consumer_8partitions_1000000rows_1024batch_Nonecompression: 
                        time:   [60.487 ms 60.780 ms 61.057 ms]
Benchmarking shuffle/stream/1producer_8consumer_8partitions_1000000rows_1024batch_Nonecompression: 
                        time:   [197.23 ms 197.69 ms 198.10 ms]
Benchmarking shuffle/stream/8producer_8consumer_8partitions_1000000rows_1024batch_Nonecompression: 
                        time:   [76.776 ms 77.825 ms 80.544 ms]

I think there should be some low hanging fruits regarding performance, and some other more complicated improvements:

  • I'm suspicious about arrow-flight cloning data unnecessarily
  • RepartitionExec might be doing something funny depending on the amount of output tasks

Activity

  1. gene-bordegaray
    added a commit that references this issue on Apr 20, 2026

    refactor / add benches for shuffle and transport shuffling (#408)

    4df1a0b

  2. Rich-T-kid commented on May 10, 2026

    @Rich-T-kid
    Rich-T-kid
    on May 10, 2026
    Collaborator

    does #408 close this ticket?
    cc @gabotechs @gene-bordegaray

  3. gene-bordegaray commented on May 10, 2026

    @gene-bordegaray
    gene-bordegaray
    on May 10, 2026
    Collaborator

    does #408 close this ticket?
    cc @gabotechs @gene-bordegaray

    No still working on this

  4. github-merge-queue
    added a commit that references this issue on May 16, 2026

    Call take arrays once per repartitioned input batch (#22159)

    1ab146a
  5. gabotechs
    added a commit that references this issue on May 18, 2026

    Call take arrays once per repartitioned input batch (apache#22159)

    6e11f6d

  6. gabotechs commented on Aug 4, 2026

    @gabotechs
    gabotechs
    on Aug 4, 2026
    CollaboratorAuthor

    One idea comes to mind:

    During network shuffles, perform two levels of RepartitionExec, with two different salts in the hashing expressions:

         ▲        ▲       ▲              ▲       ▲        ▲           ▲        ▲        ▲     
     │        │       │              │       │        │           │        │        │     
     │        │       │              │       │        │           │        │        │     
┌────┴────────┴───────┴─────┐  ┌─────┴───────┴────────┴────┐ ┌────┴────────┴────────┴────┐
│      NetworkShuffle       │  │      NetworkShuffle       │ │      NetworkShuffle       │
│┌─────────────────────────┐│  │┌─────────────────────────┐│ │┌─────────────────────────┐│
││  RepartitionExec::Hash  ││  ││  RepartitionExec::Hash  ││ ││  RepartitionExec::Hash  ││
││  [...exprs, local_salt],││  ││  [...exprs, local_salt],││ ││  [...exprs, local_salt],││
││    target_partitions,   ││  ││    target_partitions,   ││ ││    target_partitions,   ││
│└─────────────────────────┘│  │└─────────────────────────┘│ │└─────────────────────────┘│
└───────────────────────────┘  └───────────────────────────┘ └───────────────────────────┘
              ▲                              ▲                             ▲              
              │                              │                             │              
              │                              │                             │              
              │                              │                             │              
              │                              │                             │              
              │                              │                             │              
              └─────────────────┬────────────┴─────────────┬───────────────┘              
                                │                          │                              
                                │                          │                              
                                │                          │                              
                                │                          │                              
                                │                          │                              
                                │                          │                              
                   ┌────────────────────────┐ ┌────────────────────────┐                  
                   │ RepartitionExec::Hash  │ │ RepartitionExec::Hash  │                  
                   │[...exprs, remote_salt],│ │[...exprs, remote_salt],│                  
                   │  output_worker_count,  │ │  output_worker_count,  │                  
                   └────────────────────────┘ └────────────────────────┘
    1. First remote shuffle, that fanouts to 1 partition per worker, and that enriches the hashing expressions with a constant expression that acts like a "salt".
    2. Local re-shuffle, that fanouts to target_partitions partitions locally inside each NetworkShuffleExec, re-establishing local parallelism.

    Before jumping to this, I'd double check if there's any precedence in other engines about this.

  7. gabotechs commented on Sep 3, 2026

    @gabotechs
    gabotechs
    on Sep 3, 2026
    CollaboratorAuthor

    Here's another more complete diagram:

    Now

    ┌──────────────────────────┐          ┌──────────────────────────┐
│    NetworkShuffleExec    │          │    NetworkShuffleExec    │
│                          │          │                          │
│         task_i=0         │   ...    │        task_i=T-1        │
│   output_partitions=P    │          │   output_partitions=P    │
└──────────────────────────┘          └──────────────────────────┘
        ▲  ▲  ▲                                     ▲ ▲ ▲
        │  │  │                                     │ │ │
        │  │  └─         ────────┬────────         ─┘ │ │
        │  └────P streams──────┐ │ ┌──────P streams───┘ │
        └───────         ───┐  │ │ │ ┌────         ─────┘
                            │  │ │ │ │
                   ┌──────────────────────────┐
                   │  RepartitionExec(Hash)   │
                   │                          │
                   │  output_partitions=P*T   │
                   │   hash_expr=[...exprs]   │
                   └──────────────────────────┘

    Future

    ┌──────────────────────────────────────┐           ┌──────────────────────────────────────┐
│          NetworkShuffleExec          │           │          NetworkShuffleExec          │
│                                      │           │                                      │
│               task_i=0               │           │              task_i=T-1              │
│         output_partitions=P          │           │         output_partitions=P          │
│┌────────────────────────────────────┐│           │┌────────────────────────────────────┐│
││       RepartitionExec(Hash)        ││   ...     ││       RepartitionExec(Hash)        ││
││         input_partitions=1         ││           ││         input_partitions=1         ││
││        output_partitions=P         ││           ││        output_partitions=P         ││
││  hash_expr=[...exprs, local_salt]  ││           ││  hash_expr=[...exprs, local_salt]  ││
│└────────────────────────────────────┘│           │└────────────────────────────────────┘│
└──────────────────────────────────────┘           └──────────────────────────────────────┘
                    ▲                                                  ▲
                    │                                                  │
                    └─────1 streams──────────┬──────────1 streams──────┘
                                             │
                                             │
                                             │
                          ┌────────────────────────────────────┐
                          │       RepartitionExec(Hash)        │
                          │                                    │
                          │        output_partitions=T         │
                          │ hash_expr=[...exprs, remote_salt]  │
                          └────────────────────────────────────┘

  8. Rich-T-kid commented on Sep 5, 2026

    @Rich-T-kid
    Rich-T-kid
    on Sep 5, 2026
    Collaborator

    🤔 if we are going to reparation at the destination nodes do we still need to do it at the source?

    for example (I don't know how to make those cool asci diagrams 😆 )

          a           b
      ^          ^
            c

    why cant C send a stream of data directly to a & b and have the data record batches be Repartitioned there?

    edit: on second thought I think we need to do this to get groups keys together on the same node/partition.

  9. gabotechs
    closed this as completedin #717on Sep 23, 2026
  10. VedantMadane
    added a commit that references this issue on Sep 23, 2026

    feat: Reduce producer partition fan-out in NetworkShuffleExec (datafu…

    db6a843
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions