Skip to content

Commit

Permalink
Decoding FT.SEARCH reply can be disabled at field level (#3926)
Browse files Browse the repository at this point in the history 
* Decoding FT.SEARCH reply can be disabled at field level

* added javadoc in Document

* make the pr not breaking
  • Loading branch information
@sazzad16
sazzad16 authored Aug 27, 2024
1 parent 11274de commit 81c9f32
Show file tree
Hide file tree
Showing 5 changed files with 200 additions and 69 deletions.
32 changes: 20 additions & 12 deletions src/main/java/redis/clients/jedis/CommandObjects.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3384,19 +3384,20 @@ public final CommandObject<String> ftDropIndexDD(String indexName) {

public final CommandObject<SearchResult> ftSearch(String indexName, String query) {
return new CommandObject<>(checkAndRoundRobinSearchCommand(SearchCommand.SEARCH, indexName).add(query),
getSearchResultBuilder(() -> new SearchResultBuilder(true, false, true)));
getSearchResultBuilder(null, () -> new SearchResultBuilder(true, false, true)));
}

public final CommandObject<SearchResult> ftSearch(String indexName, String query, FTSearchParams params) {
return new CommandObject<>(checkAndRoundRobinSearchCommand(SearchCommand.SEARCH, indexName)
.add(query).addParams(params.dialectOptional(searchDialect.get())),
getSearchResultBuilder(() -> new SearchResultBuilder(!params.getNoContent(), params.getWithScores(), true)));
getSearchResultBuilder(params.getReturnFieldDecodeMap(), () -> new SearchResultBuilder(
!params.getNoContent(), params.getWithScores(), true, params.getReturnFieldDecodeMap())));
}

public final CommandObject<SearchResult> ftSearch(String indexName, Query query) {
return new CommandObject<>(checkAndRoundRobinSearchCommand(SearchCommand.SEARCH, indexName)
.addParams(query.dialectOptional(searchDialect.get())), getSearchResultBuilder(() ->
new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), true)));
.addParams(query.dialectOptional(searchDialect.get())), getSearchResultBuilder(null,
() -> new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), true)));
}

@Deprecated
Expand All @@ -3405,8 +3406,8 @@ public final CommandObject<SearchResult> ftSearch(byte[] indexName, Query query)
throw new UnsupportedOperationException("binary ft.search is not implemented with resp3.");
}
return new CommandObject<>(checkAndRoundRobinSearchCommand(commandArguments(SearchCommand.SEARCH), indexName)
.addParams(query.dialectOptional(searchDialect.get())), getSearchResultBuilder(() ->
new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), false)));
.addParams(query.dialectOptional(searchDialect.get())), getSearchResultBuilder(null,
() -> new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), false)));
}

public final CommandObject<String> ftExplain(String indexName, Query query) {
Expand Down Expand Up @@ -3449,20 +3450,27 @@ public final CommandObject<Map.Entry<SearchResult, Map<String, Object>>> ftProfi
String indexName, FTProfileParams profileParams, Query query) {
return new CommandObject<>(checkAndRoundRobinSearchCommand(SearchCommand.PROFILE, indexName)
.add(SearchKeyword.SEARCH).addParams(profileParams).add(SearchKeyword.QUERY)
.addParams(query.dialectOptional(searchDialect.get())), new SearchProfileResponseBuilder<>(
getSearchResultBuilder(() -> new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), true))));
.addParams(query.dialectOptional(searchDialect.get())),
new SearchProfileResponseBuilder<>(getSearchResultBuilder(null,
() -> new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), true))));
}

public final CommandObject<Map.Entry<SearchResult, Map<String, Object>>> ftProfileSearch(
String indexName, FTProfileParams profileParams, String query, FTSearchParams searchParams) {
return new CommandObject<>(checkAndRoundRobinSearchCommand(SearchCommand.PROFILE, indexName)
.add(SearchKeyword.SEARCH).addParams(profileParams).add(SearchKeyword.QUERY).add(query)
.addParams(searchParams.dialectOptional(searchDialect.get())), new SearchProfileResponseBuilder<>(
getSearchResultBuilder(() -> new SearchResultBuilder(!searchParams.getNoContent(), searchParams.getWithScores(), true))));
.addParams(searchParams.dialectOptional(searchDialect.get())),
new SearchProfileResponseBuilder<>(getSearchResultBuilder(searchParams.getReturnFieldDecodeMap(),
() -> new SearchResultBuilder(!searchParams.getNoContent(), searchParams.getWithScores(), true,
searchParams.getReturnFieldDecodeMap()))));
}

