batch-stream.js 583 B

12345678910111213141516171819202122232425262728293031323334
  1. const { Transform } = require('stream');
  2. function createBatchStream(batchSize) {
  3. let batchBuffer = [];
  4. return new Transform({
  5. // object mode
  6. objectMode: true,
  7. transform(doc, encoding, callback) {
  8. batchBuffer.push(doc);
  9. if (batchBuffer.length >= batchSize) {
  10. this.push(batchBuffer);
  11. // reset buffer
  12. batchBuffer = [];
  13. }
  14. callback();
  15. },
  16. final(callback) {
  17. if (batchBuffer.length > 0) {
  18. this.push(batchBuffer);
  19. }
  20. callback();
  21. },
  22. });
  23. }
  24. module.exports = {
  25. createBatchStream,
  26. };