Skip to content

Clean up query contexts #12633

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gianm merged 2 commits into from
Jun 15, 2022
Merged

Clean up query contexts #12633

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
341117f
Clean up query contexts
paul-rogers Jun 10, 2022
069a6b6
Revisions from review comments
paul-rogers Jun 14, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 2 additions & 8 deletions processing/src/main/java/org/apache/druid/query/BaseQuery.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,10 +35,10 @@
import org.joda.time.Interval;

import javax.annotation.Nullable;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.TreeMap;

/**
*
Expand Down Expand Up @@ -210,13 +210,7 @@ public static Map<String, Object> computeOverriddenContext(
final Map<String, Object> overrides
)
{
Map<String, Object> overridden = new TreeMap<>();
if (context != null) {
overridden.putAll(context);
}
overridden.putAll(overrides);

return overridden;
return QueryContexts.override(context, overrides);
}

/**
Expand Down
2 changes: 1 addition & 1 deletion processing/src/main/java/org/apache/druid/query/FinalizeResultsQueryRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ public Sequence<T> run(final QueryPlus<T> queryPlus, ResponseContext responseCon
final MetricManipulationFn metricManipulationFn;

if (shouldFinalize) {
queryToRun = query.withOverriddenContext(ImmutableMap.of("finalize", false));
queryToRun = query.withOverriddenContext(ImmutableMap.of(QueryContexts.FINALIZE_KEY, false));
metricManipulationFn = MetricManipulatorFns.finalizing();
} else {
queryToRun = query;
Expand Down
37 changes: 4 additions & 33 deletions processing/src/main/java/org/apache/druid/query/QueryContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,8 @@

package org.apache.druid.query;

import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.Numbers;

import javax.annotation.Nullable;

import java.util.Collections;
import java.util.Map;
import java.util.Objects;
Expand Down Expand Up @@ -162,47 +160,20 @@ public boolean getAsBoolean(
final boolean defaultValue
)
{
final Object value = get(parameter);
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return Boolean.parseBoolean((String) value);
} else if (value instanceof Boolean) {
return (Boolean) value;
} else {
throw new IAE("Expected parameter[%s] to be boolean", parameter);
}
return QueryContexts.getAsBoolean(parameter, get(parameter), defaultValue);
}

public int getAsInt(
final String parameter,
final int defaultValue
)
{
final Object value = get(parameter);
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return Numbers.parseInt(value);
} else if (value instanceof Number) {
return ((Number) value).intValue();
} else {
throw new IAE("Expected parameter[%s] to be integer", parameter);
}
return QueryContexts.getAsInt(parameter, get(parameter), defaultValue);
}

public long getAsLong(final String parameter, final long defaultValue)
{
final Object value = get(parameter);
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 long", parameter);
}
return QueryContexts.getAsLong(parameter, get(parameter), defaultValue);
}

public Map<String, Object> getMergedParams()
Expand Down
133 changes: 116 additions & 17 deletions processing/src/main/java/org/apache/druid/query/QueryContexts.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
import org.apache.druid.segment.QueryableIndexStorageAdapter;

import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.TimeUnit;

@PublicApi
Expand Down Expand Up @@ -71,6 +72,12 @@ public class QueryContexts
public static final String BROKER_SERVICE_NAME = "brokerService";
public static final String IN_SUB_QUERY_THRESHOLD_KEY = "inSubQueryThreshold";
public static final String TIME_BOUNDARY_PLANNING_KEY = "enableTimeBoundaryPlanning";
public static final String POPULATE_CACHE_KEY = "populateCache";
public static final String POPULATE_RESULT_LEVEL_CACHE_KEY = "populateResultLevelCache";
public static final String USE_RESULT_LEVEL_CACHE_KEY = "useResultLevelCache";
public static final String SERIALIZE_DATE_TIME_AS_LONG_KEY = "serializeDateTimeAsLong";
public static final String SERIALIZE_DATE_TIME_AS_LONG_INNER_KEY = "serializeDateTimeAsLongInner";
public static final String UNCOVERED_INTERVALS_LIMIT_KEY = "uncoveredIntervalsLimit";

public static final boolean DEFAULT_BY_SEGMENT = false;
public static final boolean DEFAULT_POPULATE_CACHE = true;
Expand Down Expand Up @@ -158,7 +165,7 @@ public static <T> boolean isPopulateCache(Query<T> query)

public static <T> boolean isPopulateCache(Query<T> query, boolean defaultValue)
{
return parseBoolean(query, "populateCache", defaultValue);
return parseBoolean(query, POPULATE_CACHE_KEY, defaultValue);
}

public static <T> boolean isUseCache(Query<T> query)
Expand All @@ -178,7 +185,7 @@ public static <T> boolean isPopulateResultLevelCache(Query<T> query)

public static <T> boolean isPopulateResultLevelCache(Query<T> query, boolean defaultValue)
{
return parseBoolean(query, "populateResultLevelCache", defaultValue);
return parseBoolean(query, POPULATE_RESULT_LEVEL_CACHE_KEY, defaultValue);
}

public static <T> boolean isUseResultLevelCache(Query<T> query)
Expand All @@ -188,22 +195,22 @@ public static <T> boolean isUseResultLevelCache(Query<T> query)

public static <T> boolean isUseResultLevelCache(Query<T> query, boolean defaultValue)
{
return parseBoolean(query, "useResultLevelCache", defaultValue);
return parseBoolean(query, USE_RESULT_LEVEL_CACHE_KEY, defaultValue);
}

public static <T> boolean isFinalize(Query<T> query, boolean defaultValue)

{
return parseBoolean(query, FINALIZE_KEY, defaultValue);
}

public static <T> boolean isSerializeDateTimeAsLong(Query<T> query, boolean defaultValue)
{
return parseBoolean(query, "serializeDateTimeAsLong", defaultValue);
return parseBoolean(query, SERIALIZE_DATE_TIME_AS_LONG_KEY, defaultValue);
}

public static <T> boolean isSerializeDateTimeAsLongInner(Query<T> query, boolean defaultValue)
{
return parseBoolean(query, "serializeDateTimeAsLongInner", defaultValue);
return parseBoolean(query, SERIALIZE_DATE_TIME_AS_LONG_INNER_KEY, defaultValue);
Copy link
Copy Markdown
Contributor

@gianm gianm Jun 13, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any reason for parseBoolean and getAsBoolean to remain different? Now that we have getAsBoolean, shall it replace parseBoolean?

(Btw, is the behavior the same? I haven't checked)

Copy link
Copy Markdown
Contributor Author

@paul-rogers paul-rogers Jun 14, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. I didn't really pay attention to the parseFoo methods. Those methods appear to only handle the value-as-string case, where as the getAsFoo method handle maps with values of the expected type. Went ahead and changed the parseFoo methods to use getAsFoo so we're consistent.

}

public static <T> Vectorize getVectorize(Query<T> query)
Expand Down Expand Up @@ -248,7 +255,7 @@ public static <T> int getUncoveredIntervalsLimit(Query<T> query)

public static <T> int getUncoveredIntervalsLimit(Query<T> query, int defaultValue)
{
return parseInt(query, "uncoveredIntervalsLimit", defaultValue);
return parseInt(query, UNCOVERED_INTERVALS_LIMIT_KEY, defaultValue);
Copy link
Copy Markdown
Contributor

@gianm gianm Jun 13, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any reason for parseInt and getAsInt to remain different? Now that we have getAsInt, shall it replace parseInt?

(Btw, is the behavior the same? I haven't checked)

}

public static <T> int getPriority(Query<T> query)
Expand Down Expand Up @@ -450,32 +457,124 @@ public static String getBrokerServiceName(Map<String, Object> queryContext)

static <T> long parseLong(Query<T> query, String key, long defaultValue)
{
final Object val = query.getContextValue(key);
return val == null ? defaultValue : Numbers.parseLong(val);
return getAsLong(key, query.getContextValue(key), defaultValue);
}

@SuppressWarnings("unused")
static <T> long parseLong(Map<String, Object> context, String key, long defaultValue)
{
return getAsLong(key, context.get(key), defaultValue);
}

static <T> int parseInt(Query<T> query, String key, int defaultValue)
{
final Object val = query.getContextValue(key);
return val == null ? defaultValue : Numbers.parseInt(val);
return getAsInt(key, query.getContextValue(key), defaultValue);
}

static int parseInt(Map<String, Object> context, String key, int defaultValue)
{
final Object val = context.get(key);
return val == null ? defaultValue : Numbers.parseInt(val);
return getAsInt(key, context.get(key), defaultValue);
}

static <T> boolean parseBoolean(Query<T> query, String key, boolean defaultValue)
{
final Object val = query.getContextValue(key);
return val == null ? defaultValue : Numbers.parseBoolean(val);
return getAsBoolean(key, query.getContextValue(key), defaultValue);
}

static boolean parseBoolean(Map<String, Object> context, String key, boolean defaultValue)
{
final Object val = context.get(key);
return val == null ? defaultValue : Numbers.parseBoolean(val);
return getAsBoolean(key, context.get(key), defaultValue);
}

public static String getAsString(
final String parameter,
final Object value,
final String defaultValue
)
{
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return (String) value;
} else {
throw new IAE("Expected parameter [%s] to be String", parameter);
}
}

/**
* 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
)
{
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return Boolean.parseBoolean((String) value);
} else if (value instanceof Boolean) {
return (Boolean) value;
} else {
throw new IAE("Expected parameter [%s] to be a boolean", 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.
*/
public static int getAsInt(
final String parameter,
final Object value,
final int defaultValue
)
{
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return Numbers.parseInt(value);
} else if (value instanceof Number) {
return ((Number) value).intValue();
} else {
throw new IAE("Expected parameter [%s] to be an integer", parameter);
}
}

