Skip to content

Commit

Permalink
[NOID] Fixes #4110: Add support for the Reciprocal rank fusion in the…
Browse files Browse the repository at this point in the history 
… Elastic procedures (#4155) (#4331)
  • Loading branch information
@vga91
vga91 authored Jan 21, 2025
1 parent 0a5127f commit e759105
Show file tree
Hide file tree
Showing 5 changed files with 277 additions and 5 deletions.
136 changes: 136 additions & 0 deletions docs/asciidoc/modules/ROOT/pages/database-integration/elasticsearch.adoc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -268,4 +268,140 @@ CALL apoc.es.put($host,'<indexName>', null, null, null, null, { version: 'EIGHT'

=== Results

Results are stream of map in value.

== Reciprocal Rank Fusion (RRF)

RRF can be performed from Neo4j using ES. For further details, read the https://www.elastic.co/guide/en/elasticsearch/reference/current/rrf.html[official documentation].
Note that this API is supported since version 8.14.x of Elastic.

Here an example using Neo4j with ES.

=== Step 1 - Mapping creation

[source, cypher]
----
CALL apoc.es.put($host, 'example-index', null, null, null,
{
"mappings": {
"properties": {
"text": {
"type": "text"
},
"vector": {
"type": "dense_vector",
"dims": 1,
"index": true,
"similarity": "l2_norm"
},
"integer": {
"type": "integer"
}
}
}
}, $config)
----

==== Results

Results are stream of map in value.

=== Step 2 - Put documents

[source, cypher]
----
CALL apoc.es.put($host, 'example-index/_doc/1', null, null, null,
{
"text" : "rrf",
"vector" : [5],
"integer": 1
}, $config)
CALL apoc.es.put($host, 'example-index/_doc/2', null, null, null,
{
"text" : "rrf rrf",
"vector" : [4],
"integer": 2
}, $config)
CALL apoc.es.put($host, 'example-index/_doc/3', null, null, null,
{
"text" : "rrf rrf rrf",
"vector" : [3],
"integer": 1
}, $config)
CALL apoc.es.put($host, 'example-index/_doc/4', null, null, null,
{
"text" : "rrf rrf rrf rrf",
"integer": 2
}, $config)
CALL apoc.es.put($host, 'example-index/_doc/5', null, null, null,
{
"vector" : [0],
"integer": 1
}, $config)
----

==== Results

Results are stream of map in value.

=== Step 3 - Refresh index

[source, cypher]
----
CALL apoc.es.post($host, 'example-index/_refresh', null, null, '', $config)
----

==== Results

Results are stream of map in value.

=== Step 4 - Perform search using rrf retriever

[source, cypher]
----
CALL apoc.es.getRaw($host,'example-index/_search',
{
"retriever": {
"rrf": {
"retrievers": [
{
"standard": {
"query": {
"term": {
"text": "rrf"
}
}
}
},
{
"knn": {
"field": "vector",
"query_vector": [3],
"k": 5,
"num_candidates": 5
}
}
],
"window_size": 5,
"rank_constant": 1
}
},
"size": 3,
"aggs": {
"int_count": {
"terms": {
"field": "integer"
}
}
}
}
,$config) yield value
----

==== Results

Results are stream of map in value.
2 changes: 1 addition & 1 deletion docs/asciidoc/modules/ROOT/pages/overview/apoc.es/apoc.es.post.adoc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ apoc.es.post(host-or-key,index-or-null,type-or-null,query-or-null,payload-or-nul

[source]
----
apoc.es.post(host :: STRING?, index :: STRING?, type :: STRING?, query :: ANY?, payload = {} :: MAP?, config = {} :: MAP?) :: (value :: MAP?)
apoc.es.post(host :: STRING?, index :: STRING?, type :: STRING?, query :: ANY?, payload = {} :: ANY?, config = {} :: MAP?) :: (value :: MAP?)
----

== Input parameters
Expand Down
138 changes: 136 additions & 2 deletions full-it/src/test/java/apoc/full/it/es/ElasticVersionEightTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,11 @@
import static org.junit.Assert.assertTrue;

import apoc.es.ElasticSearchHandler;
import apoc.util.JsonUtil;
import apoc.util.TestUtil;
import apoc.util.Util;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.nimbusds.jose.util.Pair;
import java.util.List;
import java.util.Map;
import java.util.Set;
Expand All @@ -22,10 +25,13 @@ public static void setUp() throws Exception {
Map.of("headers", basicAuthHeader, VERSION_KEY, ElasticSearchHandler.Version.EIGHT.name());
Map<String, Object> params = Util.map("index", ES_INDEX, "id", ES_ID, "type", ES_TYPE, "config", config);

String tag = "8.12.1";
String tag = "8.14.3";
Map<String, String> envMap = Map.of(
"xpack.security.http.ssl.enabled", "false",
"cluster.routing.allocation.disk.threshold_enabled", "false");
"cluster.routing.allocation.disk.threshold_enabled", "false",
"xpack.license.self_generated.type",
"trial" // To avoid error "current license is non-compliant for [Reciprocal Rank Fusion (RRF)]"
);

getElasticContainer(tag, envMap, params);
}
Expand Down Expand Up @@ -74,6 +80,134 @@ public void testSearchWithQueryAsPayloadAndWithoutIndex() {
this::searchQueryPayloadAssertions);
}

@Test
public void testSearchRRF() throws JsonProcessingException {
String payload = "{\n" + " \"mappings\": {\n"
+ " \"properties\": {\n"
+ " \"text\": {\n"
+ " \"type\": \"text\"\n"
+ " },\n"
+ " \"vector\": {\n"
+ " \"type\": \"dense_vector\",\n"
+ " \"dims\": 1,\n"
+ " \"index\": true,\n"
+ " \"similarity\": \"l2_norm\"\n"
+ " },\n"
+ " \"integer\": {\n"
+ " \"type\": \"integer\"\n"
+ " }\n"
+ " }\n"
+ " }\n"
+ "}";

setPayload(payload, paramsWithBasicAuth);
TestUtil.testCall(
db,
"CALL apoc.es.put($host, 'example-index', null, null, null, $payload, $config)",
paramsWithBasicAuth,
r -> {
Object actual = ((Map) r.get("value")).get("index");
assertEquals("example-index", actual);
});

assertPutForRRF();

paramsWithBasicAuth.remove("payload");
TestUtil.testCall(
db,
"CALL apoc.es.post($host, 'example-index/_refresh', null, null, '', $config)",
paramsWithBasicAuth,
r -> {
Object actual = ((Map) ((Map) r.get("value")).get("_shards")).get("successful");
assertEquals(1L, actual);
});

payload = " {\n" + " \"retriever\": {\n"
+ " \"rrf\": {\n"
+ " \"retrievers\": [\n"
+ " {\n"
+ " \"standard\": {\n"
+ " \"query\": {\n"
+ " \"term\": {\n"
+ " \"text\": \"rrf\"\n"
+ " }\n"
+ " }\n"
+ " }\n"
+ " },\n"
+ " {\n"
+ " \"knn\": {\n"
+ " \"field\": \"vector\",\n"
+ " \"query_vector\": [3],\n"
+ " \"k\": 5,\n"
+ " \"num_candidates\": 5\n"
+ " }\n"
+ " }\n"
+ " ],\n"
+ " \"window_size\": 5,\n"
+ " \"rank_constant\": 1\n"
+ " }\n"
+ " },\n"
+ " \"size\": 3,\n"
+ " \"aggs\": {\n"
+ " \"int_count\": {\n"
+ " \"terms\": {\n"
+ " \"field\": \"integer\"\n"
+ " }\n"
+ " }\n"
+ " }\n"
+ " }";

setPayload(payload, paramsWithBasicAuth);
TestUtil.testCall(
db,
"CALL apoc.es.getRaw($host,'example-index/_search',$payload,$config) yield value",
paramsWithBasicAuth,
r -> {
Object result = ((Map) ((Map) ((Map) r.get("value")).get("hits")).get("total")).get("value");
assertEquals(5L, result);
});

TestUtil.testCall(
db, "CALL apoc.es.delete($host,'example-index',null,null,null,$config)", paramsWithBasicAuth, r -> {
boolean acknowledged = ((boolean) ((Map) r.get("value")).get("acknowledged"));
assertTrue(acknowledged);
});

paramsWithBasicAuth.put("index", ES_INDEX);
}

private void assertPutForRRF() {
List<Pair<String, String>> payloads = List.of(
Pair.of("example-index/_doc/1", "{ \"text\" : \"rrf\", \"vector\" : [5], \"integer\": 1 }"),
Pair.of("example-index/_doc/2", "{ \"text\" : \"rrf rrf\", \"vector\" : [4], \"integer\": 2 }"),
Pair.of("example-index/_doc/3", "{ \"text\" : \"rrf rrf rrf\", \"vector\" : [3], \"integer\": 1 }"),
Pair.of("example-index/_doc/4", "{ \"text\" : \"rrf rrf rrf rrf\", \"integer\": 2 }"),
Pair.of("example-index/_doc/5", "{ \"vector\" : [0], \"integer\": 1 }"));

payloads.forEach(payload -> {
try {
Map mapPayload = JsonUtil.OBJECT_MAPPER.readValue(payload.getRight(), Map.class);
paramsWithBasicAuth.put("payload", mapPayload);
paramsWithBasicAuth.put("index", payload.getLeft());
TestUtil.testCall(
db,
"CALL apoc.es.put($host, $index, null, null, null, $payload, $config)",
paramsWithBasicAuth,
r -> {
Object actual = ((Map) r.get("value")).get("result");
assertEquals("created", actual);
});
} catch (JsonProcessingException e) {
throw new RuntimeException(e);
}
});
}

private void setPayload(String payload, Map<String, Object> params) throws JsonProcessingException {
Map<String, Object> mapPayload = JsonUtil.OBJECT_MAPPER.readValue(payload, Map.class);
params.put("payload", mapPayload);
}

private void searchQueryPayloadAssertions(Map<String, Object> r) {
List<Map> values = (List<Map>) extractValueFromResponse(r, "$.hits.hits");
assertEquals(3, values.size());
Expand Down
4 changes: 2 additions & 2 deletions full/src/main/java/apoc/es/ElasticSearch.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ public Stream<MapResult> post(
@Name("index") String index,
@Name("type") String type,
@Name("query") Object query,
@Name(value = "payload", defaultValue = "{}") Map<String, Object> payload,
@Name(value = "payload", defaultValue = "{}") Object payload,
@Name(value = "config", defaultValue = "{}") Map<String, Object> config) {
if (payload == null) {
payload = Collections.emptyMap();
Expand All @@ -123,7 +123,7 @@ public Stream<MapResult> put(
@Name("type") String type,
@Name("id") String id,
@Name("query") Object query,
@Name(value = "payload", defaultValue = "{}") Map<String, Object> payload,
@Name(value = "payload", defaultValue = "{}") Object payload,
@Name(value = "config", defaultValue = "{}") Map<String, Object> config) {
if (payload == null) {
payload = Collections.emptyMap();
Expand Down
2 changes: 2 additions & 0 deletions full/src/test/java/apoc/mongodb/MongoDBTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
import java.util.stream.Collectors;
import org.bson.types.ObjectId;
import org.junit.BeforeClass;
import org.junit.Ignore;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
Expand Down Expand Up @@ -103,6 +104,7 @@ public void shouldExtractObjectIdsAsMaps() {
assertFalse("should not have an exception", hasException);
}

@Ignore("flaky")
@Test
public void testObjectIdToStringMapping() {
boolean hasException = false;
Expand Down

0 comments on commit e759105

Please sign in to comment.