Merge pull request #26 from scramjetorg/feature/batch

Introduce "DataStream.batch()"

Author: f1ames

src/streams/base-stream.ts

@@ -3,5 +3,6 @@ export interface BaseStream<T extends any> {
     map<U, W extends any[]>(callback: TransformFunction<T, U, W>, ...args: W): BaseStream<U>;
     flatMap<U, W extends any[]>(callback: TransformFunction<T, AnyIterable<U>, W>, ...args: W): BaseStream<U>;
     filter<W extends any[]>(callback: TransformFunction<T, Boolean, W>, ...args: W): BaseStream<T>;
-    reduce<U = T>(callback: (previousValue: U, currentChunk: T) => MaybePromise<U>, initial?: U): Promise<U>
+    batch<W extends any[] = []>(callback: TransformFunction<T, Boolean, W>, ...args: W): BaseStream<T[]>;
+    reduce<U = T>(callback: (previousValue: U, currentChunk: T) => MaybePromise<U>, initial?: U): Promise<U>;
 }

src/streams/data-stream.ts

@@ -74,6 +74,45 @@ export class DataStream<T> implements BaseStream<T>, AsyncIterable<T> {
         return this.asNewFlattenedStream(this.map<AnyIterable<U>, W>(callback, ...args));
     }
 
+    batch<W extends any[] = []>(callback: TransformFunction<T, Boolean, W>, ...args: W): DataStream<T[]> {
+        let currentBatch: T[] = [];
+        let aggregator: TransformFunction<T, T[], W>;
+
+        if (isAsyncFunction(callback)) {
+            aggregator = async (chunk: T, ...args1: W): Promise<T[]> => {
+                currentBatch.push(chunk);
+
+                let result: T[] = [];
+
+                if (await callback(chunk, ...args1)) {
+                    result = [...currentBatch];
+                    currentBatch = [];
+                }
+
+                return result;
+            };
+        } else {
+            aggregator = (chunk: T, ...args1: W): T[] => {
+                currentBatch.push(chunk);
+
+                let result: T[] = [];
+
+                if (callback(chunk, ...args1)) {
+                    result = [...currentBatch];
+                    currentBatch = [];
+                }
+
+                return result;
+            };
+        }
+
+        const onEnd = () => {
+            return { yield: currentBatch.length > 0, value: currentBatch };
+        };
+
+        return this.asNewStream(this.map<T[], W>(aggregator, ...args).filter(chunk => chunk.length > 0), onEnd);
+    }
+
     async reduce<U = T>(callback: (previousValue: U, currentChunk: T) => MaybePromise<U>, initial?: U): Promise<U> {
         // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Array/Reduce#parameters
         //
@@ -152,6 +191,25 @@ export class DataStream<T> implements BaseStream<T>, AsyncIterable<T> {
         return reducer as Reducer<T, U>;
     }
 
+    protected asNewStream<U, W extends DataStream<U>>(
+        fromStream: W,
+        onEndYield?: () => { yield: boolean, value?: U }
+    ): DataStream<U> {
+        return DataStream.from((async function * (stream){
+            for await (const chunk of stream) {
+                yield chunk;
+            }
+
+            if (onEndYield) {
+                const yieldValue = onEndYield();
+
+                if (yieldValue.yield) {
+                    yield yieldValue.value as U;
+                }
+            }
+        })(fromStream));
+    }
+
     protected asNewFlattenedStream<U, W extends DataStream<AnyIterable<U>>>(
         fromStream: W,
         onEndYield?: () => { yield: boolean, value?: U }

test/streams/data/batch.spec.ts

@@ -0,0 +1,45 @@
+import test from "ava";
+import { DataStream } from "../../../src/streams/data-stream";
+import { deferReturn } from "../../helpers/utils";
+
+test("DataStream batch can make sentences from words", async (t) => {
+    const result = await DataStream
+        .from(["foo", "bar.", "baz", "bax", ".", "foo"])
+        .batch(chunk => chunk.endsWith("."))
+        .toArray();
+
+    t.deepEqual(result, [["foo", "bar."], ["baz", "bax", "."], ["foo"]]);
+});
+
+test("DataStream batch can make sentences from words (async)", async (t) => {
+    const result = await DataStream
+        .from(["foo", "bar.", "baz", "bax", ".", "foo"])
+        .batch(async (chunk) => deferReturn(5, chunk.endsWith(".")))
+        .toArray();
+
+    t.deepEqual(result, [["foo", "bar."], ["baz", "bax", "."], ["foo"]]);
+});
+
+test("DataStream batch can bu used to batch by amount (via variadic arg counter)", async (t) => {
+    const result = await DataStream
+        .from([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+        .batch((chunk, counter) => { counter.i++; return counter.i % 3 === 0; }, { i: 0 })
+        .toArray();
+
+    t.deepEqual(result, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]);
+});
+
+test("DataStream batch does not deep copy chunks", async (t) => {
+    const input = [{ id: 0, data: "foo" }, { id: 1, data: "bar" }, { id: 2, data: "baz" }, { id: 3, data: "bax" }];
+    const result = await DataStream
+        .from(input)
+        .batch((chunk) => chunk.id % 2 !== 0)
+        .toArray();
+
+    t.deepEqual(result, [[{ id: 0, data: "foo" }, { id: 1, data: "bar" }], [{ id: 2, data: "baz" }, { id: 3, data: "bax" }]]);
+
+    input[0].data = "changed1";
+    input[3].data = "changed2";
+
+    t.deepEqual(result, [[{ id: 0, data: "changed1" }, { id: 1, data: "bar" }], [{ id: 2, data: "baz" }, { id: 3, data: "changed2" }]]);
+});