private Builder<SearchResult> getSearchResultBuilder(Supplier<Builder<SearchResult>> resp2) {
if (protocol == RedisProtocol.RESP3) return SearchResult.SEARCH_RESULT_BUILDER;
private Builder<SearchResult> getSearchResultBuilder(
Map<String, Boolean> isReturnFieldDecode, Supplier<Builder<SearchResult>> resp2) {
if (protocol == RedisProtocol.RESP3) {
return isReturnFieldDecode == null ? SearchResult.SEARCH_RESULT_BUILDER
: new SearchResult.PerFieldDecoderSearchResultBuilder(isReturnFieldDecode);
}
return resp2.get();
}

Expand Down
101 changes: 76 additions & 25 deletions src/main/java/redis/clients/jedis/search/Document.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import redis.clients.jedis.Builder;
import redis.clients.jedis.BuilderFactory;
import redis.clients.jedis.util.KeyValue;
Expand Down Expand Up @@ -50,24 +51,6 @@ public Iterable<Map.Entry<String, Object>> getProperties() {
return fields.entrySet();
}

public static Document load(String id, double score, byte[] payload, List<byte[]> fields) {
return Document.load(id, score, fields, true);
}

public static Document load(String id, double score, List<byte[]> fields, boolean decode) {
Document ret = new Document(id, score);
if (fields != null) {
for (int i = 0; i < fields.size(); i += 2) {
byte[] rawKey = fields.get(i);
byte[] rawValue = fields.get(i + 1);
String key = SafeEncoder.encode(rawKey);
Object value = rawValue == null ? null : decode ? SafeEncoder.encode(rawValue) : rawValue;
ret.set(key, value);
}
}
return ret;
}

/**
* @return the document's id
*/
Expand Down Expand Up @@ -102,16 +85,22 @@ public Object get(String key) {
*/
public String getString(String key) {
Object value = fields.get(key);
if (value instanceof String) {
if (value == null) {
return null;
} else if (value instanceof String) {
return (String) value;
} else if (value instanceof byte[]) {
return SafeEncoder.encode((byte[]) value);
} else {
return String.valueOf(value);
}
return value instanceof byte[] ? SafeEncoder.encode((byte[]) value) : value.toString();
}

public boolean hasProperty(String key) {
return fields.containsKey(key);
}

// TODO: private ??
public Document set(String key, Object value) {
fields.put(key, value);
return this;
Expand All @@ -122,7 +111,9 @@ public Document set(String key, Object value) {
*
* @param score new score to set
* @return the document itself
* @deprecated
*/
@Deprecated
public Document setScore(float score) {
this.score = (double) score;
return this;
Expand All @@ -134,13 +125,58 @@ public String toString() {
", properties:" + this.getProperties();
}

static Builder<Document> SEARCH_DOCUMENT = new Builder<Document>() {
/// RESP2 -->
public static Document load(String id, double score, byte[] payload, List<byte[]> fields) {
return Document.load(id, score, fields, true);
}

public static Document load(String id, double score, List<byte[]> fields, boolean decode) {
return load(id, score, fields, decode, null);
}

/**
* Parse document object from FT.SEARCH reply.
* @param id
* @param score
* @param fields
* @param decode
* @param isFieldDecode checked only if {@code decode=true}
* @return document
*/
public static Document load(String id, double score, List<byte[]> fields, boolean decode,
Map<String, Boolean> isFieldDecode) {
Document ret = new Document(id, score);
if (fields != null) {
for (int i = 0; i < fields.size(); i += 2) {
byte[] rawKey = fields.get(i);
byte[] rawValue = fields.get(i + 1);
String key = SafeEncoder.encode(rawKey);
Object value = rawValue == null ? null
: (decode && (isFieldDecode == null || !Boolean.FALSE.equals(isFieldDecode.get(key))))
? SafeEncoder.encode(rawValue) : rawValue;
ret.set(key, value);
}
}
return ret;
}
/// <-- RESP2

/// RESP3 -->
// TODO: final
static Builder<Document> SEARCH_DOCUMENT = new PerFieldDecoderDocumentBuilder((Map) null);

static final class PerFieldDecoderDocumentBuilder extends Builder<Document> {

private static final String ID_STR = "id";
private static final String SCORE_STR = "score";
// private static final String FIELDS_STR = "fields";
private static final String FIELDS_STR = "extra_attributes";

private final Map<String, Boolean> isFieldDecode;

public PerFieldDecoderDocumentBuilder(Map<String, Boolean> isFieldDecode) {
this.isFieldDecode = isFieldDecode != null ? isFieldDecode : Collections.emptyMap();
}

@Override
public Document build(Object data) {
List<KeyValue> list = (List<KeyValue>) data;
Expand All @@ -157,13 +193,28 @@ public Document build(Object data) {
score = BuilderFactory.DOUBLE.build(kv.getValue());
break;
case FIELDS_STR:
fields = BuilderFactory.ENCODED_OBJECT_MAP.build(kv.getValue());
fields = makeFieldsMap(isFieldDecode, kv.getValue());
break;
}
}
// assert id != null;
// if (fields == null) fields = Collections.emptyMap();
return new Document(id, score, fields);
}
};

private static Map<String, Object> makeFieldsMap(Map<String, Boolean> isDecode, Object data) {
if (data == null) return null;

final List<KeyValue> list = (List) data;

Map<String, Object> map = new HashMap<>(list.size(), 1f);
list.stream().filter((kv) -> (kv != null && kv.getKey() != null && kv.getValue() != null))
.forEach((kv) -> {
String key = BuilderFactory.STRING.build(kv.getKey());
map.put(key,
(Boolean.FALSE.equals(isDecode.get(key)) ? BuilderFactory.RAW_OBJECT
: BuilderFactory.AGGRESSIVE_ENCODED_OBJECT).build(kv.getValue()));
});
return map;
}
/// <-- RESP3
}
62 changes: 37 additions & 25 deletions src/main/java/redis/clients/jedis/search/FTSearchParams.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,7 @@ public class FTSearchParams implements IParams {
private final List<IParams> filters = new LinkedList<>();
private Collection<String> inKeys;
private Collection<String> inFields;
private Collection<String> returnFields;
private Collection<FieldName> returnFieldNames;
private Collection<FieldName> returnFieldsNames;
private boolean summarize;
private SummarizeParams summarizeParams;
private boolean highlight;
Expand All @@ -43,6 +42,9 @@ public class FTSearchParams implements IParams {
private Map<String, Object> params;
private Integer dialect;

/// non command parameters
private Map<String, Boolean> returnFieldDecodeMap = null;

public FTSearchParams() {
}

Expand Down Expand Up @@ -78,17 +80,15 @@ public void addParams(CommandArguments args) {
args.add(INFIELDS).add(inFields.size()).addObjects(inFields);
}

if (returnFieldNames != null && !returnFieldNames.isEmpty()) {
if (returnFieldsNames != null && !returnFieldsNames.isEmpty()) {
args.add(RETURN);
LazyRawable returnCountObject = new LazyRawable();
args.add(returnCountObject); // holding a place for setting the total count later.
int returnCount = 0;
for (FieldName fn : returnFieldNames) {
for (FieldName fn : returnFieldsNames) {
returnCount += fn.addCommandArguments(args);
}
returnCountObject.setRaw(Protocol.toByteArray(returnCount));
} else if (returnFields != null && !returnFields.isEmpty()) {
args.add(RETURN).add(returnFields.size()).addObjects(returnFields);
}

if (summarizeParams != null) {
Expand Down Expand Up @@ -256,41 +256,46 @@ public FTSearchParams inFields(Collection<String> fields) {
* @return the query object itself
*/
public FTSearchParams returnFields(String... fields) {
if (returnFieldNames != null) {
Arrays.stream(fields).forEach(f -> returnFieldNames.add(FieldName.of(f)));
} else {
if (returnFields == null) {
returnFields = new ArrayList<>();
}
Arrays.stream(fields).forEach(f -> returnFields.add(f));
if (returnFieldsNames == null) {
returnFieldsNames = new ArrayList<>();
}
Arrays.stream(fields).forEach(f -> returnFieldsNames.add(FieldName.of(f)));
return this;
}

public FTSearchParams returnField(FieldName field) {
initReturnFieldNames();
returnFieldNames.add(field);
return this;
return returnFields(Collections.singleton(field));
}

public FTSearchParams returnFields(FieldName... fields) {
return returnFields(Arrays.asList(fields));
}

public FTSearchParams returnFields(Collection<FieldName> fields) {
initReturnFieldNames();
returnFieldNames.addAll(fields);
if (returnFieldsNames == null) {
returnFieldsNames = new ArrayList<>();
}
returnFieldsNames.addAll(fields);
return this;
}

private void initReturnFieldNames() {
if (returnFieldNames == null) {
returnFieldNames = new ArrayList<>();
}
if (returnFields != null) {
returnFields.forEach(f -> returnFieldNames.add(FieldName.of(f)));
returnFields = null;
public FTSearchParams returnField(String field, boolean decode) {
returnFields(field);
addReturnFieldDecode(field, decode);
return this;
}

public FTSearchParams returnField(FieldName field, boolean decode) {
returnFields(field);
addReturnFieldDecode(field.getAttribute() != null ? field.getAttribute() : field.getName(), decode);
return this;
}

private void addReturnFieldDecode(String returnName, boolean decode) {
if (returnFieldDecodeMap == null) {
returnFieldDecodeMap = new HashMap<>();
}
returnFieldDecodeMap.put(returnName, decode);
}

public FTSearchParams summarize() {
Expand Down Expand Up @@ -436,14 +441,21 @@ public FTSearchParams dialectOptional(int dialect) {
return this;
}

@Internal
public boolean getNoContent() {
return noContent;
}

@Internal
public boolean getWithScores() {
return withScores;
}

@Internal
public Map<String, Boolean> getReturnFieldDecodeMap() {
return returnFieldDecodeMap;
}

/**
* NumericFilter wraps a range filter on a numeric field. It can be inclusive or exclusive
*/
Expand Down
Loading

0 comments on commit 81c9f32

Please sign in to comment.