From 26e3a7d3d6ac815f6e0fc529dcce721ec215a160 Mon Sep 17 00:00:00 2001 From: Ardis Lu Date: Fri, 17 Jan 2025 23:52:29 -0800 Subject: [PATCH] Add maxBatchSize option Although this functionality could be implemented by end-users by returning null in the transformation function until a batch is accumulated, baking it into the transform stream simplifies this common use case. --- examples/README.md | 4 +++- examples/data/ex7-in.csv | 13 +++++++++++ examples/data/ex7-out.csv | 13 +++++++++++ examples/ex7.js | 33 ++++++++++++++++++++++++++++ src/core.js | 45 +++++++++++++++++++++++++++++++++------ test/core.test.js | 45 +++++++++++++++++++++++++++++++++++++++ 6 files changed, 146 insertions(+), 7 deletions(-) create mode 100644 examples/data/ex7-in.csv create mode 100644 examples/data/ex7-out.csv create mode 100644 examples/ex7.js diff --git a/examples/README.md b/examples/README.md index ec07240..c27eddd 100644 --- a/examples/README.md +++ b/examples/README.md @@ -20,4 +20,6 @@ node ./examples/ex1.js [Example 5: Handling an irregularly-shaped CSV file](./ex5.js) -[Example 6: Handling a flaky transformation (basic retry strategy)](./ex6.js) \ No newline at end of file +[Example 6: Handling a flaky transformation (basic retry strategy)](./ex6.js) + +[Example 7: Processing rows in batches (multi-row input)](./ex7.js) \ No newline at end of file diff --git a/examples/data/ex7-in.csv b/examples/data/ex7-in.csv new file mode 100644 index 0000000..0edcd62 --- /dev/null +++ b/examples/data/ex7-in.csv @@ -0,0 +1,13 @@ +columnA,ColumnB +1,1 +2,2 +3,3 +4,4 +5,5 +6,6 +7,7 +8,8 +9,9 +10,10 +11,11 +12,12 diff --git a/examples/data/ex7-out.csv b/examples/data/ex7-out.csv new file mode 100644 index 0000000..7e00f84 --- /dev/null +++ b/examples/data/ex7-out.csv @@ -0,0 +1,13 @@ +batch number,sum of batch,columnA,ColumnB +1,30,1,1 +1,30,2,2 +1,30,3,3 +1,30,4,4 +1,30,5,5 +2,80,6,6 +2,80,7,7 +2,80,8,8 +2,80,9,9 +2,80,10,10 +3,46,11,11 +3,46,12,12 diff --git a/examples/ex7.js b/examples/ex7.js new file mode 100644 index 0000000..81c4c2b --- /dev/null +++ b/examples/ex7.js @@ -0,0 +1,33 @@ +// ## Example 7: Processing rows in batches (multi-row input) +// +// Set the `maxBatchSize` option to a number greater than `1` to collect rows into batches before passing the whole +// batch to the transformation function. +// +// This functionality is useful if you need to collect data from multiple rows before processing (e.g., performing an +// expensive network request where is more efficient to batch multiple rows into one request). + +import { createCSVReadableStream, createCSVTransformStream, createCSVWritableStream } from '../src/index.js'; + +let firstChunk = true; +let batchNumber = 0; +function process(batch) { + if (firstChunk) { // Header row is always passed by itself, regardless of maxBatchSize + firstChunk = false; + return ['batch number', 'sum of batch', ...batch]; // Note the header row is Array, NOT Array> + } + + // Some logic that requires data across multiple rows + let sum = 0; + for (const row of batch) { + for (const field of row) { + sum += Number(field); + } + } + + batchNumber++; + return batch.map(row => [batchNumber, sum, ...row]); +} + +await createCSVReadableStream(new URL('./data/ex7-in.csv', import.meta.url)) + .pipeThrough(createCSVTransformStream(process, { includeHeaders: true, maxBatchSize: 5 })) + .pipeTo(createCSVWritableStream(new URL('./data/ex7-out.csv', import.meta.url))); \ No newline at end of file diff --git a/src/core.js b/src/core.js index f25b3cd..fdb79fd 100644 --- a/src/core.js +++ b/src/core.js @@ -128,8 +128,8 @@ export function createCSVReadableStream(path) { } /** - * A row of input CSV data. - * @typedef {Array|string} TransformationInput + * A row or rows of input CSV data. + * @typedef {Array>|Array|string} TransformationInput */ /** @@ -138,7 +138,7 @@ export function createCSVReadableStream(path) { */ /** - * A function to process a row of CSV data from `createCSVReadableStream`. + * A function to process a row or rows of CSV data from `createCSVReadableStream`. * @callback TransformationFunction * @param {TransformationInput} row * @returns {TransformationOutput} @@ -158,20 +158,27 @@ export function createCSVReadableStream(path) { * @property {null|function(TransformationInput,Error,TransformationFunction):TransformationOutput} [onError=null] Set to a function to catch * errors thrown by the transformation function. The input data, the error that was thrown, and the transformation function itself will be passed * to the `onError` function. The default value is `null` (errors will not be caught). + * @property {Number} [maxBatchSize=1] The maximum number of rows that will be passed to the transformation function per function call (greedy). + * The default value is `1`. */ /** * Create a new [TransformStream](https://developer.mozilla.org/en-US/docs/Web/API/TransformStream) to process CSV data. * - * @param {TransformationFunction} fn A function to process a row of CSV data from `createCSVReadableStream`. + * @param {TransformationFunction} fn A function to process a row or rows of CSV data from `createCSVReadableStream`. * - * If `options.rawInput` is `false` (default), the input will be a `Array` representing the CSV row. If `options.rawInput` is `true`, - * the input will be a JSON `string` representing the CSV row. + * If `options.rawInput` is `false` (default), the input will be a `Array` representing a CSV row. If `options.rawInput` is `true`, + * the input will be a JSON `string` representing a CSV row. * * If `options.rawOutput` is `false` (default), the expected output is a `Array` which will be serialized with `JSON.stringify()` and enqueued, * or a `Array>` which will have each of its sub-arrays serialized and enqueued separately. If `options.rawOutput` is `true`, the raw * output will be sent the next stream unmodified. * + * If `options.maxBatchSize` is a number greater than `1`, the input will represent multiple CSV rows as an `Array>` (if + * `options.rawInput` is `false`) or a `string` (if `options.rawInput` is `true`). The input will still be a 2-D array even if there is only + * 1 row in the batch (except for the header row). The header row is always processed by itself and passed as a `Array`, regardless of + * the `options.maxBatchSize` value. + * * Return `null` to consume an input row without emitting an output row. * @param {CreateCSVTransformStreamOptions} options Object containing flags to configure the stream logic. * @returns {TransformStream} A `TransformStream` where each chunk is one row of the CSV file, after transformations applied by `fn()`. @@ -181,6 +188,7 @@ export function createCSVTransformStream(fn, options = {}) { options.rawInput ??= false; options.rawOutput ??= false; options.onError ??= null; + options.maxBatchSize ??= 1; async function wrappedFn(row) { if (options.onError === null) { return await fn(row); } @@ -188,6 +196,7 @@ export function createCSVTransformStream(fn, options = {}) { catch (e) { return await options.onError(row, e, fn); } } + const batch = []; let firstChunk = true; return new TransformStream({ async transform(chunk, controller) { @@ -197,6 +206,13 @@ export function createCSVTransformStream(fn, options = {}) { firstChunk = false; out = options.includeHeaders ? await wrappedFn(row) : row; } + else if (options.maxBatchSize > 1) { + batch.push(row); + if (batch.length === options.maxBatchSize) { + out = await wrappedFn(batch); + batch.length = 0; + } + } else { out = await wrappedFn(row); } @@ -214,6 +230,23 @@ export function createCSVTransformStream(fn, options = {}) { else { controller.enqueue(JSON.stringify(out)); } + }, + async flush(controller) { + if (batch.length > 0) { // CSV ended before batch got filled + const out = await wrappedFn(batch); + if (out === null || out === undefined) { // Input row is consumed without emitting any output row + return; + } + if (options.rawOutput || typeof out === 'string') { + controller.enqueue(out); + } + else if (Array.isArray(out) && Array.isArray(out[0])) { // Multiple rows returned, enqueue each row separately + out.forEach(row => controller.enqueue(JSON.stringify(row))); + } + else { + controller.enqueue(JSON.stringify(out)); + } + } } }); } diff --git a/test/core.test.js b/test/core.test.js index 35ae2dd..3cf8140 100644 --- a/test/core.test.js +++ b/test/core.test.js @@ -272,6 +272,51 @@ suite('createCSVTransformStream', { concurrency: true }, () => { { name: 'Error', message: 'inner' } ); }); + test('creates one batch when maxBatchSize is greater than CSV row length', { concurrency: true }, async () => { + await createCSVMockStream([ + ['columnA', 'columnB'], + ['1', '2'], + ['3', '4'], + ['5', '6'] + ]) + .pipeThrough(createCSVTransformStream(b => b.map(r => r.map(f => Number(f) * 2)), { maxBatchSize: 100 })) + .pipeTo(csvStreamEqualWritable([ + ['columnA', 'columnB'], + [2, 4], + [6, 8], + [10, 12] + ])); + }); + test('creates one batch when maxBatchSize is equal to CSV row length', { concurrency: true }, async () => { + await createCSVMockStream([ + ['columnA', 'columnB'], + ['1', '2'], + ['3', '4'], + ['5', '6'] + ]) + .pipeThrough(createCSVTransformStream(b => b.map(r => r.map(f => Number(f) * 2)), { maxBatchSize: 3 })) + .pipeTo(csvStreamEqualWritable([ + ['columnA', 'columnB'], + [2, 4], + [6, 8], + [10, 12] + ])); + }); + test('creates multiple batches when maxBatchSize is less than CSV row length', { concurrency: true }, async () => { + await createCSVMockStream([ + ['columnA', 'columnB'], + ['1', '2'], + ['3', '4'], + ['5', '6'] + ]) + .pipeThrough(createCSVTransformStream(b => b.map(r => r.map(f => Number(f) * 2)), { maxBatchSize: 2 })) + .pipeTo(csvStreamEqualWritable([ + ['columnA', 'columnB'], + [2, 4], + [6, 8], + [10, 12] + ])); + }); }); suite('createCSVWritableStream', { concurrency: true }, () => { -- 2.47.3