How can I consume an iterable in batches (equally sized chunks)?

Question

I am often using batch() in Python. Is there some alternative in JavaScript since ES6 which has iterators and generator functions?

Answer 1

I had to write one for myself, which I'm sharing here for me and the others to find here easily:

// subsequently yield iterators of given `size`
// these have to be fully consumed
function* batches(iterable, size) {
  const it = iterable[Symbol.iterator]();
  while (true) {
    // this is for the case when batch ends at the end of iterable
    // (we don't want to yield empty batch)
    let {value, done} = it.next();
    if (done) return value;

    yield function*() {
      yield value;
      for (let curr = 1; curr < size; curr++) {
        ({value, done} = it.next());
        if (done) return;

        yield value;
      }
    }();
    if (done) return value;
  }
}

It yields generators, not Arrays for example. You have to fully consume each batch before calling next() on it again.

Answer 2

Came here looking to see what other people had suggested. Here's the version I wrote in TypeScript initially before looking at this post.

async function* batch<T>(iterable: AsyncIterableIterator<T>, batchSize: number) {
  let items: T[] = [];
  for await (const item of iterable) {
    items.push(item);
    if (items.length >= batchSize) {
      yield items;
      items = []
    }
  }
  if (items.length !== 0) {
    yield items;
  }
}

This allows you to consume an iterable in batches as shown below.

async function doYourThing<T>(iterable: AsyncIterableIterator<T>) {
  const itemsPerBatch = 5
  const batchedIterable = batch<T>(iterable, itemsPerBatch)
  for await (const items of batchedIterable) {
    await someOperation(items)
  }
}

In my case, this allowed me to use bulkOps in Mongo a little more easily as demonstrated below.

import { MongoClient, ObjectID } from 'mongodb';
import { batch } from './batch';

const config = {
  mongoUri: 'mongodb://localhost:27017/test?replicaSet=rs0',
};

interface Doc {
  readonly _id: ObjectID;
  readonly test: number;
}

async function main() {
  const client = await MongoClient.connect(config.mongoUri);
  const db = client.db('test');
  const coll = db.collection<Doc>('test');
  await coll.deleteMany({});
  console.log('Deleted test docs');

  const testDocs = new Array(4).fill(null).map(() => ({ test: 1 }));
  await coll.insertMany(testDocs);
  console.log('Inserted test docs');

  const cursor = coll.find().batchSize(5);
  for await (const docs of batch<Doc>(cursor as any, 5)) {
    const bulkOp = coll.initializeUnorderedBulkOp();
    docs.forEach((doc) => {
      bulkOp.find({ _id: doc._id }).updateOne({ test: 2 });
    });
    console.log('Updating', docs.length, 'test docs');
    await bulkOp.execute();
  }
  console.log('Updated test docs');
}

main()
  .catch(console.error)
  .then(() => process.exit());