Cutting Shuffle Out of Big Joins

Shuffle is the backbone of distributed data processing—it’s what makes joins, GroupByKey, and reduce operations possible. But it’s also one of the most expensive steps in any pipeline, moving key–value pairs across disk and network. Sort Merge Bucket (SMB) flips that trade-off: do the expensive partitioning work once, up front, when writing data, so later joins can be done by merging already-sorted files instead of shuffling everything again.

At Spotify, that idea went from concept to production in a big way. Andrea Nardelli first explored SMB in a 2018 master’s thesis, and Spotify turned it into a generalized Scio module. The payoff: when joining frequently used datasets on a known key—say, user events with user metadata on a user ID—both sides are pre-bucketed and pre-sorted by that key. At read time, shuffle reduces to merging sorted files from matching buckets, sparing the costly disk and network I/O of moving records around.

Most of Spotify’s pipelines run on Scio, its Scala API for Apache Beam, on Google Cloud Dataflow. The SMB implementation itself lives in Java, closer to the Beam SDK, with Scala syntactic sugar exposed in Scio. The design breaks into two main transforms and a few supporting abstractions.

How SMB Writes and Reads

The two poles of the SMB design are the write transform and the read transform. They work together to guarantee the bucket-and-sort invariant that makes shuffle-free joins possible.

Writing Bucketed, Sorted Data

SortedBucketSink takes a PCollection and writes it out in SMB format. It extracts keys from each record, assigns bucket IDs according to logic in BucketMetadata, groups records by ID, sorts values within a bucket, and writes files-named per bucket using the associated FileOperations instance.

Alongside those data files, the sink also emits a JSON metadata file into the output directory. That file captures everything needed to read the source later: the total number of buckets, the hashing scheme, and the method for pulling the key from each record—for Avro, that’s the specific GenericRecord field name.

Merging Buckets on Read

SortedBucketSource reads from one or more SMB-formatted sources sharing the same key and hashing scheme. It opens the corresponding bucket files from each source, merges them while maintaining sorted order, and emits results as CoGbkResult objects per key group—the same class Beam uses for regular cogroup operations, so downstream code can extract per-source results with the correct typed accessors.

Supporting Abstractions

Two helper classes fill in the details. FileOperations handles reading and writing individual bucket files. Beam’s existing file I/Os won’t work here—they operate at the PCollection level, which masks the fine-grained control SMB needs over element locality and order. Instead, SMB file operations happen at the BoundedSource and ParDo level. Current formats include Avro, BigQuery TableRow JSON, and TensorFlow TFRecord/Example, with Parquet on the roadmap.

BucketMetadata encapsulates keying and bucketing; it stores the key field and type, the number of buckets, shards, and the hash function. Serialized alongside the data when writing, it’s deserialized at read time to verify source compatibility before a join starts.

Refinements for Production Scale

Adoption at Spotify over the past year and a half has driven a steady stream of additions to the original design, matching the quirks of real data pipelines.

Date partitioning. Event data at Spotify lives in hourly or daily partitions on GCS, and ingestion often means reading many partitions at once—say, for a seven-day stream count. Standard file I/O in Beam handles this naturally via wildcard file patterns, but SMB’s read API requires a directory rather than a pattern because it needs to reach the directory’s metadata.json as well as the record files. It also has to align bucket files across partitions and sources while correctly assigning all partitions of a source to the same TupleTag. The API now accepts one or more directories per source to make that work.

Sharding. Hash-based bucketing usually spreads records evenly, but skewed key spaces can blow up a single bucket and cause OOM errors during sorting. Users can specify a number of shards to split each bucket file further. Bucket assignment gets an extra random value in the range [0, numShards), generated per bundle. That value is orthogonal to the bucket ID, so even sprawling key groups get spread across files—each page still in sorted order, so merging at read time stays simple.

Parallelism. Since the bucket count in an SMB sink is always a power of two, users choose how deeply to parallelize joins between sources with different bucket counts. With Source A at 4 buckets and Source B at 2, three strategies are available:

  • Minimum parallelism (“Merge Greatest Buckets”)—two readers, each handling two buckets from A and one from B. Because bucket IDs are calculated as a hash mod the bucket count, the merged key spaces are provably overlapping.
  • Maximum parallelism (“Least Bucket Replication”)—four readers, each pairing one bucket from A and one from B. After merging each key group, a reader rehashes keys modulo the largest bucket count to avoid emitting duplicates. It’s faster in parallel, but it pays for that throughput with recomputation.
  • Auto parallelism—a reader count between minimum and maximum, calculated from a runtime split size supplied by the runner.

