add virtual columns to search query cache key (#12907)

* add virtual columns to search query cache key
This commit is contained in:
Clint Wylie 2022-08-17 20:26:01 -07:00 committed by GitHub
parent 770358dc34
commit 7fb1153bba
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 113 additions and 73 deletions

View File

@ -113,48 +113,30 @@ public class CacheKeyBuilder
}
}
private static byte[] stringCollectionToByteArray(Collection<String> input, boolean preserveOrder)
private static byte[] stringCollectionToByteArray(
@Nullable Collection<String> input,
boolean preserveOrder
)
{
return collectionToByteArray(
input,
new Function<String, byte[]>()
{
@Override
public byte[] apply(@Nullable String input)
{
return StringUtils.toUtf8WithNullToEmpty(input);
}
},
STRING_SEPARATOR,
preserveOrder
);
return collectionToByteArray(input, StringUtils::toUtf8WithNullToEmpty, STRING_SEPARATOR, preserveOrder);
}
private static byte[] cacheableCollectionToByteArray(Collection<? extends Cacheable> input, boolean preserveOrder)
private static byte[] cacheableCollectionToByteArray(
@Nullable Collection<? extends Cacheable> input,
boolean preserveOrder
)
{
return collectionToByteArray(
input,
new Function<Cacheable, byte[]>()
{
@Override
public byte[] apply(@Nullable Cacheable input)
{
return input == null ? EMPTY_BYTES : input.getCacheKey();
}
},
EMPTY_BYTES,
preserveOrder
);
return collectionToByteArray(input, CacheKeyBuilder::cacheableToByteArray, EMPTY_BYTES, preserveOrder);
}
private static <T> byte[] collectionToByteArray(
Collection<? extends T> collection,
@Nullable Collection<? extends T> collection,
Function<T, byte[]> serializeFunction,
byte[] separator,
boolean preserveOrder
)
{
if (collection.size() > 0) {
if (collection != null && collection.size() > 0) {
List<byte[]> byteArrayList = Lists.newArrayListWithCapacity(collection.size());
int totalByteLength = 0;
for (T eachItem : collection) {
@ -184,6 +166,7 @@ public class CacheKeyBuilder
}
}
private final List<Item> items = new ArrayList<>();
private final byte id;
private int size;

View File

@ -26,7 +26,6 @@ import com.google.common.base.Functions;
import com.google.common.base.Preconditions;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.primitives.Ints;
import com.google.inject.Inject;
import org.apache.druid.java.util.common.DateTimes;
import org.apache.druid.java.util.common.IAE;
@ -42,12 +41,11 @@ import org.apache.druid.query.QueryToolChest;
import org.apache.druid.query.Result;
import org.apache.druid.query.ResultGranularTimestampComparator;
import org.apache.druid.query.aggregation.MetricManipulationFn;
import org.apache.druid.query.cache.CacheKeyBuilder;
import org.apache.druid.query.context.ResponseContext;
import org.apache.druid.query.dimension.DimensionSpec;
import org.apache.druid.query.filter.DimFilter;
import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
@ -147,41 +145,14 @@ public class SearchQueryQueryToolChest extends QueryToolChest<Result<SearchResul
@Override
public byte[] computeCacheKey(SearchQuery query)
{
final DimFilter dimFilter = query.getDimensionsFilter();
final byte[] filterBytes = dimFilter == null ? new byte[]{} : dimFilter.getCacheKey();
final byte[] querySpecBytes = query.getQuery().getCacheKey();
final byte[] granularityBytes = query.getGranularity().getCacheKey();
final List<DimensionSpec> dimensionSpecs =
query.getDimensions() != null ? query.getDimensions() : Collections.emptyList();
final byte[][] dimensionsBytes = new byte[dimensionSpecs.size()][];
int dimensionsBytesSize = 0;
int index = 0;
for (DimensionSpec dimensionSpec : dimensionSpecs) {
dimensionsBytes[index] = dimensionSpec.getCacheKey();
dimensionsBytesSize += dimensionsBytes[index].length;
++index;
}
final byte[] sortSpecBytes = query.getSort().getCacheKey();
final ByteBuffer queryCacheKey = ByteBuffer
.allocate(
1 + 4 + granularityBytes.length + filterBytes.length +
querySpecBytes.length + dimensionsBytesSize + sortSpecBytes.length
)
.put(SEARCH_QUERY)
.put(Ints.toByteArray(query.getLimit()))
.put(granularityBytes)
.put(filterBytes)
.put(querySpecBytes)
.put(sortSpecBytes);
for (byte[] bytes : dimensionsBytes) {
queryCacheKey.put(bytes);
}
return queryCacheKey.array();
return new CacheKeyBuilder(SEARCH_QUERY).appendInt(query.getLimit())
.appendCacheable(query.getGranularity())
.appendCacheable(query.getFilter())
.appendCacheable(query.getQuery())
.appendCacheable(query.getSort())
.appendCacheables(query.getDimensions())
.appendCacheable(query.getVirtualColumns())
.build();
}
@Override

View File

@ -22,6 +22,7 @@ package org.apache.druid.query.search;
import com.fasterxml.jackson.annotation.JsonSubTypes;
import com.fasterxml.jackson.annotation.JsonTypeInfo;
import org.apache.druid.annotations.SubclassesMustOverrideEqualsAndHashCode;
import org.apache.druid.java.util.common.Cacheable;
import javax.annotation.Nullable;
@ -36,9 +37,7 @@ import javax.annotation.Nullable;
@JsonSubTypes.Type(name = "all", value = AllSearchQuerySpec.class)
})
@SubclassesMustOverrideEqualsAndHashCode
public interface SearchQuerySpec
public interface SearchQuerySpec extends Cacheable
{
boolean accept(@Nullable String dimVal);
byte[] getCacheKey();
}

