fix(components): SlidingOperator when windowSize is greater than data…

…set #26
GenSpectrum · Apr 4, 2024 · 2b9f951 · 2b9f951
1 parent 57aafc3
commit 2b9f951
Show file tree

Hide file tree

Showing 6 changed files with 70 additions and 54 deletions.
diff --git a/components/src/operator/FillMissingOperator.ts b/components/src/operator/FillMissingOperator.ts
@@ -1,16 +1,16 @@
 import { Operator } from './Operator';
 import { Dataset } from './Dataset';
 
-export class FillMissingOperator<S, K extends keyof S> implements Operator<S> {
+export class FillMissingOperator<Data, KeyToFill extends keyof Data> implements Operator<Data> {
     constructor(
-        private child: Operator<S>,
-        private keyField: K,
-        private getMinMaxFn: (values: Iterable<S[K]>) => [S[K], S[K]] | null,
-        private getAllRequiredKeysFn: (min: S[K], max: S[K]) => S[K][],
-        private defaultValueFn: (key: S[K]) => S,
+        private child: Operator<Data>,
+        private keyField: KeyToFill,
+        private getMinMaxFn: (values: Iterable<Data[KeyToFill]>) => [Data[KeyToFill], Data[KeyToFill]] | null,
+        private getAllRequiredKeysFn: (min: Data[KeyToFill], max: Data[KeyToFill]) => Data[KeyToFill][],
+        private defaultValueFn: (key: Data[KeyToFill]) => Data,
     ) {}
 
-    async evaluate(lapis: string, signal?: AbortSignal): Promise<Dataset<S>> {
+    async evaluate(lapis: string, signal?: AbortSignal): Promise<Dataset<Data>> {
         const childEvaluated = await this.child.evaluate(lapis, signal);
         const existingKeys = new Set(childEvaluated.content.map((row) => row[this.keyField]));
         const minMax = this.getMinMaxFn(existingKeys);

diff --git a/components/src/operator/GroupByAndSumOperator.ts b/components/src/operator/GroupByAndSumOperator.ts
@@ -2,25 +2,25 @@ import { GroupByOperator } from './GroupByOperator';
 import { Operator } from './Operator';
 import { NumberFields } from '../utils/type-utils';
 
-type Result<T, K extends keyof T, C extends keyof T> = {
-    [P in K | C]: P extends K ? T[K] : number;
+type Result<Data, KeyToGroupBy extends keyof Data, KeyToSumBy extends keyof Data> = {
+    [P in KeyToGroupBy | KeyToSumBy]: P extends KeyToGroupBy ? Data[KeyToGroupBy] : number;
 };
 
-export class GroupByAndSumOperator<T, K extends keyof T, C extends NumberFields<T>> extends GroupByOperator<
-    T,
-    Result<T, K, C>,
-    K
-> {
-    constructor(child: Operator<T>, groupByField: K, sumField: C) {
-        super(child, groupByField, (values: T[]) => {
+export class GroupByAndSumOperator<
+    Data,
+    KeyToGroupBy extends keyof Data,
+    KeySoSumBy extends NumberFields<Data>,
+> extends GroupByOperator<Data, Result<Data, KeyToGroupBy, KeySoSumBy>, KeyToGroupBy> {
+    constructor(child: Operator<Data>, groupByField: KeyToGroupBy, sumField: KeySoSumBy) {
+        super(child, groupByField, (values: Data[]) => {
             let n = 0;
             for (const value of values) {
                 n += value[sumField] as number;
             }
             return {
                 [groupByField]: values[0][groupByField],
                 [sumField]: n,
-            } as Result<T, K, C>;
+            } as Result<Data, KeyToGroupBy, KeySoSumBy>;
         });
     }
 }
diff --git a/components/src/operator/GroupByOperator.ts b/components/src/operator/GroupByOperator.ts
@@ -1,24 +1,26 @@
 import { Operator } from './Operator';
 import { Dataset } from './Dataset';
 
-export class GroupByOperator<T, S, K extends keyof T> implements Operator<S> {
+export class GroupByOperator<Data, AggregationResult, KeyToGroupBy extends keyof Data>
+    implements Operator<AggregationResult>
+{
     constructor(
-        private child: Operator<T>,
-        private field: K,
-        private aggregate: (values: T[]) => S,
+        private child: Operator<Data>,
+        private field: KeyToGroupBy,
+        private aggregate: (values: Data[]) => AggregationResult,
     ) {}
 
-    async evaluate(lapis: string, signal?: AbortSignal): Promise<Dataset<S>> {
+    async evaluate(lapis: string, signal?: AbortSignal): Promise<Dataset<AggregationResult>> {
         const childEvaluated = await this.child.evaluate(lapis, signal);
-        const grouped = new Map<T[K], T[]>();
+        const grouped = new Map<Data[KeyToGroupBy], Data[]>();
         for (const row of childEvaluated.content) {
             const key = row[this.field];
             if (!grouped.has(key)) {
                 grouped.set(key, []);
             }
             grouped.get(key)!.push(row);
         }
-        const result = new Array<S>();
+        const result = new Array<AggregationResult>();
         for (const [, values] of grouped) {
             result.push(this.aggregate(values));
         }

diff --git a/components/src/operator/SlidingOperator.spec.ts b/components/src/operator/SlidingOperator.spec.ts
@@ -3,31 +3,49 @@ import { SlidingOperator } from './SlidingOperator';
 import { expectEqualAfterSorting } from '../utils/test-utils';
 import { describe, it } from 'vitest';
 
+const mockOperator = new MockOperator([
+    { id: 1, value: 1 },
+    { id: 2, value: 2 },
+    { id: 3, value: 3 },
+    { id: 4, value: 4 },
+    { id: 5, value: 5 },
+]);
 describe('SlidingOperator', () => {
     it('should slide the values', async () => {
-        const child = new MockOperator([
-            { id: 1, value: 1 },
-            { id: 2, value: 2 },
-            { id: 3, value: 3 },
-            { id: 4, value: 4 },
-            { id: 5, value: 5 },
-        ]);
-        const query = new SlidingOperator(child, 3, (values) => {
-            let sum = 0;
-            for (const { value } of values) {
-                sum += value;
-            }
-            return { id: values[1].id, sum };
-        });
-        const result = await query.evaluate('lapis');
-        await expectEqualAfterSorting(
+        const underTest = getSlidingOperatorWithWindowSize(3);
+
+        const result = await underTest.evaluate('lapis');
+
+        expectEqualAfterSorting(
             result.content,
             [
                 { id: 2, sum: 6 },
                 { id: 3, sum: 9 },
                 { id: 4, sum: 12 },
             ],
-            (a, b) => a.id - b.id,
+            sortById,
         );
     });
+
+    it('should return single value when window size is greater than number of entries', async () => {
+        const underTest = getSlidingOperatorWithWindowSize(999);
+
+        const result = await underTest.evaluate('lapis');
+
+        expectEqualAfterSorting(result.content, [{ id: 2, sum: 15 }], sortById);
+    });
+
+    function getSlidingOperatorWithWindowSize(windowSize: number) {
+        return new SlidingOperator(mockOperator, windowSize, (values) => {
+            let sum = 0;
+            for (const { value } of values) {
+                sum += value;
+            }
+            return { id: values[1].id, sum };
+        });
+    }
+
+    function sortById(a: { id: number }, b: { id: number }) {
+        return a.id - b.id;
+    }
 });
diff --git a/components/src/operator/SlidingOperator.ts b/components/src/operator/SlidingOperator.ts
@@ -1,21 +1,21 @@
 import { Operator } from './Operator';
-import { Dataset } from './Dataset';
 
-export class SlidingOperator<T, S> implements Operator<S> {
+export class SlidingOperator<Data, AggregationResult> implements Operator<AggregationResult> {
     constructor(
-        private child: Operator<T>,
+        private child: Operator<Data>,
         private windowSize: number,
-        private aggregate: (values: T[]) => S,
+        private aggregate: (values: Data[]) => AggregationResult,
     ) {
         if (windowSize < 1) {
             throw new Error('Window size must be at least 1');
         }
     }
 
-    async evaluate(lapis: string, signal?: AbortSignal): Promise<Dataset<S>> {
+    async evaluate(lapis: string, signal?: AbortSignal) {
         const childEvaluated = await this.child.evaluate(lapis, signal);
-        const content = new Array<S>();
-        for (let i = 0; i < childEvaluated.content.length - this.windowSize + 1; i++) {
+        const content = new Array<AggregationResult>();
+        const numberOfWindows = Math.max(childEvaluated.content.length - this.windowSize, 0) + 1;
+        for (let i = 0; i < numberOfWindows; i++) {
             content.push(this.aggregate(childEvaluated.content.slice(i, i + this.windowSize)));
         }
         return { content };

diff --git a/components/src/query/queryPrevalenceOverTime.ts b/components/src/query/queryPrevalenceOverTime.ts
@@ -4,7 +4,6 @@ import { MapOperator } from '../operator/MapOperator';
 import { GroupByAndSumOperator } from '../operator/GroupByAndSumOperator';
 import { FillMissingOperator } from '../operator/FillMissingOperator';
 import { SortOperator } from '../operator/SortOperator';
-import { Operator } from '../operator/Operator';
 import { SlidingOperator } from '../operator/SlidingOperator';
 import { DivisionOperator } from '../operator/DivisionOperator';
 import { compareTemporal, generateAllInRange, getMinMaxTemporal, Temporal, TemporalCache } from '../utils/temporal';
@@ -66,11 +65,8 @@ function fetchAndPrepare(filter: LapisFilter, granularity: TemporalGranularity,
         (key) => ({ dateRange: key, count: 0 }),
     );
     const sortData = new SortOperator(fillData, dateRangeCompare);
-    let smoothData: Operator<{ dateRange: Temporal | null; count: number }> = sortData;
-    if (smoothingWindow >= 1) {
-        smoothData = new SlidingOperator(sortData, smoothingWindow, averageSmoothing);
-    }
-    return smoothData;
+
+    return smoothingWindow >= 1 ? new SlidingOperator(sortData, smoothingWindow, averageSmoothing) : sortData;
 }
 
 function mapDateToGranularityRange(d: { date: string | null; count: number }, granularity: TemporalGranularity) {