From d3e15f3e17a877e45a77558b7f3dcec69e0d16d4 Mon Sep 17 00:00:00 2001 From: Vedant Madane <6527493+VedantMadane@users.noreply.github.com> Date: Tue, 3 Feb 2026 23:39:51 +0530 Subject: [PATCH] fix(node): allow bigint[] for takeRowIds (#2916) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary This PR changes takeRowIds to accept bigint[] instead of number[], matching the type of _rowid returned by withRowId(). ## Problem When retrieving row IDs using \withRowId()\ and querying them back with takeRowIds(), users get an error because: 1. _rowid values are returned as JavaScript bigint 2. takeRowIds() expected number[] 3. NAPI failed to convert: Error: Failed to convert napi value BigInt into rust type i64 ## Reproduction \\\js import lancedb from '@lancedb/lancedb'; const db = await lancedb.connect('memory://'); const table = await db.createTable('test', [{ id: 1, vector: [1.0, 2.0] }]); const results = await table.query().withRowId().toArray(); const rowIds = results.map(row => row._rowid); console.log('types:', rowIds.map(id => typeof id)); // ['bigint'] await table.takeRowIds(rowIds).toArray(); // ❌ Error before fix \\\ ## Solution - Updated TypeScript signature from takeRowIds(rowIds: number[]) to takeRowIds(rowIds: bigint[]) - Updated Rust NAPI binding to accept Vec and convert using get_u64() Fixes #2722 --------- Co-authored-by: Will Jones --- docs/src/js/classes/Table.md | 5 ++- nodejs/__test__/table.test.ts | 60 +++++++++++++++++++++++++++++++++++ nodejs/lancedb/table.ts | 26 +++++++++++++-- nodejs/src/table.rs | 20 ++++++++---- 4 files changed, 100 insertions(+), 11 deletions(-) diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index fd7902f06..67a94bf6f 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -705,8 +705,11 @@ Create a query that returns a subset of the rows in the table. #### Parameters -* **rowIds**: `number`[] +* **rowIds**: readonly (`number` \| `bigint`)[] The row ids of the rows to return. + Row ids returned by `withRowId()` are `bigint`, so `bigint[]` is supported. + For convenience / backwards compatibility, `number[]` is also accepted (for + small row ids that fit in a safe integer). #### Returns diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 51e3a048f..33b83ff0e 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -312,6 +312,66 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( expect(res.getChild("id")?.toJSON()).toEqual([2, 3]); }); + it("should support takeRowIds with bigint array", async () => { + await table.add([{ id: 1 }, { id: 2 }, { id: 3 }]); + // Get actual row IDs using withRowId() + const allRows = await table.query().withRowId().toArray(); + const rowIds = allRows.map((row) => row._rowid) as bigint[]; + + // Verify row IDs are bigint + expect(typeof rowIds[0]).toBe("bigint"); + + // Use takeRowIds with bigint array (the main use case from issue #2722) + const res = await table.takeRowIds([rowIds[0], rowIds[2]]).toArray(); + expect(res.map((r) => r.id)).toEqual([1, 3]); + }); + + it("should support takeRowIds with number array for backwards compatibility", async () => { + await table.add([{ id: 1 }, { id: 2 }, { id: 3 }]); + // Small row IDs can be passed as numbers + const res = await table.takeRowIds([0, 2]).toArray(); + expect(res.map((r) => r.id)).toEqual([1, 3]); + }); + + it("should support takeRowIds with mixed bigint and number array", async () => { + await table.add([{ id: 1 }, { id: 2 }, { id: 3 }]); + // Mixed array of bigint and number + const res = await table.takeRowIds([0n, 1, 2n]).toArray(); + expect(res.map((r) => r.id)).toEqual([1, 2, 3]); + }); + + it("should throw for non-integer number in takeRowIds", () => { + expect(() => table.takeRowIds([1.5])).toThrow( + "Row id must be an integer (or bigint)", + ); + expect(() => table.takeRowIds([0, 1.1, 2])).toThrow( + "Row id must be an integer (or bigint)", + ); + }); + + it("should throw for negative number in takeRowIds", () => { + expect(() => table.takeRowIds([-1])).toThrow("Row id cannot be negative"); + expect(() => table.takeRowIds([0, -5, 2])).toThrow( + "Row id cannot be negative", + ); + }); + + it("should throw for unsafe large number in takeRowIds", () => { + // Number.MAX_SAFE_INTEGER + 1 is not safe + const unsafeNumber = Number.MAX_SAFE_INTEGER + 1; + expect(() => table.takeRowIds([unsafeNumber])).toThrow( + "Row id is too large for number; use bigint instead", + ); + }); + + it("should reject negative bigint in takeRowIds", async () => { + await table.add([{ id: 1 }]); + // Negative bigint should be rejected by the Rust layer + expect(() => { + table.takeRowIds([-1n]); + }).toThrow("Row id cannot be negative"); + }); + it("should return the table as an instance of an arrow table", async () => { const arrowTbl = await table.toArrow(); expect(arrowTbl).toBeInstanceOf(ArrowTable); diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 43389893c..c5a14bdc4 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -347,9 +347,13 @@ export abstract class Table { /** * Create a query that returns a subset of the rows in the table. * @param rowIds The row ids of the rows to return. + * + * Row ids returned by `withRowId()` are `bigint`, so `bigint[]` is supported. + * For convenience / backwards compatibility, `number[]` is also accepted (for + * small row ids that fit in a safe integer). * @returns A builder that can be used to parameterize the query. */ - abstract takeRowIds(rowIds: number[]): TakeQuery; + abstract takeRowIds(rowIds: readonly (bigint | number)[]): TakeQuery; /** * Create a search query to find the nearest neighbors @@ -686,8 +690,24 @@ export class LocalTable extends Table { return new TakeQuery(this.inner.takeOffsets(offsets)); } - takeRowIds(rowIds: number[]): TakeQuery { - return new TakeQuery(this.inner.takeRowIds(rowIds)); + takeRowIds(rowIds: readonly (bigint | number)[]): TakeQuery { + const ids = rowIds.map((id) => { + if (typeof id === "bigint") { + return id; + } + if (!Number.isInteger(id)) { + throw new Error("Row id must be an integer (or bigint)"); + } + if (id < 0) { + throw new Error("Row id cannot be negative"); + } + if (!Number.isSafeInteger(id)) { + throw new Error("Row id is too large for number; use bigint instead"); + } + return BigInt(id); + }); + + return new TakeQuery(this.inner.takeRowIds(ids)); } query(): Query { diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index aa2bc51e2..253182e70 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -208,18 +208,24 @@ impl Table { } #[napi(catch_unwind)] - pub fn take_row_ids(&self, row_ids: Vec) -> napi::Result { + pub fn take_row_ids(&self, row_ids: Vec) -> napi::Result { Ok(TakeQuery::new( self.inner_ref()?.take_row_ids( row_ids .into_iter() - .map(|o| { - u64::try_from(o).map_err(|e| { - napi::Error::from_reason(format!( - "Failed to convert row id to u64: {}", - e + .map(|id| { + let (negative, value, lossless) = id.get_u64(); + if negative { + Err(napi::Error::from_reason( + "Row id cannot be negative".to_string(), )) - }) + } else if !lossless { + Err(napi::Error::from_reason( + "Row id is too large to fit in u64".to_string(), + )) + } else { + Ok(value) + } }) .collect::>>()?, ),