Limit number of record per shard in beam pipeline AvroIO.write()

There's no direct option to limit the number of records per shard, but a similar result can be achieved by setting the number of shards to 1 and then limiting the number of records per file.

One option is to use FinishBundle in a DoFn to limit the number of elements, and then write them, you can see an example of this approach with Python here.

In addition, I think that a similar behavior can be achieved with GroupIntoBatches by creating batches of maximum a a determined number of records, and then use those results to write the files.

Note that both approaches will likely affect the performance of the pipeline, as the need to track the number of records will limit the parallelism.

Here's a code that you can use as a reference. I used the WordCount quickstart and I modified it to write the Avro files. The example gives you a KV<String, Long> with the words counted, so I used it as if it were similar to your KV<item1,item2>.

static void runWordCount(WordCountOptions options) {
    Pipeline p = Pipeline.create(options);

    // Reading a file
    p.apply("ReadLines", TextIO.read().from(options.getInputFile()))

        .apply(new CountWords()) // Original data - KV<item1, item2>
        .apply(WithKeys.of("1")) // Add artificial key - KV<1, KV<item1, item2>>
        .apply(GroupIntoBatches.ofSize(1000)) // Group with max size of records per shard -> KV<1, [KV<item1, item2>, KV<item1, item2>, KV<item1, item2>]>
        .apply(ParDo.of(new SplitForFiles())) // Split to write in files -  KV<String, GenericRecord>
            .setCoder(KvCoder.of(StringUtf8Coder.of(), AvroCoder.of(GenericRecord.class, getSchema())))      
        // Write dynamically using FileIO
        .apply(FileIO.<String, KV<String, GenericRecord>> writeDynamic()
                        .by(elem -> elem.getKey()) // dest = key in the KV element
                        .via(Contextful.fn(elem -> elem.getValue()),  // Get the Generic Record
                                Contextful.fn(dest -> AvroIO.<GenericRecord>sink(getSchema())))
                        .withNumShards(1) // Limit number of shards
                        .to(options.getOutput())
                        .withNaming(key -> FileIO.Write.defaultNaming("results-", key + ".avro"))
                        .withDestinationCoder(StringUtf8Coder.of()));
        
    p.run().waitUntilFinish();
  }

Here's the method that will be used to SplitForFiles where it is defined that items1 and items2 ends up with the same key prefix, and different suffix, which will later be used to write the files:

static class SplitForFiles extends DoFn<KV<String, Iterable<KV<String, Long>>>, KV<String, GenericRecord>> {
  
    @ProcessElement
    public void processElement(@Element KV<String, Iterable<KV<String, Long>>> element, OutputReceiver<KV<String, GenericRecord>> receiver) {

        String key = getCurrentTimeStamp();
        Schema schema = getSchema();

        for (KV<String,Long> record : element.getValue()) {

            GenericRecord item1 = new GenericData.Record(schema);
            item1.put("value", record.getKey()) ;
            receiver.output(KV.of(key + "-01", item1));

            GenericRecord item2 = new GenericData.Record(schema);
            item2.put("value", record.getValue().toString()) ;
            receiver.output(KV.of(key + "-02", item2));
        }

    }
  }

The key uses the custom method getCurrentTimeStamp() which retrieves a string with the timestamp up to miliseconds. This means that if the unlikely event of two keys are generated at the same milisecond, the files might have more record that you were expecting. If this limit is critical, then I suggest to change the generation of the key, you can use unique identifiers, or even combine it with the approach mentioned above of using FinishBundle.

At the end of the pipeline, you can see that I used FileIO to write Avro files with dynamic destinations.

This was a simplified example writing Avro files with the same schema, but you can also use it as a reference to write more complex pipelines, for examble, writing to files with different schemas.