/**
* Get the value of a parameter as an {@code long}. The parameter is expected
* to be {@code null}, a string or a {@code Number} object.
*/
public static long getAsLong(
final String parameter,
final Object value,
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);
}
}

public static Map<String, Object> override(
final Map<String, Object> context,
final Map<String, Object> overrides
)
{
Map<String, Object> overridden = new TreeMap<>();
if (context != null) {
overridden.putAll(context);
}
overridden.putAll(overrides);

return overridden;
}

private QueryContexts()
Expand Down
3 changes: 2 additions & 1 deletion processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV1.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import org.apache.druid.java.util.common.guava.Sequence;
import org.apache.druid.java.util.common.guava.Sequences;
import org.apache.druid.query.GroupByMergedQueryRunner;
import org.apache.druid.query.QueryContexts;
import org.apache.druid.query.QueryPlus;
import org.apache.druid.query.QueryProcessingPool;
import org.apache.druid.query.QueryRunner;
Expand Down Expand Up @@ -115,7 +116,7 @@ public Sequence<ResultRow> mergeResults(
.overrideContext(
ImmutableMap.<String, Object>builder()
.put(GroupByQueryConfig.CTX_KEY_STRATEGY, GroupByStrategySelector.STRATEGY_V1)
.put("finalize", false)
.put(QueryContexts.FINALIZE_KEY, false)

// Always request array result rows when passing the query down.
.put(GroupByQueryConfig.CTX_KEY_ARRAY_RESULT_ROWS, true)
Expand Down
2 changes: 1 addition & 1 deletion processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,7 @@ public Sequence<ResultRow> mergeResults(

// Set up downstream context.
final ImmutableMap.Builder<String, Object> context = ImmutableMap.builder();
context.put("finalize", false);
context.put(QueryContexts.FINALIZE_KEY, false);
context.put(GroupByQueryConfig.CTX_KEY_STRATEGY, GroupByStrategySelector.STRATEGY_V2);
context.put(CTX_KEY_OUTERMOST, false);

Expand Down
Loading