This is an automated email from the ASF dual-hosted git repository. raulcd pushed a commit to branch maint-16.x.x in repository https://gitbox.apache.org/repos/asf/arrow.git
commit 3d690b7e2ac5061e4b60119ba42080a3c2e1ee5d Author: Dominik Moritz <[email protected]> AuthorDate: Tue Apr 16 16:19:07 2024 -0400 GH-39131: [JS] Add at() for array like types (#40730) Simpler version of https://github.com/apache/arrow/pull/40712 that preserves `get`. * GitHub Issue: #39131 --- js/src/recordbatch.ts | 17 +++++++++++++---- js/src/table.ts | 11 ++++++++++- js/src/util/vector.ts | 14 +++----------- js/src/vector.ts | 10 +++++++++- js/test/unit/vector/bool-vector-tests.ts | 7 ++++--- js/test/unit/vector/numeric-vector-tests.ts | 6 ++++-- js/test/unit/vector/vector-tests.ts | 11 +++++++---- 7 files changed, 50 insertions(+), 26 deletions(-) diff --git a/js/src/recordbatch.ts b/js/src/recordbatch.ts index b9061c8b9b..33dbe9e59c 100644 --- a/js/src/recordbatch.ts +++ b/js/src/recordbatch.ts @@ -20,6 +20,7 @@ import { Table } from './table.js'; import { Vector } from './vector.js'; import { Schema, Field } from './schema.js'; import { DataType, Struct, Null, TypeMap } from './type.js'; +import { wrapIndex } from './util/vector.js'; import { instance as getVisitor } from './visitor/get.js'; import { instance as setVisitor } from './visitor/set.js'; @@ -116,7 +117,7 @@ export class RecordBatch<T extends TypeMap = any> { } /** - * Check whether an element is null. + * Check whether an row is null. * @param index The index at which to read the validity bitmap. */ public isValid(index: number) { @@ -125,15 +126,23 @@ export class RecordBatch<T extends TypeMap = any> { /** * Get a row by position. - * @param index The index of the element to read. + * @param index The index of the row to read. */ public get(index: number) { return getVisitor.visit(this.data, index); } + /** + * Get a row value by position. + * @param index The index of the row to read. A negative index will count back from the last row. + */ + public at(index: number) { + return this.get(wrapIndex(index, this.numRows)); + } + /** * Set a row by position. - * @param index The index of the element to write. + * @param index The index of the row to write. * @param value The value to set. */ public set(index: number, value: Struct<T>['TValue']) { @@ -175,7 +184,7 @@ export class RecordBatch<T extends TypeMap = any> { /** * Return a zero-copy sub-section of this RecordBatch. * @param start The beginning of the specified portion of the RecordBatch. - * @param end The end of the specified portion of the RecordBatch. This is exclusive of the element at the index 'end'. + * @param end The end of the specified portion of the RecordBatch. This is exclusive of the row at the index 'end'. */ public slice(begin?: number, end?: number): RecordBatch<T> { const [slice] = new Vector([this.data]).slice(begin, end).data; diff --git a/js/src/table.ts b/js/src/table.ts index d7a6617530..2aab2b7ec2 100644 --- a/js/src/table.ts +++ b/js/src/table.ts @@ -40,7 +40,7 @@ import { instance as indexOfVisitor } from './visitor/indexof.js'; import { instance as iteratorVisitor } from './visitor/iterator.js'; import { DataProps } from './data.js'; -import { clampRange } from './util/vector.js'; +import { clampRange, wrapIndex } from './util/vector.js'; import { ArrayDataType, BigIntArray, TypedArray, TypedArrayDataType } from './interfaces.js'; import { RecordBatch, _InternalEmptyPlaceholderRecordBatch } from './recordbatch.js'; @@ -196,6 +196,15 @@ export class Table<T extends TypeMap = any> { // @ts-ignore public get(index: number): Struct<T>['TValue'] | null { return null; } + /** + * Get an element value by position. + * @param index The index of the element to read. A negative index will count back from the last element. + */ + // @ts-ignore + public at(index: number): Struct<T>['TValue'] | null { + return this.get(wrapIndex(index, this.numRows)); + } + /** * Set an element value by position. * diff --git a/js/src/util/vector.ts b/js/src/util/vector.ts index 179b17a39f..9414aef68c 100644 --- a/js/src/util/vector.ts +++ b/js/src/util/vector.ts @@ -23,19 +23,8 @@ import { compareArrayLike } from '../util/buffer.js'; /** @ignore */ type RangeLike = { length: number; stride?: number }; /** @ignore */ -type ClampThen<T extends RangeLike> = (source: T, index: number) => any; -/** @ignore */ type ClampRangeThen<T extends RangeLike> = (source: T, offset: number, length: number) => any; -export function clampIndex<T extends RangeLike>(source: T, index: number): number; -export function clampIndex<T extends RangeLike, N extends ClampThen<T> = ClampThen<T>>(source: T, index: number, then: N): ReturnType<N>; -/** @ignore */ -export function clampIndex<T extends RangeLike, N extends ClampThen<T> = ClampThen<T>>(source: T, index: number, then?: N) { - const length = source.length; - const adjust = index > -1 ? index : (length + (index % length)); - return then ? then(source, adjust) : adjust; -} - /** @ignore */ let tmp: number; export function clampRange<T extends RangeLike>(source: T, begin: number | undefined, end: number | undefined): [number, number]; @@ -60,6 +49,9 @@ export function clampRange<T extends RangeLike, N extends ClampRangeThen<T> = Cl return then ? then(source, lhs, rhs) : [lhs, rhs]; } +/** @ignore */ +export const wrapIndex = (index: number, len: number) => index < 0 ? (len + index) : index; + const isNaNFast = (value: any) => value !== value; /** @ignore */ diff --git a/js/src/vector.ts b/js/src/vector.ts index 1b0d9a0579..aeaa1c1342 100644 --- a/js/src/vector.ts +++ b/js/src/vector.ts @@ -16,7 +16,7 @@ // under the License. import { Type } from './enum.js'; -import { clampRange } from './util/vector.js'; +import { clampRange, wrapIndex } from './util/vector.js'; import { DataType, strideForType } from './type.js'; import { Data, makeData, DataProps } from './data.js'; import { BigIntArray, TypedArray, TypedArrayDataType } from './interfaces.js'; @@ -177,6 +177,14 @@ export class Vector<T extends DataType = any> { // @ts-ignore public get(index: number): T['TValue'] | null { return null; } + /** + * Get an element value by position. + * @param index The index of the element to read. A negative index will count back from the last element. + */ + public at(index: number): T['TValue'] | null { + return this.get(wrapIndex(index, this.length)); + } + /** * Set an element value by position. * @param index The index of the element to write. diff --git a/js/test/unit/vector/bool-vector-tests.ts b/js/test/unit/vector/bool-vector-tests.ts index afbaef01a1..7b6ed3e589 100644 --- a/js/test/unit/vector/bool-vector-tests.ts +++ b/js/test/unit/vector/bool-vector-tests.ts @@ -25,9 +25,10 @@ describe(`BoolVector`, () => { const n = values.length; const vector = newBoolVector(n, new Uint8Array([27, 0, 0, 0, 0, 0, 0, 0])); test(`gets expected values`, () => { - let i = -1; - while (++i < n) { + for (let i = 0; i < values.length; i++) { expect(vector.get(i)).toEqual(values[i]); + expect(vector.at(i)).toEqual(values.at(i)); + expect(vector.at(-i)).toEqual(values.at(-i)); } }); test(`iterates expected values`, () => { @@ -53,7 +54,7 @@ describe(`BoolVector`, () => { const expected2 = [true, true, true, true, true, false, false, false]; const expected3 = [true, true, false, false, false, false, true, true]; function validate(expected: boolean[]) { - for (let i = -1; ++i < n;) { + for (let i = 0; i < n; i++) { expect(v.get(i)).toEqual(expected[i]); } } diff --git a/js/test/unit/vector/numeric-vector-tests.ts b/js/test/unit/vector/numeric-vector-tests.ts index 032af796de..e5c1789068 100644 --- a/js/test/unit/vector/numeric-vector-tests.ts +++ b/js/test/unit/vector/numeric-vector-tests.ts @@ -337,10 +337,12 @@ function testAndValidateVector<T extends Int | Float>(vector: Vector<T>, typed: function gets_expected_values<T extends Int | Float>(vector: Vector<T>, typed: T['TArray'], values: any[] = [...typed]) { test(`gets expected values`, () => { expect.hasAssertions(); - let i = -1, n = vector.length; + let i = -1; try { - while (++i < n) { + while (++i < vector.length) { expect(vector.get(i)).toEqual(values[i]); + expect(vector.at(i)).toEqual(values.at(i)); + expect(vector.at(-i)).toEqual(values.at(-i)); } } catch (e) { throw new Error(`${i}: ${e}`); } }); diff --git a/js/test/unit/vector/vector-tests.ts b/js/test/unit/vector/vector-tests.ts index a10d7c757c..881bf987b5 100644 --- a/js/test/unit/vector/vector-tests.ts +++ b/js/test/unit/vector/vector-tests.ts @@ -265,8 +265,10 @@ describe(`ListVector`, () => { }); test(`get value`, () => { - for (const [i, value] of values.entries()) { - expect(vector.get(i)!.toJSON()).toEqual(value); + for (let i = 0; i < values.length; i++) { + expect(vector.get(i)!.toJSON()).toEqual(values[i]); + expect(vector.at(i)!.toJSON()).toEqual(values.at(i)); + expect(vector.at(-i)!.toJSON()).toEqual(values.at(-i)); } }); }); @@ -308,9 +310,10 @@ function basicVectorTests(vector: Vector, values: any[], extras: any[]) { const n = values.length; test(`gets expected values`, () => { - let i = -1; - while (++i < n) { + for (let i = 0; i < values.length; i++) { expect(vector.get(i)).toEqual(values[i]); + expect(vector.at(i)).toEqual(values.at(i)); + expect(vector.at(-i)).toEqual(values.at(-i)); } }); test(`iterates expected values`, () => {
