Skip to content

Instantly share code, notes, and snippets.

@chaitanyamannem
Created May 11, 2016 07:19
Show Gist options
  • Save chaitanyamannem/2655c74fd551fef9229477d33f2bcb26 to your computer and use it in GitHub Desktop.
Save chaitanyamannem/2655c74fd551fef9229477d33f2bcb26 to your computer and use it in GitHub Desktop.
Implementation of solr custom function class which extends ValueSource
package com.solrExample.solr4;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.logging.Logger;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.analysis.util.TokenFilterFactory;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.queries.function.FunctionValues;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.queries.function.docvalues.IntDocValues;
import org.apache.lucene.queryparser.classic.QueryParser;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopScoreDocCollector;
import org.apache.lucene.util.Version;
import org.apache.solr.common.SolrException;
/**
*
* Created May 10, 2016 11:18:12 AM
*
* @author chaitanya
*/
public class AvailabilityValueSource extends ValueSource {
protected String startDateTime;
protected String endDateTime;
protected String fieldName;
public AvailabilityValueSource(String startDateTime, String endDateTime,
String fieldName) {
if (startDateTime == null || endDateTime == null || fieldName == null) {
throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
"One or more inputs missing for petAvail function");
}
this.startDateTime = startDateTime;
this.endDateTime = endDateTime;
this.fieldName = fieldName;
}
@Override
public String description() {
return "";
}
@Override
public boolean equals(Object otherObj) {
if (otherObj == null || !(otherObj instanceof AvailabilityValueSource))
return false;
AvailabilityValueSource other = (AvailabilityValueSource) otherObj;
return this.startDateTime.equals(other.startDateTime)
&& this.endDateTime.equals(other.endDateTime)
&& this.fieldName.equals(other.fieldName);
}
protected String name() {
return "AvailabilityValueSource";
}
@Override
public FunctionValues getValues(Map context, AtomicReaderContext reader)
throws IOException {
final Map<String, String> availableUsers =
(Map<String, String>) context.get("availableUsers");
return new IntDocValues(this) {
@Override
public int intVal(int doc) {
if (availableUsers != null) {
String isAvailable = availableUsers.get(String.valueOf(doc));
// if current doc is found in availableUsers map then his availability
// matches return 1 which can be used for sort
if (isAvailable != null) {
return 1;
}
}
return 0;
}
};
}
@Override
public int hashCode() {
return (int) startDateTime.hashCode() + endDateTime.hashCode() + fieldName.hashCode();
}
@Override
public void createWeight(Map context, IndexSearcher searcher) {
Map<String, String> availableUsers = new HashMap<String, String>();
try {
TokenFilterFactory[] factory = new TokenFilterFactory[0];
Analyzer analyzer = new StandardAnalyzer(Version.LUCENE_41);
// build availability query string using startDateTime, endDateTime and availabilityType
//Here we are taking the start date and end date and building the query based
//on how it is in the solr, eg: +availble:0506 +available:0507
//Here available is the solr filed name
String parsedDateTokens =
SplitDateIntoTokensParser.getAvailabilitySlots(startDateTime, endDateTime,
availabilityType);
Query q =
new QueryParser(Version.LUCENE_41, fieldName, analyzer).parse(parsedDateTokens);
TopScoreDocCollector collector =
TopScoreDocCollector.create(searcher.getIndexReader().numDocs(), true);
searcher.search(q, collector);
ScoreDoc[] hits = collector.topDocs().scoreDocs;
// If availaiblity matches add those docs to map with value 1
for (int i = 0; i < hits.length; ++i) {
int docId = hits[i].doc;
Document doc = searcher.doc(docId);
availableProviders.put(String.valueOf(docId), String.valueOf(1));
}
} catch (Exception e) {
logger.severe("EXception while creating availability map" + e);
}
context.put("availableUsers", availableUsers);
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment