CustomFeature.java

package com.example.solr.ltr.feature;

import java.io.IOException;
import java.util.LinkedHashMap;
import java.util.Map;

import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.ReaderUtil;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.TermContext;
import org.apache.lucene.index.TermState;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.TermQuery;
import org.apache.solr.ltr.feature.Feature;
import org.apache.solr.ltr.feature.FeatureException;
import org.apache.solr.request.SolrQueryRequest;

public class CustomFeature extends Feature {
	public CustomFeature(String name, Map<String, Object> params) {
		super(name, params);
	}

	@Override
	protected void validate() throws FeatureException {

	}

	@Override
	public FeatureWeight createWeight(IndexSearcher searcher, boolean needsScores, SolrQueryRequest request, Query originalQuery,
			Map<String, String[]> efi) throws IOException {
		return new CustomWeight(this, searcher, request, (TermQuery) originalQuery, efi);
	}

	@Override
	public LinkedHashMap<String, Object> paramsToMap() {
		return null;
	}

	public class CustomWeight extends Feature.FeatureWeight {
		private final Term term;
		private final TermContext termStates;

		public CustomWeight(Query featureQuery, IndexSearcher indexSearcher, SolrQueryRequest solrRequest,
				TermQuery originalQuery, Map<String, String[]> efi) throws IOException {
			super(featureQuery, indexSearcher, solrRequest, originalQuery, efi);
			term = originalQuery.getTerm();
			termStates = TermContext.build(indexSearcher.getTopReaderContext(), term);
		}

		@Override
		public float getValueForNormalization() throws IOException {
			return 0;
		}

		@Override
		public void normalize(float norm, float boost) {

		}

		@Override
		public FeatureScorer scorer(LeafReaderContext context) throws IOException {
			TermsEnum termsEnum = getTermsEnum(context);
			if (termsEnum == null) {
				return null;
			}
			PostingsEnum postingsEnum = termsEnum.postings(null, PostingsEnum.FREQS);
			return new TFScorer(this, postingsEnum);
		}

		protected TermsEnum getTermsEnum(LeafReaderContext context) throws IOException {
			assert termStates.wasBuiltFor(ReaderUtil.getTopLevelContext(
					context)) : "The top-reader used to create Weight is not the same as the current reader's top-reader ("
							+ ReaderUtil.getTopLevelContext(context);
			final TermState state = termStates.get(context.ord);
			if (state == null) {
				if (context.reader().docFreq(term) != 0) {
					throw new RuntimeException("Term is exist but TermState is not found. term=" + term);
				}
				return null;
			}
			final TermsEnum termsEnum = context.reader().terms(term.field()).iterator();
			termsEnum.seekExact(term.bytes(), state);
			return termsEnum;
		}

		public class TFScorer extends FeatureScorer {
			private final PostingsEnum postingsEnum;

			public TFScorer(FeatureWeight weight, PostingsEnum postingsEnum) {
				super(weight, postingsEnum);
				this.postingsEnum = postingsEnum;
			}

			@Override
			public float score() throws IOException {
				return postingsEnum.freq();
			}
		}
	}
}