/*
* Copyright 2015, The Querydsl Team (http://www.querydsl.com/team)
*
* Licensed 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 com.querydsl.lucene4;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import javax.annotation.Nullable;
import org.apache.lucene.document.Document;
import org.apache.lucene.queries.ChainedFilter;
import org.apache.lucene.sandbox.queries.DuplicateFilter;
import org.apache.lucene.search.*;
import org.apache.lucene.search.Query;
import com.google.common.base.CharMatcher;
import com.google.common.base.Function;
import com.google.common.collect.ImmutableList;
import com.mysema.commons.lang.CloseableIterator;
import com.mysema.commons.lang.EmptyCloseableIterator;
import com.mysema.commons.lang.IteratorAdapter;
import com.querydsl.core.*;
import com.querydsl.core.support.QueryMixin;
import com.querydsl.core.types.OrderSpecifier;
import com.querydsl.core.types.ParamExpression;
import com.querydsl.core.types.Path;
import com.querydsl.core.types.Predicate;
/**
* AbstractLuceneQuery is an abstract super class for Lucene query implementations
*
* @author tiwe
*
* @param <T> projection type
* @param <Q> concrete subtype of querydsl
*/
public abstract class AbstractLuceneQuery<T,Q extends AbstractLuceneQuery<T,Q>> implements SimpleQuery<Q>, Fetchable<T> {
private final QueryMixin<Q> queryMixin;
private final IndexSearcher searcher;
private final LuceneSerializer serializer;
private final Function<Document, T> transformer;
@Nullable
private Set<String> fieldsToLoad;
private List<Filter> filters = ImmutableList.of();
@Nullable
private Filter filter;
@Nullable
private Sort querySort;
@SuppressWarnings("unchecked")
public AbstractLuceneQuery(LuceneSerializer serializer, IndexSearcher searcher,
Function<Document, T> transformer) {
queryMixin = new QueryMixin<Q>((Q) this, new DefaultQueryMetadata());
this.serializer = serializer;
this.searcher = searcher;
this.transformer = transformer;
}
public AbstractLuceneQuery(IndexSearcher searcher, Function<Document, T> transformer) {
this(LuceneSerializer.DEFAULT, searcher, transformer);
}
private long innerCount() {
try {
final int maxDoc = searcher.getIndexReader().maxDoc();
if (maxDoc == 0) {
return 0;
}
TotalHitCountCollector collector = new TotalHitCountCollector();
searcher.search(createQuery(), getFilter(), collector);
return collector.getTotalHits();
} catch (IOException e) {
throw new QueryException(e);
} catch (IllegalArgumentException e) {
throw new QueryException(e);
}
}
@Override
public long fetchCount() {
return innerCount();
}
protected Query createQuery() {
if (queryMixin.getMetadata().getWhere() == null) {
return new MatchAllDocsQuery();
}
return serializer.toQuery(queryMixin.getMetadata().getWhere(), queryMixin.getMetadata());
}
/**
* Create a filter for constraints defined in this querydsl
*
* @return filter
*/
public Filter asFilter() {
return new QueryWrapperFilter(createQuery());
}
@Override
public Q distinct() {
throw new UnsupportedOperationException("use distinct(path) instead");
}
/**
* Add a DuplicateFilter for the field of the given property path
*
* @param property distinct property
* @return the current object
*/
public Q distinct(Path<?> property) {
return filter(new DuplicateFilter(serializer.toField(property)));
}
/**
* Apply the given Lucene filter to the search results
*
* @param filter filter
* @return the current object
*/
@SuppressWarnings("unchecked")
public Q filter(Filter filter) {
if (filters.isEmpty()) {
this.filter = filter;
filters = ImmutableList.of(filter);
} else {
this.filter = null;
if (filters.size() == 1) {
filters = new ArrayList<Filter>();
}
filters.add(filter);
}
return (Q) this;
}
private Filter getFilter() {
if (filter == null && !filters.isEmpty()) {
filter = new ChainedFilter(filters.toArray(new Filter[filters.size()]));
}
return filter;
}
@Override
public Q limit(long limit) {
return queryMixin.limit(limit);
}
@Override
public CloseableIterator<T> iterate() {
final QueryMetadata metadata = queryMixin.getMetadata();
final List<OrderSpecifier<?>> orderBys = metadata.getOrderBy();
final Integer queryLimit = metadata.getModifiers().getLimitAsInteger();
final Integer queryOffset = metadata.getModifiers().getOffsetAsInteger();
Sort sort = querySort;
int limit;
final int offset = queryOffset != null ? queryOffset : 0;
try {
limit = maxDoc();
if (limit == 0) {
return new EmptyCloseableIterator<T>();
}
} catch (IOException e) {
throw new QueryException(e);
} catch (IllegalArgumentException e) {
throw new QueryException(e);
}
if (queryLimit != null && queryLimit < limit) {
limit = queryLimit;
}
if (sort == null && !orderBys.isEmpty()) {
sort = serializer.toSort(orderBys);
}
try {
ScoreDoc[] scoreDocs;
int sumOfLimitAndOffset = limit + offset;
if (sumOfLimitAndOffset < 1) {
throw new QueryException("The given limit (" + limit + ") and offset (" + offset + ") cause an integer overflow.");
}
if (sort != null) {
scoreDocs = searcher.search(createQuery(), getFilter(), sumOfLimitAndOffset, sort, false, false).scoreDocs;
} else {
scoreDocs = searcher.search(createQuery(), getFilter(), sumOfLimitAndOffset, Sort.INDEXORDER, false, false).scoreDocs;
}
if (offset < scoreDocs.length) {
return new ResultIterator<T>(scoreDocs, offset, searcher, fieldsToLoad, transformer);
}
return new EmptyCloseableIterator<T>();
} catch (final IOException e) {
throw new QueryException(e);
}
}
private List<T> innerList() {
return new IteratorAdapter<T>(iterate()).asList();
}
@Override
public List<T> fetch() {
return innerList();
}
/**
* Set the given fields to load
*
* @param fieldsToLoad fields to load
* @return the current object
*/
@SuppressWarnings("unchecked")
public Q load(Set<String> fieldsToLoad) {
this.fieldsToLoad = fieldsToLoad;
return (Q) this;
}
/**
* Load only the fields of the given paths
*
* @param paths fields to load
* @return the current object
*/
@SuppressWarnings("unchecked")
public Q load(Path<?>... paths) {
Set<String> fields = new HashSet<String>();
for (Path<?> path : paths) {
fields.add(serializer.toField(path));
}
this.fieldsToLoad = fields;
return (Q) this;
}
@Override
public QueryResults<T> fetchResults() {
List<T> documents = innerList();
/*
* TODO Get rid of fetchCount(). It could be implemented by iterating the
* fetch results in fetch* from n to m.
*/
return new QueryResults<T>(documents, queryMixin.getMetadata().getModifiers(), innerCount());
}
@Override
public Q offset(long offset) {
return queryMixin.offset(offset);
}
public Q orderBy(OrderSpecifier<?> o) {
return queryMixin.orderBy(o);
}
@Override
public Q orderBy(OrderSpecifier<?>... o) {
return queryMixin.orderBy(o);
}
@Override
public Q restrict(QueryModifiers modifiers) {
return queryMixin.restrict(modifiers);
}
@Override
public <P> Q set(ParamExpression<P> param, P value) {
return queryMixin.set(param, value);
}
@SuppressWarnings("unchecked")
public Q sort(Sort sort) {
this.querySort = sort;
return (Q) this;
}
@Nullable
private T oneResult(boolean unique) {
try {
int maxDoc = maxDoc();
if (maxDoc == 0) {
return null;
}
final ScoreDoc[] scoreDocs = searcher.search(createQuery(), getFilter(), maxDoc, Sort.INDEXORDER, false, false).scoreDocs;
int index = 0;
QueryModifiers modifiers = queryMixin.getMetadata().getModifiers();
Long offset = modifiers.getOffset();
if (offset != null) {
index = offset.intValue();
}
Long limit = modifiers.getLimit();
if (unique && (limit == null ? scoreDocs.length - index > 1 :
limit > 1 && scoreDocs.length > 1)) {
throw new NonUniqueResultException("Unique result requested, but " + scoreDocs.length + " found.");
} else if (scoreDocs.length > index) {
Document document;
if (fieldsToLoad != null) {
document = searcher.doc(scoreDocs[index].doc, fieldsToLoad);
} else {
document = searcher.doc(scoreDocs[index].doc);
}
return transformer.apply(document);
} else {
return null;
}
} catch (IOException e) {
throw new QueryException(e);
} catch (IllegalArgumentException e) {
throw new QueryException(e);
}
}
@Override
public T fetchFirst() {
return oneResult(false);
}
@Override
public T fetchOne() {
return oneResult(true);
}
public Q where(Predicate e) {
return queryMixin.where(e);
}
@Override
public Q where(Predicate... e) {
return queryMixin.where(e);
}
@Override
public String toString() {
String str = createQuery().toString();
return CharMatcher.JAVA_ISO_CONTROL.replaceFrom(str, '_');
}
private int maxDoc() throws IOException {
return searcher.getIndexReader().maxDoc();
}
}