diff --git a/extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQueryQueryToolChest.java b/extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQueryQueryToolChest.java index 34b556669fb..39e2f19b1a1 100644 --- a/extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQueryQueryToolChest.java +++ b/extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQueryQueryToolChest.java @@ -73,6 +73,13 @@ public class MaterializedViewQueryQueryToolChest extends QueryToolChest return warehouse.getToolChest(realQuery).makePreComputeManipulatorFn(realQuery, fn); } + @Override + public Function makePostComputeManipulatorFn(Query query, MetricManipulationFn fn) + { + Query realQuery = getRealQuery(query); + return warehouse.getToolChest(realQuery).makePostComputeManipulatorFn(realQuery, fn); + } + @Override public TypeReference getResultTypeReference() { diff --git a/extensions-contrib/materialized-view-selection/src/test/java/org/apache/druid/query/materializedview/MaterializedViewQueryQueryToolChestTest.java b/extensions-contrib/materialized-view-selection/src/test/java/org/apache/druid/query/materializedview/MaterializedViewQueryQueryToolChestTest.java new file mode 100644 index 00000000000..fae69e61db2 --- /dev/null +++ b/extensions-contrib/materialized-view-selection/src/test/java/org/apache/druid/query/materializedview/MaterializedViewQueryQueryToolChestTest.java @@ -0,0 +1,92 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.query.materializedview; + +import com.google.common.base.Function; +import com.google.common.collect.ImmutableMap; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.query.Druids; +import org.apache.druid.query.MapQueryToolChestWarehouse; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryRunnerTestHelper; +import org.apache.druid.query.QueryToolChest; +import org.apache.druid.query.Result; +import org.apache.druid.query.aggregation.AggregatorFactory; +import org.apache.druid.query.aggregation.MetricManipulationFn; +import org.apache.druid.query.timeseries.TimeseriesQuery; +import org.apache.druid.query.timeseries.TimeseriesQueryQueryToolChest; +import org.apache.druid.query.timeseries.TimeseriesResultValue; +import org.junit.Assert; +import org.junit.Test; + +import java.util.Map; + +public class MaterializedViewQueryQueryToolChestTest +{ + @Test + public void testMakePostComputeManipulatorFn() + { + TimeseriesQuery realQuery = Druids.newTimeseriesQueryBuilder() + .dataSource(QueryRunnerTestHelper.dataSource) + .granularity(QueryRunnerTestHelper.dayGran) + .intervals(QueryRunnerTestHelper.fullOnInterval) + .aggregators(QueryRunnerTestHelper.rowsCount) + .descending(true) + .build(); + MaterializedViewQuery materializedViewQuery = new MaterializedViewQuery(realQuery, null); + + QueryToolChest materializedViewQueryQueryToolChest = + new MaterializedViewQueryQueryToolChest(new MapQueryToolChestWarehouse( + ImmutableMap., QueryToolChest>builder() + .put( + TimeseriesQuery.class, + new TimeseriesQueryQueryToolChest( + QueryRunnerTestHelper.NoopIntervalChunkingQueryRunnerDecorator() + ) + ) + .build() + )); + + Function postFn = + materializedViewQueryQueryToolChest.makePostComputeManipulatorFn( + materializedViewQuery, + new MetricManipulationFn() { + @Override + public Object manipulate(AggregatorFactory factory, Object object) + { + return "metricvalue1"; + } + }); + + + + Result result = new Result<>(DateTimes.nowUtc(), new TimeseriesResultValue(ImmutableMap.builder() + .put("dim1", "dimvalue1") + .build())); + + Result postResult = (Result) postFn.apply(result); + Map postResultMap = postResult.getValue().getBaseObject(); + + Assert.assertEquals(postResult.getTimestamp(), result.getTimestamp()); + Assert.assertEquals(postResultMap.size(), 2); + Assert.assertEquals(postResultMap.get(QueryRunnerTestHelper.rowsCount.getName()), "metricvalue1"); + Assert.assertEquals(postResultMap.get("dim1"), "dimvalue1"); + } +}