From 4b97c745502eec28e7c3be3b642659968a7f10ad Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 2 Sep 2022 22:54:29 +0800 Subject: [PATCH 1/9] Use safe conversion methods --- .../MaterializedViewQuery.java | 53 ++---- .../org/apache/druid/query/BaseQuery.java | 7 - .../java/org/apache/druid/query/Query.java | 121 ++++++++++++- .../org/apache/druid/query/QueryContext.java | 29 +++- .../org/apache/druid/query/QueryContexts.java | 163 ++++++++++-------- .../druid/query/groupby/GroupByQuery.java | 2 +- .../query/groupby/GroupByQueryConfig.java | 16 +- .../query/groupby/GroupByQueryEngine.java | 2 +- .../query/groupby/GroupByQueryHelper.java | 4 +- .../epinephelinae/GroupByQueryEngineV2.java | 2 +- .../groupby/orderby/DefaultLimitSpec.java | 4 +- .../groupby/strategy/GroupByStrategyV2.java | 4 +- .../apache/druid/query/scan/ScanQuery.java | 4 +- .../druid/query/search/SearchQueryConfig.java | 2 +- .../druid/query/select/SelectQuery.java | 19 -- .../TimeBoundaryQueryQueryToolChest.java | 4 +- .../query/timeseries/TimeseriesQuery.java | 2 +- .../query/topn/TopNQueryQueryToolChest.java | 2 +- .../apache/druid/query/QueryContextTest.java | 71 +++++--- .../DataSourceMetadataQueryTest.java | 8 +- .../timeboundary/TimeBoundaryQueryTest.java | 16 +- .../druid/client/DirectDruidClient.java | 2 +- .../druid/client/JsonParserIterator.java | 2 +- .../appenderator/SinkQuerySegmentWalker.java | 2 +- .../server/ClientQuerySegmentWalker.java | 2 +- .../scheduling/HiLoQueryLaningStrategy.java | 4 +- ...ingClusteredClientCacheKeyManagerTest.java | 7 +- .../client/CachingClusteredClientTest.java | 2 +- .../druid/client/JsonParserIteratorTest.java | 2 +- .../SetAndVerifyContextQueryRunnerTest.java | 8 +- 30 files changed, 352 insertions(+), 214 deletions(-) diff --git a/extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQuery.java b/extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQuery.java index cb98d6b4b29c..73c74dba68ce 100644 --- a/extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQuery.java +++ b/extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQuery.java @@ -24,7 +24,6 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.base.Preconditions; import com.google.common.collect.Ordering; -import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.granularity.Granularity; import org.apache.druid.query.BaseQuery; import org.apache.druid.query.DataSource; @@ -47,20 +46,20 @@ import java.util.Objects; /** - * MaterializedViewQuery helps to do materialized view selection automatically. - * + * MaterializedViewQuery helps to do materialized view selection automatically. + * * Each MaterializedViewQuery contains a real query which type can be topn, timeseries or groupBy. * The real query will be optimized based on its dataSources and intervals. It will be converted into one or more * sub-queries, in which dataSources and intervals are replaced by derived dataSources and related sub-intervals. - * + * * Derived dataSources always have less dimensions, but contains all dimensions which real query required. */ -public class MaterializedViewQuery implements Query +public class MaterializedViewQuery implements Query { public static final String TYPE = "view"; private final Query query; private final DataSourceOptimizer optimizer; - + @JsonCreator public MaterializedViewQuery( @JsonProperty("query") Query query, @@ -74,24 +73,24 @@ public MaterializedViewQuery( this.query = query; this.optimizer = optimizer; } - + @JsonProperty("query") public Query getQuery() { return query; } - + public DataSourceOptimizer getOptimizer() { return optimizer; } - + @Override public DataSource getDataSource() { return query.getDataSource(); } - + @Override public boolean hasFilters() { @@ -111,14 +110,14 @@ public String getType() } @Override - public QueryRunner getRunner(QuerySegmentWalker walker) + public QueryRunner getRunner(QuerySegmentWalker walker) { return ((BaseQuery) query).getQuerySegmentSpec().lookup(this, walker); } @Override public List getIntervals() - + { return query.getIntervals(); } @@ -153,30 +152,6 @@ public QueryContext getQueryContext() return query.getQueryContext(); } - @Override - public ContextType getContextValue(String key) - { - return (ContextType) query.getContextValue(key); - } - - @Override - public ContextType getContextValue(String key, ContextType defaultValue) - { - return (ContextType) query.getContextValue(key, defaultValue); - } - - @Override - public boolean getContextBoolean(String key, boolean defaultValue) - { - return query.getContextBoolean(key, defaultValue); - } - - @Override - public HumanReadableBytes getContextHumanReadableBytes(String key, HumanReadableBytes defaultValue) - { - return query.getContextHumanReadableBytes(key, defaultValue); - } - @Override public boolean isDescending() { @@ -190,13 +165,13 @@ public Ordering getResultOrdering() } @Override - public MaterializedViewQuery withOverriddenContext(Map contextOverride) + public MaterializedViewQuery withOverriddenContext(Map contextOverride) { return new MaterializedViewQuery(query.withOverriddenContext(contextOverride), optimizer); } @Override - public MaterializedViewQuery withQuerySegmentSpec(QuerySegmentSpec spec) + public MaterializedViewQuery withQuerySegmentSpec(QuerySegmentSpec spec) { return new MaterializedViewQuery(query.withQuerySegmentSpec(spec), optimizer); } @@ -227,7 +202,7 @@ public String getSubQueryId() } @Override - public MaterializedViewQuery withDataSource(DataSource dataSource) + public MaterializedViewQuery withDataSource(DataSource dataSource) { return new MaterializedViewQuery(query.withDataSource(dataSource), optimizer); } diff --git a/processing/src/main/java/org/apache/druid/query/BaseQuery.java b/processing/src/main/java/org/apache/druid/query/BaseQuery.java index c581dfcfc374..12a24fb1d67f 100644 --- a/processing/src/main/java/org/apache/druid/query/BaseQuery.java +++ b/processing/src/main/java/org/apache/druid/query/BaseQuery.java @@ -187,13 +187,6 @@ public ContextType getContextValue(String key) return (ContextType) context.get(key); } - @Override - public ContextType getContextValue(String key, ContextType defaultValue) - { - ContextType retVal = getContextValue(key); - return retVal == null ? defaultValue : retVal; - } - @Override public boolean getContextBoolean(String key, boolean defaultValue) { diff --git a/processing/src/main/java/org/apache/druid/query/Query.java b/processing/src/main/java/org/apache/druid/query/Query.java index 9db2f0ca006d..4f6337c20f74 100644 --- a/processing/src/main/java/org/apache/druid/query/Query.java +++ b/processing/src/main/java/org/apache/druid/query/Query.java @@ -124,18 +124,129 @@ default QueryContext getQueryContext() return null; } - ContextType getContextValue(String key); + default boolean containsContextValue(String key) + { + if (getQueryContext() != null) { + return getQueryContext().containsKey(key); + } else { + return false; + } + } - ContextType getContextValue(String key, ContextType defaultValue); + /** + * Get context value and cast it to ContextType. + * NOTE: If value is not type of ContextType, {@link java.lang.ClassCastException} is raised. + * + * For safe conversion, please use following methods: + * {@link #getContextBoolean(String, boolean)} + * {@link #getContextAsString(String)} + * {@link #getContextAsInt(String)} + * {@link #getContextAsLong(String)} + * {@link #getContextAsFloat(String, float)} + * {@link #getContextAsEnum(String, Class, Enum)} + * {@link #getContextHumanReadableBytes(String, HumanReadableBytes)} + */ + @Nullable + default ContextType getContextValue(String key) + { + if (getQueryContext() != null) { + //noinspection unchecked + return (ContextType) getQueryContext().get(key); + } else { + return null; + } + } - boolean getContextBoolean(String key, boolean defaultValue); + @Nullable + default String getContextAsString(String key) + { + if (getQueryContext() != null) { + return getQueryContext().getAsString(key); + } else { + return null; + } + } + + default String getContextAsString(String key, String defaultValue) + { + if (getQueryContext() != null) { + return getQueryContext().getAsString(key, defaultValue); + } else { + return defaultValue; + } + } + + @Nullable + default Integer getContextAsInt(String key) + { + if (getQueryContext() != null) { + return getQueryContext().getAsInt(key); + } else { + return null; + } + } + + default int getContextAsInt(String key, int defaultValue) + { + if (getQueryContext() != null) { + return getQueryContext().getAsInt(key, defaultValue); + } else { + return defaultValue; + } + } + + @Nullable + default Long getContextAsLong(String key) + { + if (getQueryContext() != null) { + return getQueryContext().getAsLong(key); + } else { + return null; + } + } + + default long getContextAsLong(String key, long defaultValue) + { + if (getQueryContext() != null) { + return getQueryContext().getAsLong(key, defaultValue); + } else { + return defaultValue; + } + } + + default float getContextAsFloat(String key, float defaultValue) + { + if (getQueryContext() != null) { + return getQueryContext().getAsFloat(key, defaultValue); + } else { + return defaultValue; + } + } + + default > E getContextAsEnum(String key, Class clazz, E defaultValue) + { + if (getQueryContext() != null) { + return getQueryContext().getAsEnum(key, clazz, defaultValue); + } else { + return defaultValue; + } + } + + default boolean getContextBoolean(String key, boolean defaultValue) + { + if (getQueryContext() != null) { + return getQueryContext().getAsBoolean(key, defaultValue); + } else { + return defaultValue; + } + } /** * Returns {@link HumanReadableBytes} for a specified context key. If the context is null or the key doesn't exist * a caller specified default value is returned. A default implementation is provided since Query is an extension * point. Extensions can choose to rely on this default to retain compatibility with core Druid. * - * @param key The context key value being looked up + * @param key The context key value being looked up * @param defaultValue The default to return if the key value doesn't exist or the context is null. * @return {@link HumanReadableBytes} */ @@ -204,7 +315,7 @@ default Query withSqlQueryId(String sqlQueryId) @Nullable default String getSqlQueryId() { - return getContextValue(BaseQuery.SQL_QUERY_ID); + return getContextAsString(BaseQuery.SQL_QUERY_ID); } /** diff --git a/processing/src/main/java/org/apache/druid/query/QueryContext.java b/processing/src/main/java/org/apache/druid/query/QueryContext.java index 29f967429b74..52291ceda952 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContext.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContext.java @@ -168,7 +168,14 @@ public Object getOrDefault(String key, Object defaultValue) @Nullable public String getAsString(String key) { - return (String) get(key); + Object val = get(key); + return val == null ? null : val.toString(); + } + + public String getAsString(String key, String defaultValue) + { + Object val = get(key); + return val == null ? defaultValue : val.toString(); } public boolean getAsBoolean( @@ -179,6 +186,11 @@ public boolean getAsBoolean( return QueryContexts.getAsBoolean(parameter, get(parameter), defaultValue); } + public Integer getAsInt(final String parameter) + { + return QueryContexts.getAsInt(parameter, get(parameter)); + } + public int getAsInt( final String parameter, final int defaultValue @@ -187,6 +199,11 @@ public int getAsInt( return QueryContexts.getAsInt(parameter, get(parameter), defaultValue); } + public Long getAsLong(final String parameter) + { + return QueryContexts.getAsLong(parameter, get(parameter)); + } + public long getAsLong(final String parameter, final long defaultValue) { return QueryContexts.getAsLong(parameter, get(parameter), defaultValue); @@ -197,6 +214,16 @@ public HumanReadableBytes getAsHumanReadableBytes(final String parameter, final return QueryContexts.getAsHumanReadableBytes(parameter, get(parameter), defaultValue); } + public float getAsFloat(final String parameter, final float defaultValue) + { + return QueryContexts.getAsFloat(parameter, get(parameter), defaultValue); + } + + public > E getAsEnum(String parameter, Class clazz, E defaultValue) + { + return QueryContexts.getAsEnum(parameter, get(parameter), clazz, defaultValue); + } + public Map getMergedParams() { if (mergedParams == null) { diff --git a/processing/src/main/java/org/apache/druid/query/QueryContexts.java b/processing/src/main/java/org/apache/druid/query/QueryContexts.java index 416760cb43b8..a1bd87e1c9cf 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContexts.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContexts.java @@ -31,6 +31,7 @@ import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.segment.QueryableIndexStorageAdapter; +import javax.annotation.Nullable; import java.util.Map; import java.util.TreeMap; import java.util.concurrent.TimeUnit; @@ -156,7 +157,7 @@ public static boolean isBySegment(Query query) public static boolean isBySegment(Query query, boolean defaultValue) { - return parseBoolean(query, BY_SEGMENT_KEY, defaultValue); + return query.getContextBoolean(BY_SEGMENT_KEY, defaultValue); } public static boolean isPopulateCache(Query query) @@ -166,7 +167,7 @@ public static boolean isPopulateCache(Query query) public static boolean isPopulateCache(Query query, boolean defaultValue) { - return parseBoolean(query, POPULATE_CACHE_KEY, defaultValue); + return query.getContextBoolean(POPULATE_CACHE_KEY, defaultValue); } public static boolean isUseCache(Query query) @@ -176,7 +177,7 @@ public static boolean isUseCache(Query query) public static boolean isUseCache(Query query, boolean defaultValue) { - return parseBoolean(query, USE_CACHE_KEY, defaultValue); + return query.getContextBoolean(USE_CACHE_KEY, defaultValue); } public static boolean isPopulateResultLevelCache(Query query) @@ -186,7 +187,7 @@ public static boolean isPopulateResultLevelCache(Query query) public static boolean isPopulateResultLevelCache(Query query, boolean defaultValue) { - return parseBoolean(query, POPULATE_RESULT_LEVEL_CACHE_KEY, defaultValue); + return query.getContextBoolean(POPULATE_RESULT_LEVEL_CACHE_KEY, defaultValue); } public static boolean isUseResultLevelCache(Query query) @@ -196,22 +197,23 @@ public static boolean isUseResultLevelCache(Query query) public static boolean isUseResultLevelCache(Query query, boolean defaultValue) { - return parseBoolean(query, USE_RESULT_LEVEL_CACHE_KEY, defaultValue); + return query.getContextBoolean(USE_RESULT_LEVEL_CACHE_KEY, defaultValue); } + public static boolean isFinalize(Query query, boolean defaultValue) { - return parseBoolean(query, FINALIZE_KEY, defaultValue); + return query.getContextBoolean(FINALIZE_KEY, defaultValue); } public static boolean isSerializeDateTimeAsLong(Query query, boolean defaultValue) { - return parseBoolean(query, SERIALIZE_DATE_TIME_AS_LONG_KEY, defaultValue); + return query.getContextBoolean(SERIALIZE_DATE_TIME_AS_LONG_KEY, defaultValue); } public static boolean isSerializeDateTimeAsLongInner(Query query, boolean defaultValue) { - return parseBoolean(query, SERIALIZE_DATE_TIME_AS_LONG_INNER_KEY, defaultValue); + return query.getContextBoolean(SERIALIZE_DATE_TIME_AS_LONG_INNER_KEY, defaultValue); } public static Vectorize getVectorize(Query query) @@ -221,7 +223,7 @@ public static Vectorize getVectorize(Query query) public static Vectorize getVectorize(Query query, Vectorize defaultValue) { - return parseEnum(query, VECTORIZE_KEY, Vectorize.class, defaultValue); + return query.getContextAsEnum(VECTORIZE_KEY, Vectorize.class, defaultValue); } public static Vectorize getVectorizeVirtualColumns(Query query) @@ -231,7 +233,7 @@ public static Vectorize getVectorizeVirtualColumns(Query query) public static Vectorize getVectorizeVirtualColumns(Query query, Vectorize defaultValue) { - return parseEnum(query, VECTORIZE_VIRTUAL_COLUMNS_KEY, Vectorize.class, defaultValue); + return query.getContextAsEnum(VECTORIZE_VIRTUAL_COLUMNS_KEY, Vectorize.class, defaultValue); } public static int getVectorSize(Query query) @@ -241,12 +243,12 @@ public static int getVectorSize(Query query) public static int getVectorSize(Query query, int defaultSize) { - return parseInt(query, VECTOR_SIZE_KEY, defaultSize); + return query.getContextAsInt(VECTOR_SIZE_KEY, defaultSize); } public static int getMaxSubqueryRows(Query query, int defaultSize) { - return parseInt(query, MAX_SUBQUERY_ROWS_KEY, defaultSize); + return query.getContextAsInt(MAX_SUBQUERY_ROWS_KEY, defaultSize); } public static int getUncoveredIntervalsLimit(Query query) @@ -256,7 +258,7 @@ public static int getUncoveredIntervalsLimit(Query query) public static int getUncoveredIntervalsLimit(Query query, int defaultValue) { - return parseInt(query, UNCOVERED_INTERVALS_LIMIT_KEY, defaultValue); + return query.getContextAsInt(UNCOVERED_INTERVALS_LIMIT_KEY, defaultValue); } public static int getPriority(Query query) @@ -266,38 +268,37 @@ public static int getPriority(Query query) public static int getPriority(Query query, int defaultValue) { - return parseInt(query, PRIORITY_KEY, defaultValue); + return query.getContextAsInt(PRIORITY_KEY, defaultValue); } public static String getLane(Query query) { - return (String) query.getContextValue(LANE_KEY); + return query.getContextAsString(LANE_KEY); } public static boolean getEnableParallelMerges(Query query) { - return parseBoolean(query, BROKER_PARALLEL_MERGE_KEY, DEFAULT_ENABLE_PARALLEL_MERGE); + return query.getContextBoolean(BROKER_PARALLEL_MERGE_KEY, DEFAULT_ENABLE_PARALLEL_MERGE); } public static int getParallelMergeInitialYieldRows(Query query, int defaultValue) { - return parseInt(query, BROKER_PARALLEL_MERGE_INITIAL_YIELD_ROWS_KEY, defaultValue); + return query.getContextAsInt(BROKER_PARALLEL_MERGE_INITIAL_YIELD_ROWS_KEY, defaultValue); } public static int getParallelMergeSmallBatchRows(Query query, int defaultValue) { - return parseInt(query, BROKER_PARALLEL_MERGE_SMALL_BATCH_ROWS_KEY, defaultValue); + return query.getContextAsInt(BROKER_PARALLEL_MERGE_SMALL_BATCH_ROWS_KEY, defaultValue); } public static int getParallelMergeParallelism(Query query, int defaultValue) { - return parseInt(query, BROKER_PARALLELISM, defaultValue); + return query.getContextAsInt(BROKER_PARALLELISM, defaultValue); } public static boolean getEnableJoinFilterRewriteValueColumnFilters(Query query) { - return parseBoolean( - query, + return query.getContextBoolean( JOIN_FILTER_REWRITE_VALUE_COLUMN_FILTERS_ENABLE_KEY, DEFAULT_ENABLE_JOIN_FILTER_REWRITE_VALUE_COLUMN_FILTERS ); @@ -305,8 +306,7 @@ public static boolean getEnableJoinFilterRewriteValueColumnFilters(Query public static boolean getEnableRewriteJoinToFilter(Query query) { - return parseBoolean( - query, + return query.getContextBoolean( REWRITE_JOIN_TO_FILTER_ENABLE_KEY, DEFAULT_ENABLE_REWRITE_JOIN_TO_FILTER ); @@ -314,32 +314,32 @@ public static boolean getEnableRewriteJoinToFilter(Query query) public static long getJoinFilterRewriteMaxSize(Query query) { - return parseLong(query, JOIN_FILTER_REWRITE_MAX_SIZE_KEY, DEFAULT_ENABLE_JOIN_FILTER_REWRITE_MAX_SIZE); + return query.getContextAsLong(JOIN_FILTER_REWRITE_MAX_SIZE_KEY, DEFAULT_ENABLE_JOIN_FILTER_REWRITE_MAX_SIZE); } public static boolean getEnableJoinFilterPushDown(Query query) { - return parseBoolean(query, JOIN_FILTER_PUSH_DOWN_KEY, DEFAULT_ENABLE_JOIN_FILTER_PUSH_DOWN); + return query.getContextBoolean(JOIN_FILTER_PUSH_DOWN_KEY, DEFAULT_ENABLE_JOIN_FILTER_PUSH_DOWN); } public static boolean getEnableJoinFilterRewrite(Query query) { - return parseBoolean(query, JOIN_FILTER_REWRITE_ENABLE_KEY, DEFAULT_ENABLE_JOIN_FILTER_REWRITE); + return query.getContextBoolean(JOIN_FILTER_REWRITE_ENABLE_KEY, DEFAULT_ENABLE_JOIN_FILTER_REWRITE); } - public static boolean getEnableJoinLeftScanDirect(Map context) + public static boolean getEnableJoinLeftScanDirect(Map context) { return parseBoolean(context, SQL_JOIN_LEFT_SCAN_DIRECT, DEFAULT_ENABLE_SQL_JOIN_LEFT_SCAN_DIRECT); } public static boolean isSecondaryPartitionPruningEnabled(Query query) { - return parseBoolean(query, SECONDARY_PARTITION_PRUNING_KEY, DEFAULT_SECONDARY_PARTITION_PRUNING); + return query.getContextBoolean(SECONDARY_PARTITION_PRUNING_KEY, DEFAULT_SECONDARY_PARTITION_PRUNING); } public static boolean isDebug(Query query) { - return parseBoolean(query, ENABLE_DEBUG, DEFAULT_ENABLE_DEBUG); + return query.getContextBoolean(ENABLE_DEBUG, DEFAULT_ENABLE_DEBUG); } public static boolean isDebug(Map queryContext) @@ -364,11 +364,10 @@ public static boolean isTimeBoundaryPlanningEnabled(Map queryCon public static Query withMaxScatterGatherBytes(Query query, long maxScatterGatherBytesLimit) { - Object obj = query.getContextValue(MAX_SCATTER_GATHER_BYTES_KEY); - if (obj == null) { + Long curr = query.getContextAsLong(MAX_SCATTER_GATHER_BYTES_KEY); + if (curr == null) { return query.withOverriddenContext(ImmutableMap.of(MAX_SCATTER_GATHER_BYTES_KEY, maxScatterGatherBytesLimit)); } else { - long curr = ((Number) obj).longValue(); if (curr > maxScatterGatherBytesLimit) { throw new IAE( "configured [%s = %s] is more than enforced limit of [%s].", @@ -399,12 +398,12 @@ public static Query verifyMaxQueryTimeout(Query query, long maxQueryTi public static long getMaxQueuedBytes(Query query, long defaultValue) { - return parseLong(query, MAX_QUEUED_BYTES_KEY, defaultValue); + return query.getContextAsLong(MAX_QUEUED_BYTES_KEY, defaultValue); } public static long getMaxScatterGatherBytes(Query query) { - return parseLong(query, MAX_SCATTER_GATHER_BYTES_KEY, Long.MAX_VALUE); + return query.getContextAsLong(MAX_SCATTER_GATHER_BYTES_KEY, Long.MAX_VALUE); } public static boolean hasTimeout(Query query) @@ -420,7 +419,7 @@ public static long getTimeout(Query query) public static long getTimeout(Query query, long defaultTimeout) { try { - final long timeout = parseLong(query, TIMEOUT_KEY, defaultTimeout); + final long timeout = query.getContextAsLong(TIMEOUT_KEY, defaultTimeout); Preconditions.checkState(timeout >= 0, "Timeout must be a non negative value, but was [%s]", timeout); return timeout; } @@ -441,14 +440,14 @@ public static Query withDefaultTimeout(Query query, long defaultTimeou static long getDefaultTimeout(Query query) { - final long defaultTimeout = parseLong(query, DEFAULT_TIMEOUT_KEY, DEFAULT_TIMEOUT_MILLIS); + final long defaultTimeout = query.getContextAsLong(DEFAULT_TIMEOUT_KEY, DEFAULT_TIMEOUT_MILLIS); Preconditions.checkState(defaultTimeout >= 0, "Timeout must be a non negative value, but was [%s]", defaultTimeout); return defaultTimeout; } public static int getNumRetriesOnMissingSegments(Query query, int defaultValue) { - return query.getContextValue(NUM_RETRIES_ON_MISSING_SEGMENTS_KEY, defaultValue); + return query.getContextAsInt(NUM_RETRIES_ON_MISSING_SEGMENTS_KEY, defaultValue); } public static boolean allowReturnPartialResults(Query query, boolean defaultValue) @@ -461,32 +460,17 @@ public static String getBrokerServiceName(Map queryContext) return queryContext == null ? null : (String) queryContext.get(BROKER_SERVICE_NAME); } - static long parseLong(Query query, String key, long defaultValue) - { - return getAsLong(key, query.getContextValue(key), defaultValue); - } - @SuppressWarnings("unused") static long parseLong(Map context, String key, long defaultValue) { return getAsLong(key, context.get(key), defaultValue); } - static int parseInt(Query query, String key, int defaultValue) - { - return getAsInt(key, query.getContextValue(key), defaultValue); - } - static int parseInt(Map context, String key, int defaultValue) { return getAsInt(key, context.get(key), defaultValue); } - static boolean parseBoolean(Query query, String key, boolean defaultValue) - { - return getAsBoolean(key, query.getContextValue(key), defaultValue); - } - static boolean parseBoolean(Map context, String key, boolean defaultValue) { return getAsBoolean(key, context.get(key), defaultValue); @@ -528,6 +512,24 @@ public static boolean getAsBoolean( } } + @Nullable + public static Integer getAsInt(String parameter, Object value) + { + if (value == null) { + return null; + } else if (value instanceof Number) { + return ((Number) value).intValue(); + } else if (value instanceof String) { + try { + return Numbers.parseInt(value); + } + catch (NumberFormatException ignored) { + } + } + + throw new IAE("Expected parameter [%s] to be an integer", parameter); + } + /** * Get the value of a parameter as an {@code int}. The parameter is expected * to be {@code null}, a string or a {@code Number} object. @@ -537,16 +539,26 @@ public static int getAsInt( final Object value, final int defaultValue ) + { + Integer val = getAsInt(parameter, value); + return val == null ? defaultValue : val; + } + + @Nullable + public static Long getAsLong(String parameter, Object value) { if (value == null) { - return defaultValue; - } else if (value instanceof String) { - return Numbers.parseInt(value); + return null; } else if (value instanceof Number) { - return ((Number) value).intValue(); - } else { - throw new IAE("Expected parameter [%s] to be an integer", parameter); + return ((Number) value).longValue(); + } else if (value instanceof String) { + try { + return Numbers.parseLong(value); + } + catch (NumberFormatException ignored) { + } } + throw new IAE("Expected parameter [%s] to be a long", parameter); } /** @@ -556,17 +568,11 @@ public static int getAsInt( public static long getAsLong( final String parameter, final Object value, - final long defaultValue) + final long defaultValue + ) { - if (value == null) { - return defaultValue; - } else if (value instanceof String) { - return Numbers.parseLong(value); - } else if (value instanceof Number) { - return ((Number) value).longValue(); - } else { - throw new IAE("Expected parameter [%s] to be a long", parameter); - } + Long val = getAsLong(parameter, value); + return val == null ? defaultValue : val; } public static HumanReadableBytes getAsHumanReadableBytes( @@ -586,6 +592,22 @@ public static HumanReadableBytes getAsHumanReadableBytes( } } + public static float getAsFloat(String parameter, Object value, float defaultValue) + { + if (null == value) { + return defaultValue; + } else if (value instanceof Number) { + return ((Number) value).floatValue(); + } else if (value instanceof String) { + try { + return Float.parseFloat((String) value); + } + catch (NumberFormatException ignored) { + } + } + throw new IAE("Expected parameter [%s] to be a float", parameter); + } + public static Map override( final Map context, final Map overrides @@ -604,9 +626,8 @@ private QueryContexts() { } - static > E parseEnum(Query query, String key, Class clazz, E defaultValue) + public static > E getAsEnum(String parameter, Object val, Class clazz, E defaultValue) { - Object val = query.getContextValue(key); if (val == null) { return defaultValue; } @@ -615,7 +636,7 @@ static > E parseEnum(Query query, String key, Class c } else if (val instanceof Boolean) { return Enum.valueOf(clazz, StringUtils.toUpperCase(String.valueOf(val))); } else { - throw new ISE("Unknown type [%s]. Cannot parse!", val.getClass()); + throw new ISE("Expected parameter [%s] must be type of [%s], actual type is [%s]. ", parameter, val.getClass()); } } } diff --git a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQuery.java b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQuery.java index 17a30f732f59..0a67fc052d57 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQuery.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQuery.java @@ -748,7 +748,7 @@ private int compareDims(List dimensions, ResultRow lhs, ResultRow @Nullable private DateTime computeUniversalTimestamp() { - final String timestampStringFromContext = getContextValue(CTX_KEY_FUDGE_TIMESTAMP, ""); + final String timestampStringFromContext = getContextAsString(CTX_KEY_FUDGE_TIMESTAMP, ""); final Granularity granularity = getGranularity(); if (!timestampStringFromContext.isEmpty()) { diff --git a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java index 6125577ac306..8468f851325d 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java @@ -335,25 +335,25 @@ public boolean isMultiValueUnnestingEnabled() public GroupByQueryConfig withOverrides(final GroupByQuery query) { final GroupByQueryConfig newConfig = new GroupByQueryConfig(); - newConfig.defaultStrategy = query.getContextValue(CTX_KEY_STRATEGY, getDefaultStrategy()); + newConfig.defaultStrategy = query.getContextAsString(CTX_KEY_STRATEGY, getDefaultStrategy()); newConfig.singleThreaded = query.getContextBoolean(CTX_KEY_IS_SINGLE_THREADED, isSingleThreaded()); newConfig.maxIntermediateRows = Math.min( - query.getContextValue(CTX_KEY_MAX_INTERMEDIATE_ROWS, getMaxIntermediateRows()), + query.getContextAsInt(CTX_KEY_MAX_INTERMEDIATE_ROWS, getMaxIntermediateRows()), getMaxIntermediateRows() ); newConfig.maxResults = Math.min( - query.getContextValue(CTX_KEY_MAX_RESULTS, getMaxResults()), + query.getContextAsInt(CTX_KEY_MAX_RESULTS, getMaxResults()), getMaxResults() ); newConfig.bufferGrouperMaxSize = Math.min( - query.getContextValue(CTX_KEY_BUFFER_GROUPER_MAX_SIZE, getBufferGrouperMaxSize()), + query.getContextAsInt(CTX_KEY_BUFFER_GROUPER_MAX_SIZE, getBufferGrouperMaxSize()), getBufferGrouperMaxSize() ); - newConfig.bufferGrouperMaxLoadFactor = query.getContextValue( + newConfig.bufferGrouperMaxLoadFactor = query.getContextAsFloat( CTX_KEY_BUFFER_GROUPER_MAX_LOAD_FACTOR, getBufferGrouperMaxLoadFactor() ); - newConfig.bufferGrouperInitialBuckets = query.getContextValue( + newConfig.bufferGrouperInitialBuckets = query.getContextAsInt( CTX_KEY_BUFFER_GROUPER_INITIAL_BUCKETS, getBufferGrouperInitialBuckets() ); @@ -378,11 +378,11 @@ public GroupByQueryConfig withOverrides(final GroupByQuery query) CTX_KEY_FORCE_PUSH_DOWN_NESTED_QUERY, isForcePushDownNestedQuery() ); - newConfig.intermediateCombineDegree = query.getContextValue( + newConfig.intermediateCombineDegree = query.getContextAsInt( CTX_KEY_INTERMEDIATE_COMBINE_DEGREE, getIntermediateCombineDegree() ); - newConfig.numParallelCombineThreads = query.getContextValue( + newConfig.numParallelCombineThreads = query.getContextAsInt( CTX_KEY_NUM_PARALLEL_COMBINE_THREADS, getNumParallelCombineThreads() ); diff --git a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryEngine.java b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryEngine.java index 611d0084061d..d971e1bcce2c 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryEngine.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryEngine.java @@ -96,7 +96,7 @@ public Sequence process( "Null storage adapter found. Probably trying to issue a query against a segment being memory unmapped." ); } - if (!query.getContextValue(GroupByQueryConfig.CTX_KEY_ENABLE_MULTI_VALUE_UNNESTING, true)) { + if (!query.getContextBoolean(GroupByQueryConfig.CTX_KEY_ENABLE_MULTI_VALUE_UNNESTING, true)) { throw new UOE( "GroupBy v1 does not support %s as false. Set %s to true or use groupBy v2", GroupByQueryConfig.CTX_KEY_ENABLE_MULTI_VALUE_UNNESTING, diff --git a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryHelper.java b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryHelper.java index eaeeadafa87f..4fd84c9b62b3 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryHelper.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryHelper.java @@ -100,7 +100,7 @@ public String apply(DimensionSpec input) ); final IncrementalIndex index; - final boolean sortResults = query.getContextValue(CTX_KEY_SORT_RESULTS, true); + final boolean sortResults = query.getContextBoolean(CTX_KEY_SORT_RESULTS, true); // All groupBy dimensions are strings, for now. final List dimensionSchemas = new ArrayList<>(); @@ -118,7 +118,7 @@ public String apply(DimensionSpec input) final AppendableIndexBuilder indexBuilder; - if (query.getContextValue("useOffheap", false)) { + if (query.getContextBoolean("useOffheap", false)) { throw new UnsupportedOperationException( "The 'useOffheap' option is no longer available for groupBy v1. Please move to the newer groupBy engine, " + "which always operates off-heap, by removing any custom 'druid.query.groupBy.defaultStrategy' runtime " diff --git a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java index e2cea67b8da2..cc994a05616e 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java @@ -141,7 +141,7 @@ public static Sequence process( try { final String fudgeTimestampString = NullHandling.emptyToNullIfNeeded( - query.getContextValue(GroupByStrategyV2.CTX_KEY_FUDGE_TIMESTAMP, null) + query.getContextAsString(GroupByStrategyV2.CTX_KEY_FUDGE_TIMESTAMP) ); final DateTime fudgeTimestamp = fudgeTimestampString == null diff --git a/processing/src/main/java/org/apache/druid/query/groupby/orderby/DefaultLimitSpec.java b/processing/src/main/java/org/apache/druid/query/groupby/orderby/DefaultLimitSpec.java index 335b220b74db..c6e44b6f89b7 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/orderby/DefaultLimitSpec.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/orderby/DefaultLimitSpec.java @@ -232,9 +232,9 @@ public Function, Sequence> build(final GroupByQue } if (!sortingNeeded) { - String timestampField = query.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD); + String timestampField = query.getContextAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD); if (timestampField != null && !timestampField.isEmpty()) { - int timestampResultFieldIndex = query.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX); + int timestampResultFieldIndex = query.getContextAsInt(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX); sortingNeeded = query.getContextSortByDimsFirst() ? timestampResultFieldIndex != query.getDimensions().size() - 1 : timestampResultFieldIndex != 0; diff --git a/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java b/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java index e6adfa4bee12..bc8f411e54d3 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java @@ -221,7 +221,7 @@ public Sequence mergeResults( Granularity granularity = query.getGranularity(); List dimensionSpecs = query.getDimensions(); // the CTX_TIMESTAMP_RESULT_FIELD is set in DruidQuery.java - final String timestampResultField = query.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD); + final String timestampResultField = query.getContextAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD); final boolean hasTimestampResultField = (timestampResultField != null && !timestampResultField.isEmpty()) && query.getContextBoolean(CTX_KEY_OUTERMOST, true) && !query.isApplyLimitPushDown(); @@ -258,7 +258,7 @@ public Sequence mergeResults( granularity = timestampResultFieldGranularity; // when timestampResultField is the last dimension, should set sortByDimsFirst=true, // otherwise the downstream is sorted by row's timestamp first which makes the final ordering not as expected - timestampResultFieldIndex = query.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX); + timestampResultFieldIndex = query.getContextAsInt(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX); if (!query.getContextSortByDimsFirst() && timestampResultFieldIndex == query.getDimensions().size() - 1) { context.put(GroupByQuery.CTX_KEY_SORT_BY_DIMS_FIRST, true); } diff --git a/processing/src/main/java/org/apache/druid/query/scan/ScanQuery.java b/processing/src/main/java/org/apache/druid/query/scan/ScanQuery.java index 8e0a6a1c259e..3fdf849e1f79 100644 --- a/processing/src/main/java/org/apache/druid/query/scan/ScanQuery.java +++ b/processing/src/main/java/org/apache/druid/query/scan/ScanQuery.java @@ -264,7 +264,7 @@ public static void verifyOrderByForNativeExecution(final ScanQuery query) private Integer validateAndGetMaxRowsQueuedForOrdering() { final Integer maxRowsQueuedForOrdering = - getContextValue(ScanQueryConfig.CTX_KEY_MAX_ROWS_QUEUED_FOR_ORDERING, null); + getContextAsInt(ScanQueryConfig.CTX_KEY_MAX_ROWS_QUEUED_FOR_ORDERING); Preconditions.checkArgument( maxRowsQueuedForOrdering == null || maxRowsQueuedForOrdering > 0, "maxRowsQueuedForOrdering must be greater than 0" @@ -275,7 +275,7 @@ private Integer validateAndGetMaxRowsQueuedForOrdering() private Integer validateAndGetMaxSegmentPartitionsOrderedInMemory() { final Integer maxSegmentPartitionsOrderedInMemory = - getContextValue(ScanQueryConfig.CTX_KEY_MAX_SEGMENT_PARTITIONS_FOR_ORDERING, null); + getContextAsInt(ScanQueryConfig.CTX_KEY_MAX_SEGMENT_PARTITIONS_FOR_ORDERING); Preconditions.checkArgument( maxSegmentPartitionsOrderedInMemory == null || maxSegmentPartitionsOrderedInMemory > 0, "maxRowsQueuedForOrdering must be greater than 0" diff --git a/processing/src/main/java/org/apache/druid/query/search/SearchQueryConfig.java b/processing/src/main/java/org/apache/druid/query/search/SearchQueryConfig.java index c45793c98517..318901caf437 100644 --- a/processing/src/main/java/org/apache/druid/query/search/SearchQueryConfig.java +++ b/processing/src/main/java/org/apache/druid/query/search/SearchQueryConfig.java @@ -55,7 +55,7 @@ public SearchQueryConfig withOverrides(final SearchQuery query) { final SearchQueryConfig newConfig = new SearchQueryConfig(); newConfig.maxSearchLimit = query.getLimit(); - newConfig.searchStrategy = query.getContextValue(CTX_KEY_STRATEGY, searchStrategy); + newConfig.searchStrategy = query.getContextAsString(CTX_KEY_STRATEGY, searchStrategy); return newConfig; } } diff --git a/processing/src/main/java/org/apache/druid/query/select/SelectQuery.java b/processing/src/main/java/org/apache/druid/query/select/SelectQuery.java index f2895f0f9bc0..ad4ad3616356 100644 --- a/processing/src/main/java/org/apache/druid/query/select/SelectQuery.java +++ b/processing/src/main/java/org/apache/druid/query/select/SelectQuery.java @@ -21,7 +21,6 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.google.common.collect.Ordering; -import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.granularity.Granularity; import org.apache.druid.query.DataSource; import org.apache.druid.query.Query; @@ -123,24 +122,6 @@ public ContextType getContextValue(String key) throw new RuntimeException(REMOVED_ERROR_MESSAGE); } - @Override - public ContextType getContextValue(String key, ContextType defaultValue) - { - throw new RuntimeException(REMOVED_ERROR_MESSAGE); - } - - @Override - public boolean getContextBoolean(String key, boolean defaultValue) - { - throw new RuntimeException(REMOVED_ERROR_MESSAGE); - } - - @Override - public HumanReadableBytes getContextHumanReadableBytes(String key, HumanReadableBytes defaultValue) - { - throw new RuntimeException(REMOVED_ERROR_MESSAGE); - } - @Override public boolean isDescending() { diff --git a/processing/src/main/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java b/processing/src/main/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java index 9dc0859f4e29..861e3c1bb565 100644 --- a/processing/src/main/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java +++ b/processing/src/main/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java @@ -233,8 +233,8 @@ public RowSignature resultArraySignature(TimeBoundaryQuery query) if (query.isMinTime() || query.isMaxTime()) { RowSignature.Builder builder = RowSignature.builder(); String outputName = query.isMinTime() ? - query.getContextValue(TimeBoundaryQuery.MIN_TIME_ARRAY_OUTPUT_NAME, TimeBoundaryQuery.MIN_TIME) : - query.getContextValue(TimeBoundaryQuery.MAX_TIME_ARRAY_OUTPUT_NAME, TimeBoundaryQuery.MAX_TIME); + query.getContextAsString(TimeBoundaryQuery.MIN_TIME_ARRAY_OUTPUT_NAME, TimeBoundaryQuery.MIN_TIME) : + query.getContextAsString(TimeBoundaryQuery.MAX_TIME_ARRAY_OUTPUT_NAME, TimeBoundaryQuery.MAX_TIME); return builder.add(outputName, ColumnType.LONG).build(); } return super.resultArraySignature(query); diff --git a/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQuery.java b/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQuery.java index 1e7300588c88..65dffef6e0d3 100644 --- a/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQuery.java +++ b/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQuery.java @@ -159,7 +159,7 @@ public boolean isGrandTotal() public String getTimestampResultField() { - return getContextValue(CTX_TIMESTAMP_RESULT_FIELD); + return getContextAsString(CTX_TIMESTAMP_RESULT_FIELD); } public boolean isSkipEmptyBuckets() diff --git a/processing/src/main/java/org/apache/druid/query/topn/TopNQueryQueryToolChest.java b/processing/src/main/java/org/apache/druid/query/topn/TopNQueryQueryToolChest.java index 3b634bc5fe1d..87855bf84f01 100644 --- a/processing/src/main/java/org/apache/druid/query/topn/TopNQueryQueryToolChest.java +++ b/processing/src/main/java/org/apache/druid/query/topn/TopNQueryQueryToolChest.java @@ -574,7 +574,7 @@ public Sequence> run( } final TopNQuery query = (TopNQuery) input; - final int minTopNThreshold = query.getContextValue("minTopNThreshold", config.getMinTopNThreshold()); + final int minTopNThreshold = query.getContextAsInt("minTopNThreshold", config.getMinTopNThreshold()); if (query.getThreshold() > minTopNThreshold) { return runner.run(queryPlus, responseContext); } diff --git a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java index b3f794b1c68a..7d4b945bc1d1 100644 --- a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java +++ b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java @@ -80,11 +80,13 @@ public void testIsEmpty() public void testGetString() { final QueryContext context = new QueryContext( - ImmutableMap.of("key", "val") + ImmutableMap.of("key", "val", + "key2", 2) ); Assert.assertEquals("val", context.get("key")); Assert.assertEquals("val", context.getAsString("key")); + Assert.assertEquals("2", context.getAsString("key2")); Assert.assertNull(context.getAsString("non-exist")); } @@ -109,13 +111,16 @@ public void testGetInt() final QueryContext context = new QueryContext( ImmutableMap.of( "key1", "100", - "key2", 100 + "key2", 100, + "key3", "abc" ) ); Assert.assertEquals(100, context.getAsInt("key1", 0)); Assert.assertEquals(100, context.getAsInt("key2", 0)); Assert.assertEquals(0, context.getAsInt("non-exist", 0)); + + Assert.assertThrows(IAE.class, () -> context.getAsInt("key3", 5)); } @Test @@ -124,24 +129,57 @@ public void testGetLong() final QueryContext context = new QueryContext( ImmutableMap.of( "key1", "100", - "key2", 100 + "key2", 100, + "key3", "abc" ) ); Assert.assertEquals(100L, context.getAsLong("key1", 0)); Assert.assertEquals(100L, context.getAsLong("key2", 0)); Assert.assertEquals(0L, context.getAsLong("non-exist", 0)); + + Assert.assertThrows(IAE.class, () -> context.getAsLong("key3", 5)); } @Test - public void testGetHumanReadableBytes() + public void testGetFloat() { final QueryContext context = new QueryContext( ImmutableMap.of( - "maxOnDiskStorage", "500M" + "f1", "500", + "f2", 500, + "f3", 500.1, + "f4", "ab" ) ); - Assert.assertEquals(500_000_000, context.getAsHumanReadableBytes("maxOnDiskStorage", HumanReadableBytes.ZERO).getBytes()); + + Assert.assertEquals(0, Float.compare(500, context.getAsFloat("f1", 100))); + Assert.assertEquals(0, Float.compare(500, context.getAsFloat("f2", 100))); + Assert.assertEquals(0, Float.compare(500.1f, context.getAsFloat("f3", 100))); + + Assert.assertThrows(IAE.class, () -> context.getAsLong("f4", 5)); + } + + @Test + public void testGetHumanReadableBytes() + { + final QueryContext context = new QueryContext( + ImmutableMap.builder() + .put("m1", 500_000_000) + .put("m2", "500M") + .put("m3", "500Mi") + .put("m4", "500MiB") + .put("m5", "500000000") + .put("m6", "abc") + .build() + ); + Assert.assertEquals(500_000_000, context.getAsHumanReadableBytes("m1", HumanReadableBytes.ZERO).getBytes()); + Assert.assertEquals(500_000_000, context.getAsHumanReadableBytes("m2", HumanReadableBytes.ZERO).getBytes()); + Assert.assertEquals(500 * 1024 * 1024L, context.getAsHumanReadableBytes("m3", HumanReadableBytes.ZERO).getBytes()); + Assert.assertEquals(500 * 1024 * 1024L, context.getAsHumanReadableBytes("m4", HumanReadableBytes.ZERO).getBytes()); + Assert.assertEquals(500_000_000, context.getAsHumanReadableBytes("m5", HumanReadableBytes.ZERO).getBytes()); + + Assert.assertThrows(IAE.class, () -> context.getAsHumanReadableBytes("m6", HumanReadableBytes.ZERO)); } @Test @@ -302,12 +340,12 @@ public void testLegacyReturnsLegacy() public void testNonLegacyIsNotLegacyContext() { Query timeseries = Druids.newTimeseriesQueryBuilder() - .dataSource("test") - .intervals("2015-01-02/2015-01-03") - .granularity(Granularities.DAY) - .aggregators(Collections.singletonList(new CountAggregatorFactory("theCount"))) - .context(ImmutableMap.of("foo", "bar")) - .build(); + .dataSource("test") + .intervals("2015-01-02/2015-01-03") + .granularity(Granularities.DAY) + .aggregators(Collections.singletonList(new CountAggregatorFactory("theCount"))) + .context(ImmutableMap.of("foo", "bar")) + .build(); Assert.assertNotNull(timeseries.getQueryContext()); } @@ -463,15 +501,6 @@ public Query withOverriddenContext(Map contextOverride) return new LegacyContextQuery(contextOverride); } - @Override - public Object getContextValue(String key, Object defaultValue) - { - if (!context.containsKey(key)) { - return defaultValue; - } - return context.get(key); - } - @Override public Object getContextValue(String key) { diff --git a/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java b/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java index 97958818c9db..1473e39c4604 100644 --- a/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java +++ b/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java @@ -102,10 +102,10 @@ public void testContextSerde() throws Exception ), Query.class ); - Assert.assertEquals((Integer) 1, serdeQuery.getContextValue(QueryContexts.PRIORITY_KEY)); - Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.USE_CACHE_KEY)); - Assert.assertEquals("true", serdeQuery.getContextValue(QueryContexts.POPULATE_CACHE_KEY)); - Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.FINALIZE_KEY)); + Assert.assertEquals((Integer) 1, serdeQuery.getContextAsInt(QueryContexts.PRIORITY_KEY)); + Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.USE_CACHE_KEY, false)); + Assert.assertEquals("true", serdeQuery.getContextAsString(QueryContexts.POPULATE_CACHE_KEY)); + Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.FINALIZE_KEY, false)); Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.USE_CACHE_KEY, false)); Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.POPULATE_CACHE_KEY, false)); Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.FINALIZE_KEY, false)); diff --git a/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java b/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java index 868e83392db9..40d1d3fda985 100644 --- a/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java +++ b/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java @@ -79,10 +79,10 @@ public void testContextSerde() throws Exception ); - Assert.assertEquals(new Integer(1), serdeQuery.getContextValue(QueryContexts.PRIORITY_KEY)); - Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.USE_CACHE_KEY)); - Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.POPULATE_CACHE_KEY)); - Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.FINALIZE_KEY)); + Assert.assertEquals(new Integer(1), serdeQuery.getContextAsInt(QueryContexts.PRIORITY_KEY)); + Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.USE_CACHE_KEY, false)); + Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.POPULATE_CACHE_KEY, false)); + Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.FINALIZE_KEY, false)); } @Test @@ -117,9 +117,9 @@ public void testContextSerde2() throws Exception ); - Assert.assertEquals("1", serdeQuery.getContextValue(QueryContexts.PRIORITY_KEY)); - Assert.assertEquals("true", serdeQuery.getContextValue(QueryContexts.USE_CACHE_KEY)); - Assert.assertEquals("true", serdeQuery.getContextValue(QueryContexts.POPULATE_CACHE_KEY)); - Assert.assertEquals("true", serdeQuery.getContextValue(QueryContexts.FINALIZE_KEY)); + Assert.assertEquals("1", serdeQuery.getContextAsString(QueryContexts.PRIORITY_KEY)); + Assert.assertEquals("true", serdeQuery.getContextAsString(QueryContexts.USE_CACHE_KEY)); + Assert.assertEquals("true", serdeQuery.getContextAsString(QueryContexts.POPULATE_CACHE_KEY)); + Assert.assertEquals("true", serdeQuery.getContextAsString(QueryContexts.FINALIZE_KEY)); } } diff --git a/server/src/main/java/org/apache/druid/client/DirectDruidClient.java b/server/src/main/java/org/apache/druid/client/DirectDruidClient.java index c7556b03cc99..7fa3364e56fa 100644 --- a/server/src/main/java/org/apache/druid/client/DirectDruidClient.java +++ b/server/src/main/java/org/apache/druid/client/DirectDruidClient.java @@ -163,7 +163,7 @@ public Sequence run(final QueryPlus queryPlus, final ResponseContext conte log.debug("Querying queryId[%s] url[%s]", query.getId(), url); final long requestStartTimeNs = System.nanoTime(); - final long timeoutAt = query.getContextValue(QUERY_FAIL_TIME); + final long timeoutAt = query.getContextAsLong(QUERY_FAIL_TIME); final long maxScatterGatherBytes = QueryContexts.getMaxScatterGatherBytes(query); final AtomicLong totalBytesGathered = context.getTotalBytes(); final long maxQueuedBytes = QueryContexts.getMaxQueuedBytes(query, 0); diff --git a/server/src/main/java/org/apache/druid/client/JsonParserIterator.java b/server/src/main/java/org/apache/druid/client/JsonParserIterator.java index 35532276e93f..d10ee099ae4e 100644 --- a/server/src/main/java/org/apache/druid/client/JsonParserIterator.java +++ b/server/src/main/java/org/apache/druid/client/JsonParserIterator.java @@ -75,7 +75,7 @@ public JsonParserIterator( this.future = future; this.url = url; if (query != null) { - this.timeoutAt = query.getContextValue(DirectDruidClient.QUERY_FAIL_TIME, -1L); + this.timeoutAt = query.getContextAsLong(DirectDruidClient.QUERY_FAIL_TIME, -1L); this.queryId = query.getId(); } else { this.timeoutAt = -1; diff --git a/server/src/main/java/org/apache/druid/segment/realtime/appenderator/SinkQuerySegmentWalker.java b/server/src/main/java/org/apache/druid/segment/realtime/appenderator/SinkQuerySegmentWalker.java index 53783fcbdb39..a0fb767eaad7 100644 --- a/server/src/main/java/org/apache/druid/segment/realtime/appenderator/SinkQuerySegmentWalker.java +++ b/server/src/main/java/org/apache/druid/segment/realtime/appenderator/SinkQuerySegmentWalker.java @@ -161,7 +161,7 @@ public QueryRunner getQueryRunnerForSegments(final Query query, final } final QueryToolChest> toolChest = factory.getToolchest(); - final boolean skipIncrementalSegment = query.getContextValue(CONTEXT_SKIP_INCREMENTAL_SEGMENT, false); + final boolean skipIncrementalSegment = query.getContextBoolean(CONTEXT_SKIP_INCREMENTAL_SEGMENT, false); final AtomicLong cpuTimeAccumulator = new AtomicLong(0L); // Make sure this query type can handle the subquery, if present. diff --git a/server/src/main/java/org/apache/druid/server/ClientQuerySegmentWalker.java b/server/src/main/java/org/apache/druid/server/ClientQuerySegmentWalker.java index 136953e8719e..1220be5e2b6d 100644 --- a/server/src/main/java/org/apache/druid/server/ClientQuerySegmentWalker.java +++ b/server/src/main/java/org/apache/druid/server/ClientQuerySegmentWalker.java @@ -431,7 +431,7 @@ private QueryRunner decorateClusterRunner(Query query, QueryRunner .emitCPUTimeMetric(emitter) .postProcess( objectMapper.convertValue( - query.getContextValue("postProcessing"), + query.getContextAsString("postProcessing"), new TypeReference>() {} ) ) diff --git a/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java b/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java index 3af5d10ac9ac..36f2057b385f 100644 --- a/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java +++ b/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java @@ -70,10 +70,10 @@ public Optional computeLane(QueryPlus query, Set keyManager = makeKeyManager(); Set selectors = ImmutableSet.of( @@ -271,7 +272,7 @@ public void testSegmentQueryCacheKey_joinWithSupportedCaching() public void testSegmentQueryCacheKey_noCachingIfBySegment() { reset(query); - expect(query.getContextValue(QueryContexts.BY_SEGMENT_KEY)).andReturn(true).anyTimes(); + expect(query.getContextBoolean(QueryContexts.BY_SEGMENT_KEY, DEFAULT_BY_SEGMENT)).andReturn(true).anyTimes(); replayAll(); byte[] cacheKey = makeKeyManager().computeSegmentLevelQueryCacheKey(); Assert.assertNull(cacheKey); diff --git a/server/src/test/java/org/apache/druid/client/CachingClusteredClientTest.java b/server/src/test/java/org/apache/druid/client/CachingClusteredClientTest.java index c7f05747fa9b..434ab49b7ea0 100644 --- a/server/src/test/java/org/apache/druid/client/CachingClusteredClientTest.java +++ b/server/src/test/java/org/apache/druid/client/CachingClusteredClientTest.java @@ -2301,7 +2301,7 @@ public Iterable>> apply(@Nullable Integer input) Assert.assertEquals(true, capturedQuery.getContextValue(QueryContexts.BY_SEGMENT_KEY)); } else { Assert.assertTrue( - capturedQuery.getContextValue(QueryContexts.BY_SEGMENT_KEY) == null || + !capturedQuery.containsContextValue(QueryContexts.BY_SEGMENT_KEY) || capturedQuery.getContextValue(QueryContexts.BY_SEGMENT_KEY).equals(false) ); } diff --git a/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java b/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java index d0e154d08a16..a1bc9ddb6576 100644 --- a/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java +++ b/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java @@ -307,7 +307,7 @@ private Query mockQuery(String queryId, long timeoutAt) { Query query = Mockito.mock(Query.class); Mockito.when(query.getId()).thenReturn(queryId); - Mockito.when(query.getContextValue(ArgumentMatchers.eq(DirectDruidClient.QUERY_FAIL_TIME), ArgumentMatchers.eq(-1L))) + Mockito.when(query.getContextAsLong(ArgumentMatchers.eq(DirectDruidClient.QUERY_FAIL_TIME), ArgumentMatchers.eq(-1L))) .thenReturn(timeoutAt); return query; } diff --git a/server/src/test/java/org/apache/druid/server/SetAndVerifyContextQueryRunnerTest.java b/server/src/test/java/org/apache/druid/server/SetAndVerifyContextQueryRunnerTest.java index bffadfe5c558..fccfd0500d75 100644 --- a/server/src/test/java/org/apache/druid/server/SetAndVerifyContextQueryRunnerTest.java +++ b/server/src/test/java/org/apache/druid/server/SetAndVerifyContextQueryRunnerTest.java @@ -58,7 +58,7 @@ public void testTimeoutIsUsedIfTimeoutIsNonZero() throws InterruptedException // time + 1 at the time the method was called // this means that after sleeping for 1 millis, the fail time should be less than the current time when checking Assert.assertTrue( - System.currentTimeMillis() > (Long) transformed.getContextValue(DirectDruidClient.QUERY_FAIL_TIME) + System.currentTimeMillis() > transformed.getContextAsLong(DirectDruidClient.QUERY_FAIL_TIME) ); } @@ -85,7 +85,7 @@ public long getDefaultQueryTimeout() Query transformed = queryRunner.withTimeoutAndMaxScatterGatherBytes(query, defaultConfig); // timeout is not set, default timeout has been set to long.max, make sure timeout is still in the future - Assert.assertEquals((Long) Long.MAX_VALUE, transformed.getContextValue(DirectDruidClient.QUERY_FAIL_TIME)); + Assert.assertEquals((Long) Long.MAX_VALUE, transformed.getContextAsLong(DirectDruidClient.QUERY_FAIL_TIME)); } @Test @@ -107,7 +107,7 @@ public void testTimeoutZeroIsNotImmediateTimeoutDefaultServersideMax() // timeout is set to 0, so withTimeoutAndMaxScatterGatherBytes should set QUERY_FAIL_TIME to be the current // time + max query timeout at the time the method was called // since default is long max, expect long max since current time would overflow - Assert.assertEquals((Long) Long.MAX_VALUE, transformed.getContextValue(DirectDruidClient.QUERY_FAIL_TIME)); + Assert.assertEquals((Long) Long.MAX_VALUE, transformed.getContextAsLong(DirectDruidClient.QUERY_FAIL_TIME)); } @Test @@ -137,7 +137,7 @@ public long getMaxQueryTimeout() // time + max query timeout at the time the method was called // this means that the fail time should be greater than the current time when checking Assert.assertTrue( - System.currentTimeMillis() < (Long) transformed.getContextValue(DirectDruidClient.QUERY_FAIL_TIME) + System.currentTimeMillis() < (Long) transformed.getContextAsLong(DirectDruidClient.QUERY_FAIL_TIME) ); } } From 134484867e5ca5e1efd8e331cc098aad7317c170 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Sep 2022 11:55:23 +0800 Subject: [PATCH 2/9] Rename method --- .../query/materializedview/MaterializedViewQueryTest.java | 2 +- .../src/main/java/org/apache/druid/query/BaseQuery.java | 2 +- processing/src/main/java/org/apache/druid/query/Query.java | 4 ++-- .../org/apache/druid/query/groupby/GroupByQueryConfig.java | 2 +- .../test/java/org/apache/druid/query/QueryContextTest.java | 5 ++--- 5 files changed, 7 insertions(+), 8 deletions(-) diff --git a/extensions-contrib/materialized-view-selection/src/test/java/org/apache/druid/query/materializedview/MaterializedViewQueryTest.java b/extensions-contrib/materialized-view-selection/src/test/java/org/apache/druid/query/materializedview/MaterializedViewQueryTest.java index 1a55cdd02666..13dfe567cd22 100644 --- a/extensions-contrib/materialized-view-selection/src/test/java/org/apache/druid/query/materializedview/MaterializedViewQueryTest.java +++ b/extensions-contrib/materialized-view-selection/src/test/java/org/apache/druid/query/materializedview/MaterializedViewQueryTest.java @@ -121,7 +121,7 @@ public void testGetContextHumanReadableBytes() .postAggregators(QueryRunnerTestHelper.ADD_ROWS_INDEX_CONSTANT) .build(); MaterializedViewQuery query = new MaterializedViewQuery(topNQuery, optimizer); - Assert.assertEquals(20_000_000, query.getContextHumanReadableBytes("maxOnDiskStorage", HumanReadableBytes.ZERO).getBytes()); + Assert.assertEquals(20_000_000, query.getContextAsHumanReadableBytes("maxOnDiskStorage", HumanReadableBytes.ZERO).getBytes()); } } diff --git a/processing/src/main/java/org/apache/druid/query/BaseQuery.java b/processing/src/main/java/org/apache/druid/query/BaseQuery.java index 12a24fb1d67f..514220b2506b 100644 --- a/processing/src/main/java/org/apache/druid/query/BaseQuery.java +++ b/processing/src/main/java/org/apache/druid/query/BaseQuery.java @@ -194,7 +194,7 @@ public boolean getContextBoolean(String key, boolean defaultValue) } @Override - public HumanReadableBytes getContextHumanReadableBytes(String key, HumanReadableBytes defaultValue) + public HumanReadableBytes getContextAsHumanReadableBytes(String key, HumanReadableBytes defaultValue) { return context.getAsHumanReadableBytes(key, defaultValue); } diff --git a/processing/src/main/java/org/apache/druid/query/Query.java b/processing/src/main/java/org/apache/druid/query/Query.java index 4f6337c20f74..6c172e89bb6f 100644 --- a/processing/src/main/java/org/apache/druid/query/Query.java +++ b/processing/src/main/java/org/apache/druid/query/Query.java @@ -144,7 +144,7 @@ default boolean containsContextValue(String key) * {@link #getContextAsLong(String)} * {@link #getContextAsFloat(String, float)} * {@link #getContextAsEnum(String, Class, Enum)} - * {@link #getContextHumanReadableBytes(String, HumanReadableBytes)} + * {@link #getContextAsHumanReadableBytes(String, HumanReadableBytes)} */ @Nullable default ContextType getContextValue(String key) @@ -250,7 +250,7 @@ default boolean getContextBoolean(String key, boolean defaultValue) * @param defaultValue The default to return if the key value doesn't exist or the context is null. * @return {@link HumanReadableBytes} */ - default HumanReadableBytes getContextHumanReadableBytes(String key, HumanReadableBytes defaultValue) + default HumanReadableBytes getContextAsHumanReadableBytes(String key, HumanReadableBytes defaultValue) { if (null != getQueryContext()) { return getQueryContext().getAsHumanReadableBytes(key, defaultValue); diff --git a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java index 8468f851325d..1c7e03a6e87c 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java @@ -362,7 +362,7 @@ public GroupByQueryConfig withOverrides(final GroupByQuery query) // choose a default value lower than the max allowed when the context key is missing in the client query. newConfig.maxOnDiskStorage = HumanReadableBytes.valueOf( Math.min( - query.getContextHumanReadableBytes(CTX_KEY_MAX_ON_DISK_STORAGE, getDefaultOnDiskStorage()).getBytes(), + query.getContextAsHumanReadableBytes(CTX_KEY_MAX_ON_DISK_STORAGE, getDefaultOnDiskStorage()).getBytes(), getMaxOnDiskStorage().getBytes() ) ); diff --git a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java index 7d4b945bc1d1..77d70c66412f 100644 --- a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java +++ b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java @@ -26,7 +26,6 @@ import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.IAE; import org.apache.druid.java.util.common.Intervals; -import org.apache.druid.java.util.common.Numbers; import org.apache.druid.java.util.common.granularity.Granularities; import org.apache.druid.java.util.common.granularity.Granularity; import org.apache.druid.query.aggregation.CountAggregatorFactory; @@ -428,14 +427,14 @@ public boolean getContextBoolean(String key, boolean defaultValue) } @Override - public HumanReadableBytes getContextHumanReadableBytes(String key, HumanReadableBytes defaultValue) + public HumanReadableBytes getContextAsHumanReadableBytes(String key, HumanReadableBytes defaultValue) { if (null == context || !context.containsKey(key)) { return defaultValue; } Object value = context.get(key); if (value instanceof Number) { - return HumanReadableBytes.valueOf(Numbers.parseLong(value)); + return HumanReadableBytes.valueOf(((Number) value).longValue()); } else if (value instanceof String) { return new HumanReadableBytes((String) value); } else { From 63041bf518ee35c4be05aa2066cb02f9538af237 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Sep 2022 12:15:56 +0800 Subject: [PATCH 3/9] Add getContextAsBoolean --- .../java/org/apache/druid/query/Query.java | 10 +++++++ .../org/apache/druid/query/QueryContext.java | 6 +++++ .../org/apache/druid/query/QueryContexts.java | 26 +++++++++++++------ .../DataSourceMetadataQueryTest.java | 4 +-- 4 files changed, 36 insertions(+), 10 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/query/Query.java b/processing/src/main/java/org/apache/druid/query/Query.java index 6c172e89bb6f..296c964ffcfa 100644 --- a/processing/src/main/java/org/apache/druid/query/Query.java +++ b/processing/src/main/java/org/apache/druid/query/Query.java @@ -232,6 +232,16 @@ default > E getContextAsEnum(String key, Class clazz, E def } } + @Nullable + default Boolean getContextAsBoolean(String key) + { + if (getQueryContext() != null) { + return getQueryContext().getAsBoolean(key); + } else { + return null; + } + } + default boolean getContextBoolean(String key, boolean defaultValue) { if (getQueryContext() != null) { diff --git a/processing/src/main/java/org/apache/druid/query/QueryContext.java b/processing/src/main/java/org/apache/druid/query/QueryContext.java index 52291ceda952..3269c928bfa0 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContext.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContext.java @@ -178,6 +178,12 @@ public String getAsString(String key, String defaultValue) return val == null ? defaultValue : val.toString(); } + @Nullable + public Boolean getAsBoolean(String parameter) + { + return QueryContexts.getAsBoolean(parameter, get(parameter)); + } + public boolean getAsBoolean( final String parameter, final boolean defaultValue diff --git a/processing/src/main/java/org/apache/druid/query/QueryContexts.java b/processing/src/main/java/org/apache/druid/query/QueryContexts.java index a1bd87e1c9cf..9b127b2044c5 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContexts.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContexts.java @@ -491,18 +491,14 @@ public static String getAsString( } } - /** - * Get the value of a parameter as a {@code boolean}. The parameter is expected - * to be {@code null}, a string or a {@code Boolean} object. - */ - public static boolean getAsBoolean( + @Nullable + public static Boolean getAsBoolean( final String parameter, - final Object value, - final boolean defaultValue + final Object value ) { if (value == null) { - return defaultValue; + return null; } else if (value instanceof String) { return Boolean.parseBoolean((String) value); } else if (value instanceof Boolean) { @@ -512,6 +508,20 @@ public static boolean getAsBoolean( } } + /** + * Get the value of a parameter as a {@code boolean}. The parameter is expected + * to be {@code null}, a string or a {@code Boolean} object. + */ + public static boolean getAsBoolean( + final String parameter, + final Object value, + final boolean defaultValue + ) + { + Boolean val = getAsBoolean(parameter, value); + return val == null ? defaultValue : val; + } + @Nullable public static Integer getAsInt(String parameter, Object value) { diff --git a/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java b/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java index 1473e39c4604..94a48d896f48 100644 --- a/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java +++ b/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java @@ -103,9 +103,9 @@ public void testContextSerde() throws Exception ); Assert.assertEquals((Integer) 1, serdeQuery.getContextAsInt(QueryContexts.PRIORITY_KEY)); - Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.USE_CACHE_KEY, false)); + Assert.assertEquals(true, serdeQuery.getContextAsBoolean(QueryContexts.USE_CACHE_KEY)); Assert.assertEquals("true", serdeQuery.getContextAsString(QueryContexts.POPULATE_CACHE_KEY)); - Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.FINALIZE_KEY, false)); + Assert.assertEquals(true, serdeQuery.getContextAsBoolean(QueryContexts.FINALIZE_KEY)); Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.USE_CACHE_KEY, false)); Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.POPULATE_CACHE_KEY, false)); Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.FINALIZE_KEY, false)); From 803ea0a897a4bb41b3296449a78d1a9adb16c398 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Sep 2022 12:17:15 +0800 Subject: [PATCH 4/9] Update test case --- .../druid/query/timeboundary/TimeBoundaryQueryTest.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java b/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java index 40d1d3fda985..46e7eb780f66 100644 --- a/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java +++ b/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java @@ -78,11 +78,10 @@ public void testContextSerde() throws Exception ), TimeBoundaryQuery.class ); - Assert.assertEquals(new Integer(1), serdeQuery.getContextAsInt(QueryContexts.PRIORITY_KEY)); - Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.USE_CACHE_KEY, false)); - Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.POPULATE_CACHE_KEY, false)); - Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.FINALIZE_KEY, false)); + Assert.assertEquals(true, serdeQuery.getContextAsBoolean(QueryContexts.USE_CACHE_KEY)); + Assert.assertEquals(true, serdeQuery.getContextAsBoolean(QueryContexts.POPULATE_CACHE_KEY)); + Assert.assertEquals(true, serdeQuery.getContextAsBoolean(QueryContexts.FINALIZE_KEY)); } @Test From 9ae638efeda601d1623e6d4373246e34500ebc69 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Sep 2022 12:26:41 +0800 Subject: [PATCH 5/9] Remove generic from getContextValue --- .../apache/druid/msq/exec/ControllerImpl.java | 4 ++-- .../druid/msq/querykit/QueryKitUtils.java | 2 +- .../druid/msq/querykit/scan/ScanQueryKit.java | 2 +- .../org/apache/druid/query/BaseQuery.java | 6 ------ .../java/org/apache/druid/query/Query.java | 19 +++---------------- .../groupby/strategy/GroupByStrategyV2.java | 2 +- .../druid/query/select/SelectQuery.java | 6 ------ .../scheduling/HiLoQueryLaningStrategy.java | 2 +- .../client/CachingClusteredClientTest.java | 6 +++--- 9 files changed, 12 insertions(+), 37 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java index 3aa346c38463..1b3bdb1aeae1 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java @@ -1454,7 +1454,7 @@ private static GranularitySpec makeGranularitySpecForIngestion( ) { if (isRollupQuery) { - final String queryGranularity = query.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY, ""); + final String queryGranularity = query.getContextAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY, ""); if (timeIsGroupByDimension((GroupByQuery) query, columnMappings) && !queryGranularity.isEmpty()) { return new ArbitraryGranularitySpec( @@ -1483,7 +1483,7 @@ private static boolean timeIsGroupByDimension(GroupByQuery groupByQuery, ColumnM { if (columnMappings.hasOutputColumn(ColumnHolder.TIME_COLUMN_NAME)) { final String queryTimeColumn = columnMappings.getQueryColumnForOutputColumn(ColumnHolder.TIME_COLUMN_NAME); - return queryTimeColumn.equals(groupByQuery.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD)); + return queryTimeColumn.equals(groupByQuery.getContextAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD)); } else { return false; } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitUtils.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitUtils.java index 2e5cd0c34c9d..6d587448f3b2 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitUtils.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitUtils.java @@ -191,7 +191,7 @@ public static RowSignature sortableSignature( public static VirtualColumn makeSegmentGranularityVirtualColumn(final Query query) { final Granularity segmentGranularity = QueryKitUtils.getSegmentGranularityFromContext(query.getContext()); - final String timeColumnName = query.getContextValue(QueryKitUtils.CTX_TIME_COLUMN_NAME); + final String timeColumnName = query.getContextAsString(QueryKitUtils.CTX_TIME_COLUMN_NAME); if (timeColumnName == null || Granularities.ALL.equals(segmentGranularity)) { return null; diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java index c8341a1b4766..b72968258670 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java @@ -57,7 +57,7 @@ public static RowSignature getAndValidateSignature(final ScanQuery scanQuery, fi { RowSignature scanSignature; try { - final String s = scanQuery.getContextValue(DruidQuery.CTX_SCAN_SIGNATURE); + final String s = scanQuery.getContextAsString(DruidQuery.CTX_SCAN_SIGNATURE); scanSignature = jsonMapper.readValue(s, RowSignature.class); } catch (JsonProcessingException e) { diff --git a/processing/src/main/java/org/apache/druid/query/BaseQuery.java b/processing/src/main/java/org/apache/druid/query/BaseQuery.java index 514220b2506b..a4c1a999a812 100644 --- a/processing/src/main/java/org/apache/druid/query/BaseQuery.java +++ b/processing/src/main/java/org/apache/druid/query/BaseQuery.java @@ -181,12 +181,6 @@ public QueryContext getQueryContext() return context; } - @Override - public ContextType getContextValue(String key) - { - return (ContextType) context.get(key); - } - @Override public boolean getContextBoolean(String key, boolean defaultValue) { diff --git a/processing/src/main/java/org/apache/druid/query/Query.java b/processing/src/main/java/org/apache/druid/query/Query.java index 296c964ffcfa..cf971afcf7ed 100644 --- a/processing/src/main/java/org/apache/druid/query/Query.java +++ b/processing/src/main/java/org/apache/druid/query/Query.java @@ -124,20 +124,8 @@ default QueryContext getQueryContext() return null; } - default boolean containsContextValue(String key) - { - if (getQueryContext() != null) { - return getQueryContext().containsKey(key); - } else { - return false; - } - } - /** - * Get context value and cast it to ContextType. - * NOTE: If value is not type of ContextType, {@link java.lang.ClassCastException} is raised. - * - * For safe conversion, please use following methods: + * Get context value. It's recommended to use following methods instead * {@link #getContextBoolean(String, boolean)} * {@link #getContextAsString(String)} * {@link #getContextAsInt(String)} @@ -147,11 +135,10 @@ default boolean containsContextValue(String key) * {@link #getContextAsHumanReadableBytes(String, HumanReadableBytes)} */ @Nullable - default ContextType getContextValue(String key) + default Object getContextValue(String key) { if (getQueryContext() != null) { - //noinspection unchecked - return (ContextType) getQueryContext().get(key); + return getQueryContext().get(key); } else { return null; } diff --git a/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java b/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java index bc8f411e54d3..9d10f29e0948 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java @@ -249,7 +249,7 @@ public Sequence mergeResults( // the granularity and dimensions are slightly different. // now, part of the query plan logic is handled in GroupByStrategyV2, not only in DruidQuery.toGroupByQuery() final Granularity timestampResultFieldGranularity - = query.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY); + = (Granularity) query.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY); dimensionSpecs = query.getDimensions() .stream() diff --git a/processing/src/main/java/org/apache/druid/query/select/SelectQuery.java b/processing/src/main/java/org/apache/druid/query/select/SelectQuery.java index ad4ad3616356..e6a9ec197fd3 100644 --- a/processing/src/main/java/org/apache/druid/query/select/SelectQuery.java +++ b/processing/src/main/java/org/apache/druid/query/select/SelectQuery.java @@ -116,12 +116,6 @@ public QueryContext getQueryContext() throw new RuntimeException(REMOVED_ERROR_MESSAGE); } - @Override - public ContextType getContextValue(String key) - { - throw new RuntimeException(REMOVED_ERROR_MESSAGE); - } - @Override public boolean isDescending() { diff --git a/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java b/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java index 36f2057b385f..ecb432d3fb71 100644 --- a/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java +++ b/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java @@ -70,7 +70,7 @@ public Optional computeLane(QueryPlus query, Set>> apply(@Nullable Integer input) QueryPlus capturedQueryPlus = (QueryPlus) queryCapture.getValue(); Query capturedQuery = capturedQueryPlus.getQuery(); if (expectBySegment) { - Assert.assertEquals(true, capturedQuery.getContextValue(QueryContexts.BY_SEGMENT_KEY)); + Assert.assertEquals(true, capturedQuery.getContextAsBoolean(QueryContexts.BY_SEGMENT_KEY)); } else { Assert.assertTrue( - !capturedQuery.containsContextValue(QueryContexts.BY_SEGMENT_KEY) || - capturedQuery.getContextValue(QueryContexts.BY_SEGMENT_KEY).equals(false) + capturedQuery.getContextValue(QueryContexts.BY_SEGMENT_KEY) == null || + capturedQuery.getContextAsBoolean(QueryContexts.BY_SEGMENT_KEY).equals(false) ); } } From 93db08253d83ea093c82527f3829a19a18c2c757 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Sep 2022 14:24:39 +0800 Subject: [PATCH 6/9] Update catch-handler --- .../org/apache/druid/query/QueryContexts.java | 26 ++++++++++++------- .../druid/sql/http/SqlResourceTest.java | 2 +- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/query/QueryContexts.java b/processing/src/main/java/org/apache/druid/query/QueryContexts.java index 9b127b2044c5..a64858a878e5 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContexts.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContexts.java @@ -423,7 +423,7 @@ public static long getTimeout(Query query, long defaultTimeout) Preconditions.checkState(timeout >= 0, "Timeout must be a non negative value, but was [%s]", timeout); return timeout; } - catch (NumberFormatException e) { + catch (IAE e) { throw new BadQueryContextException(e); } } @@ -487,7 +487,7 @@ public static String getAsString( } else if (value instanceof String) { return (String) value; } else { - throw new IAE("Expected parameter [%s] to be String", parameter); + throw new IAE("Expected parameter [%s] to be a String, but got [%s]", parameter, value.getClass().getName()); } } @@ -504,7 +504,7 @@ public static Boolean getAsBoolean( } else if (value instanceof Boolean) { return (Boolean) value; } else { - throw new IAE("Expected parameter [%s] to be a boolean", parameter); + throw new IAE("Expected parameter [%s] to be a Boolean, but got [%s]", parameter, value.getClass().getName()); } } @@ -534,10 +534,11 @@ public static Integer getAsInt(String parameter, Object value) return Numbers.parseInt(value); } catch (NumberFormatException ignored) { + throw new IAE("Expected parameter [%s] in integer format, but got [%s]", parameter, value); } } - throw new IAE("Expected parameter [%s] to be an integer", parameter); + throw new IAE("Expected parameter [%s] to be an Integer, but got [%s]", parameter, value.getClass().getName()); } /** @@ -566,9 +567,10 @@ public static Long getAsLong(String parameter, Object value) return Numbers.parseLong(value); } catch (NumberFormatException ignored) { + throw new IAE("Expected parameter [%s] in long format, but got [%s]", parameter, value); } } - throw new IAE("Expected parameter [%s] to be a long", parameter); + throw new IAE("Expected parameter [%s] to be a Long, but got [%s]", parameter, value.getClass().getName()); } /** @@ -596,10 +598,15 @@ public static HumanReadableBytes getAsHumanReadableBytes( } else if (value instanceof Number) { return HumanReadableBytes.valueOf(Numbers.parseLong(value)); } else if (value instanceof String) { - return new HumanReadableBytes((String) value); - } else { - throw new IAE("Expected parameter [%s] to be in human readable format", parameter); + try { + return HumanReadableBytes.valueOf(HumanReadableBytes.parse((String) value)); + } + catch (IAE e) { + throw new IAE("Expected parameter [%s] in human readable format, but got [%s]", parameter, value); + } } + + throw new IAE("Expected parameter [%s] to be a human readable number, but got [%s]", parameter, value.getClass().getName()); } public static float getAsFloat(String parameter, Object value, float defaultValue) @@ -613,9 +620,10 @@ public static float getAsFloat(String parameter, Object value, float defaultValu return Float.parseFloat((String) value); } catch (NumberFormatException ignored) { + throw new IAE("Expected parameter [%s] in float format, but got [%s]", parameter, value); } } - throw new IAE("Expected parameter [%s] to be a float", parameter); + throw new IAE("Expected parameter [%s] to be a Float, but got [%s]", parameter, value.getClass().getName()); } public static Map override( diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index 4a16c344c122..4d5f240ef99a 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -1705,7 +1705,7 @@ public void testQueryContextException() throws Exception Assert.assertNotNull(queryContextException); Assert.assertEquals(BadQueryContextException.ERROR_CODE, queryContextException.getErrorCode()); Assert.assertEquals(BadQueryContextException.ERROR_CLASS, queryContextException.getErrorClass()); - Assert.assertTrue(queryContextException.getMessage().contains("For input string: \"2000'\"")); + Assert.assertTrue(queryContextException.getMessage().contains("2000'")); checkSqlRequestLog(false); Assert.assertTrue(lifecycleManager.getAll(sqlQueryId).isEmpty()); } From 5bfe7e74b1e0f5843b33d13e7beb65e66f09a259 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Mon, 5 Sep 2022 13:13:42 +0800 Subject: [PATCH 7/9] Add test --- .../org/apache/druid/query/QueryContexts.java | 26 ++++++--- .../apache/druid/query/QueryContextsTest.java | 53 ++++++++++++++++--- 2 files changed, 66 insertions(+), 13 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/query/QueryContexts.java b/processing/src/main/java/org/apache/druid/query/QueryContexts.java index a64858a878e5..c51ff8509b9d 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContexts.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContexts.java @@ -649,12 +649,26 @@ public static > E getAsEnum(String parameter, Object val, Clas if (val == null) { return defaultValue; } - if (val instanceof String) { - return Enum.valueOf(clazz, StringUtils.toUpperCase((String) val)); - } else if (val instanceof Boolean) { - return Enum.valueOf(clazz, StringUtils.toUpperCase(String.valueOf(val))); - } else { - throw new ISE("Expected parameter [%s] must be type of [%s], actual type is [%s]. ", parameter, val.getClass()); + + try { + if (val instanceof String) { + return Enum.valueOf(clazz, StringUtils.toUpperCase((String) val)); + } else if (val instanceof Boolean) { + return Enum.valueOf(clazz, StringUtils.toUpperCase(String.valueOf(val))); + } } + catch (IllegalArgumentException e) { + throw new IAE("Expected parameter [%s] must be value of enum [%s], but got [%s].", + parameter, + clazz.getName(), + val.toString()); + } + + throw new ISE( + "Expected parameter [%s] must be type of [%s], actual type is [%s].", + parameter, + clazz.getName(), + val.getClass() + ); } } diff --git a/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java b/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java index 3d34e2d0e40f..22478ff64d4b 100644 --- a/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java +++ b/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java @@ -137,15 +137,19 @@ public void testDefaultSegmentPruning() @Test public void testDefaultInSubQueryThreshold() { - Assert.assertEquals(QueryContexts.DEFAULT_IN_SUB_QUERY_THRESHOLD, - QueryContexts.getInSubQueryThreshold(ImmutableMap.of())); + Assert.assertEquals( + QueryContexts.DEFAULT_IN_SUB_QUERY_THRESHOLD, + QueryContexts.getInSubQueryThreshold(ImmutableMap.of()) + ); } @Test public void testDefaultPlanTimeBoundarySql() { - Assert.assertEquals(QueryContexts.DEFAULT_ENABLE_TIME_BOUNDARY_PLANNING, - QueryContexts.isTimeBoundaryPlanningEnabled(ImmutableMap.of())); + Assert.assertEquals( + QueryContexts.DEFAULT_ENABLE_TIME_BOUNDARY_PLANNING, + QueryContexts.isTimeBoundaryPlanningEnabled(ImmutableMap.of()) + ); } @Test @@ -279,8 +283,43 @@ public void testGetAs() @Test public void testGetAsHumanReadableBytes() { - Assert.assertEquals(new HumanReadableBytes("500M").getBytes(), QueryContexts.getAsHumanReadableBytes("maxOnDiskStorage", 500_000_000, HumanReadableBytes.ZERO).getBytes()); - Assert.assertEquals(new HumanReadableBytes("500M").getBytes(), QueryContexts.getAsHumanReadableBytes("maxOnDiskStorage", "500000000", HumanReadableBytes.ZERO).getBytes()); - Assert.assertEquals(new HumanReadableBytes("500M").getBytes(), QueryContexts.getAsHumanReadableBytes("maxOnDiskStorage", "500M", HumanReadableBytes.ZERO).getBytes()); + Assert.assertEquals( + new HumanReadableBytes("500M").getBytes(), + QueryContexts.getAsHumanReadableBytes("maxOnDiskStorage", 500_000_000, HumanReadableBytes.ZERO) + .getBytes() + ); + Assert.assertEquals( + new HumanReadableBytes("500M").getBytes(), + QueryContexts.getAsHumanReadableBytes("maxOnDiskStorage", "500000000", HumanReadableBytes.ZERO) + .getBytes() + ); + Assert.assertEquals( + new HumanReadableBytes("500M").getBytes(), + QueryContexts.getAsHumanReadableBytes("maxOnDiskStorage", "500M", HumanReadableBytes.ZERO) + .getBytes() + ); + } + + @Test + public void testGetEnum() + { + Query query = new TestQuery( + new TableDataSource("test"), + new MultipleIntervalSegmentSpec(ImmutableList.of(Intervals.of("0/100"))), + false, + ImmutableMap.of("e1", "FORCE", + "e2", "INVALID_ENUM" + ) + ); + + Assert.assertEquals( + QueryContexts.Vectorize.FORCE, + query.getContextAsEnum("e1", QueryContexts.Vectorize.class, QueryContexts.Vectorize.FALSE) + ); + + Assert.assertThrows( + IAE.class, + () -> query.getContextAsEnum("e2", QueryContexts.Vectorize.class, QueryContexts.Vectorize.FALSE) + ); } } From 73202dedc05c2687212ca0756409e7c5f65c7aef Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Sep 2022 15:32:44 +0800 Subject: [PATCH 8/9] Resolve comments --- .../java/org/apache/druid/query/Query.java | 79 ++++++++++--------- .../org/apache/druid/query/QueryContext.java | 36 ++++----- .../org/apache/druid/query/QueryContexts.java | 48 +++++------ .../groupby/strategy/GroupByStrategyV2.java | 2 +- 4 files changed, 84 insertions(+), 81 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/query/Query.java b/processing/src/main/java/org/apache/druid/query/Query.java index cf971afcf7ed..c3d9628e969c 100644 --- a/processing/src/main/java/org/apache/druid/query/Query.java +++ b/processing/src/main/java/org/apache/druid/query/Query.java @@ -125,7 +125,10 @@ default QueryContext getQueryContext() } /** - * Get context value. It's recommended to use following methods instead + * Get context value and cast to ContextType in an unsafe way. + * + * For safe conversion, it's recommended to use following methods instead + * * {@link #getContextBoolean(String, boolean)} * {@link #getContextAsString(String)} * {@link #getContextAsInt(String)} @@ -135,106 +138,106 @@ default QueryContext getQueryContext() * {@link #getContextAsHumanReadableBytes(String, HumanReadableBytes)} */ @Nullable - default Object getContextValue(String key) + default ContextType getContextValue(String key) { - if (getQueryContext() != null) { - return getQueryContext().get(key); - } else { + if (getQueryContext() == null) { return null; + } else { + return (ContextType) getQueryContext().get(key); } } @Nullable default String getContextAsString(String key) { - if (getQueryContext() != null) { - return getQueryContext().getAsString(key); - } else { + if (getQueryContext() == null) { return null; + } else { + return getQueryContext().getAsString(key); } } default String getContextAsString(String key, String defaultValue) { - if (getQueryContext() != null) { - return getQueryContext().getAsString(key, defaultValue); - } else { + if (getQueryContext() == null) { return defaultValue; + } else { + return getQueryContext().getAsString(key, defaultValue); } } @Nullable default Integer getContextAsInt(String key) { - if (getQueryContext() != null) { - return getQueryContext().getAsInt(key); - } else { + if (getQueryContext() == null) { return null; + } else { + return getQueryContext().getAsInt(key); } } default int getContextAsInt(String key, int defaultValue) { - if (getQueryContext() != null) { - return getQueryContext().getAsInt(key, defaultValue); - } else { + if (getQueryContext() == null) { return defaultValue; + } else { + return getQueryContext().getAsInt(key, defaultValue); } } @Nullable default Long getContextAsLong(String key) { - if (getQueryContext() != null) { - return getQueryContext().getAsLong(key); - } else { + if (getQueryContext() == null) { return null; + } else { + return getQueryContext().getAsLong(key); } } default long getContextAsLong(String key, long defaultValue) { - if (getQueryContext() != null) { - return getQueryContext().getAsLong(key, defaultValue); - } else { + if (getQueryContext() == null) { return defaultValue; + } else { + return getQueryContext().getAsLong(key, defaultValue); } } default float getContextAsFloat(String key, float defaultValue) { - if (getQueryContext() != null) { - return getQueryContext().getAsFloat(key, defaultValue); - } else { + if (getQueryContext() == null) { return defaultValue; + } else { + return getQueryContext().getAsFloat(key, defaultValue); } } default > E getContextAsEnum(String key, Class clazz, E defaultValue) { - if (getQueryContext() != null) { - return getQueryContext().getAsEnum(key, clazz, defaultValue); - } else { + if (getQueryContext() == null) { return defaultValue; + } else { + return getQueryContext().getAsEnum(key, clazz, defaultValue); } } @Nullable default Boolean getContextAsBoolean(String key) { - if (getQueryContext() != null) { - return getQueryContext().getAsBoolean(key); - } else { + if (getQueryContext() == null) { return null; + } else { + return getQueryContext().getAsBoolean(key); } } default boolean getContextBoolean(String key, boolean defaultValue) { - if (getQueryContext() != null) { - return getQueryContext().getAsBoolean(key, defaultValue); - } else { + if (getQueryContext() == null) { return defaultValue; + } else { + return getQueryContext().getAsBoolean(key, defaultValue); } } @@ -249,10 +252,10 @@ default boolean getContextBoolean(String key, boolean defaultValue) */ default HumanReadableBytes getContextAsHumanReadableBytes(String key, HumanReadableBytes defaultValue) { - if (null != getQueryContext()) { - return getQueryContext().getAsHumanReadableBytes(key, defaultValue); - } else { + if (getQueryContext() == null) { return defaultValue; + } else { + return getQueryContext().getAsHumanReadableBytes(key, defaultValue); } } diff --git a/processing/src/main/java/org/apache/druid/query/QueryContext.java b/processing/src/main/java/org/apache/druid/query/QueryContext.java index 3269c928bfa0..f902bddb2d97 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContext.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContext.java @@ -179,55 +179,55 @@ public String getAsString(String key, String defaultValue) } @Nullable - public Boolean getAsBoolean(String parameter) + public Boolean getAsBoolean(String key) { - return QueryContexts.getAsBoolean(parameter, get(parameter)); + return QueryContexts.getAsBoolean(key, get(key)); } public boolean getAsBoolean( - final String parameter, + final String key, final boolean defaultValue ) { - return QueryContexts.getAsBoolean(parameter, get(parameter), defaultValue); + return QueryContexts.getAsBoolean(key, get(key), defaultValue); } - public Integer getAsInt(final String parameter) + public Integer getAsInt(final String key) { - return QueryContexts.getAsInt(parameter, get(parameter)); + return QueryContexts.getAsInt(key, get(key)); } public int getAsInt( - final String parameter, + final String key, final int defaultValue ) { - return QueryContexts.getAsInt(parameter, get(parameter), defaultValue); + return QueryContexts.getAsInt(key, get(key), defaultValue); } - public Long getAsLong(final String parameter) + public Long getAsLong(final String key) { - return QueryContexts.getAsLong(parameter, get(parameter)); + return QueryContexts.getAsLong(key, get(key)); } - public long getAsLong(final String parameter, final long defaultValue) + public long getAsLong(final String key, final long defaultValue) { - return QueryContexts.getAsLong(parameter, get(parameter), defaultValue); + return QueryContexts.getAsLong(key, get(key), defaultValue); } - public HumanReadableBytes getAsHumanReadableBytes(final String parameter, final HumanReadableBytes defaultValue) + public HumanReadableBytes getAsHumanReadableBytes(final String key, final HumanReadableBytes defaultValue) { - return QueryContexts.getAsHumanReadableBytes(parameter, get(parameter), defaultValue); + return QueryContexts.getAsHumanReadableBytes(key, get(key), defaultValue); } - public float getAsFloat(final String parameter, final float defaultValue) + public float getAsFloat(final String key, final float defaultValue) { - return QueryContexts.getAsFloat(parameter, get(parameter), defaultValue); + return QueryContexts.getAsFloat(key, get(key), defaultValue); } - public > E getAsEnum(String parameter, Class clazz, E defaultValue) + public > E getAsEnum(String key, Class clazz, E defaultValue) { - return QueryContexts.getAsEnum(parameter, get(parameter), clazz, defaultValue); + return QueryContexts.getAsEnum(key, get(key), clazz, defaultValue); } public Map getMergedParams() diff --git a/processing/src/main/java/org/apache/druid/query/QueryContexts.java b/processing/src/main/java/org/apache/druid/query/QueryContexts.java index c51ff8509b9d..22f5b2960066 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContexts.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContexts.java @@ -477,7 +477,7 @@ static boolean parseBoolean(Map context, String key, boolean def } public static String getAsString( - final String parameter, + final String key, final Object value, final String defaultValue ) @@ -487,7 +487,7 @@ public static String getAsString( } else if (value instanceof String) { return (String) value; } else { - throw new IAE("Expected parameter [%s] to be a String, but got [%s]", parameter, value.getClass().getName()); + throw new IAE("Expected key [%s] to be a String, but got [%s]", key, value.getClass().getName()); } } @@ -513,17 +513,17 @@ public static Boolean getAsBoolean( * to be {@code null}, a string or a {@code Boolean} object. */ public static boolean getAsBoolean( - final String parameter, + final String key, final Object value, final boolean defaultValue ) { - Boolean val = getAsBoolean(parameter, value); + Boolean val = getAsBoolean(key, value); return val == null ? defaultValue : val; } @Nullable - public static Integer getAsInt(String parameter, Object value) + public static Integer getAsInt(String key, Object value) { if (value == null) { return null; @@ -534,11 +534,11 @@ public static Integer getAsInt(String parameter, Object value) return Numbers.parseInt(value); } catch (NumberFormatException ignored) { - throw new IAE("Expected parameter [%s] in integer format, but got [%s]", parameter, value); + throw new IAE("Expected key [%s] in integer format, but got [%s]", key, value); } } - throw new IAE("Expected parameter [%s] to be an Integer, but got [%s]", parameter, value.getClass().getName()); + throw new IAE("Expected key [%s] to be an Integer, but got [%s]", key, value.getClass().getName()); } /** @@ -546,17 +546,17 @@ public static Integer getAsInt(String parameter, Object value) * to be {@code null}, a string or a {@code Number} object. */ public static int getAsInt( - final String parameter, + final String ke, final Object value, final int defaultValue ) { - Integer val = getAsInt(parameter, value); + Integer val = getAsInt(ke, value); return val == null ? defaultValue : val; } @Nullable - public static Long getAsLong(String parameter, Object value) + public static Long getAsLong(String key, Object value) { if (value == null) { return null; @@ -567,10 +567,10 @@ public static Long getAsLong(String parameter, Object value) return Numbers.parseLong(value); } catch (NumberFormatException ignored) { - throw new IAE("Expected parameter [%s] in long format, but got [%s]", parameter, value); + throw new IAE("Expected key [%s] in long format, but got [%s]", key, value); } } - throw new IAE("Expected parameter [%s] to be a Long, but got [%s]", parameter, value.getClass().getName()); + throw new IAE("Expected key [%s] to be a Long, but got [%s]", key, value.getClass().getName()); } /** @@ -578,12 +578,12 @@ public static Long getAsLong(String parameter, Object value) * to be {@code null}, a string or a {@code Number} object. */ public static long getAsLong( - final String parameter, + final String key, final Object value, final long defaultValue ) { - Long val = getAsLong(parameter, value); + Long val = getAsLong(key, value); return val == null ? defaultValue : val; } @@ -602,14 +602,14 @@ public static HumanReadableBytes getAsHumanReadableBytes( return HumanReadableBytes.valueOf(HumanReadableBytes.parse((String) value)); } catch (IAE e) { - throw new IAE("Expected parameter [%s] in human readable format, but got [%s]", parameter, value); + throw new IAE("Expected key [%s] in human readable format, but got [%s]", parameter, value); } } - throw new IAE("Expected parameter [%s] to be a human readable number, but got [%s]", parameter, value.getClass().getName()); + throw new IAE("Expected key [%s] to be a human readable number, but got [%s]", parameter, value.getClass().getName()); } - public static float getAsFloat(String parameter, Object value, float defaultValue) + public static float getAsFloat(String key, Object value, float defaultValue) { if (null == value) { return defaultValue; @@ -620,10 +620,10 @@ public static float getAsFloat(String parameter, Object value, float defaultValu return Float.parseFloat((String) value); } catch (NumberFormatException ignored) { - throw new IAE("Expected parameter [%s] in float format, but got [%s]", parameter, value); + throw new IAE("Expected key [%s] in float format, but got [%s]", key, value); } } - throw new IAE("Expected parameter [%s] to be a Float, but got [%s]", parameter, value.getClass().getName()); + throw new IAE("Expected key [%s] to be a Float, but got [%s]", key, value.getClass().getName()); } public static Map override( @@ -644,7 +644,7 @@ private QueryContexts() { } - public static > E getAsEnum(String parameter, Object val, Class clazz, E defaultValue) + public static > E getAsEnum(String key, Object val, Class clazz, E defaultValue) { if (val == null) { return defaultValue; @@ -658,15 +658,15 @@ public static > E getAsEnum(String parameter, Object val, Clas } } catch (IllegalArgumentException e) { - throw new IAE("Expected parameter [%s] must be value of enum [%s], but got [%s].", - parameter, + throw new IAE("Expected key [%s] must be value of enum [%s], but got [%s].", + key, clazz.getName(), val.toString()); } throw new ISE( - "Expected parameter [%s] must be type of [%s], actual type is [%s].", - parameter, + "Expected key [%s] must be type of [%s], actual type is [%s].", + key, clazz.getName(), val.getClass() ); diff --git a/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java b/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java index 9d10f29e0948..bc8f411e54d3 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java @@ -249,7 +249,7 @@ public Sequence mergeResults( // the granularity and dimensions are slightly different. // now, part of the query plan logic is handled in GroupByStrategyV2, not only in DruidQuery.toGroupByQuery() final Granularity timestampResultFieldGranularity - = (Granularity) query.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY); + = query.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY); dimensionSpecs = query.getDimensions() .stream() From 8f6656371d2f14bcb380b843fc14c513b8b6f0b8 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Mon, 12 Sep 2022 12:05:02 +0800 Subject: [PATCH 9/9] Replace 'getContextXXX' to 'getQueryContext().getAsXXXX' --- .../apache/druid/msq/exec/ControllerImpl.java | 4 +- .../druid/msq/querykit/QueryKitUtils.java | 2 +- .../druid/msq/querykit/scan/ScanQueryKit.java | 2 +- .../java/org/apache/druid/query/Query.java | 101 ++---------------- .../org/apache/druid/query/QueryContexts.java | 34 +++--- .../druid/query/groupby/GroupByQuery.java | 2 +- .../query/groupby/GroupByQueryConfig.java | 18 ++-- .../epinephelinae/GroupByQueryEngineV2.java | 2 +- .../groupby/orderby/DefaultLimitSpec.java | 4 +- .../groupby/strategy/GroupByStrategyV2.java | 4 +- .../apache/druid/query/scan/ScanQuery.java | 4 +- .../druid/query/search/SearchQueryConfig.java | 2 +- .../TimeBoundaryQueryQueryToolChest.java | 4 +- .../query/timeseries/TimeseriesQuery.java | 2 +- .../query/topn/TopNQueryQueryToolChest.java | 2 +- .../apache/druid/query/QueryContextsTest.java | 4 +- .../DataSourceMetadataQueryTest.java | 8 +- .../timeboundary/TimeBoundaryQueryTest.java | 16 +-- .../druid/client/DirectDruidClient.java | 2 +- .../druid/client/JsonParserIterator.java | 2 +- .../server/ClientQuerySegmentWalker.java | 2 +- .../scheduling/HiLoQueryLaningStrategy.java | 2 +- .../client/CachingClusteredClientTest.java | 4 +- .../druid/client/JsonParserIteratorTest.java | 3 +- .../SetAndVerifyContextQueryRunnerTest.java | 8 +- 25 files changed, 77 insertions(+), 161 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java index 1b3bdb1aeae1..b6414c3db6e8 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java @@ -1454,7 +1454,7 @@ private static GranularitySpec makeGranularitySpecForIngestion( ) { if (isRollupQuery) { - final String queryGranularity = query.getContextAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY, ""); + final String queryGranularity = query.getQueryContext().getAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY, ""); if (timeIsGroupByDimension((GroupByQuery) query, columnMappings) && !queryGranularity.isEmpty()) { return new ArbitraryGranularitySpec( @@ -1483,7 +1483,7 @@ private static boolean timeIsGroupByDimension(GroupByQuery groupByQuery, ColumnM { if (columnMappings.hasOutputColumn(ColumnHolder.TIME_COLUMN_NAME)) { final String queryTimeColumn = columnMappings.getQueryColumnForOutputColumn(ColumnHolder.TIME_COLUMN_NAME); - return queryTimeColumn.equals(groupByQuery.getContextAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD)); + return queryTimeColumn.equals(groupByQuery.getQueryContext().getAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD)); } else { return false; } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitUtils.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitUtils.java index 6d587448f3b2..8630fec754d2 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitUtils.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitUtils.java @@ -191,7 +191,7 @@ public static RowSignature sortableSignature( public static VirtualColumn makeSegmentGranularityVirtualColumn(final Query query) { final Granularity segmentGranularity = QueryKitUtils.getSegmentGranularityFromContext(query.getContext()); - final String timeColumnName = query.getContextAsString(QueryKitUtils.CTX_TIME_COLUMN_NAME); + final String timeColumnName = query.getQueryContext().getAsString(QueryKitUtils.CTX_TIME_COLUMN_NAME); if (timeColumnName == null || Granularities.ALL.equals(segmentGranularity)) { return null; diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java index b72968258670..edd553b2b046 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java @@ -57,7 +57,7 @@ public static RowSignature getAndValidateSignature(final ScanQuery scanQuery, fi { RowSignature scanSignature; try { - final String s = scanQuery.getContextAsString(DruidQuery.CTX_SCAN_SIGNATURE); + final String s = scanQuery.getQueryContext().getAsString(DruidQuery.CTX_SCAN_SIGNATURE); scanSignature = jsonMapper.readValue(s, RowSignature.class); } catch (JsonProcessingException e) { diff --git a/processing/src/main/java/org/apache/druid/query/Query.java b/processing/src/main/java/org/apache/druid/query/Query.java index c3d9628e969c..5662b988c238 100644 --- a/processing/src/main/java/org/apache/druid/query/Query.java +++ b/processing/src/main/java/org/apache/druid/query/Query.java @@ -129,13 +129,13 @@ default QueryContext getQueryContext() * * For safe conversion, it's recommended to use following methods instead * - * {@link #getContextBoolean(String, boolean)} - * {@link #getContextAsString(String)} - * {@link #getContextAsInt(String)} - * {@link #getContextAsLong(String)} - * {@link #getContextAsFloat(String, float)} - * {@link #getContextAsEnum(String, Class, Enum)} - * {@link #getContextAsHumanReadableBytes(String, HumanReadableBytes)} + * {@link QueryContext#getAsBoolean(String)} + * {@link QueryContext#getAsString(String)} + * {@link QueryContext#getAsInt(String)} + * {@link QueryContext#getAsLong(String)} + * {@link QueryContext#getAsFloat(String, float)} + * {@link QueryContext#getAsEnum(String, Class, Enum)} + * {@link QueryContext#getAsHumanReadableBytes(String, HumanReadableBytes)} */ @Nullable default ContextType getContextValue(String key) @@ -147,91 +147,6 @@ default ContextType getContextValue(String key) } } - @Nullable - default String getContextAsString(String key) - { - if (getQueryContext() == null) { - return null; - } else { - return getQueryContext().getAsString(key); - } - } - - default String getContextAsString(String key, String defaultValue) - { - if (getQueryContext() == null) { - return defaultValue; - } else { - return getQueryContext().getAsString(key, defaultValue); - } - } - - @Nullable - default Integer getContextAsInt(String key) - { - if (getQueryContext() == null) { - return null; - } else { - return getQueryContext().getAsInt(key); - } - } - - default int getContextAsInt(String key, int defaultValue) - { - if (getQueryContext() == null) { - return defaultValue; - } else { - return getQueryContext().getAsInt(key, defaultValue); - } - } - - @Nullable - default Long getContextAsLong(String key) - { - if (getQueryContext() == null) { - return null; - } else { - return getQueryContext().getAsLong(key); - } - } - - default long getContextAsLong(String key, long defaultValue) - { - if (getQueryContext() == null) { - return defaultValue; - } else { - return getQueryContext().getAsLong(key, defaultValue); - } - } - - default float getContextAsFloat(String key, float defaultValue) - { - if (getQueryContext() == null) { - return defaultValue; - } else { - return getQueryContext().getAsFloat(key, defaultValue); - } - } - - default > E getContextAsEnum(String key, Class clazz, E defaultValue) - { - if (getQueryContext() == null) { - return defaultValue; - } else { - return getQueryContext().getAsEnum(key, clazz, defaultValue); - } - } - - @Nullable - default Boolean getContextAsBoolean(String key) - { - if (getQueryContext() == null) { - return null; - } else { - return getQueryContext().getAsBoolean(key); - } - } - default boolean getContextBoolean(String key, boolean defaultValue) { if (getQueryContext() == null) { @@ -315,7 +230,7 @@ default Query withSqlQueryId(String sqlQueryId) @Nullable default String getSqlQueryId() { - return getContextAsString(BaseQuery.SQL_QUERY_ID); + return getQueryContext().getAsString(BaseQuery.SQL_QUERY_ID); } /** diff --git a/processing/src/main/java/org/apache/druid/query/QueryContexts.java b/processing/src/main/java/org/apache/druid/query/QueryContexts.java index 22f5b2960066..6138979faca7 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContexts.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContexts.java @@ -223,7 +223,7 @@ public static Vectorize getVectorize(Query query) public static Vectorize getVectorize(Query query, Vectorize defaultValue) { - return query.getContextAsEnum(VECTORIZE_KEY, Vectorize.class, defaultValue); + return query.getQueryContext().getAsEnum(VECTORIZE_KEY, Vectorize.class, defaultValue); } public static Vectorize getVectorizeVirtualColumns(Query query) @@ -233,7 +233,7 @@ public static Vectorize getVectorizeVirtualColumns(Query query) public static Vectorize getVectorizeVirtualColumns(Query query, Vectorize defaultValue) { - return query.getContextAsEnum(VECTORIZE_VIRTUAL_COLUMNS_KEY, Vectorize.class, defaultValue); + return query.getQueryContext().getAsEnum(VECTORIZE_VIRTUAL_COLUMNS_KEY, Vectorize.class, defaultValue); } public static int getVectorSize(Query query) @@ -243,12 +243,12 @@ public static int getVectorSize(Query query) public static int getVectorSize(Query query, int defaultSize) { - return query.getContextAsInt(VECTOR_SIZE_KEY, defaultSize); + return query.getQueryContext().getAsInt(VECTOR_SIZE_KEY, defaultSize); } public static int getMaxSubqueryRows(Query query, int defaultSize) { - return query.getContextAsInt(MAX_SUBQUERY_ROWS_KEY, defaultSize); + return query.getQueryContext().getAsInt(MAX_SUBQUERY_ROWS_KEY, defaultSize); } public static int getUncoveredIntervalsLimit(Query query) @@ -258,7 +258,7 @@ public static int getUncoveredIntervalsLimit(Query query) public static int getUncoveredIntervalsLimit(Query query, int defaultValue) { - return query.getContextAsInt(UNCOVERED_INTERVALS_LIMIT_KEY, defaultValue); + return query.getQueryContext().getAsInt(UNCOVERED_INTERVALS_LIMIT_KEY, defaultValue); } public static int getPriority(Query query) @@ -268,12 +268,12 @@ public static int getPriority(Query query) public static int getPriority(Query query, int defaultValue) { - return query.getContextAsInt(PRIORITY_KEY, defaultValue); + return query.getQueryContext().getAsInt(PRIORITY_KEY, defaultValue); } public static String getLane(Query query) { - return query.getContextAsString(LANE_KEY); + return query.getQueryContext().getAsString(LANE_KEY); } public static boolean getEnableParallelMerges(Query query) @@ -283,17 +283,17 @@ public static boolean getEnableParallelMerges(Query query) public static int getParallelMergeInitialYieldRows(Query query, int defaultValue) { - return query.getContextAsInt(BROKER_PARALLEL_MERGE_INITIAL_YIELD_ROWS_KEY, defaultValue); + return query.getQueryContext().getAsInt(BROKER_PARALLEL_MERGE_INITIAL_YIELD_ROWS_KEY, defaultValue); } public static int getParallelMergeSmallBatchRows(Query query, int defaultValue) { - return query.getContextAsInt(BROKER_PARALLEL_MERGE_SMALL_BATCH_ROWS_KEY, defaultValue); + return query.getQueryContext().getAsInt(BROKER_PARALLEL_MERGE_SMALL_BATCH_ROWS_KEY, defaultValue); } public static int getParallelMergeParallelism(Query query, int defaultValue) { - return query.getContextAsInt(BROKER_PARALLELISM, defaultValue); + return query.getQueryContext().getAsInt(BROKER_PARALLELISM, defaultValue); } public static boolean getEnableJoinFilterRewriteValueColumnFilters(Query query) @@ -314,7 +314,7 @@ public static boolean getEnableRewriteJoinToFilter(Query query) public static long getJoinFilterRewriteMaxSize(Query query) { - return query.getContextAsLong(JOIN_FILTER_REWRITE_MAX_SIZE_KEY, DEFAULT_ENABLE_JOIN_FILTER_REWRITE_MAX_SIZE); + return query.getQueryContext().getAsLong(JOIN_FILTER_REWRITE_MAX_SIZE_KEY, DEFAULT_ENABLE_JOIN_FILTER_REWRITE_MAX_SIZE); } public static boolean getEnableJoinFilterPushDown(Query query) @@ -364,7 +364,7 @@ public static boolean isTimeBoundaryPlanningEnabled(Map queryCon public static Query withMaxScatterGatherBytes(Query query, long maxScatterGatherBytesLimit) { - Long curr = query.getContextAsLong(MAX_SCATTER_GATHER_BYTES_KEY); + Long curr = query.getQueryContext().getAsLong(MAX_SCATTER_GATHER_BYTES_KEY); if (curr == null) { return query.withOverriddenContext(ImmutableMap.of(MAX_SCATTER_GATHER_BYTES_KEY, maxScatterGatherBytesLimit)); } else { @@ -398,12 +398,12 @@ public static Query verifyMaxQueryTimeout(Query query, long maxQueryTi public static long getMaxQueuedBytes(Query query, long defaultValue) { - return query.getContextAsLong(MAX_QUEUED_BYTES_KEY, defaultValue); + return query.getQueryContext().getAsLong(MAX_QUEUED_BYTES_KEY, defaultValue); } public static long getMaxScatterGatherBytes(Query query) { - return query.getContextAsLong(MAX_SCATTER_GATHER_BYTES_KEY, Long.MAX_VALUE); + return query.getQueryContext().getAsLong(MAX_SCATTER_GATHER_BYTES_KEY, Long.MAX_VALUE); } public static boolean hasTimeout(Query query) @@ -419,7 +419,7 @@ public static long getTimeout(Query query) public static long getTimeout(Query query, long defaultTimeout) { try { - final long timeout = query.getContextAsLong(TIMEOUT_KEY, defaultTimeout); + final long timeout = query.getQueryContext().getAsLong(TIMEOUT_KEY, defaultTimeout); Preconditions.checkState(timeout >= 0, "Timeout must be a non negative value, but was [%s]", timeout); return timeout; } @@ -440,14 +440,14 @@ public static Query withDefaultTimeout(Query query, long defaultTimeou static long getDefaultTimeout(Query query) { - final long defaultTimeout = query.getContextAsLong(DEFAULT_TIMEOUT_KEY, DEFAULT_TIMEOUT_MILLIS); + final long defaultTimeout = query.getQueryContext().getAsLong(DEFAULT_TIMEOUT_KEY, DEFAULT_TIMEOUT_MILLIS); Preconditions.checkState(defaultTimeout >= 0, "Timeout must be a non negative value, but was [%s]", defaultTimeout); return defaultTimeout; } public static int getNumRetriesOnMissingSegments(Query query, int defaultValue) { - return query.getContextAsInt(NUM_RETRIES_ON_MISSING_SEGMENTS_KEY, defaultValue); + return query.getQueryContext().getAsInt(NUM_RETRIES_ON_MISSING_SEGMENTS_KEY, defaultValue); } public static boolean allowReturnPartialResults(Query query, boolean defaultValue) diff --git a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQuery.java b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQuery.java index 0a67fc052d57..31d4b031e664 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQuery.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQuery.java @@ -748,7 +748,7 @@ private int compareDims(List dimensions, ResultRow lhs, ResultRow @Nullable private DateTime computeUniversalTimestamp() { - final String timestampStringFromContext = getContextAsString(CTX_KEY_FUDGE_TIMESTAMP, ""); + final String timestampStringFromContext = getQueryContext().getAsString(CTX_KEY_FUDGE_TIMESTAMP, ""); final Granularity granularity = getGranularity(); if (!timestampStringFromContext.isEmpty()) { diff --git a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java index 1c7e03a6e87c..ac018b942c0b 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryConfig.java @@ -335,25 +335,25 @@ public boolean isMultiValueUnnestingEnabled() public GroupByQueryConfig withOverrides(final GroupByQuery query) { final GroupByQueryConfig newConfig = new GroupByQueryConfig(); - newConfig.defaultStrategy = query.getContextAsString(CTX_KEY_STRATEGY, getDefaultStrategy()); - newConfig.singleThreaded = query.getContextBoolean(CTX_KEY_IS_SINGLE_THREADED, isSingleThreaded()); + newConfig.defaultStrategy = query.getQueryContext().getAsString(CTX_KEY_STRATEGY, getDefaultStrategy()); + newConfig.singleThreaded = query.getQueryContext().getAsBoolean(CTX_KEY_IS_SINGLE_THREADED, isSingleThreaded()); newConfig.maxIntermediateRows = Math.min( - query.getContextAsInt(CTX_KEY_MAX_INTERMEDIATE_ROWS, getMaxIntermediateRows()), + query.getQueryContext().getAsInt(CTX_KEY_MAX_INTERMEDIATE_ROWS, getMaxIntermediateRows()), getMaxIntermediateRows() ); newConfig.maxResults = Math.min( - query.getContextAsInt(CTX_KEY_MAX_RESULTS, getMaxResults()), + query.getQueryContext().getAsInt(CTX_KEY_MAX_RESULTS, getMaxResults()), getMaxResults() ); newConfig.bufferGrouperMaxSize = Math.min( - query.getContextAsInt(CTX_KEY_BUFFER_GROUPER_MAX_SIZE, getBufferGrouperMaxSize()), + query.getQueryContext().getAsInt(CTX_KEY_BUFFER_GROUPER_MAX_SIZE, getBufferGrouperMaxSize()), getBufferGrouperMaxSize() ); - newConfig.bufferGrouperMaxLoadFactor = query.getContextAsFloat( + newConfig.bufferGrouperMaxLoadFactor = query.getQueryContext().getAsFloat( CTX_KEY_BUFFER_GROUPER_MAX_LOAD_FACTOR, getBufferGrouperMaxLoadFactor() ); - newConfig.bufferGrouperInitialBuckets = query.getContextAsInt( + newConfig.bufferGrouperInitialBuckets = query.getQueryContext().getAsInt( CTX_KEY_BUFFER_GROUPER_INITIAL_BUCKETS, getBufferGrouperInitialBuckets() ); @@ -378,11 +378,11 @@ public GroupByQueryConfig withOverrides(final GroupByQuery query) CTX_KEY_FORCE_PUSH_DOWN_NESTED_QUERY, isForcePushDownNestedQuery() ); - newConfig.intermediateCombineDegree = query.getContextAsInt( + newConfig.intermediateCombineDegree = query.getQueryContext().getAsInt( CTX_KEY_INTERMEDIATE_COMBINE_DEGREE, getIntermediateCombineDegree() ); - newConfig.numParallelCombineThreads = query.getContextAsInt( + newConfig.numParallelCombineThreads = query.getQueryContext().getAsInt( CTX_KEY_NUM_PARALLEL_COMBINE_THREADS, getNumParallelCombineThreads() ); diff --git a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java index cc994a05616e..96055d521d50 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java @@ -141,7 +141,7 @@ public static Sequence process( try { final String fudgeTimestampString = NullHandling.emptyToNullIfNeeded( - query.getContextAsString(GroupByStrategyV2.CTX_KEY_FUDGE_TIMESTAMP) + query.getQueryContext().getAsString(GroupByStrategyV2.CTX_KEY_FUDGE_TIMESTAMP) ); final DateTime fudgeTimestamp = fudgeTimestampString == null diff --git a/processing/src/main/java/org/apache/druid/query/groupby/orderby/DefaultLimitSpec.java b/processing/src/main/java/org/apache/druid/query/groupby/orderby/DefaultLimitSpec.java index c6e44b6f89b7..33650cd5d9d1 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/orderby/DefaultLimitSpec.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/orderby/DefaultLimitSpec.java @@ -232,9 +232,9 @@ public Function, Sequence> build(final GroupByQue } if (!sortingNeeded) { - String timestampField = query.getContextAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD); + String timestampField = query.getQueryContext().getAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD); if (timestampField != null && !timestampField.isEmpty()) { - int timestampResultFieldIndex = query.getContextAsInt(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX); + int timestampResultFieldIndex = query.getQueryContext().getAsInt(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX); sortingNeeded = query.getContextSortByDimsFirst() ? timestampResultFieldIndex != query.getDimensions().size() - 1 : timestampResultFieldIndex != 0; diff --git a/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java b/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java index bc8f411e54d3..060127880122 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java @@ -221,7 +221,7 @@ public Sequence mergeResults( Granularity granularity = query.getGranularity(); List dimensionSpecs = query.getDimensions(); // the CTX_TIMESTAMP_RESULT_FIELD is set in DruidQuery.java - final String timestampResultField = query.getContextAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD); + final String timestampResultField = query.getQueryContext().getAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD); final boolean hasTimestampResultField = (timestampResultField != null && !timestampResultField.isEmpty()) && query.getContextBoolean(CTX_KEY_OUTERMOST, true) && !query.isApplyLimitPushDown(); @@ -258,7 +258,7 @@ public Sequence mergeResults( granularity = timestampResultFieldGranularity; // when timestampResultField is the last dimension, should set sortByDimsFirst=true, // otherwise the downstream is sorted by row's timestamp first which makes the final ordering not as expected - timestampResultFieldIndex = query.getContextAsInt(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX); + timestampResultFieldIndex = query.getQueryContext().getAsInt(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX); if (!query.getContextSortByDimsFirst() && timestampResultFieldIndex == query.getDimensions().size() - 1) { context.put(GroupByQuery.CTX_KEY_SORT_BY_DIMS_FIRST, true); } diff --git a/processing/src/main/java/org/apache/druid/query/scan/ScanQuery.java b/processing/src/main/java/org/apache/druid/query/scan/ScanQuery.java index 3fdf849e1f79..41273907623e 100644 --- a/processing/src/main/java/org/apache/druid/query/scan/ScanQuery.java +++ b/processing/src/main/java/org/apache/druid/query/scan/ScanQuery.java @@ -264,7 +264,7 @@ public static void verifyOrderByForNativeExecution(final ScanQuery query) private Integer validateAndGetMaxRowsQueuedForOrdering() { final Integer maxRowsQueuedForOrdering = - getContextAsInt(ScanQueryConfig.CTX_KEY_MAX_ROWS_QUEUED_FOR_ORDERING); + getQueryContext().getAsInt(ScanQueryConfig.CTX_KEY_MAX_ROWS_QUEUED_FOR_ORDERING); Preconditions.checkArgument( maxRowsQueuedForOrdering == null || maxRowsQueuedForOrdering > 0, "maxRowsQueuedForOrdering must be greater than 0" @@ -275,7 +275,7 @@ private Integer validateAndGetMaxRowsQueuedForOrdering() private Integer validateAndGetMaxSegmentPartitionsOrderedInMemory() { final Integer maxSegmentPartitionsOrderedInMemory = - getContextAsInt(ScanQueryConfig.CTX_KEY_MAX_SEGMENT_PARTITIONS_FOR_ORDERING); + getQueryContext().getAsInt(ScanQueryConfig.CTX_KEY_MAX_SEGMENT_PARTITIONS_FOR_ORDERING); Preconditions.checkArgument( maxSegmentPartitionsOrderedInMemory == null || maxSegmentPartitionsOrderedInMemory > 0, "maxRowsQueuedForOrdering must be greater than 0" diff --git a/processing/src/main/java/org/apache/druid/query/search/SearchQueryConfig.java b/processing/src/main/java/org/apache/druid/query/search/SearchQueryConfig.java index 318901caf437..c7b6f201f88c 100644 --- a/processing/src/main/java/org/apache/druid/query/search/SearchQueryConfig.java +++ b/processing/src/main/java/org/apache/druid/query/search/SearchQueryConfig.java @@ -55,7 +55,7 @@ public SearchQueryConfig withOverrides(final SearchQuery query) { final SearchQueryConfig newConfig = new SearchQueryConfig(); newConfig.maxSearchLimit = query.getLimit(); - newConfig.searchStrategy = query.getContextAsString(CTX_KEY_STRATEGY, searchStrategy); + newConfig.searchStrategy = query.getQueryContext().getAsString(CTX_KEY_STRATEGY, searchStrategy); return newConfig; } } diff --git a/processing/src/main/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java b/processing/src/main/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java index 861e3c1bb565..a5e58d094616 100644 --- a/processing/src/main/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java +++ b/processing/src/main/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java @@ -233,8 +233,8 @@ public RowSignature resultArraySignature(TimeBoundaryQuery query) if (query.isMinTime() || query.isMaxTime()) { RowSignature.Builder builder = RowSignature.builder(); String outputName = query.isMinTime() ? - query.getContextAsString(TimeBoundaryQuery.MIN_TIME_ARRAY_OUTPUT_NAME, TimeBoundaryQuery.MIN_TIME) : - query.getContextAsString(TimeBoundaryQuery.MAX_TIME_ARRAY_OUTPUT_NAME, TimeBoundaryQuery.MAX_TIME); + query.getQueryContext().getAsString(TimeBoundaryQuery.MIN_TIME_ARRAY_OUTPUT_NAME, TimeBoundaryQuery.MIN_TIME) : + query.getQueryContext().getAsString(TimeBoundaryQuery.MAX_TIME_ARRAY_OUTPUT_NAME, TimeBoundaryQuery.MAX_TIME); return builder.add(outputName, ColumnType.LONG).build(); } return super.resultArraySignature(query); diff --git a/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQuery.java b/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQuery.java index 65dffef6e0d3..291428f5c2a3 100644 --- a/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQuery.java +++ b/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQuery.java @@ -159,7 +159,7 @@ public boolean isGrandTotal() public String getTimestampResultField() { - return getContextAsString(CTX_TIMESTAMP_RESULT_FIELD); + return getQueryContext().getAsString(CTX_TIMESTAMP_RESULT_FIELD); } public boolean isSkipEmptyBuckets() diff --git a/processing/src/main/java/org/apache/druid/query/topn/TopNQueryQueryToolChest.java b/processing/src/main/java/org/apache/druid/query/topn/TopNQueryQueryToolChest.java index 87855bf84f01..f6fa421719da 100644 --- a/processing/src/main/java/org/apache/druid/query/topn/TopNQueryQueryToolChest.java +++ b/processing/src/main/java/org/apache/druid/query/topn/TopNQueryQueryToolChest.java @@ -574,7 +574,7 @@ public Sequence> run( } final TopNQuery query = (TopNQuery) input; - final int minTopNThreshold = query.getContextAsInt("minTopNThreshold", config.getMinTopNThreshold()); + final int minTopNThreshold = query.getQueryContext().getAsInt("minTopNThreshold", config.getMinTopNThreshold()); if (query.getThreshold() > minTopNThreshold) { return runner.run(queryPlus, responseContext); } diff --git a/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java b/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java index 22478ff64d4b..7431022240d3 100644 --- a/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java +++ b/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java @@ -314,12 +314,12 @@ public void testGetEnum() Assert.assertEquals( QueryContexts.Vectorize.FORCE, - query.getContextAsEnum("e1", QueryContexts.Vectorize.class, QueryContexts.Vectorize.FALSE) + query.getQueryContext().getAsEnum("e1", QueryContexts.Vectorize.class, QueryContexts.Vectorize.FALSE) ); Assert.assertThrows( IAE.class, - () -> query.getContextAsEnum("e2", QueryContexts.Vectorize.class, QueryContexts.Vectorize.FALSE) + () -> query.getQueryContext().getAsEnum("e2", QueryContexts.Vectorize.class, QueryContexts.Vectorize.FALSE) ); } } diff --git a/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java b/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java index 94a48d896f48..0b9717919d82 100644 --- a/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java +++ b/processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java @@ -102,10 +102,10 @@ public void testContextSerde() throws Exception ), Query.class ); - Assert.assertEquals((Integer) 1, serdeQuery.getContextAsInt(QueryContexts.PRIORITY_KEY)); - Assert.assertEquals(true, serdeQuery.getContextAsBoolean(QueryContexts.USE_CACHE_KEY)); - Assert.assertEquals("true", serdeQuery.getContextAsString(QueryContexts.POPULATE_CACHE_KEY)); - Assert.assertEquals(true, serdeQuery.getContextAsBoolean(QueryContexts.FINALIZE_KEY)); + Assert.assertEquals((Integer) 1, serdeQuery.getQueryContext().getAsInt(QueryContexts.PRIORITY_KEY)); + Assert.assertEquals(true, serdeQuery.getQueryContext().getAsBoolean(QueryContexts.USE_CACHE_KEY)); + Assert.assertEquals("true", serdeQuery.getQueryContext().getAsString(QueryContexts.POPULATE_CACHE_KEY)); + Assert.assertEquals(true, serdeQuery.getQueryContext().getAsBoolean(QueryContexts.FINALIZE_KEY)); Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.USE_CACHE_KEY, false)); Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.POPULATE_CACHE_KEY, false)); Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.FINALIZE_KEY, false)); diff --git a/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java b/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java index 46e7eb780f66..aaa293f4355c 100644 --- a/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java +++ b/processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java @@ -78,10 +78,10 @@ public void testContextSerde() throws Exception ), TimeBoundaryQuery.class ); - Assert.assertEquals(new Integer(1), serdeQuery.getContextAsInt(QueryContexts.PRIORITY_KEY)); - Assert.assertEquals(true, serdeQuery.getContextAsBoolean(QueryContexts.USE_CACHE_KEY)); - Assert.assertEquals(true, serdeQuery.getContextAsBoolean(QueryContexts.POPULATE_CACHE_KEY)); - Assert.assertEquals(true, serdeQuery.getContextAsBoolean(QueryContexts.FINALIZE_KEY)); + Assert.assertEquals(new Integer(1), serdeQuery.getQueryContext().getAsInt(QueryContexts.PRIORITY_KEY)); + Assert.assertEquals(true, serdeQuery.getQueryContext().getAsBoolean(QueryContexts.USE_CACHE_KEY)); + Assert.assertEquals(true, serdeQuery.getQueryContext().getAsBoolean(QueryContexts.POPULATE_CACHE_KEY)); + Assert.assertEquals(true, serdeQuery.getQueryContext().getAsBoolean(QueryContexts.FINALIZE_KEY)); } @Test @@ -116,9 +116,9 @@ public void testContextSerde2() throws Exception ); - Assert.assertEquals("1", serdeQuery.getContextAsString(QueryContexts.PRIORITY_KEY)); - Assert.assertEquals("true", serdeQuery.getContextAsString(QueryContexts.USE_CACHE_KEY)); - Assert.assertEquals("true", serdeQuery.getContextAsString(QueryContexts.POPULATE_CACHE_KEY)); - Assert.assertEquals("true", serdeQuery.getContextAsString(QueryContexts.FINALIZE_KEY)); + Assert.assertEquals("1", serdeQuery.getQueryContext().getAsString(QueryContexts.PRIORITY_KEY)); + Assert.assertEquals("true", serdeQuery.getQueryContext().getAsString(QueryContexts.USE_CACHE_KEY)); + Assert.assertEquals("true", serdeQuery.getQueryContext().getAsString(QueryContexts.POPULATE_CACHE_KEY)); + Assert.assertEquals("true", serdeQuery.getQueryContext().getAsString(QueryContexts.FINALIZE_KEY)); } } diff --git a/server/src/main/java/org/apache/druid/client/DirectDruidClient.java b/server/src/main/java/org/apache/druid/client/DirectDruidClient.java index 7fa3364e56fa..9c244bbdc731 100644 --- a/server/src/main/java/org/apache/druid/client/DirectDruidClient.java +++ b/server/src/main/java/org/apache/druid/client/DirectDruidClient.java @@ -163,7 +163,7 @@ public Sequence run(final QueryPlus queryPlus, final ResponseContext conte log.debug("Querying queryId[%s] url[%s]", query.getId(), url); final long requestStartTimeNs = System.nanoTime(); - final long timeoutAt = query.getContextAsLong(QUERY_FAIL_TIME); + final long timeoutAt = query.getQueryContext().getAsLong(QUERY_FAIL_TIME); final long maxScatterGatherBytes = QueryContexts.getMaxScatterGatherBytes(query); final AtomicLong totalBytesGathered = context.getTotalBytes(); final long maxQueuedBytes = QueryContexts.getMaxQueuedBytes(query, 0); diff --git a/server/src/main/java/org/apache/druid/client/JsonParserIterator.java b/server/src/main/java/org/apache/druid/client/JsonParserIterator.java index d10ee099ae4e..42834b0fbaf6 100644 --- a/server/src/main/java/org/apache/druid/client/JsonParserIterator.java +++ b/server/src/main/java/org/apache/druid/client/JsonParserIterator.java @@ -75,7 +75,7 @@ public JsonParserIterator( this.future = future; this.url = url; if (query != null) { - this.timeoutAt = query.getContextAsLong(DirectDruidClient.QUERY_FAIL_TIME, -1L); + this.timeoutAt = query.getQueryContext().getAsLong(DirectDruidClient.QUERY_FAIL_TIME, -1L); this.queryId = query.getId(); } else { this.timeoutAt = -1; diff --git a/server/src/main/java/org/apache/druid/server/ClientQuerySegmentWalker.java b/server/src/main/java/org/apache/druid/server/ClientQuerySegmentWalker.java index 1220be5e2b6d..05cc77f1cd75 100644 --- a/server/src/main/java/org/apache/druid/server/ClientQuerySegmentWalker.java +++ b/server/src/main/java/org/apache/druid/server/ClientQuerySegmentWalker.java @@ -431,7 +431,7 @@ private QueryRunner decorateClusterRunner(Query query, QueryRunner .emitCPUTimeMetric(emitter) .postProcess( objectMapper.convertValue( - query.getContextAsString("postProcessing"), + query.getQueryContext().getAsString("postProcessing"), new TypeReference>() {} ) ) diff --git a/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java b/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java index ecb432d3fb71..cb365d82c598 100644 --- a/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java +++ b/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java @@ -73,7 +73,7 @@ public Optional computeLane(QueryPlus query, Set>> apply(@Nullable Integer input) QueryPlus capturedQueryPlus = (QueryPlus) queryCapture.getValue(); Query capturedQuery = capturedQueryPlus.getQuery(); if (expectBySegment) { - Assert.assertEquals(true, capturedQuery.getContextAsBoolean(QueryContexts.BY_SEGMENT_KEY)); + Assert.assertEquals(true, capturedQuery.getQueryContext().getAsBoolean(QueryContexts.BY_SEGMENT_KEY)); } else { Assert.assertTrue( capturedQuery.getContextValue(QueryContexts.BY_SEGMENT_KEY) == null || - capturedQuery.getContextAsBoolean(QueryContexts.BY_SEGMENT_KEY).equals(false) + capturedQuery.getQueryContext().getAsBoolean(QueryContexts.BY_SEGMENT_KEY).equals(false) ); } } diff --git a/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java b/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java index a1bc9ddb6576..472fa6c96f2d 100644 --- a/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java +++ b/server/src/test/java/org/apache/druid/client/JsonParserIteratorTest.java @@ -307,7 +307,8 @@ private Query mockQuery(String queryId, long timeoutAt) { Query query = Mockito.mock(Query.class); Mockito.when(query.getId()).thenReturn(queryId); - Mockito.when(query.getContextAsLong(ArgumentMatchers.eq(DirectDruidClient.QUERY_FAIL_TIME), ArgumentMatchers.eq(-1L))) + Mockito.when(query.getQueryContext().getAsLong(ArgumentMatchers.eq(DirectDruidClient.QUERY_FAIL_TIME), + ArgumentMatchers.eq(-1L))) .thenReturn(timeoutAt); return query; } diff --git a/server/src/test/java/org/apache/druid/server/SetAndVerifyContextQueryRunnerTest.java b/server/src/test/java/org/apache/druid/server/SetAndVerifyContextQueryRunnerTest.java index fccfd0500d75..59853bc7c5d9 100644 --- a/server/src/test/java/org/apache/druid/server/SetAndVerifyContextQueryRunnerTest.java +++ b/server/src/test/java/org/apache/druid/server/SetAndVerifyContextQueryRunnerTest.java @@ -58,7 +58,7 @@ public void testTimeoutIsUsedIfTimeoutIsNonZero() throws InterruptedException // time + 1 at the time the method was called // this means that after sleeping for 1 millis, the fail time should be less than the current time when checking Assert.assertTrue( - System.currentTimeMillis() > transformed.getContextAsLong(DirectDruidClient.QUERY_FAIL_TIME) + System.currentTimeMillis() > transformed.getQueryContext().getAsLong(DirectDruidClient.QUERY_FAIL_TIME) ); } @@ -85,7 +85,7 @@ public long getDefaultQueryTimeout() Query transformed = queryRunner.withTimeoutAndMaxScatterGatherBytes(query, defaultConfig); // timeout is not set, default timeout has been set to long.max, make sure timeout is still in the future - Assert.assertEquals((Long) Long.MAX_VALUE, transformed.getContextAsLong(DirectDruidClient.QUERY_FAIL_TIME)); + Assert.assertEquals((Long) Long.MAX_VALUE, transformed.getQueryContext().getAsLong(DirectDruidClient.QUERY_FAIL_TIME)); } @Test @@ -107,7 +107,7 @@ public void testTimeoutZeroIsNotImmediateTimeoutDefaultServersideMax() // timeout is set to 0, so withTimeoutAndMaxScatterGatherBytes should set QUERY_FAIL_TIME to be the current // time + max query timeout at the time the method was called // since default is long max, expect long max since current time would overflow - Assert.assertEquals((Long) Long.MAX_VALUE, transformed.getContextAsLong(DirectDruidClient.QUERY_FAIL_TIME)); + Assert.assertEquals((Long) Long.MAX_VALUE, transformed.getQueryContext().getAsLong(DirectDruidClient.QUERY_FAIL_TIME)); } @Test @@ -137,7 +137,7 @@ public long getMaxQueryTimeout() // time + max query timeout at the time the method was called // this means that the fail time should be greater than the current time when checking Assert.assertTrue( - System.currentTimeMillis() < (Long) transformed.getContextAsLong(DirectDruidClient.QUERY_FAIL_TIME) + System.currentTimeMillis() < (Long) transformed.getQueryContext().getAsLong(DirectDruidClient.QUERY_FAIL_TIME) ); } }