SortedBucketTransform. Enrichment pipelines typically read a dataset, join it, and write it out again. SMB has a dedicated transform for that pattern that reads, transforms, and rewrites using the same key and bucketing scheme per bucket on the same worker. Since the key doesn’t change, it’s certain that transformed elements from input bucket M belong to output bucket M—in the same sorted order they arrived in—so no reshuffle and no bucket recomputation is needed.

External sort improvements. Contributing back to Beam’s external sorter extension, Spotify’s team swapped out the Hadoop sequence file for native file I/O, lifted the prior 2GB memory cap, and reduced coder overhead and disk usage in the process.

Rolling Out SMB to Core Datasets

Adoption had to start with data producers. A handful of core datasets serve as the single source of truth for business domains like streaming activity and user metadata, and the team worked with their maintainers to convert a full year of history into SMB format.

The move turned out to be mostly a drop-in swap from the vanilla Avro sink. Compared to running with the existing sharding option—which already demanded a full shuffle—the SMB sink didn’t add any noticeable vCPU, vRAM, or wall-clock time. A few details required attention:

  • An agreement to treat user_id as a hexadecimal string for both bucket and sort keys, ensuring one consistent key semantic across all SMB datasets.
  • A shift to DEFLATE compression at level 6, matching the Scio Avro sink default. A nice side effect of the bucketed, sorted layout: storage dropped ~50% thanks to better compression from grouping similar records.
  • Backwards compatibility at the file level. SMB output files carry a distinct bucket-X-shard-Y filename pattern, but the records and schemas inside are unchanged—so existing consumers keep reading them, just without the join speedup.

Putting SMB to Work

Building on the groundwork from the Wrapped 2019 campaign, the team set out to rework the pipeline for Wrapped 2020. The architecture from the prior year was reusable, but its reliance on a large Bigtable cluster as the listening history source was both costly and cumbersome; the cluster had to be scaled up significantly just to handle the Wrapped workload. The goal was to move away from Bigtable entirely.

The new approach leaned on sorted, mergeable (SMB) data sources keyed by user_id. This shift also addressed a new challenge: complex filtering and aggregation of streams based on contextual information. Joining a large contextual dataset to full-year listening histories would have been exorbitantly expensive with conventional joins, so the team used SMB to sidestep that join altogether.

Wrapped 2020 required reading from three primary sources: streaming activity, user metadata, and streaming context. In the old system, five years of listening history sat in Bigtable, already keyed by user_id. The new design read all three SMB sources in the same keyed format, then aggregated a year of data per key to produce each user's Wrapped result.

Aggregating by Partition

A practical complication emerged from the source data's partitioning: one source is partitioned hourly while the other two are daily. Reading a full year of data in one job would have caused an impractical number of concurrent reads from the hourly partition. The solution was a two-stage process. First, smaller jobs computed weekly (or daily) aggregates of play counts, msPlayed, and related metrics per user. Then, a single job merged those smaller partitions into one dataset holding a full year's worth of data.

This is where SMB really paid off. The team used sortMergeTransform to combine the three sources, each already keyed by user_id, and wrote the results — play counts, play context, and so on — back out in SMB format. The final aggregation step used sortMergeGroupByKey to read all weekly SMB partitions, merge a year of data, and emit the output for the downstream Wrapped calculations.

A particularly useful property of this design is flexibility: the aggregate job can accept any mixture of weekly and daily partitions. That made job scheduling far more manageable.

Measured Results

The change delivered substantial cost savings. By leveraging SMB, the team joined roughly 1PB of data without conventional shuffle operations or Bigtable. They estimate a 50% reduction in Dataflow costs compared to previous years' Bigtable-based pipeline. Additionally, the Bigtable cluster no longer needed to be scaled to two or three times its normal capacity — peak usage had reached around 1,500 nodes in prior years.

The win was clear: Wrapped 2020 was delivered more cost-effectively than any prior campaign, and the SMB-based pipeline eliminated the need for the expensive workaround that Bigtable had become.

Future Direction

The success of SMB opens the door to migrating more workflows off Bigtable. Remaining work includes handling edge cases such as data skew, composite keys, and additional file formats. The team sees SMB as a foundational tool for making previously unfeasible joins routine and cost-effective.