add DataFrame filter and count ops

Author: Brian Hulette

js/perf/index.js

@@ -48,6 +48,7 @@ for (let {name, buffers, tests} of require('./table_config')) {
     const dfIteratorCountSuite = new Benchmark.Suite(`DataFrame Iterator Count "${name}"`, { async: true });
     const dfDirectCountSuite = new Benchmark.Suite(`DataFrame Direct Count "${name}"`, { async: true });
     const dfScanCountSuite = new Benchmark.Suite(`DataFrame Scan Count "${name}"`, { async: true });
+    const dfFilterCountSuite = new Benchmark.Suite(`DataFrame Filter Scan Count "${name}"`, { async: true });
     const vectorCountSuite = new Benchmark.Suite(`Vector Count "${name}"`, { async: true });
     const table = Table.from(buffers);
 
@@ -58,10 +59,11 @@ for (let {name, buffers, tests} of require('./table_config')) {
         dfIteratorCountSuite.add(createDataFrameIteratorCountTest(table, test.col, test.test, test.value))
         dfDirectCountSuite.add(createDataFrameDirectCountTest(table, test.col, test.test, test.value))
         dfScanCountSuite.add(createDataFrameScanCountTest(table, test.col, test.test, test.value))
+        dfFilterCountSuite.add(createDataFrameFilterCountTest(table, test.col, test.test, test.value))
         vectorCountSuite.add(createVectorCountTest(table.columns[test.col], test.test, test.value))
     }
 
-    suites.push(tableIteratorSuite, tableCountSuite, dfIteratorSuite, dfIteratorCountSuite, dfDirectCountSuite, dfScanCountSuite, vectorCountSuite)
+    suites.push(tableIteratorSuite, tableCountSuite, dfIteratorSuite, dfIteratorCountSuite, dfDirectCountSuite, dfScanCountSuite, dfFilterCountSuite, vectorCountSuite)
 }
 
 console.log('Running apache-arrow performance tests...\n');
@@ -275,6 +277,25 @@ function createDataFrameScanCountTest(table, column, test, value) {
     };
 }
 
