org.apache.druid.query.search.SearchQueryRunner Maven / Gradle / Ivy
Show all versions of druid-processing Show documentation
/*
* 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.druid.query.search;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import it.unimi.dsi.fastutil.objects.Object2IntRBTreeMap;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.guava.FunctionalIterable;
import org.apache.druid.java.util.common.guava.Sequence;
import org.apache.druid.java.util.common.guava.Sequences;
import org.apache.druid.query.Query;
import org.apache.druid.query.QueryPlus;
import org.apache.druid.query.QueryRunner;
import org.apache.druid.query.Result;
import org.apache.druid.query.context.ResponseContext;
import org.apache.druid.query.dimension.ColumnSelectorStrategy;
import org.apache.druid.query.dimension.ColumnSelectorStrategyFactory;
import org.apache.druid.segment.BaseDoubleColumnValueSelector;
import org.apache.druid.segment.BaseFloatColumnValueSelector;
import org.apache.druid.segment.BaseLongColumnValueSelector;
import org.apache.druid.segment.ColumnValueSelector;
import org.apache.druid.segment.DimensionSelector;
import org.apache.druid.segment.Segment;
import org.apache.druid.segment.column.ColumnCapabilities;
import org.apache.druid.segment.data.IndexedInts;
import java.util.List;
/**
*/
public class SearchQueryRunner implements QueryRunner>
{
public static final SearchColumnSelectorStrategyFactory SEARCH_COLUMN_SELECTOR_STRATEGY_FACTORY
= new SearchColumnSelectorStrategyFactory();
private final Segment segment;
private final SearchStrategySelector strategySelector;
public SearchQueryRunner(Segment segment, SearchStrategySelector strategySelector)
{
this.segment = segment;
this.strategySelector = strategySelector;
}
private static class SearchColumnSelectorStrategyFactory
implements ColumnSelectorStrategyFactory
{
@Override
public SearchColumnSelectorStrategy makeColumnSelectorStrategy(
ColumnCapabilities capabilities,
ColumnValueSelector selector
)
{
switch (capabilities.getType()) {
case STRING:
return new StringSearchColumnSelectorStrategy();
case LONG:
return new LongSearchColumnSelectorStrategy();
case FLOAT:
return new FloatSearchColumnSelectorStrategy();
case DOUBLE:
return new DoubleSearchColumnSelectorStrategy();
default:
throw new IAE("Cannot create query type helper from invalid type [%s]", capabilities.asTypeString());
}
}
}
public interface SearchColumnSelectorStrategy extends ColumnSelectorStrategy
{
/**
* Read the current row from dimSelector and update the search result set.
*
* For each row value:
* 1. Check if searchQuerySpec accept()s the value
* 2. If so, add the value to the result set and increment the counter for that value
* 3. If the size of the result set reaches the limit after adding a value, return early.
*
* @param outputName Output name for this dimension in the search query being served
* @param dimSelector Dimension value selector
* @param searchQuerySpec Spec for the search query
* @param set The result set of the search query
* @param limit The limit of the search query
*/
void updateSearchResultSet(
String outputName,
ValueSelectorType dimSelector,
SearchQuerySpec searchQuerySpec,
int limit,
Object2IntRBTreeMap set
);
}
public static class StringSearchColumnSelectorStrategy implements SearchColumnSelectorStrategy
{
@Override
public void updateSearchResultSet(
String outputName,
DimensionSelector selector,
SearchQuerySpec searchQuerySpec,
int limit,
final Object2IntRBTreeMap set
)
{
if (!DimensionSelector.isNilSelector(selector)) {
final IndexedInts row = selector.getRow();
for (int i = 0, rowSize = row.size(); i < rowSize; ++i) {
final String dimVal = selector.lookupName(row.get(i));
if (searchQuerySpec.accept(dimVal)) {
set.addTo(new SearchHit(outputName, dimVal), 1);
if (set.size() >= limit) {
return;
}
}
}
}
}
}
public static class LongSearchColumnSelectorStrategy
implements SearchColumnSelectorStrategy
{
@Override
public void updateSearchResultSet(
String outputName,
BaseLongColumnValueSelector selector,
SearchQuerySpec searchQuerySpec,
int limit,
Object2IntRBTreeMap set
)
{
if (selector != null) {
final String dimVal = selector.isNull() ? null : String.valueOf(selector.getLong());
if (searchQuerySpec.accept(dimVal)) {
set.addTo(new SearchHit(outputName, dimVal), 1);
}
}
}
}
public static class FloatSearchColumnSelectorStrategy
implements SearchColumnSelectorStrategy
{
@Override
public void updateSearchResultSet(
String outputName,
BaseFloatColumnValueSelector selector,
SearchQuerySpec searchQuerySpec,
int limit,
Object2IntRBTreeMap set
)
{
if (selector != null) {
final String dimVal = selector.isNull() ? null : String.valueOf(selector.getFloat());
if (searchQuerySpec.accept(dimVal)) {
set.addTo(new SearchHit(outputName, dimVal), 1);
}
}
}
}
public static class DoubleSearchColumnSelectorStrategy
implements SearchColumnSelectorStrategy
{
@Override
public void updateSearchResultSet(
String outputName,
BaseDoubleColumnValueSelector selector,
SearchQuerySpec searchQuerySpec,
int limit,
Object2IntRBTreeMap set
)
{
if (selector != null) {
final String dimVal = selector.isNull() ? null : String.valueOf(selector.getDouble());
if (searchQuerySpec.accept(dimVal)) {
set.addTo(new SearchHit(outputName, dimVal), 1);
}
}
}
}
@Override
public Sequence> run(
final QueryPlus> queryPlus,
ResponseContext responseContext
)
{
Query> input = queryPlus.getQuery();
if (!(input instanceof SearchQuery)) {
throw new ISE("Got a [%s] which isn't a %s", input.getClass(), SearchQuery.class);
}
final SearchQuery query = (SearchQuery) input;
final List plan = strategySelector.strategize(query).getExecutionPlan(query, segment);
final Object2IntRBTreeMap retVal = new Object2IntRBTreeMap<>(query.getSort().getComparator());
retVal.defaultReturnValue(0);
int remain = query.getLimit();
for (final SearchQueryExecutor executor : plan) {
retVal.putAll(executor.execute(remain));
remain -= retVal.size();
}
return makeReturnResult(segment, query.getLimit(), retVal);
}
private static Sequence> makeReturnResult(
Segment segment,
int limit,
Object2IntRBTreeMap retVal
)
{
Iterable source = Iterables.transform(
retVal.object2IntEntrySet(),
input -> {
SearchHit hit = input.getKey();
return new SearchHit(hit.getDimension(), hit.getValue(), input.getIntValue());
}
);
return Sequences.simple(
ImmutableList.of(
new Result<>(
segment.getDataInterval().getStart(),
new SearchResultValue(
Lists.newArrayList(new FunctionalIterable<>(source).limit(limit))
)
)
)
);
}
}