View File

@ -21,13 +21,14 @@ package org.apache.druid.query.search;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.apache.druid.java.util.common.Cacheable;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.query.ordering.StringComparator;
import org.apache.druid.query.ordering.StringComparators;
import java.util.Comparator;
public class SearchSortSpec
public class SearchSortSpec implements Cacheable
{
public static final StringComparator DEFAULT_ORDERING = StringComparators.LEXICOGRAPHIC;
@ -66,6 +67,7 @@ public class SearchSortSpec
};
}
@Override
public byte[] getCacheKey()
{
return ordering.getCacheKey();

View File

@ -144,6 +144,7 @@ public class NestedFieldVirtualColumn implements VirtualColumn
{
final String partsString = NestedPathFinder.toNormalizedJsonPath(parts);
return new CacheKeyBuilder(VirtualColumnCacheHelper.CACHE_TYPE_ID_USER_DEFINED).appendString("nested-field")
.appendString(outputName)
.appendString(columnName)
.appendString(partsString)
.appendBoolean(processFromRaw)

View File

@ -29,12 +29,18 @@ import org.apache.druid.query.CacheStrategy;
import org.apache.druid.query.Druids;
import org.apache.druid.query.Result;
import org.apache.druid.query.TableDataSource;
import org.apache.druid.query.expression.TestExprMacroTable;
import org.apache.druid.query.spec.MultipleIntervalSegmentSpec;
import org.apache.druid.segment.VirtualColumns;
import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.virtual.ExpressionVirtualColumn;
import org.apache.druid.testing.InitializedNullHandlingTest;
import org.junit.Assert;
import org.junit.Test;
public class SearchQueryQueryToolChestTest
import java.util.Arrays;
public class SearchQueryQueryToolChestTest extends InitializedNullHandlingTest
{
@Test
@ -75,4 +81,82 @@ public class SearchQueryQueryToolChestTest
Assert.assertEquals(result, fromCacheResult);
}
@Test
public void testCacheStrategyVirtualColumns()
{
SearchQueryQueryToolChest toolChest = new SearchQueryQueryToolChest(null, null);
SearchQuery query1 = new SearchQuery(
new TableDataSource("dummy"),
null,
Granularities.ALL,
1,
new MultipleIntervalSegmentSpec(ImmutableList.of(Intervals.of("2015-01-01/2015-01-02"))),
ImmutableList.of(Druids.DIMENSION_IDENTITY.apply("v0")),
VirtualColumns.create(
ImmutableList.of(
new ExpressionVirtualColumn("v0", "concat(dim1, 'foo')", ColumnType.STRING, TestExprMacroTable.INSTANCE)
)
),
new FragmentSearchQuerySpec(ImmutableList.of("a", "b")),
null,
null
);
SearchQuery query2 = new SearchQuery(
new TableDataSource("dummy"),
null,
Granularities.ALL,
1,
new MultipleIntervalSegmentSpec(ImmutableList.of(Intervals.of("2015-01-01/2015-01-02"))),
ImmutableList.of(Druids.DIMENSION_IDENTITY.apply("v0")),
VirtualColumns.create(
ImmutableList.of(
new ExpressionVirtualColumn("v0", "concat(dim2, 'foo')", ColumnType.STRING, TestExprMacroTable.INSTANCE)
)
),
new FragmentSearchQuerySpec(ImmutableList.of("a", "b")),
null,
null
);
SearchQuery query3 = new SearchQuery(
new TableDataSource("dummy"),
null,
Granularities.ALL,
1,
new MultipleIntervalSegmentSpec(ImmutableList.of(Intervals.of("2015-01-01/2015-01-02"))),
ImmutableList.of(Druids.DIMENSION_IDENTITY.apply("v0")),
VirtualColumns.create(
ImmutableList.of(
new ExpressionVirtualColumn("v0", "concat(dim1, 'foo')", ColumnType.STRING, TestExprMacroTable.INSTANCE)
)
),
new FragmentSearchQuerySpec(ImmutableList.of("a", "b")),
null,
null
);
Assert.assertArrayEquals(
toolChest.getCacheStrategy(query1).computeCacheKey(query1),
toolChest.getCacheStrategy(query1).computeCacheKey(query1)
);
Assert.assertArrayEquals(
toolChest.getCacheStrategy(query2).computeCacheKey(query2),
toolChest.getCacheStrategy(query2).computeCacheKey(query2)
);
Assert.assertArrayEquals(
toolChest.getCacheStrategy(query1).computeCacheKey(query1),
toolChest.getCacheStrategy(query3).computeCacheKey(query3)
);
Assert.assertFalse(
Arrays.equals(
toolChest.getCacheStrategy(query1).computeCacheKey(query1),
toolChest.getCacheStrategy(query2).computeCacheKey(query2)
)
);
}
}