+function createDataFrameFilterCountTest(table, column, test, value) {
+    let df = DataFrame.from(table);
+    if (test == 'gteq') {
+        df = df.filter((idx, cols)=>cols[column].get(idx) >= value);
+    } else if (test == 'eq') {
+        df = df.filter((idx, cols)=>cols[column].get(idx) == value);
+    } else {
+        throw new Error(`Unrecognized test "${test}"`);
+    }
+
+    return {
+        async: true,
+        name: `name: '${table.columns[column].name}', length: ${table.length}, type: ${table.columns[column].type}, test: ${test}, value: ${value}`,
+        fn() {
+            df.count();
+        }
+    };
+}
+
 function createDataFrameIteratorCountTest(table, column, test, value) {
     let df = DataFrame.from(table);
 

js/src/dataframe/dataframe.ts

@@ -2,10 +2,18 @@ import { Vector } from "../vector/vector";
 import { StructVector } from "../vector/struct";
 import { VirtualVector } from "../vector/virtual";
 
+export type NextFunc = (idx: number, cols: Vector[]) => void;
+export type PredicateFunc = (idx: number, cols: Vector[]) => boolean;
+
 export abstract class DataFrame {
+    constructor(readonly lengths: Uint32Array) {}
     public abstract columns: Vector<any>[];
     public abstract getBatch(batch: number): Vector[];
-    public abstract scan(next: (idx: number, cols: Vector[])=>void): void;
+    public abstract scan(next: NextFunc): void;
+    public filter(predicate: PredicateFunc): DataFrame {
+        return new FilteredDataFrame(this, predicate);
+    }
+
     static from(table: Vector<any>): DataFrame {
         // There are two types of Vectors we might want to make into
         // a ChunkedDataFrame:
@@ -31,23 +39,26 @@ export abstract class DataFrame {
             return new SimpleDataFrame([table]);
         }
     }
+
+    count(): number {
+        return this.lengths.reduce((acc, val) => acc + val);
+    }
 }
 
 class SimpleDataFrame extends DataFrame {
     readonly lengths: Uint32Array;
     constructor(public columns: Vector<any>[]) {
-        super();
+        super(new Uint32Array([0, columns[0].length]));
         if (!this.columns.slice(1).every((v) => v.length === this.columns[0].length)) {
             throw new Error("Attempted to create a DataFrame with un-aligned vectors");
         }
-        this.lengths = new Uint32Array([0, this.columns[0].length]);
     }
 
     public getBatch() {
         return this.columns;
     }
 
-    public scan(next: (idx: number, cols: Vector[])=>void) {
+    public scan(next: NextFunc) {
         for (let idx = -1; ++idx < this.lengths[1];) {
             next(idx, this.columns)
         }
@@ -62,24 +73,16 @@ class SimpleDataFrame extends DataFrame {
 
 class ChunkedDataFrame extends DataFrame {
     public columns: Vector<any>[];
-    readonly lengths: Uint32Array;
     constructor(private virtuals: VirtualVector<any>[]) {
-        super();
-        const offsets = virtuals[0].offsets;
-        if (!this.virtuals.slice(1).every((v) => v.aligned(virtuals[0]))) {
-            throw new Error("Attempted to create a DataFrame with un-aligned vectors");
-        }
-        this.lengths = new Uint32Array(offsets.length);
-        offsets.forEach((offset, i) => {
-            this.lengths[i] = offsets[i+1] - offset;;
-        });
+        super(ChunkedDataFrame.getLengths(virtuals));
+        this.virtuals = virtuals;
     }
 
     getBatch(batch: number): Vector[] {
         return this.virtuals.map((virt) => virt.vectors[batch]);
     }
 
-    scan(next: (idx: number, cols: Vector[])=>void) {
+    scan(next: NextFunc) {
         for (let batch = -1; ++batch < this.lengths.length;) {
             const length = this.lengths[batch];
 
@@ -106,4 +109,69 @@ class ChunkedDataFrame extends DataFrame {
             }
         }
     }
+
+    private static getLengths(virtuals: VirtualVector<any>[]): Uint32Array {
+        if (!virtuals.slice(1).every((v) => v.aligned(virtuals[0]))) {
+            throw new Error("Attempted to create a DataFrame with un-aligned vectors");
+        }
+        return new Uint32Array(virtuals[0].vectors.map((v)=>v.length));
+    }
+}
+
+class FilteredDataFrame extends DataFrame {
+    public columns: Vector<any>[];
+    constructor (readonly parent: DataFrame, private predicate: PredicateFunc) {
+        super(parent.lengths);
+    }
+
+    getBatch(batch: number): Vector[] {
+        return this.parent.getBatch(batch);
+    };
+
+    scan(next: NextFunc) {
+        // inlined version of this:
+        // this.parent.scan((idx, columns) => {
+        //     if (this.predicate(idx, columns)) next(idx, columns);
+        // });
+        for (let batch = -1; ++batch < this.parent.lengths.length;) {
+            const length = this.parent.lengths[batch];
+
+            // load batches
+            const columns = this.parent.getBatch(batch);
+
+            // yield all indices
+            for (let idx = -1; ++idx < length;) {
+                if (this.predicate(idx, columns)) next(idx, columns);
+            }
+        }
+    }
+
+    count(): number {
+        // inlined version of this:
+        // let sum = 0;
+        // this.parent.scan((idx, columns) => {
+        //     if (this.predicate(idx, columns)) ++sum;
+        // });
+        // return sum;
+        let sum = 0;
+        for (let batch = -1; ++batch < this.parent.lengths.length;) {
+            const length = this.parent.lengths[batch];
+
+            // load batches
+            const columns = this.parent.getBatch(batch);
+
+            // yield all indices
+            for (let idx = -1; ++idx < length;) {
+                if (this.predicate(idx, columns)) ++sum;
+            }
+        }
+        return sum;
+    }
+
+    filter(predicate: PredicateFunc): DataFrame {
+        return new FilteredDataFrame(
+            this.parent,
+            (idx, cols) => this.predicate(idx, cols) && predicate(idx, cols)
+        );
+    }
 }