Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

community[patch]: add ?| (arrayContains) filter on metadata to PGVector search #5381

Merged
merged 8 commits into from
May 26, 2024
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,8 @@ const pgvectorStore = await PGVectorStore.initialize(
);

await pgvectorStore.addDocuments([
{ pageContent: "what's this", metadata: { a: 2 } },
{ pageContent: "Cat drinks milk", metadata: { a: 1 } },
{ pageContent: "what's this", metadata: { a: 2, b: ["tag1", "tag2"] } },
{ pageContent: "Cat drinks milk", metadata: { a: 1, b: ["tag2"] } },
]);

const results = await pgvectorStore.similaritySearch("water", 1);
Expand Down Expand Up @@ -84,4 +84,17 @@ console.log(results4);
[ Document { pageContent: 'what's this', metadata: { a: 2 } } ]
*/

// Filtering using arrayContains (?|) is supported
const results5 = await pgvectorStore.similaritySearch("water", 1, {
b: {
arrayContains: ["tag1"],
},
});

console.log(results5);

/*
[ Document { pageContent: "what's this", metadata: { a: 2, b: ['tag1', 'tag2'] } } } ]
*/

await pgvectorStore.end();
12 changes: 12 additions & 0 deletions libs/langchain-community/src/vectorstores/pgvector.ts
Original file line number Diff line number Diff line change
Expand Up @@ -501,6 +501,18 @@ export class PGVectorStore extends VectorStore {
parameters.push(..._value.in);
paramCount += _value.in.length;
}
if (Array.isArray(_value.arrayContains)) {
const placeholders = _value.arrayContains
.map(
(_: unknown, index: number) => `$${currentParamCount + index + 1}`
)
.join(",");
whereClauses.push(
`${this.metadataColumnName}->'${key}' ?| array[${placeholders}]`
);
parameters.push(..._value.arrayContains);
paramCount += _value.arrayContains.length;
}
} else {
paramCount += 1;
whereClauses.push(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,44 @@ describe("PGVectorStore", () => {
expect(result3.length).toEqual(3);
});

test("PGvector supports arrayContains (?|) in metadata filter ", async () => {
const documents = [
{ pageContent: "Lorem Ipsum", metadata: { a: ["tag1", "tag2"] } },
{ pageContent: "Lorem Ipsum", metadata: { a: ["tag2"] } },
{ pageContent: "Lorem Ipsum", metadata: { a: ["tag1"] } },
];

await pgvectorVectorStore.addDocuments(documents);

const result = await pgvectorVectorStore.similaritySearch("hello", 2, {
a: {
arrayContains: ["tag1"],
},
});

expect(result.length).toEqual(2);
expect(result).toEqual([
{ pageContent: "Lorem Ipsum", metadata: { a: ["tag1", "tag2"] } },
{ pageContent: "Lorem Ipsum", metadata: { a: ["tag1"] } },
]);

const result2 = await pgvectorVectorStore.similaritySearch("hello", 2, {
a: {
arrayContains: ["tag2"],
},
});
expect(result2.length).toEqual(2);
expect(result2).toEqual([
{ pageContent: "Lorem Ipsum", metadata: { a: ["tag1", "tag2"] } },
{ pageContent: "Lorem Ipsum", metadata: { a: ["tag2"] } },
]);

const result3 = await pgvectorVectorStore.similaritySearch("hello", 3);

expect(result3.length).toEqual(3);
expect(result3).toEqual(documents);
});

test("PGvector can delete document by id", async () => {
try {
const documents = [
Expand Down
Loading