From f3b47e6a68a8020652b46692de164b7d498275b0 Mon Sep 17 00:00:00 2001 From: "Lucas.Capistrant" Date: Mon, 25 Jul 2022 15:08:48 -0500 Subject: [PATCH 1/9] Introduce defaultOnDiskStorage config for groupBy --- docs/configuration/index.md | 3 +- .../query/groupby/GroupByQueryConfig.java | 24 +++++++++- .../query/groupby/GroupByQueryConfigTest.java | 47 +++++++++++++++++++ 3 files changed, 72 insertions(+), 2 deletions(-) diff --git a/docs/configuration/index.md b/docs/configuration/index.md index 2246727ad288..32ccb8f2261f 100644 --- a/docs/configuration/index.md +++ b/docs/configuration/index.md @@ -2079,6 +2079,7 @@ Supported runtime properties: |`druid.query.groupBy.maxSelectorDictionarySize`|Maximum amount of heap space (approximately) to use for per-segment string dictionaries. See [groupBy memory tuning and resource limits](../querying/groupbyquery.md#memory-tuning-and-resource-limits) for details.|100000000| |`druid.query.groupBy.maxMergingDictionarySize`|Maximum amount of heap space (approximately) to use for per-query string dictionaries. When the dictionary exceeds this size, a spill to disk will be triggered. See [groupBy memory tuning and resource limits](../querying/groupbyquery.md#memory-tuning-and-resource-limits) for details.|100000000| |`druid.query.groupBy.maxOnDiskStorage`|Maximum amount of disk space to use, per-query, for spilling result sets to disk when either the merging buffer or the dictionary fills up. Queries that exceed this limit will fail. Set to zero to disable disk spilling.|0 (disabled)| +|`druid.query.groupBy.defaultOnDiskStorage`|Default amount of disk space to use, per-query, for spilling the result sets to disk when either the merging buffer or the dictionary fills up. Set to zero to disable disk spilling for queries who don't override `maxOnDiskStorage` in their context.|`druid.query.groupBy.maxOnDiskStorage`| Supported query contexts: @@ -2086,7 +2087,7 @@ Supported query contexts: |---|-----------| |`maxSelectorDictionarySize`|Can be used to lower the value of `druid.query.groupBy.maxMergingDictionarySize` for this query.| |`maxMergingDictionarySize`|Can be used to lower the value of `druid.query.groupBy.maxMergingDictionarySize` for this query.| -|`maxOnDiskStorage`|Can be used to lower the value of `druid.query.groupBy.maxOnDiskStorage` for this query.| +|`maxOnDiskStorage`|Can be used to set `maxOnDiskStorage` to a value between 0 and `druid.query.groupBy.maxOnDiskStorage` for this query. If this query context override exceeds `druid.query.groupBy.maxOnDiskStorage`, the query will use `druid.query.groupBy.maxOnDiskStorage`. Omitting this from the query context will cause the query to use `druid.query.groupBy.defaultOnDiskStorage` for `maxOnDiskStorage`| ### Advanced configurations 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 be4fb985a45c..f0004f18c667 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 @@ -67,6 +67,8 @@ public class GroupByQueryConfig private static final long MIN_AUTOMATIC_DICTIONARY_SIZE = 1; private static final long MAX_AUTOMATIC_DICTIONARY_SIZE = 1_000_000_000; + + @JsonProperty private String defaultStrategy = GroupByStrategySelector.STRATEGY_V2; @@ -102,6 +104,9 @@ public class GroupByQueryConfig // Max on-disk temporary storage, per-query; when exceeded, the query fails private long maxOnDiskStorage = 0L; + @JsonProperty + private long defaultOnDiskStorage = Integer.MAX_VALUE; + @JsonProperty private boolean forcePushDownLimit = false; @@ -263,6 +268,19 @@ public long getMaxOnDiskStorage() return maxOnDiskStorage; } + /** + * Mirror maxOnDiskStorage if defaultOnDiskStorage's default is not overridden by cluster operator. + * + * This mirroring is done to maintain continuity in behavior between Druid versions. If an operator wants to use + * defaultOnDiskStorage, they have to explicitly override it. + * + * @return The working value for defaultOnDiskStorage + */ + public long getDefaultOnDiskStorage() + { + return defaultOnDiskStorage == Integer.MAX_VALUE ? getMaxOnDiskStorage() : defaultOnDiskStorage; + } + public boolean isForcePushDownLimit() { return forcePushDownLimit; @@ -338,8 +356,11 @@ public GroupByQueryConfig withOverrides(final GroupByQuery query) CTX_KEY_BUFFER_GROUPER_INITIAL_BUCKETS, getBufferGrouperInitialBuckets() ); + // If the client overrides do not provide "maxOnDiskStorage" context key, the server side "defaultOnDiskStorage" + // value is used in the calculation of the newConfig value of maxOnDiskStorage. This allows the operator to + // choose a default value lower than the max allowed when the context key is missing in the client query. newConfig.maxOnDiskStorage = Math.min( - ((Number) query.getContextValue(CTX_KEY_MAX_ON_DISK_STORAGE, getMaxOnDiskStorage())).longValue(), + ((Number) query.getContextValue(CTX_KEY_MAX_ON_DISK_STORAGE, getDefaultOnDiskStorage())).longValue(), getMaxOnDiskStorage() ); newConfig.maxSelectorDictionarySize = maxSelectorDictionarySize; // No overrides @@ -384,6 +405,7 @@ public String toString() ", bufferGrouperInitialBuckets=" + bufferGrouperInitialBuckets + ", maxMergingDictionarySize=" + maxMergingDictionarySize + ", maxOnDiskStorage=" + maxOnDiskStorage + + ", defaultOnDiskStorage=" + getDefaultOnDiskStorage() + // use the getter because of special behavior for mirroring maxOnDiskStorage if defaultOnDiskStorage not explicitly set. ", forcePushDownLimit=" + forcePushDownLimit + ", forceHashAggregation=" + forceHashAggregation + ", intermediateCombineDegree=" + intermediateCombineDegree + diff --git a/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java b/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java index 869bd75765d3..3806b60ddc6c 100644 --- a/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java +++ b/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java @@ -55,6 +55,7 @@ public void testSerde() Assert.assertEquals(2, config.getMaxIntermediateRows()); Assert.assertEquals(3, config.getMaxResults()); Assert.assertEquals(4, config.getMaxOnDiskStorage()); + Assert.assertEquals(4, config.getDefaultOnDiskStorage()); Assert.assertEquals(5, config.getConfiguredMaxSelectorDictionarySize()); Assert.assertEquals(6_000_000, config.getConfiguredMaxMergingDictionarySize()); Assert.assertEquals(7.0, config.getBufferGrouperMaxLoadFactor(), 0.0); @@ -166,4 +167,50 @@ public void testNonAutomaticSelectorDictionarySize() Assert.assertEquals(100, config.getConfiguredMaxSelectorDictionarySize()); Assert.assertEquals(100, config.getActualMaxSelectorDictionarySize(1_000_000_000, 2)); } + + /** + * Tests that the defaultOnDiskStorage value is used when applying override context that is lacking maxOnDiskStorage. + */ + @Test + public void testUseDefaultOnDiskStorage() + { + final GroupByQueryConfig config = MAPPER.convertValue( + ImmutableMap.of( + "maxOnDiskStorage", "10", + "defaultOnDiskStorage", "5" + ), + GroupByQueryConfig.class + ); + final GroupByQueryConfig config2 = config.withOverrides( + GroupByQuery.builder() + .setDataSource("test") + .setInterval(Intervals.of("2000/P1D")) + .setGranularity(Granularities.ALL) + .setContext(ImmutableMap.builder().build()) + .build() + ); + Assert.assertEquals(5L, config2.getMaxOnDiskStorage()); + } + + @Test + public void testUseMaxOnDiskStorageWhenClientOverrideIsTooLarge() + { + final GroupByQueryConfig config = MAPPER.convertValue( + ImmutableMap.of("maxOnDiskStorage", "10"), + GroupByQueryConfig.class + ); + final GroupByQueryConfig config2 = config.withOverrides( + GroupByQuery.builder() + .setDataSource("test") + .setInterval(Intervals.of("2000/P1D")) + .setGranularity(Granularities.ALL) + .setContext( + ImmutableMap.builder() + .put("maxOnDiskStorage", 500) + .build() + ) + .build() + ); + Assert.assertEquals(10L, config2.getMaxOnDiskStorage()); + } } From 5eb988c950b2b83cdaaedd6e43c67bc15ad4c21f Mon Sep 17 00:00:00 2001 From: "Lucas.Capistrant" Date: Thu, 28 Jul 2022 11:11:15 -0500 Subject: [PATCH 2/9] add debug log to groupby query config --- .../org/apache/druid/query/groupby/GroupByQueryConfig.java | 5 +++++ 1 file changed, 5 insertions(+) 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 f0004f18c667..f4b95457a23b 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 @@ -21,6 +21,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import org.apache.druid.java.util.common.HumanReadableBytes; +import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.query.DruidProcessingConfig; import org.apache.druid.query.QueryContexts; import org.apache.druid.query.groupby.strategy.GroupByStrategySelector; @@ -31,6 +32,8 @@ */ public class GroupByQueryConfig { + private static final Logger logger = new Logger(GroupByQueryConfig.class); + public static final long AUTOMATIC = 0; public static final String CTX_KEY_STRATEGY = "groupByStrategy"; @@ -389,6 +392,8 @@ public GroupByQueryConfig withOverrides(final GroupByQuery query) CTX_KEY_ENABLE_MULTI_VALUE_UNNESTING, isMultiValueUnnestingEnabled() ); + + logger.debug(newConfig.toString()); return newConfig; } From 03301286756513bde45218e7e7a97d82f2ed6b47 Mon Sep 17 00:00:00 2001 From: "Lucas.Capistrant" Date: Mon, 1 Aug 2022 13:09:48 -0500 Subject: [PATCH 3/9] Apply config change suggestion from review --- docs/configuration/index.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/configuration/index.md b/docs/configuration/index.md index 32ccb8f2261f..e19a6b3c05e7 100644 --- a/docs/configuration/index.md +++ b/docs/configuration/index.md @@ -2079,7 +2079,7 @@ Supported runtime properties: |`druid.query.groupBy.maxSelectorDictionarySize`|Maximum amount of heap space (approximately) to use for per-segment string dictionaries. See [groupBy memory tuning and resource limits](../querying/groupbyquery.md#memory-tuning-and-resource-limits) for details.|100000000| |`druid.query.groupBy.maxMergingDictionarySize`|Maximum amount of heap space (approximately) to use for per-query string dictionaries. When the dictionary exceeds this size, a spill to disk will be triggered. See [groupBy memory tuning and resource limits](../querying/groupbyquery.md#memory-tuning-and-resource-limits) for details.|100000000| |`druid.query.groupBy.maxOnDiskStorage`|Maximum amount of disk space to use, per-query, for spilling result sets to disk when either the merging buffer or the dictionary fills up. Queries that exceed this limit will fail. Set to zero to disable disk spilling.|0 (disabled)| -|`druid.query.groupBy.defaultOnDiskStorage`|Default amount of disk space to use, per-query, for spilling the result sets to disk when either the merging buffer or the dictionary fills up. Set to zero to disable disk spilling for queries who don't override `maxOnDiskStorage` in their context.|`druid.query.groupBy.maxOnDiskStorage`| +|`druid.query.groupBy.defaultOnDiskStorage`|Default amount of disk space to use, per-query, for spilling the result sets to disk when either the merging buffer or the dictionary fills up. Set to zero to disable disk spilling for queries which don't override `maxOnDiskStorage` in their context.|`druid.query.groupBy.maxOnDiskStorage`| Supported query contexts: From 9afffcfb2cd96e2db96ddb99dadd131c673a93c0 Mon Sep 17 00:00:00 2001 From: "Lucas.Capistrant" Date: Mon, 1 Aug 2022 13:10:25 -0500 Subject: [PATCH 4/9] Remove accidental new lines --- .../java/org/apache/druid/query/groupby/GroupByQueryConfig.java | 2 -- 1 file changed, 2 deletions(-) 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 f4b95457a23b..7778129e1376 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 @@ -70,8 +70,6 @@ public class GroupByQueryConfig private static final long MIN_AUTOMATIC_DICTIONARY_SIZE = 1; private static final long MAX_AUTOMATIC_DICTIONARY_SIZE = 1_000_000_000; - - @JsonProperty private String defaultStrategy = GroupByStrategySelector.STRATEGY_V2; From 4ea3e99704161424a26be0fbb128de10ff0b1be4 Mon Sep 17 00:00:00 2001 From: "Lucas.Capistrant" Date: Mon, 1 Aug 2022 13:17:42 -0500 Subject: [PATCH 5/9] update default value of new default disk storage config --- .../org/apache/druid/query/groupby/GroupByQueryConfig.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 7778129e1376..5dab5c5f35dd 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 @@ -106,7 +106,7 @@ public class GroupByQueryConfig private long maxOnDiskStorage = 0L; @JsonProperty - private long defaultOnDiskStorage = Integer.MAX_VALUE; + private long defaultOnDiskStorage = -1L; @JsonProperty private boolean forcePushDownLimit = false; @@ -279,7 +279,7 @@ public long getMaxOnDiskStorage() */ public long getDefaultOnDiskStorage() { - return defaultOnDiskStorage == Integer.MAX_VALUE ? getMaxOnDiskStorage() : defaultOnDiskStorage; + return defaultOnDiskStorage < 0L ? getMaxOnDiskStorage() : defaultOnDiskStorage; } public boolean isForcePushDownLimit() From a7e41478a1f8138a32177aed168037132902c452 Mon Sep 17 00:00:00 2001 From: "Lucas.Capistrant" Date: Mon, 1 Aug 2022 13:18:03 -0500 Subject: [PATCH 6/9] update debug log to have more descriptive text --- .../java/org/apache/druid/query/groupby/GroupByQueryConfig.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 5dab5c5f35dd..5d2754cebcd9 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 @@ -391,7 +391,7 @@ public GroupByQueryConfig withOverrides(final GroupByQuery query) isMultiValueUnnestingEnabled() ); - logger.debug(newConfig.toString()); + logger.debug("Override config for GroupBy query %s - %s", query.getId(), newConfig.toString()); return newConfig; } From cb240fe8eb092a8c9c5cc198f3e84871c0c179b8 Mon Sep 17 00:00:00 2001 From: "Lucas.Capistrant" Date: Tue, 2 Aug 2022 14:29:43 -0500 Subject: [PATCH 7/9] Make maxOnDiskStorage and defaultOnDiskStorage HumanRedadableBytes --- .../GroupByTypeInterfaceBenchmark.java | 5 ++-- .../benchmark/query/GroupByBenchmark.java | 5 ++-- .../MaterializedViewQuery.java | 7 ++++++ .../org/apache/druid/query/BaseQuery.java | 7 ++++++ .../java/org/apache/druid/query/Query.java | 3 +++ .../org/apache/druid/query/QueryContext.java | 7 ++++++ .../org/apache/druid/query/QueryContexts.java | 18 +++++++++++++ .../query/groupby/GroupByQueryConfig.java | 22 ++++++++-------- .../GroupByMergingQueryRunnerV2.java | 2 +- .../epinephelinae/GroupByRowProcessor.java | 2 +- .../druid/query/select/SelectQuery.java | 7 ++++++ .../apache/druid/query/QueryContextTest.java | 19 ++++++++++++++ ...ByLimitPushDownInsufficientBufferTest.java | 5 ++-- ...roupByLimitPushDownMultiNodeMergeTest.java | 5 ++-- .../groupby/GroupByMultiSegmentTest.java | 5 ++-- .../query/groupby/GroupByQueryConfigTest.java | 25 ++++++++++--------- .../query/groupby/GroupByQueryRunnerTest.java | 11 ++++---- .../groupby/NestedQueryPushDownTest.java | 5 ++-- 18 files changed, 119 insertions(+), 41 deletions(-) diff --git a/benchmarks/src/test/java/org/apache/druid/benchmark/GroupByTypeInterfaceBenchmark.java b/benchmarks/src/test/java/org/apache/druid/benchmark/GroupByTypeInterfaceBenchmark.java index 445068927534..8e885856bac6 100644 --- a/benchmarks/src/test/java/org/apache/druid/benchmark/GroupByTypeInterfaceBenchmark.java +++ b/benchmarks/src/test/java/org/apache/druid/benchmark/GroupByTypeInterfaceBenchmark.java @@ -32,6 +32,7 @@ import org.apache.druid.data.input.InputRow; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.FileUtils; +import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.granularity.Granularities; import org.apache.druid.java.util.common.granularity.Granularity; import org.apache.druid.java.util.common.guava.Sequence; @@ -362,9 +363,9 @@ public int getBufferGrouperInitialBuckets() } @Override - public long getMaxOnDiskStorage() + public HumanReadableBytes getMaxOnDiskStorage() { - return 1_000_000_000L; + return HumanReadableBytes.valueOf(1_000_000_000L); } }; config.setSingleThreaded(false); diff --git a/benchmarks/src/test/java/org/apache/druid/benchmark/query/GroupByBenchmark.java b/benchmarks/src/test/java/org/apache/druid/benchmark/query/GroupByBenchmark.java index 0fdffa8b87b9..9dd0ea5c5eac 100644 --- a/benchmarks/src/test/java/org/apache/druid/benchmark/query/GroupByBenchmark.java +++ b/benchmarks/src/test/java/org/apache/druid/benchmark/query/GroupByBenchmark.java @@ -34,6 +34,7 @@ import org.apache.druid.common.config.NullHandling; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.FileUtils; +import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.concurrent.Execs; import org.apache.druid.java.util.common.granularity.Granularities; import org.apache.druid.java.util.common.granularity.Granularity; @@ -477,9 +478,9 @@ public int getBufferGrouperInitialBuckets() } @Override - public long getMaxOnDiskStorage() + public HumanReadableBytes getMaxOnDiskStorage() { - return 1_000_000_000L; + return HumanReadableBytes.valueOf(1_000_000_000L); } }; config.setSingleThreaded(false); 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 cfd4bc92559e..cb98d6b4b29c 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,6 +24,7 @@ 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; @@ -170,6 +171,12 @@ 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() { 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 bc45380d2658..c581dfcfc374 100644 --- a/processing/src/main/java/org/apache/druid/query/BaseQuery.java +++ b/processing/src/main/java/org/apache/druid/query/BaseQuery.java @@ -27,6 +27,7 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.Ordering; import org.apache.druid.guice.annotations.ExtensionPoint; +import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.granularity.Granularities; import org.apache.druid.java.util.common.granularity.Granularity; import org.apache.druid.java.util.common.granularity.PeriodGranularity; @@ -199,6 +200,12 @@ public boolean getContextBoolean(String key, boolean defaultValue) return context.getAsBoolean(key, defaultValue); } + @Override + public HumanReadableBytes getContextHumanReadableBytes(String key, HumanReadableBytes defaultValue) + { + return context.getAsHumanReadableBytes(key, defaultValue); + } + /** * @deprecated use {@link #computeOverriddenContext(Map, Map) computeOverriddenContext(getContext(), overrides))} * instead. This method may be removed in the next minor or major version of Druid. 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 ced91a6383ca..924cd9454e06 100644 --- a/processing/src/main/java/org/apache/druid/query/Query.java +++ b/processing/src/main/java/org/apache/druid/query/Query.java @@ -25,6 +25,7 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.Ordering; import org.apache.druid.guice.annotations.ExtensionPoint; +import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.granularity.Granularity; import org.apache.druid.query.datasourcemetadata.DataSourceMetadataQuery; import org.apache.druid.query.filter.DimFilter; @@ -129,6 +130,8 @@ default QueryContext getQueryContext() boolean getContextBoolean(String key, boolean defaultValue); + HumanReadableBytes getContextHumanReadableBytes(String key, HumanReadableBytes defaultValue); + boolean isDescending(); /** 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 20b607f784dd..af7352c8d97e 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContext.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContext.java @@ -19,6 +19,8 @@ package org.apache.druid.query; +import org.apache.druid.java.util.common.HumanReadableBytes; + import javax.annotation.Nullable; import java.util.Collections; @@ -176,6 +178,11 @@ public long getAsLong(final String parameter, final long defaultValue) return QueryContexts.getAsLong(parameter, get(parameter), defaultValue); } + public HumanReadableBytes getAsHumanReadableBytes(final String parameter, final HumanReadableBytes defaultValue) + { + return QueryContexts.getAsHumanReadableBytes(parameter, get(parameter), 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 67cb49be9150..e531869ac5dd 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContexts.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContexts.java @@ -24,6 +24,7 @@ import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; import org.apache.druid.guice.annotations.PublicApi; +import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.IAE; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.Numbers; @@ -568,6 +569,23 @@ public static long getAsLong( } } + public static HumanReadableBytes getAsHumanReadableBytes( + final String parameter, + final Object value, + final HumanReadableBytes defaultValue + ) + { + if (null == value) { + return defaultValue; + } 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); + } + } + public static Map override( final Map context, final Map overrides 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 5d2754cebcd9..6125577ac306 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 @@ -103,10 +103,10 @@ public class GroupByQueryConfig @JsonProperty // Max on-disk temporary storage, per-query; when exceeded, the query fails - private long maxOnDiskStorage = 0L; + private HumanReadableBytes maxOnDiskStorage = HumanReadableBytes.valueOf(0); @JsonProperty - private long defaultOnDiskStorage = -1L; + private HumanReadableBytes defaultOnDiskStorage = HumanReadableBytes.valueOf(-1); @JsonProperty private boolean forcePushDownLimit = false; @@ -264,7 +264,7 @@ public long getActualMaxMergingDictionarySize(final DruidProcessingConfig proces ); } - public long getMaxOnDiskStorage() + public HumanReadableBytes getMaxOnDiskStorage() { return maxOnDiskStorage; } @@ -277,9 +277,9 @@ public long getMaxOnDiskStorage() * * @return The working value for defaultOnDiskStorage */ - public long getDefaultOnDiskStorage() + public HumanReadableBytes getDefaultOnDiskStorage() { - return defaultOnDiskStorage < 0L ? getMaxOnDiskStorage() : defaultOnDiskStorage; + return defaultOnDiskStorage.getBytes() < 0L ? getMaxOnDiskStorage() : defaultOnDiskStorage; } public boolean isForcePushDownLimit() @@ -360,9 +360,11 @@ public GroupByQueryConfig withOverrides(final GroupByQuery query) // If the client overrides do not provide "maxOnDiskStorage" context key, the server side "defaultOnDiskStorage" // value is used in the calculation of the newConfig value of maxOnDiskStorage. This allows the operator to // choose a default value lower than the max allowed when the context key is missing in the client query. - newConfig.maxOnDiskStorage = Math.min( - ((Number) query.getContextValue(CTX_KEY_MAX_ON_DISK_STORAGE, getDefaultOnDiskStorage())).longValue(), - getMaxOnDiskStorage() + newConfig.maxOnDiskStorage = HumanReadableBytes.valueOf( + Math.min( + query.getContextHumanReadableBytes(CTX_KEY_MAX_ON_DISK_STORAGE, getDefaultOnDiskStorage()).getBytes(), + getMaxOnDiskStorage().getBytes() + ) ); newConfig.maxSelectorDictionarySize = maxSelectorDictionarySize; // No overrides newConfig.maxMergingDictionarySize = maxMergingDictionarySize; // No overrides @@ -407,8 +409,8 @@ public String toString() ", bufferGrouperMaxLoadFactor=" + bufferGrouperMaxLoadFactor + ", bufferGrouperInitialBuckets=" + bufferGrouperInitialBuckets + ", maxMergingDictionarySize=" + maxMergingDictionarySize + - ", maxOnDiskStorage=" + maxOnDiskStorage + - ", defaultOnDiskStorage=" + getDefaultOnDiskStorage() + // use the getter because of special behavior for mirroring maxOnDiskStorage if defaultOnDiskStorage not explicitly set. + ", maxOnDiskStorage=" + maxOnDiskStorage.getBytes() + + ", defaultOnDiskStorage=" + getDefaultOnDiskStorage().getBytes() + // use the getter because of special behavior for mirroring maxOnDiskStorage if defaultOnDiskStorage not explicitly set. ", forcePushDownLimit=" + forcePushDownLimit + ", forceHashAggregation=" + forceHashAggregation + ", intermediateCombineDegree=" + intermediateCombineDegree + diff --git a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByMergingQueryRunnerV2.java b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByMergingQueryRunnerV2.java index 16e3dd3ac7ce..3e4a92f6e4a1 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByMergingQueryRunnerV2.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByMergingQueryRunnerV2.java @@ -175,7 +175,7 @@ public CloseableGrouperIterator make() try { final LimitedTemporaryStorage temporaryStorage = new LimitedTemporaryStorage( temporaryStorageDirectory, - querySpecificConfig.getMaxOnDiskStorage() + querySpecificConfig.getMaxOnDiskStorage().getBytes() ); final ReferenceCountingResourceHolder temporaryStorageHolder = ReferenceCountingResourceHolder.fromCloseable(temporaryStorage); diff --git a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByRowProcessor.java b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByRowProcessor.java index 88c3c965442d..76125c2c504b 100644 --- a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByRowProcessor.java +++ b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByRowProcessor.java @@ -106,7 +106,7 @@ public static ResultSupplier process( final LimitedTemporaryStorage temporaryStorage = new LimitedTemporaryStorage( temporaryStorageDirectory, - querySpecificConfig.getMaxOnDiskStorage() + querySpecificConfig.getMaxOnDiskStorage().getBytes() ); closeOnExit.register(temporaryStorage); 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 e46e6e9b5661..f2895f0f9bc0 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,6 +21,7 @@ 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; @@ -134,6 +135,12 @@ 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/test/java/org/apache/druid/query/QueryContextTest.java b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java index 3654f85af175..98b6167d2ff2 100644 --- a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java +++ b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java @@ -23,7 +23,10 @@ import com.google.common.collect.Ordering; import nl.jqno.equalsverifier.EqualsVerifier; import nl.jqno.equalsverifier.Warning; +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; @@ -346,6 +349,22 @@ public boolean getContextBoolean(String key, boolean defaultValue) return (boolean) context.get(key); } + @Override + public HumanReadableBytes getContextHumanReadableBytes(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)); + } else if (value instanceof String) { + return new HumanReadableBytes((String) value); + } else { + throw new IAE("Expected parameter [%s] to be in human readable format", key); + } + } + @Override public boolean isDescending() { diff --git a/processing/src/test/java/org/apache/druid/query/groupby/GroupByLimitPushDownInsufficientBufferTest.java b/processing/src/test/java/org/apache/druid/query/groupby/GroupByLimitPushDownInsufficientBufferTest.java index 8e5d7011ed9a..0dfdac7fb4df 100644 --- a/processing/src/test/java/org/apache/druid/query/groupby/GroupByLimitPushDownInsufficientBufferTest.java +++ b/processing/src/test/java/org/apache/druid/query/groupby/GroupByLimitPushDownInsufficientBufferTest.java @@ -34,6 +34,7 @@ import org.apache.druid.data.input.impl.StringDimensionSchema; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.FileUtils; +import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.concurrent.Execs; import org.apache.druid.java.util.common.granularity.Granularities; @@ -291,9 +292,9 @@ public int getBufferGrouperInitialBuckets() } @Override - public long getMaxOnDiskStorage() + public HumanReadableBytes getMaxOnDiskStorage() { - return 1_000_000_000L; + return HumanReadableBytes.valueOf(1_000_000_000L); } }; config.setSingleThreaded(false); diff --git a/processing/src/test/java/org/apache/druid/query/groupby/GroupByLimitPushDownMultiNodeMergeTest.java b/processing/src/test/java/org/apache/druid/query/groupby/GroupByLimitPushDownMultiNodeMergeTest.java index 4c362a53c803..73685d7aa5d8 100644 --- a/processing/src/test/java/org/apache/druid/query/groupby/GroupByLimitPushDownMultiNodeMergeTest.java +++ b/processing/src/test/java/org/apache/druid/query/groupby/GroupByLimitPushDownMultiNodeMergeTest.java @@ -34,6 +34,7 @@ import org.apache.druid.data.input.impl.StringDimensionSchema; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.FileUtils; +import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.concurrent.Execs; import org.apache.druid.java.util.common.granularity.Granularities; @@ -562,9 +563,9 @@ public int getBufferGrouperInitialBuckets() } @Override - public long getMaxOnDiskStorage() + public HumanReadableBytes getMaxOnDiskStorage() { - return 1_000_000_000L; + return HumanReadableBytes.valueOf(1_000_000_000L); } }; config.setSingleThreaded(false); diff --git a/processing/src/test/java/org/apache/druid/query/groupby/GroupByMultiSegmentTest.java b/processing/src/test/java/org/apache/druid/query/groupby/GroupByMultiSegmentTest.java index 58c4abed1745..b18b4d9252c2 100644 --- a/processing/src/test/java/org/apache/druid/query/groupby/GroupByMultiSegmentTest.java +++ b/processing/src/test/java/org/apache/druid/query/groupby/GroupByMultiSegmentTest.java @@ -32,6 +32,7 @@ import org.apache.druid.data.input.impl.StringDimensionSchema; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.FileUtils; +import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.concurrent.Execs; import org.apache.druid.java.util.common.granularity.Granularities; @@ -227,9 +228,9 @@ public int getBufferGrouperInitialBuckets() } @Override - public long getMaxOnDiskStorage() + public HumanReadableBytes getMaxOnDiskStorage() { - return 1_000_000_000L; + return HumanReadableBytes.valueOf(1_000_000_000L); } }; config.setSingleThreaded(false); diff --git a/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java b/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java index 3806b60ddc6c..369fe36ec663 100644 --- a/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java +++ b/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java @@ -38,7 +38,8 @@ public class GroupByQueryConfigTest .put("bufferGrouperInitialBuckets", "1") .put("maxIntermediateRows", "2") .put("maxResults", "3") - .put("maxOnDiskStorage", "4") + .put("defaultOnDiskStorage", "1M") + .put("maxOnDiskStorage", "4M") .put("maxSelectorDictionarySize", "5") .put("maxMergingDictionarySize", "6M") .put("bufferGrouperMaxLoadFactor", "7") @@ -54,8 +55,8 @@ public void testSerde() Assert.assertEquals(1, config.getBufferGrouperInitialBuckets()); Assert.assertEquals(2, config.getMaxIntermediateRows()); Assert.assertEquals(3, config.getMaxResults()); - Assert.assertEquals(4, config.getMaxOnDiskStorage()); - Assert.assertEquals(4, config.getDefaultOnDiskStorage()); + Assert.assertEquals(4_000_000, config.getMaxOnDiskStorage().getBytes()); + Assert.assertEquals(1_000_000, config.getDefaultOnDiskStorage().getBytes()); Assert.assertEquals(5, config.getConfiguredMaxSelectorDictionarySize()); Assert.assertEquals(6_000_000, config.getConfiguredMaxMergingDictionarySize()); Assert.assertEquals(7.0, config.getBufferGrouperMaxLoadFactor(), 0.0); @@ -79,7 +80,7 @@ public void testNoOverrides() Assert.assertEquals(1, config2.getBufferGrouperInitialBuckets()); Assert.assertEquals(2, config2.getMaxIntermediateRows()); Assert.assertEquals(3, config2.getMaxResults()); - Assert.assertEquals(4, config2.getMaxOnDiskStorage()); + Assert.assertEquals(1_000_000, config2.getMaxOnDiskStorage().getBytes()); Assert.assertEquals(5, config2.getConfiguredMaxSelectorDictionarySize()); Assert.assertEquals(6_000_000, config2.getConfiguredMaxMergingDictionarySize()); Assert.assertEquals(7.0, config2.getBufferGrouperMaxLoadFactor(), 0.0); @@ -98,7 +99,7 @@ public void testOverrides() .setContext( ImmutableMap.builder() .put("groupByStrategy", "v1") - .put("maxOnDiskStorage", 0) + .put("maxOnDiskStorage", "3M") .put("maxResults", 2) .put("maxSelectorDictionarySize", 3) .put("maxMergingDictionarySize", 4) @@ -113,7 +114,7 @@ public void testOverrides() Assert.assertEquals(1, config2.getBufferGrouperInitialBuckets()); Assert.assertEquals(2, config2.getMaxIntermediateRows()); Assert.assertEquals(2, config2.getMaxResults()); - Assert.assertEquals(0, config2.getMaxOnDiskStorage()); + Assert.assertEquals(3_000_000, config2.getMaxOnDiskStorage().getBytes()); Assert.assertEquals(5 /* Can't override */, config2.getConfiguredMaxSelectorDictionarySize()); Assert.assertEquals(6_000_000 /* Can't override */, config2.getConfiguredMaxMergingDictionarySize()); Assert.assertEquals(7.0, config2.getBufferGrouperMaxLoadFactor(), 0.0); @@ -176,8 +177,8 @@ public void testUseDefaultOnDiskStorage() { final GroupByQueryConfig config = MAPPER.convertValue( ImmutableMap.of( - "maxOnDiskStorage", "10", - "defaultOnDiskStorage", "5" + "maxOnDiskStorage", "10G", + "defaultOnDiskStorage", "5G" ), GroupByQueryConfig.class ); @@ -189,14 +190,14 @@ public void testUseDefaultOnDiskStorage() .setContext(ImmutableMap.builder().build()) .build() ); - Assert.assertEquals(5L, config2.getMaxOnDiskStorage()); + Assert.assertEquals(5_000_000_000L, config2.getMaxOnDiskStorage().getBytes()); } @Test public void testUseMaxOnDiskStorageWhenClientOverrideIsTooLarge() { final GroupByQueryConfig config = MAPPER.convertValue( - ImmutableMap.of("maxOnDiskStorage", "10"), + ImmutableMap.of("maxOnDiskStorage", "500M"), GroupByQueryConfig.class ); final GroupByQueryConfig config2 = config.withOverrides( @@ -206,11 +207,11 @@ public void testUseMaxOnDiskStorageWhenClientOverrideIsTooLarge() .setGranularity(Granularities.ALL) .setContext( ImmutableMap.builder() - .put("maxOnDiskStorage", 500) + .put("maxOnDiskStorage", "1G") .build() ) .build() ); - Assert.assertEquals(10L, config2.getMaxOnDiskStorage()); + Assert.assertEquals(500_000_000, config2.getMaxOnDiskStorage().getBytes()); } } diff --git a/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryRunnerTest.java b/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryRunnerTest.java index 71274c9f7884..db1c689e8414 100644 --- a/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryRunnerTest.java +++ b/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryRunnerTest.java @@ -34,6 +34,7 @@ import org.apache.druid.data.input.Row; import org.apache.druid.data.input.Rows; import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.IAE; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.Intervals; @@ -283,9 +284,9 @@ public int getBufferGrouperMaxSize() } @Override - public long getMaxOnDiskStorage() + public HumanReadableBytes getMaxOnDiskStorage() { - return 10L * 1024 * 1024; + return HumanReadableBytes.valueOf(10L * 1024 * 1024); } @Override @@ -315,9 +316,9 @@ public long getConfiguredMaxMergingDictionarySize() } @Override - public long getMaxOnDiskStorage() + public HumanReadableBytes getMaxOnDiskStorage() { - return 10L * 1024 * 1024; + return HumanReadableBytes.valueOf(10L * 1024 * 1024); } @Override @@ -3018,7 +3019,7 @@ public void testNotEnoughDiskSpaceThroughContextOverride() List expectedResults = null; if (config.getDefaultStrategy().equals(GroupByStrategySelector.STRATEGY_V2)) { expectedException.expect(ResourceLimitExceededException.class); - if (config.getMaxOnDiskStorage() > 0) { + if (config.getMaxOnDiskStorage().getBytes() > 0) { // The error message always mentions disk if you have spilling enabled (maxOnDiskStorage > 0) expectedException.expectMessage("Not enough disk space to execute this query"); } else { diff --git a/processing/src/test/java/org/apache/druid/query/groupby/NestedQueryPushDownTest.java b/processing/src/test/java/org/apache/druid/query/groupby/NestedQueryPushDownTest.java index 81554b26725b..8592fca11e0b 100644 --- a/processing/src/test/java/org/apache/druid/query/groupby/NestedQueryPushDownTest.java +++ b/processing/src/test/java/org/apache/druid/query/groupby/NestedQueryPushDownTest.java @@ -33,6 +33,7 @@ import org.apache.druid.data.input.impl.StringDimensionSchema; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.FileUtils; +import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.concurrent.Execs; import org.apache.druid.java.util.common.granularity.Granularities; @@ -272,9 +273,9 @@ public int getBufferGrouperInitialBuckets() } @Override - public long getMaxOnDiskStorage() + public HumanReadableBytes getMaxOnDiskStorage() { - return 1_000_000_000L; + return HumanReadableBytes.valueOf(1_000_000_000L); } }; config.setSingleThreaded(false); From 7345c7c9678d22d63a1497ce54332a2a2e7a43f8 Mon Sep 17 00:00:00 2001 From: "Lucas.Capistrant" Date: Thu, 4 Aug 2022 16:21:12 -0500 Subject: [PATCH 8/9] improve test coverage --- .../MaterializedViewQueryTest.java | 35 +++++++++++++++++++ .../apache/druid/query/QueryContextTest.java | 11 ++++++ .../apache/druid/query/QueryContextsTest.java | 9 +++++ .../query/groupby/GroupByQueryConfigTest.java | 16 +++++++++ 4 files changed, 71 insertions(+) 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 1432f519ef73..1a55cdd02666 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 @@ -22,8 +22,10 @@ import com.fasterxml.jackson.databind.InjectableValues; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.jsontype.NamedType; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; import com.google.common.collect.Lists; +import org.apache.druid.java.util.common.HumanReadableBytes; import org.apache.druid.math.expr.ExprMacroTable; import org.apache.druid.query.Query; import org.apache.druid.query.QueryRunnerTestHelper; @@ -89,4 +91,37 @@ public void testQuerySerialization() throws IOException Assert.assertEquals(QueryRunnerTestHelper.ALL_GRAN, query.getGranularity()); Assert.assertEquals(QueryRunnerTestHelper.FULL_ON_INTERVAL_SPEC.getIntervals(), query.getIntervals()); } + + @Test + public void testGetContextHumanReadableBytes() + { + TopNQuery topNQuery = new TopNQueryBuilder() + .dataSource(QueryRunnerTestHelper.DATA_SOURCE) + .granularity(QueryRunnerTestHelper.ALL_GRAN) + .dimension(QueryRunnerTestHelper.MARKET_DIMENSION) + .metric(QueryRunnerTestHelper.INDEX_METRIC) + .threshold(4) + .intervals(QueryRunnerTestHelper.FULL_ON_INTERVAL_SPEC) + .aggregators( + Lists.newArrayList( + Iterables.concat( + QueryRunnerTestHelper.COMMON_DOUBLE_AGGREGATORS, + Lists.newArrayList( + new DoubleMaxAggregatorFactory("maxIndex", "index"), + new DoubleMinAggregatorFactory("minIndex", "index") + ) + ) + ) + ) + .context( + ImmutableMap.of( + "maxOnDiskStorage", "20M" + ) + ) + .postAggregators(QueryRunnerTestHelper.ADD_ROWS_INDEX_CONSTANT) + .build(); + MaterializedViewQuery query = new MaterializedViewQuery(topNQuery, optimizer); + Assert.assertEquals(20_000_000, query.getContextHumanReadableBytes("maxOnDiskStorage", HumanReadableBytes.ZERO).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 98b6167d2ff2..3b1eee0db065 100644 --- a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java +++ b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java @@ -133,6 +133,17 @@ public void testGetLong() Assert.assertEquals(0L, context.getAsLong("non-exist", 0)); } + @Test + public void testGetHumanReadableBytes() + { + final QueryContext context = new QueryContext( + ImmutableMap.of( + "maxOnDiskStorage", "500M" + ) + ); + Assert.assertEquals(500_000_000, context.getAsHumanReadableBytes("maxOnDiskStorage", HumanReadableBytes.ZERO).getBytes()); + } + @Test public void testAddSystemParamOverrideUserParam() { 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 2ee9b9363e15..3d34e2d0e40f 100644 --- a/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java +++ b/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java @@ -21,6 +21,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +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.query.spec.MultipleIntervalSegmentSpec; @@ -274,4 +275,12 @@ public void testGetAs() // Expected } } + + @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()); + } } diff --git a/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java b/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java index 369fe36ec663..0c7a57ca0306 100644 --- a/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java +++ b/processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryConfigTest.java @@ -214,4 +214,20 @@ public void testUseMaxOnDiskStorageWhenClientOverrideIsTooLarge() ); Assert.assertEquals(500_000_000, config2.getMaxOnDiskStorage().getBytes()); } + + @Test + public void testGetDefaultOnDiskStorageReturnsCorrectValue() + { + final GroupByQueryConfig config = MAPPER.convertValue( + ImmutableMap.of("maxOnDiskStorage", "500M"), + GroupByQueryConfig.class + ); + final GroupByQueryConfig config2 = MAPPER.convertValue( + ImmutableMap.of("maxOnDiskStorage", "500M", + "defaultOnDiskStorage", "100M"), + GroupByQueryConfig.class + ); + Assert.assertEquals(500_000_000, config.getDefaultOnDiskStorage().getBytes()); + Assert.assertEquals(100_000_000, config2.getDefaultOnDiskStorage().getBytes()); + } } From 986fab858ac6ebf9f74ca74375ca852e59788433 Mon Sep 17 00:00:00 2001 From: "Lucas.Capistrant" Date: Tue, 9 Aug 2022 16:15:22 -0500 Subject: [PATCH 9/9] Provide default implementation to new default method on advice of reviewer --- .../java/org/apache/druid/query/Query.java | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) 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 924cd9454e06..9db2f0ca006d 100644 --- a/processing/src/main/java/org/apache/druid/query/Query.java +++ b/processing/src/main/java/org/apache/druid/query/Query.java @@ -130,7 +130,23 @@ default QueryContext getQueryContext() boolean getContextBoolean(String key, boolean defaultValue); - HumanReadableBytes getContextHumanReadableBytes(String key, HumanReadableBytes 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 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) + { + if (null != getQueryContext()) { + return getQueryContext().getAsHumanReadableBytes(key, defaultValue); + } else { + return defaultValue; + } + } boolean isDescending();