diff --git a/js/src/Arrow.ts b/js/src/Arrow.ts index eefed2d9a3099..e58aa69865bb3 100644 --- a/js/src/Arrow.ts +++ b/js/src/Arrow.ts @@ -25,7 +25,7 @@ import * as predicate_ from './predicate'; import { Vector } from './vector'; import { RecordBatch } from './recordbatch'; import { Schema, Field, Type } from './type'; -import { Table, DataFrame, NextFunc, CountByResult } from './table'; +import { Table, DataFrame, NextFunc, BindFunc, CountByResult } from './table'; import { read, readAsync } from './ipc/reader/arrow'; export import View = vector_.View; @@ -36,7 +36,7 @@ export import TimeBitWidth = type_.TimeBitWidth; export import TypedArrayConstructor = type_.TypedArrayConstructor; export { read, readAsync }; -export { Table, DataFrame, NextFunc, CountByResult }; +export { Table, DataFrame, NextFunc, BindFunc, CountByResult }; export { Field, Schema, RecordBatch, Vector, Type }; export namespace util { diff --git a/js/src/table.ts b/js/src/table.ts index 193c947a22db7..b9fdeb99c64e4 100644 --- a/js/src/table.ts +++ b/js/src/table.ts @@ -23,11 +23,12 @@ import { isPromise, isAsyncIterable } from './util/compat'; import { Vector, DictionaryVector, IntVector, StructVector } from './vector'; import { ChunkedView } from './vector/chunked'; -export type NextFunc = (idx: number, cols: RecordBatch) => void; +export type NextFunc = (idx: number, batch: RecordBatch) => void; +export type BindFunc = (batch: RecordBatch) => void; export interface DataFrame { filter(predicate: Predicate): DataFrame; - scan(next: NextFunc): void; + scan(next: NextFunc, bind?: BindFunc): void; count(): number; countBy(col: (Col|string)): CountByResult; } @@ -128,11 +129,12 @@ export class Table implements DataFrame { public filter(predicate: Predicate): DataFrame { return new FilteredDataFrame(this.batches, predicate); } - public scan(next: NextFunc) { + public scan(next: NextFunc, bind?: BindFunc) { const batches = this.batches, numBatches = batches.length; for (let batchIndex = -1; ++batchIndex < numBatches;) { // load batches const batch = batches[batchIndex]; + if (bind) { bind(batch); } // yield all indices for (let index = -1, numRows = batch.length; ++index < numRows;) { next(index, batch); @@ -189,7 +191,7 @@ class FilteredDataFrame implements DataFrame { this.batches = batches; this.predicate = predicate; } - public scan(next: NextFunc) { + public scan(next: NextFunc, bind?: BindFunc) { // inlined version of this: // this.parent.scan((idx, columns) => { // if (this.predicate(idx, columns)) next(idx, columns); @@ -199,6 +201,7 @@ class FilteredDataFrame implements DataFrame { for (let batchIndex = -1; ++batchIndex < numBatches;) { // load batches const batch = batches[batchIndex]; + if (bind) { bind(batch); } const predicate = this.predicate.bind(batch); // yield all indices for (let index = -1, numRows = batch.length; ++index < numRows;) {