Skip to content

Commit

Permalink
Add table unit tests...
Browse files Browse the repository at this point in the history
.. also found and resolved some minor bugs (get(idx) batch length check should
be <=, various extern issues with UMD builds)
  • Loading branch information
Brian Hulette committed Jan 15, 2018
1 parent 6719147 commit 7244887
Show file tree
Hide file tree
Showing 4 changed files with 411 additions and 11 deletions.
10 changes: 10 additions & 0 deletions js/src/Arrow.externs.ts
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,16 @@ Table.prototype.toString;
Table.prototype.lengths;
/** @type {?} */
Table.prototype.batches;
/** @type {?} */
Table.prototype.countBy;
/** @type {?} */
Table.prototype.scan;
/** @type {?} */
Table.prototype.get;

let CountByResult = function() {};
/** @type {?} */
CountByResult.prototype.asJSON;

let Vector = function() {};
/** @type {?} */
Expand Down
13 changes: 8 additions & 5 deletions js/src/Arrow.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@
// specific language governing permissions and limitations
// under the License.

import { Table, TableRow } from './table';
import { lit, col } from './predicate';
import { Table, TableRow, CountByResult } from './table';
import { lit, col, Col, Value } from './predicate';
import { Vector } from './vector/vector';
import { Utf8Vector } from './vector/utf8';
import { DictionaryVector } from './vector/dictionary';
Expand Down Expand Up @@ -54,8 +54,8 @@ Table['fromAsync'] = Table.fromAsync;
BoolVector['pack'] = BoolVector.pack;

export { read, readAsync };
export { Table, TableRow };
export { lit, col };
export { Table, TableRow, CountByResult };
export { lit, col, Col, Value };
export { Vector, StructRow };
export { Uint64, Int64, Int128 };
export { NumericVectorConstructor } from './vector/numeric';
Expand Down Expand Up @@ -94,9 +94,11 @@ try {
// string indexers tell closure compiler not to rename these properties
Arrow['lit'] = lit;
Arrow['col'] = col;
Arrow['Col'] = Col;
Arrow['read'] = read;
Arrow['readAsync'] = readAsync;
Arrow['Value'] = Value;
Arrow['Table'] = Table;
Arrow['readAsync'] = readAsync;
Arrow['Vector'] = Vector;
Arrow['StructRow'] = StructRow;
Arrow['BoolVector'] = BoolVector;
Expand All @@ -120,6 +122,7 @@ try {
Arrow['Float32Vector'] = Float32Vector;
Arrow['Float64Vector'] = Float64Vector;
Arrow['DecimalVector'] = DecimalVector;
Arrow['CountByResult'] = CountByResult;
Arrow['TimestampVector'] = TimestampVector;
Arrow['DictionaryVector'] = DictionaryVector;
Arrow['FixedSizeListVector'] = FixedSizeListVector;
Expand Down
28 changes: 22 additions & 6 deletions js/src/table.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ export interface DataFrame {
filter(predicate: Predicate): DataFrame;
scan(next: NextFunc): void;
count(): number;
countBy(col: (Col|string)): Table;
countBy(col: (Col|string)): CountByResult;
}

function columnsFromBatches(batches: Vector[][]) {
Expand Down Expand Up @@ -87,7 +87,7 @@ export class Table implements DataFrame {
}
get(idx: number): TableRow {
let batch = 0;
while (idx > this.lengths[batch] && batch < this.lengths.length) {
while (idx >= this.lengths[batch] && batch < this.lengths.length) {
idx -= this.lengths[batch++];
}

Expand All @@ -114,7 +114,7 @@ export class Table implements DataFrame {
count(): number {
return this.lengths.reduce((acc, val) => acc + val);
}
countBy(count_by: (Col|string)): Table {
countBy(count_by: (Col|string)): CountByResult {
if (count_by instanceof String) {
count_by = new Col(count_by);
}
Expand Down Expand Up @@ -146,7 +146,7 @@ export class Table implements DataFrame {
}
}

return new Table({batches: [[keys, new Uint32Vector({data: counts})]]})
return new CountByResult(keys, new Uint32Vector({data: counts}))
}
*[Symbol.iterator]() {
for (let batch = -1; ++batch < this.lengths.length;) {
Expand Down Expand Up @@ -215,7 +215,7 @@ class FilteredDataFrame implements DataFrame {
);
}

countBy(count_by: (Col|string)): Table {
countBy(count_by: (Col|string)): CountByResult {
if (count_by instanceof String) {
count_by = new Col(count_by);
}
Expand Down Expand Up @@ -246,6 +246,22 @@ class FilteredDataFrame implements DataFrame {
}
}

return new Table({batches: [[keys, new Uint32Vector({data: counts})]]})
return new CountByResult(keys, new Uint32Vector({data: counts}))
}
}

export class CountByResult extends Table implements DataFrame {
constructor(readonly keys: Vector, readonly counts: Vector<number|null>) {
super({batches: [[keys, counts]]});
}

asJSON(): Object {
let result: {[key: string]: number|null} = {};

for (let i = -1; ++i < this.length;) {
result[this.keys.get(i)] = this.counts.get(i);
}

return result;
}
}
Loading

0 comments on commit 7244887

Please sign in to comment.