diff --git a/processing/src/main/java/org/apache/druid/query/topn/TimeExtractionTopNAlgorithm.java b/processing/src/main/java/org/apache/druid/query/topn/TimeExtractionTopNAlgorithm.java index 3fd8795303c9..7d4c23cee23b 100644 --- a/processing/src/main/java/org/apache/druid/query/topn/TimeExtractionTopNAlgorithm.java +++ b/processing/src/main/java/org/apache/druid/query/topn/TimeExtractionTopNAlgorithm.java @@ -22,25 +22,36 @@ import org.apache.druid.query.ColumnSelectorPlus; import org.apache.druid.query.aggregation.Aggregator; import org.apache.druid.segment.Cursor; +import org.apache.druid.segment.DimensionHandlerUtils; import org.apache.druid.segment.DimensionSelector; import org.apache.druid.segment.StorageAdapter; +import org.apache.druid.segment.column.ValueType; -import java.util.HashMap; import java.util.Map; +import java.util.function.Function; -public class TimeExtractionTopNAlgorithm extends BaseTopNAlgorithm, TopNParams> +public class TimeExtractionTopNAlgorithm extends BaseTopNAlgorithm, TopNParams> { - public static final int[] EMPTY_INTS = new int[]{}; + private static final int[] EMPTY_INTS = new int[]{}; + private final TopNQuery query; + private final Function> dimensionValueConverter; public TimeExtractionTopNAlgorithm(StorageAdapter storageAdapter, TopNQuery query) { super(storageAdapter); this.query = query; - } + // This strategy is used for ExtractionFns on the __time column. They always return STRING, so we need to convert + // from STRING to the desired output type. + this.dimensionValueConverter = DimensionHandlerUtils.converterFromTypeToType( + ValueType.STRING, + query.getDimensionSpec().getOutputType() + ); + } @Override + @SuppressWarnings("unchecked") public TopNParams makeInitParams(ColumnSelectorPlus selectorPlus, Cursor cursor) { return new TopNParams( @@ -63,13 +74,18 @@ protected int[] updateDimValSelector(int[] dimValSelector, int numProcessed, int } @Override - protected Map makeDimValAggregateStore(TopNParams params) + @SuppressWarnings("unchecked") + protected Map makeDimValAggregateStore(TopNParams params) { - return new HashMap<>(); + return params.getSelectorPlus().getColumnSelectorStrategy().makeDimExtractionAggregateStore(); } @Override - protected long scanAndAggregate(TopNParams params, int[] dimValSelector, Map aggregatesStore) + protected long scanAndAggregate( + TopNParams params, + int[] dimValSelector, + Map aggregatesStore + ) { if (params.getCardinality() < 0) { throw new UnsupportedOperationException("Cannot operate on a dimension with unknown cardinality"); @@ -80,7 +96,7 @@ protected long scanAndAggregate(TopNParams params, int[] dimValSelector, Map aggregatesStore, + Map aggregatesStore, TopNResultBuilder resultBuilder ) { - for (Map.Entry entry : aggregatesStore.entrySet()) { + for (Map.Entry entry : aggregatesStore.entrySet()) { Aggregator[] aggs = entry.getValue(); if (aggs != null) { Object[] vals = new Object[aggs.length]; @@ -124,7 +140,7 @@ protected void updateResults( } @Override - protected void closeAggregators(Map stringMap) + protected void closeAggregators(Map stringMap) { for (Aggregator[] aggregators : stringMap.values()) { for (Aggregator agg : aggregators) { diff --git a/processing/src/test/java/org/apache/druid/query/topn/TopNQueryRunnerTest.java b/processing/src/test/java/org/apache/druid/query/topn/TopNQueryRunnerTest.java index 0244372dc309..cf636ecf5585 100644 --- a/processing/src/test/java/org/apache/druid/query/topn/TopNQueryRunnerTest.java +++ b/processing/src/test/java/org/apache/druid/query/topn/TopNQueryRunnerTest.java @@ -2364,6 +2364,53 @@ public void testTopNDimExtractionToOne() TestHelper.assertExpectedResults(expectedResults, list, "Failed to match"); } + @Test + public void testTopNDimExtractionTimeToOneLong() + { + TopNQuery query = new TopNQueryBuilder() + .dataSource(QueryRunnerTestHelper.dataSource) + .granularity(QueryRunnerTestHelper.allGran) + .dimension( + new ExtractionDimensionSpec( + ColumnHolder.TIME_COLUMN_NAME, + "t", + ValueType.LONG, + new JavaScriptExtractionFn( + "function(f) { return \"42\"; }", + false, + JavaScriptConfig.getEnabledInstance() + ) + ) + ) + .metric("rows") + .threshold(10) + .intervals(QueryRunnerTestHelper.fullOnIntervalSpec) + .aggregators(commonAggregators) + .postAggregators(Collections.singletonList(QueryRunnerTestHelper.addRowsIndexConstant)) + .build(); + + List> expectedResults = Collections.singletonList( + new Result<>( + DateTimes.of("2011-01-12T00:00:00.000Z"), + new TopNResultValue( + Collections.>singletonList( + ImmutableMap.of( + "addRowsIndexConstant", 504542.5071372986D, + "index", 503332.5071372986D, + "t", 42L, + "uniques", QueryRunnerTestHelper.UNIQUES_9, + "rows", 1209L + ) + ) + ) + ) + ); + List> list = runWithMerge(query).toList(); + Assert.assertEquals(list.size(), 1); + Assert.assertEquals("Didn't merge results", list.get(0).getValue().getValue().size(), 1); + TestHelper.assertExpectedResults(expectedResults, list, "Failed to match"); + } + @Test public void testTopNCollapsingDimExtraction() {