diff --git a/js/test/unit/table-tests.ts b/js/test/unit/table-tests.ts index f346ed61056e9..d745908a48745 100644 --- a/js/test/unit/table-tests.ts +++ b/js/test/unit/table-tests.ts @@ -21,13 +21,10 @@ const { predicate, Table } = Arrow; const { col } = predicate; -describe(`Table`, () => { - test(`can create an empty table`, () => { - expect(Table.empty().length).toEqual(0); - }); - - describe(`single record batch`, () => { - const table = Table.from({ +const F32 = 0, I32 = 1, DICT = 2; +const test_data = [ + {name: `single record batch`, + table: Table.from({ 'schema': { 'fields': [ { @@ -115,79 +112,19 @@ describe(`Table`, () => { } ] }] - }); - - // Wrap floating point values in a Float32Array and take them back out to - // make sure that equality checks will pass - const values = [ - [Math.fround(-0.3), -1, 'a'], - [Math.fround(-0.2), 1, 'b'], - [Math.fround(-0.1), -1, 'c'], - [Math.fround( 0 ), 1, 'a'], - [Math.fround( 0.1), -1, 'b'], - [Math.fround( 0.2), 1, 'c'], - [Math.fround( 0.3), -1, 'a'] - ]; - test(`has the correct length`, () => { - expect(table.length).toEqual(values.length); - }); - test(`gets expected values`, () => { - for (let i = -1; ++i < values.length;) { - expect(table.get(i).toArray()).toEqual(values[i]); - } - }); - test(`iterates expected values`, () => { - let i = 0; - for (let row of table) { - expect(row.toArray()).toEqual(values[i++]); - } - }); - test(`scans expected values`, () => { - let expected_idx = 0; - table.scan((idx, batch) => { - const columns = batch.schema.fields.map((_, i) => batch.getChildAt(i)!); - expect(columns.map((c) => c.get(idx))).toEqual(values[expected_idx++]); - }); - }); - test(`count() returns the correct length`, () => { - expect(table.count()).toEqual(values.length); - }); - test(`filter on f32 >= 0 returns the correct length`, () => { - expect(table.filter(col('f32').gteq(0)).count()).toEqual(4); - }); - test(`filter on i32 <= 0 returns the correct length`, () => { - expect(table.filter(col('i32').lteq(0)).count()).toEqual(4); - }); - test(`filter on dictionary == 'a' returns the correct length`, () => { - expect(table.filter(col('dictionary').eq('a')).count()).toEqual(3); - }); - test(`countBy on dictionary returns the correct counts`, () => { - // Make sure countBy works both with and without the Col wrapper - // class - expect(table.countBy(col('dictionary')).toJSON()).toEqual({ - 'a': 3, - 'b': 2, - 'c': 2, - }); - expect(table.countBy('dictionary').toJSON()).toEqual({ - 'a': 3, - 'b': 2, - 'c': 2, - }); - }); - test(`countBy on dictionary with filter returns the correct counts`, () => { - expect(table.filter(col('i32').eq(1)).countBy('dictionary').toJSON()).toEqual({ - 'a': 1, - 'b': 1, - 'c': 1, - }); - }); - test(`countBy on non dictionary column throws error`, () => { - expect(() => { table.countBy('i32'); }).toThrow(); - }); - }); - describe(`multiple record batches`, () => { - const table = Table.from({ + }), + // Use Math.fround to coerce to float32 + values: [ + [Math.fround(-0.3), -1, 'a'], + [Math.fround(-0.2), 1, 'b'], + [Math.fround(-0.1), -1, 'c'], + [Math.fround( 0 ), 1, 'a'], + [Math.fround( 0.1), -1, 'b'], + [Math.fround( 0.2), 1, 'c'], + [Math.fround( 0.3), -1, 'a'] + ]}, + {name: `multiple record batches`, + table: Table.from({ 'schema': { 'fields': [ { @@ -319,11 +256,8 @@ describe(`Table`, () => { } ] }] - }); - - // Wrap floating point values in a Float32Array and take them back out to - // make sure that equality checks will pass - const values = [ + }), + values: [ [Math.fround(-0.3), -1, 'a'], [Math.fround(-0.2), 1, 'b'], [Math.fround(-0.1), -1, 'c'], @@ -333,63 +267,134 @@ describe(`Table`, () => { [Math.fround( 0.3), -1, 'a'], [Math.fround( 0.2), 1, 'b'], [Math.fround( 0.1), -1, 'c'], - ]; - test(`has the correct length`, () => { - expect(table.length).toEqual(values.length); - }); - test(`gets expected values`, () => { - for (let i = -1; ++i < values.length;) { - expect(table.get(i).toArray()).toEqual(values[i]); - } - }); - test(`iterates expected values`, () => { - let i = 0; - for (let row of table) { - expect(row.toArray()).toEqual(values[i++]); + ]} +] + +describe(`Table`, () => { + test(`can create an empty table`, () => { + expect(Table.empty().length).toEqual(0); + }); + test(`Table.from([]) creates an empty table`, () => { + expect(Table.from([]).length).toEqual(0); + }); + test(`Table.from() creates an empty table`, () => { + expect(Table.from().length).toEqual(0); + }); + for (let datum of test_data) { + describe(datum.name, () => { + const table = datum.table; + const values = datum.values; + + test(`has the correct length`, () => { + expect(table.length).toEqual(values.length); + }); + test(`gets expected values`, () => { + for (let i = -1; ++i < values.length;) { + expect(table.get(i).toArray()).toEqual(values[i]); + } + }); + test(`iterates expected values`, () => { + let i = 0; + for (let row of table) { + expect(row.toArray()).toEqual(values[i++]); + } + }); + test(`scans expected values`, () => { + let expected_idx = 0; + table.scan((idx, batch) => { + const columns = batch.schema.fields.map((_, i) => batch.getChildAt(i)!); + expect(columns.map((c) => c.get(idx))).toEqual(values[expected_idx++]); + }); + }); + test(`count() returns the correct length`, () => { + expect(table.count()).toEqual(values.length); + }); + const filter_tests = [ + { + name: `filter on f32 >= 0`, + filtered: table.filter(col('f32').gteq(0)), + expected: values.filter((row)=>row[F32] >= 0) + }, { + name: `filter on i32 <= 0 returns the correct length`, + filtered: table.filter(col('i32').lteq(0)), + expected: values.filter((row)=>row[I32] <= 0) + }, { + name: `filter method combines predicates (f32 >= 0 && i32 <= 0)`, + filtered: table.filter(col('i32').lteq(0)).filter(col('f32').gteq(0)), + expected: values.filter((row)=>row[I32] <= 0 && row[F32] >= 0) + }, { + name: `filter on dictionary == 'a'`, + filtered: table.filter(col('dictionary').eq('a')), + expected: values.filter((row)=>row[DICT] === 'a') + } + ] + for (let this_test of filter_tests) { + describe(`filter on f32 >= 0`, () => { + const filtered = this_test.filtered; + const expected = this_test.expected; + test(`count() returns the correct length`, () => { + expect(filtered.count()).toEqual(expected.length); + }); + test(`scans expected values`, () => { + let expected_idx = 0; + filtered.scan((idx, batch) => { + const columns = batch.schema.fields.map((_, i) => batch.getChildAt(i)!); + expect(columns.map((c) => c.get(idx))).toEqual(expected[expected_idx++]); + }); + }) + }); } - }); - test(`scans expected values`, () => { - let expected_idx = 0; - table.scan((idx, batch) => { - const columns = batch.schema.fields.map((_, i) => batch.getChildAt(i)!); - expect(columns.map((c) => c.get(idx))).toEqual(values[expected_idx++]); + test(`countBy on dictionary returns the correct counts`, () => { + // Make sure countBy works both with and without the Col wrapper + // class + let expected: {[key: string]: number} = {'a': 0, 'b': 0, 'c': 0}; + for (let row of values) { + expected[row[DICT]] += 1; + } + + expect(table.countBy(col('dictionary')).toJSON()).toEqual(expected); + expect(table.countBy('dictionary').toJSON()).toEqual(expected); }); - }); - test(`count() returns the correct length`, () => { - expect(table.count()).toEqual(values.length); - }); - test(`filter on f32 >= 0 returns the correct length`, () => { - expect(table.filter(col('f32').gteq(0)).count()).toEqual(6); - }); - test(`filter on i32 <= 0 returns the correct length`, () => { - expect(table.filter(col('i32').lteq(0)).count()).toEqual(5); - }); - test(`filter on dictionary == 'a' returns the correct length`, () => { - expect(table.filter(col('dictionary').eq('a')).count()).toEqual(3); - }); - test(`countBy on dictionary returns the correct counts`, () => { - // Make sure countBy works both with and without the Col wrapper - // class - expect(table.countBy(col('dictionary')).toJSON()).toEqual({ - 'a': 3, - 'b': 3, - 'c': 3, + test(`countBy on dictionary with filter returns the correct counts`, () => { + let expected: {[key: string]: number} = {'a': 0, 'b': 0, 'c': 0}; + for (let row of values) { + if(row[I32] === 1) { expected[row[DICT]] += 1; } + } + + expect(table.filter(col('i32').eq(1)).countBy('dictionary').toJSON()).toEqual(expected); }); - expect(table.countBy('dictionary').toJSON()).toEqual({ - 'a': 3, - 'b': 3, - 'c': 3, + test(`countBy on non dictionary column throws error`, () => { + expect(() => { table.countBy('i32'); }).toThrow(); + expect(() => { table.filter(col('dict').eq('a')).countBy('i32'); }).toThrow(); }); - }); - test(`countBy on dictionary with filter returns the correct counts`, () => { - expect(table.filter(col('i32').eq(1)).countBy(col('dictionary')).toJSON()).toEqual({ - 'a': 1, - 'b': 2, - 'c': 1, + test(`table.select() basic tests`, () => { + let selected = table.select('f32', 'dictionary'); + expect(selected.schema.fields.length).toEqual(2); + expect(selected.schema.fields[0]).toEqual(table.schema.fields[0]); + expect(selected.schema.fields[1]).toEqual(table.schema.fields[2]); + + expect(selected.length).toEqual(values.length); + let idx = 0, expected_row; + for (let row of selected) { + expected_row = values[idx++]; + expect(row.get(0)).toEqual(expected_row[F32]); + expect(row.get(1)).toEqual(expected_row[DICT]); + } + }); + test(`table.toString()`, () => { + let selected = table.select('i32', 'dictionary'); + let headers = [`"row_id"`, `"i32: Int32"`, `"dictionary: Dictionary"`] + let expected = [headers.join(' | '), ...values.map((row, idx) => { + return [`${idx}`, `${row[I32]}`, `"${row[DICT]}"`].map((str, col) => { + return leftPad(str, ' ', headers[col].length); + }).join(' | '); + })].join('\n') + '\n'; + expect(selected.toString()).toEqual(expected); }); }); - test(`countBy on non dictionary column throws error`, () => { - expect(() => { table.countBy('i32'); }).toThrow(); - }); - }); + } }); + +function leftPad(str: string, fill: string, n: number) { + return (new Array(n + 1).join(fill) + str).slice(-1 * n); +}