Skip to content

Commit

Permalink
ARROW-2159: [JS] Support custom predicates
Browse files Browse the repository at this point in the history
Adds support for defining a custom predicate with callback functions, like so:
```js
table.filter(predicate.custom(
    (idx) => {...},
    (batch) => {...}
));
```

Author: Brian Hulette <[email protected]>

Closes apache#1616 from TheNeuralBit/custom-predicate and squashes the following commits:

77f5441 [Brian Hulette] Add test, squash bug
29ecc80 [Brian Hulette] Add CustomPredicate
  • Loading branch information
Brian Hulette committed Feb 19, 2018
1 parent 8e90836 commit 77f2841
Show file tree
Hide file tree
Showing 4 changed files with 29 additions and 3 deletions.
3 changes: 2 additions & 1 deletion js/src/Arrow.externs.js
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ CountByResult.prototype.asJSON;

var col = function () {};
var lit = function () {};
var custom = function () {};

var Value = function() {};
/** @type {?} */
Expand Down Expand Up @@ -738,4 +739,4 @@ VectorVisitor.prototype.visitInterval;
/** @type {?} */
VectorVisitor.prototype.visitFixedSizeList;
/** @type {?} */
VectorVisitor.prototype.visitMap;
VectorVisitor.prototype.visitMap;
1 change: 1 addition & 0 deletions js/src/Arrow.ts
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@ export namespace view {
export namespace predicate {
export import col = predicate_.col;
export import lit = predicate_.lit;
export import custom = predicate_.custom;

export import Or = predicate_.Or;
export import Col = predicate_.Col;
Expand Down
14 changes: 14 additions & 0 deletions js/src/predicate.ts
Original file line number Diff line number Diff line change
Expand Up @@ -222,5 +222,19 @@ export class GTeq extends ComparisonPredicate {
}
}

export class CustomPredicate extends Predicate {
constructor(private next: PredicateFunc, private bind_: (batch: RecordBatch) => void) {
super();
}

bind(batch: RecordBatch) {
this.bind_(batch);
return this.next;
}
}

export function lit(v: any): Value<any> { return new Literal(v); }
export function col(n: string): Col<any> { return new Col(n); }
export function custom(next: PredicateFunc, bind: (batch: RecordBatch) => void) {
return new CustomPredicate(next, bind);
}
14 changes: 12 additions & 2 deletions js/test/unit/table-tests.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,11 @@
// specific language governing permissions and limitations
// under the License.

import Arrow from '../Arrow';
import Arrow, { RecordBatch } from '../Arrow';

const { predicate, Table } = Arrow;

const { col, lit } = predicate;
const { col, lit, custom } = predicate;

const F32 = 0, I32 = 1, DICT = 2;
const test_data = [
Expand Down Expand Up @@ -323,6 +323,7 @@ describe(`Table`, () => {
expect(table.getColumnIndex('f32')).toEqual(F32);
expect(table.getColumnIndex('dictionary')).toEqual(DICT);
});
let get_i32: (idx: number) => number, get_f32: (idx: number) => number;
const filter_tests = [
{
name: `filter on f32 >= 0`,
Expand Down Expand Up @@ -364,6 +365,15 @@ describe(`Table`, () => {
name: `filter on f32 <= i32`,
filtered: table.filter(col('f32').lteq(col('i32'))),
expected: values.filter((row) => row[F32] <= row[I32])
}, {
name: `filter on f32*i32 > 0 (custom predicate)`,
filtered: table.filter(custom(
(idx: number) => (get_f32(idx) * get_i32(idx) > 0),
(batch: RecordBatch) => {
get_f32 = col('f32').bind(batch);
get_i32 = col('i32').bind(batch);
})),
expected: values.filter((row) => (row[F32] as number) * (row[I32] as number) > 0)
}
];
for (let this_test of filter_tests) {
Expand Down

0 comments on commit 77f2841

Please sign in to comment.