Apache Beam: Partition Transform

Overview

If you want to split a PCollection into multiple PCollections based on a function, use the Partition transform.

When You Should Use the Keys Transform

When you want to split the input PCollection into multiple PCollections. You might want to partition a collection of strings based on their length. Then, truncate the longer strings into smaller sizes, but keep the shorter strings the same.

How to Use the Keys Transform

Just apply the builtin Transform to a PCollection with the following arguments:

Number of Partitions
Partition function - This function must return an integer between 0 and the number of partitions

Example: Partition based on Percentile

    // Provide an int value with the desired number of result partitions, and a
    // PartitionFn that represents the
    // partitioning function. In this example, we define the PartitionFn in-line.
    // Returns a PCollectionList
    // containing each of the resulting partitions as individual PCollection
    // objects.
    PCollection<Student> students =
        pipeline.apply(
            Create.of(
                Student.of("Amy", 88),
                Student.of("Bob", 87),
                Student.of("Chris", 49),
                Student.of("Dylan", 62),
                Student.of("Ellen", 78),
                Student.of("Francis", 53)));
    // Split students up into 10 partitions, by percentile:
    PCollectionList<Student> studentsByPercentile =
        students.apply(
            Partition.of(
                10,
                new Partition.PartitionFn<Student>() {
                  @Override
                  public int partitionFor(Student student, int numPartitions) {
                    return student.getPercentile() // 0..99
                        * numPartitions
                        / 100;
                  }
                }));

    // You can extract each partition from the PCollectionList using the get method,
    // as follows:
    PCollection<Student> fortiethPercentile = studentsByPercentile.get(4);
    fortiethPercentile.apply(ParDo.of(new LogOutput<>("Fortieth percentile: ")));
    pipeline.run();
  }

Conclusion

Check out other useful transforms from the official Apache Beam documentation.

Apache Beam: Partition Transform

Comments

Apache Beam and Google Cloud Dataflow

Apache Beam: ParDo

More from this blog

Harnessing the IP Address  Data Type in Apache Druid (2025)

Geospatial Queries in Apache Druid: A Ride-Sharing Example

Automating Google Cloud with GitHub Actions using gcloud CLI

How to Run a Scheduled Bash Script with Google Cloud Platform

Google Dataflow Optimization: Streaming Engine

Overview

When You Should Use the Keys Transform

How to Use the Keys Transform

Example: Partition based on Percentile

Conclusion

Command Palette

Comments

Apache Beam and Google Cloud Dataflow

Apache Beam: ParDo

More from this blog

Overview

When You Should Use the Keys Transform

How to Use the Keys Transform

Example: Partition based on Percentile

Conclusion