rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-08-18 03:58:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(nodejs): upgrade arrow versions (#1924)

Browse Source 
Closes #1626
This commit is contained in:
Will Jones
2024-12-09 15:37:11 -08:00
committed by GitHub
parent b70513ca72
commit a43193c99b
5 changed files with 214 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodejs/__test__/arrow.test.ts
+24 -31
View File
@@ -13,11 +13,10 @@ import { Schema } from "apache-arrow";
// See the License for the specific language governing permissions and
// limitations under the License.
import * as arrow13 from "apache-arrow-13";
import * as arrow14 from "apache-arrow-14";
import * as arrow15 from "apache-arrow-15";
import * as arrow16 from "apache-arrow-16";
import * as arrow17 from "apache-arrow-17";
import * as arrow18 from "apache-arrow-18";
import {
convertToTable,
@@ -45,22 +44,16 @@ function sampleRecords(): Array<Record<string, any>> {
},
];
}
describe.each([arrow13, arrow14, arrow15, arrow16, arrow17])(
describe.each([arrow15, arrow16, arrow17, arrow18])(
"Arrow",
(
arrow:
| typeof arrow13
| typeof arrow14
| typeof arrow15
| typeof arrow16
| typeof arrow17,
arrow: typeof arrow15 | typeof arrow16 | typeof arrow17 | typeof arrow18,
) => {
type ApacheArrow =
| typeof arrow13
| typeof arrow14
| typeof arrow15
| typeof arrow16
| typeof arrow17;
| typeof arrow17
| typeof arrow18;
const {
Schema,
Field,
@@ -498,40 +491,40 @@ describe.each([arrow13, arrow14, arrow15, arrow16, arrow17])(
describe("when using two versions of arrow", function () {
it("can still import data", async function () {
const schema = new arrow13.Schema([
new arrow13.Field("id", new arrow13.Int32()),
new arrow13.Field(
const schema = new arrow15.Schema([
new arrow15.Field("id", new arrow15.Int32()),
new arrow15.Field(
"vector",
new arrow13.FixedSizeList(
new arrow15.FixedSizeList(
1024,
new arrow13.Field("item", new arrow13.Float32(), true),
new arrow15.Field("item", new arrow15.Float32(), true),
),
),
new arrow13.Field(
new arrow15.Field(
"struct",
new arrow13.Struct([
new arrow13.Field(
new arrow15.Struct([
new arrow15.Field(
"nested",
new arrow13.Dictionary(
new arrow13.Utf8(),
new arrow13.Int32(),
new arrow15.Dictionary(
new arrow15.Utf8(),
new arrow15.Int32(),
1,
true,
),
),
new arrow13.Field(
new arrow15.Field(
"ts_with_tz",
new arrow13.TimestampNanosecond("some_tz"),
new arrow15.TimestampNanosecond("some_tz"),
),
new arrow13.Field(
new arrow15.Field(
"ts_no_tz",
new arrow13.TimestampNanosecond(null),
new arrow15.TimestampNanosecond(null),
),
]),
),
// biome-ignore lint/suspicious/noExplicitAny: skip
]) as any;
schema.metadataVersion = arrow13.MetadataVersion.V5;
schema.metadataVersion = arrow15.MetadataVersion.V5;
const table = makeArrowTable([], { schema });
const buf = await fromTableToBuffer(table);
@@ -543,13 +536,13 @@ describe.each([arrow13, arrow14, arrow15, arrow16, arrow17])(
// Deep equality gets hung up on some very minor unimportant differences
// between arrow version 13 and 15 which isn't really what we're testing for
// and so we do our own comparison that just checks name/type/nullability
function compareFields(lhs: arrow13.Field, rhs: arrow13.Field) {
function compareFields(lhs: arrow15.Field, rhs: arrow15.Field) {
expect(lhs.name).toEqual(rhs.name);
expect(lhs.nullable).toEqual(rhs.nullable);
expect(lhs.typeId).toEqual(rhs.typeId);
if ("children" in lhs.type && lhs.type.children !== null) {
const lhsChildren = lhs.type.children as arrow13.Field[];
lhsChildren.forEach((child: arrow13.Field, idx) => {
const lhsChildren = lhs.type.children as arrow15.Field[];
lhsChildren.forEach((child: arrow15.Field, idx) => {
compareFields(child, rhs.type.children[idx]);
});
}
nodejs/__test__/registry.test.ts
+133 -144
View File
@@ -12,11 +12,10 @@ import * as apiArrow from "apache-arrow";
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
import * as arrow13 from "apache-arrow-13";
import * as arrow14 from "apache-arrow-14";
import * as arrow15 from "apache-arrow-15";
import * as arrow16 from "apache-arrow-16";
import * as arrow17 from "apache-arrow-17";
import * as arrow18 from "apache-arrow-18";
import * as tmp from "tmp";
@@ -24,154 +23,144 @@ import { connect } from "../lancedb";
import { EmbeddingFunction, LanceSchema } from "../lancedb/embedding";
import { getRegistry, register } from "../lancedb/embedding/registry";
describe.each([arrow13, arrow14, arrow15, arrow16, arrow17])(
"LanceSchema",
(arrow) => {
test("should preserve input order", async () => {
const schema = LanceSchema({
id: new arrow.Int32(),
text: new arrow.Utf8(),
vector: new arrow.Float32(),
});
expect(schema.fields.map((x) => x.name)).toEqual([
"id",
"text",
"vector",
]);
describe.each([arrow15, arrow16, arrow17, arrow18])("LanceSchema", (arrow) => {
test("should preserve input order", async () => {
const schema = LanceSchema({
id: new arrow.Int32(),
text: new arrow.Utf8(),
vector: new arrow.Float32(),
});
},
);
expect(schema.fields.map((x) => x.name)).toEqual(["id", "text", "vector"]);
});
});
describe.each([arrow13, arrow14, arrow15, arrow16, arrow17])(
"Registry",
(arrow) => {
let tmpDir: tmp.DirResult;
beforeEach(() => {
tmpDir = tmp.dirSync({ unsafeCleanup: true });
});
describe.each([arrow15, arrow16, arrow17, arrow18])("Registry", (arrow) => {
let tmpDir: tmp.DirResult;
beforeEach(() => {
tmpDir = tmp.dirSync({ unsafeCleanup: true });
});
afterEach(() => {
tmpDir.removeCallback();
getRegistry().reset();
});
afterEach(() => {
tmpDir.removeCallback();
getRegistry().reset();
});
it("should register a new item to the registry", async () => {
@register("mock-embedding")
class MockEmbeddingFunction extends EmbeddingFunction<string> {
toJSON(): object {
return {
someText: "hello",
};
}
constructor() {
super();
}
ndims() {
return 3;
}
embeddingDataType() {
return new arrow.Float32() as apiArrow.Float;
}
async computeSourceEmbeddings(data: string[]) {
return data.map(() => [1, 2, 3]);
}
it("should register a new item to the registry", async () => {
@register("mock-embedding")
class MockEmbeddingFunction extends EmbeddingFunction<string> {
toJSON(): object {
return {
someText: "hello",
};
}
const func = getRegistry()
.get<MockEmbeddingFunction>("mock-embedding")!
.create();
const schema = LanceSchema({
id: new arrow.Int32(),
text: func.sourceField(new arrow.Utf8() as apiArrow.DataType),
vector: func.vectorField(),
});
const db = await connect(tmpDir.name);
const table = await db.createTable(
"test",
[
{ id: 1, text: "hello" },
{ id: 2, text: "world" },
],
{ schema },
);
const expected = [
[1, 2, 3],
[1, 2, 3],
];
const actual = await table.query().toArrow();
const vectors = actual.getChild("vector")!.toArray();
expect(JSON.parse(JSON.stringify(vectors))).toEqual(
JSON.parse(JSON.stringify(expected)),
);
});
test("should error if registering with the same name", async () => {
class MockEmbeddingFunction extends EmbeddingFunction<string> {
toJSON(): object {
return {
someText: "hello",
};
}
constructor() {
super();
}
ndims() {
return 3;
}
embeddingDataType() {
return new arrow.Float32() as apiArrow.Float;
}
async computeSourceEmbeddings(data: string[]) {
return data.map(() => [1, 2, 3]);
}
constructor() {
super();
}
register("mock-embedding")(MockEmbeddingFunction);
expect(() => register("mock-embedding")(MockEmbeddingFunction)).toThrow(
'Embedding function with alias "mock-embedding" already exists',
);
});
test("schema should contain correct metadata", async () => {
class MockEmbeddingFunction extends EmbeddingFunction<string> {
toJSON(): object {
return {
someText: "hello",
};
}
constructor() {
super();
}
ndims() {
return 3;
}
embeddingDataType() {
return new arrow.Float32() as apiArrow.Float;
}
async computeSourceEmbeddings(data: string[]) {
return data.map(() => [1, 2, 3]);
}
ndims() {
return 3;
}
const func = new MockEmbeddingFunction();
embeddingDataType() {
return new arrow.Float32() as apiArrow.Float;
}
async computeSourceEmbeddings(data: string[]) {
return data.map(() => [1, 2, 3]);
}
}
const schema = LanceSchema({
id: new arrow.Int32(),
text: func.sourceField(new arrow.Utf8() as apiArrow.DataType),
vector: func.vectorField(),
});
const expectedMetadata = new Map<string, string>([
[
"embedding_functions",
JSON.stringify([
{
sourceColumn: "text",
vectorColumn: "vector",
name: "MockEmbeddingFunction",
model: { someText: "hello" },
},
]),
],
]);
expect(schema.metadata).toEqual(expectedMetadata);
const func = getRegistry()
.get<MockEmbeddingFunction>("mock-embedding")!
.create();
const schema = LanceSchema({
id: new arrow.Int32(),
text: func.sourceField(new arrow.Utf8() as apiArrow.DataType),
vector: func.vectorField(),
});
},
);
const db = await connect(tmpDir.name);
const table = await db.createTable(
"test",
[
{ id: 1, text: "hello" },
{ id: 2, text: "world" },
],
{ schema },
);
const expected = [
[1, 2, 3],
[1, 2, 3],
];
const actual = await table.query().toArrow();
const vectors = actual.getChild("vector")!.toArray();
expect(JSON.parse(JSON.stringify(vectors))).toEqual(
JSON.parse(JSON.stringify(expected)),
);
});
test("should error if registering with the same name", async () => {
class MockEmbeddingFunction extends EmbeddingFunction<string> {
toJSON(): object {
return {
someText: "hello",
};
}
constructor() {
super();
}
ndims() {
return 3;
}
embeddingDataType() {
return new arrow.Float32() as apiArrow.Float;
}
async computeSourceEmbeddings(data: string[]) {
return data.map(() => [1, 2, 3]);
}
}
register("mock-embedding")(MockEmbeddingFunction);
expect(() => register("mock-embedding")(MockEmbeddingFunction)).toThrow(
'Embedding function with alias "mock-embedding" already exists',
);
});
test("schema should contain correct metadata", async () => {
class MockEmbeddingFunction extends EmbeddingFunction<string> {
toJSON(): object {
return {
someText: "hello",
};
}
constructor() {
super();
}
ndims() {
return 3;
}
embeddingDataType() {
return new arrow.Float32() as apiArrow.Float;
}
async computeSourceEmbeddings(data: string[]) {
return data.map(() => [1, 2, 3]);
}
}
const func = new MockEmbeddingFunction();
const schema = LanceSchema({
id: new arrow.Int32(),
text: func.sourceField(new arrow.Utf8() as apiArrow.DataType),
vector: func.vectorField(),
});
const expectedMetadata = new Map<string, string>([
[
"embedding_functions",
JSON.stringify([
{
sourceColumn: "text",
vectorColumn: "vector",
name: "MockEmbeddingFunction",
model: { someText: "hello" },
},
]),
],
]);
expect(schema.metadata).toEqual(expectedMetadata);
});
});
nodejs/__test__/table.test.ts
+5 -7
View File
@@ -16,11 +16,10 @@ import * as fs from "fs";
import * as path from "path";
import * as tmp from "tmp";
import * as arrow13 from "apache-arrow-13";
import * as arrow14 from "apache-arrow-14";
import * as arrow15 from "apache-arrow-15";
import * as arrow16 from "apache-arrow-16";
import * as arrow17 from "apache-arrow-17";
import * as arrow18 from "apache-arrow-18";
import { Table, connect } from "../lancedb";
import {
@@ -44,7 +43,7 @@ import {
} from "../lancedb/embedding";
import { Index } from "../lancedb/indices";
describe.each([arrow13, arrow14, arrow15, arrow16, arrow17])(
describe.each([arrow15, arrow16, arrow17, arrow18])(
"Given a table",
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
(arrow: any) => {
@@ -52,11 +51,10 @@ describe.each([arrow13, arrow14, arrow15, arrow16, arrow17])(
let table: Table;
const schema:
| import("apache-arrow-13").Schema
| import("apache-arrow-14").Schema
| import("apache-arrow-15").Schema
| import("apache-arrow-16").Schema
| import("apache-arrow-17").Schema = new arrow.Schema([
| import("apache-arrow-17").Schema
| import("apache-arrow-18").Schema = new arrow.Schema([
new arrow.Field("id", new arrow.Float64(), true),
]);
@@ -939,7 +937,7 @@ describe("when optimizing a dataset", () => {
});
});
describe.each([arrow13, arrow14, arrow15, arrow16, arrow17])(
describe.each([arrow15, arrow16, arrow17, arrow18])(
"when optimizing a dataset",
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
(arrow: any) => {