org.apache.lucene.queries.function.valuesource.EnumFieldSource Maven / Gradle / Ivy
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.lucene.queries.function.valuesource;
import java.io.IOException;
import java.util.Map;
import org.apache.lucene.index.DocValues;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.queries.function.FunctionValues;
import org.apache.lucene.queries.function.ValueSourceScorer;
import org.apache.lucene.queries.function.docvalues.IntDocValues;
import org.apache.lucene.search.Weight;
/**
* Obtains int field values from {@link org.apache.lucene.index.LeafReader#getNumericDocValues} and
* makes those values available as other numeric types, casting as needed. strVal of the value is
* not the int value, but its string (displayed) value
*/
public class EnumFieldSource extends FieldCacheSource {
static final Integer DEFAULT_VALUE = -1;
final Map enumIntToStringMap;
final Map enumStringToIntMap;
public EnumFieldSource(
String field,
Map enumIntToStringMap,
Map enumStringToIntMap) {
super(field);
this.enumIntToStringMap = enumIntToStringMap;
this.enumStringToIntMap = enumStringToIntMap;
}
private static Integer tryParseInt(String valueStr) {
Integer intValue = null;
try {
intValue = Integer.parseInt(valueStr);
} catch (
@SuppressWarnings("unused")
NumberFormatException e) {
}
return intValue;
}
private String intValueToStringValue(Integer intVal) {
if (intVal == null) {
return null;
}
final String enumString = enumIntToStringMap.get(intVal);
if (enumString != null) {
return enumString;
}
// can't find matching enum name - return DEFAULT_VALUE.toString()
return DEFAULT_VALUE.toString();
}
private Integer stringValueToIntValue(String stringVal) {
if (stringVal == null) {
return null;
}
Integer intValue;
final Integer enumInt = enumStringToIntMap.get(stringVal);
if (enumInt != null) // enum int found for string
return enumInt;
// enum int not found for string
intValue = tryParseInt(stringVal);
if (intValue == null) // not Integer
intValue = DEFAULT_VALUE;
final String enumString = enumIntToStringMap.get(intValue);
if (enumString != null) // has matching string
return intValue;
return DEFAULT_VALUE;
}
@Override
public String description() {
return "enum(" + field + ')';
}
@Override
public FunctionValues getValues(Map context, LeafReaderContext readerContext)
throws IOException {
final NumericDocValues arr = DocValues.getNumeric(readerContext.reader(), field);
return new IntDocValues(this) {
int lastDocID;
@Override
public int intVal(int doc) throws IOException {
if (exists(doc)) {
return (int) arr.longValue();
} else {
return 0;
}
}
@Override
public String strVal(int doc) throws IOException {
Integer intValue = intVal(doc);
return intValueToStringValue(intValue);
}
@Override
public boolean exists(int doc) throws IOException {
if (doc < lastDocID) {
throw new IllegalArgumentException(
"docs were sent out-of-order: lastDocID=" + lastDocID + " vs docID=" + doc);
}
lastDocID = doc;
int curDocID = arr.docID();
if (doc > curDocID) {
curDocID = arr.advance(doc);
}
return doc == curDocID;
}
@Override
public ValueSourceScorer getRangeScorer(
Weight weight,
LeafReaderContext readerContext,
String lowerVal,
String upperVal,
boolean includeLower,
boolean includeUpper) {
Integer lower = stringValueToIntValue(lowerVal);
Integer upper = stringValueToIntValue(upperVal);
// instead of using separate comparison functions, adjust the endpoints.
if (lower == null) {
lower = Integer.MIN_VALUE;
} else {
if (!includeLower && lower < Integer.MAX_VALUE) lower++;
}
if (upper == null) {
upper = Integer.MAX_VALUE;
} else {
if (!includeUpper && upper > Integer.MIN_VALUE) upper--;
}
final int ll = lower;
final int uu = upper;
return new ValueSourceScorer(weight, readerContext, this) {
@Override
public boolean matches(int doc) throws IOException {
if (!exists(doc)) return false;
int val = intVal(doc);
return val >= ll && val <= uu;
}
};
}
};
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
if (!super.equals(o)) return false;
EnumFieldSource that = (EnumFieldSource) o;
if (!enumIntToStringMap.equals(that.enumIntToStringMap)) return false;
if (!enumStringToIntMap.equals(that.enumStringToIntMap)) return false;
return true;
}
@Override
public int hashCode() {
int result = super.hashCode();
result = 31 * result + enumIntToStringMap.hashCode();
result = 31 * result + enumStringToIntMap.hashCode();
return result;
}
}