Examples with MetricManipulationFn org.apache.druid.query.aggregation.MetricManipulationFn used on opensource projects

Example 1 with MetricManipulationFn

use of org.apache.druid.query.aggregation.MetricManipulationFn in project druid by druid-io.

the class GroupByQueryQueryToolChest method makePostComputeManipulatorFn.

@Override
public Function<ResultRow, ResultRow> makePostComputeManipulatorFn(final GroupByQuery query, final MetricManipulationFn fn) {
    final BitSet optimizedDims = extractionsToRewrite(query);
    final Function<ResultRow, ResultRow> preCompute = makePreComputeManipulatorFn(query, fn);
    if (optimizedDims.isEmpty()) {
        return preCompute;
    }
    // If we have optimizations that can be done at this level, we apply them here
    final List<DimensionSpec> dimensions = query.getDimensions();
    final List<ExtractionFn> extractionFns = new ArrayList<>(dimensions.size());
    for (int i = 0; i < dimensions.size(); i++) {
        final DimensionSpec dimensionSpec = dimensions.get(i);
        final ExtractionFn extractionFnToAdd;
        if (optimizedDims.get(i)) {
            extractionFnToAdd = dimensionSpec.getExtractionFn();
        } else {
            extractionFnToAdd = null;
        }
        extractionFns.add(extractionFnToAdd);
    }
    final int dimensionStart = query.getResultRowDimensionStart();
    return row -> {
        // preCompute.apply(row) will either return the original row, or create a copy.
        ResultRow newRow = preCompute.apply(row);
        // noinspection ObjectEquality (if preCompute made a copy, no need to make another copy)
        if (newRow == row) {
            newRow = row.copy();
        }
        for (int i = optimizedDims.nextSetBit(0); i >= 0; i = optimizedDims.nextSetBit(i + 1)) {
            newRow.set(dimensionStart + i, extractionFns.get(i).apply(newRow.get(dimensionStart + i)));
        }
        return newRow;
    };
}

Example 2 with MetricManipulationFn

use of org.apache.druid.query.aggregation.MetricManipulationFn in project druid by druid-io.

the class MaterializedViewQueryQueryToolChestTest method testMakePostComputeManipulatorFn.

@Test
public void testMakePostComputeManipulatorFn() {
    TimeseriesQuery realQuery = Druids.newTimeseriesQueryBuilder().dataSource(QueryRunnerTestHelper.DATA_SOURCE).granularity(QueryRunnerTestHelper.DAY_GRAN).intervals(QueryRunnerTestHelper.FULL_ON_INTERVAL_SPEC).aggregators(QueryRunnerTestHelper.ROWS_COUNT).descending(true).build();
    MaterializedViewQuery materializedViewQuery = new MaterializedViewQuery(realQuery, null);
    QueryToolChest materializedViewQueryQueryToolChest = new MaterializedViewQueryQueryToolChest(new MapQueryToolChestWarehouse(ImmutableMap.<Class<? extends Query>, QueryToolChest>builder().put(TimeseriesQuery.class, new TimeseriesQueryQueryToolChest()).build()));
    Function postFn = materializedViewQueryQueryToolChest.makePostComputeManipulatorFn(materializedViewQuery, new MetricManipulationFn() {

        @Override
        public Object manipulate(AggregatorFactory factory, Object object) {
            return "metricvalue1";
        }
    });
    Result<TimeseriesResultValue> result = new Result<>(DateTimes.nowUtc(), new TimeseriesResultValue(ImmutableMap.of("dim1", "dimvalue1")));
    Result<TimeseriesResultValue> postResult = (Result<TimeseriesResultValue>) postFn.apply(result);
    Map<String, Object> postResultMap = postResult.getValue().getBaseObject();
    Assert.assertEquals(postResult.getTimestamp(), result.getTimestamp());
    Assert.assertEquals(postResultMap.size(), 2);
    Assert.assertEquals(postResultMap.get(QueryRunnerTestHelper.ROWS_COUNT.getName()), "metricvalue1");
    Assert.assertEquals(postResultMap.get("dim1"), "dimvalue1");
}

Example 3 with MetricManipulationFn

use of org.apache.druid.query.aggregation.MetricManipulationFn in project druid by druid-io.

the class FinalizeResultsQueryRunner method run.

@Override
public Sequence<T> run(final QueryPlus<T> queryPlus, ResponseContext responseContext) {
    final Query<T> query = queryPlus.getQuery();
    final boolean isBySegment = QueryContexts.isBySegment(query);
    final boolean shouldFinalize = QueryContexts.isFinalize(query, true);
    final Query<T> queryToRun;
    final Function<T, ?> finalizerFn;
    final MetricManipulationFn metricManipulationFn;
    if (shouldFinalize) {
        queryToRun = query.withOverriddenContext(ImmutableMap.of("finalize", false));
        metricManipulationFn = MetricManipulatorFns.finalizing();
    } else {
        queryToRun = query;
        metricManipulationFn = MetricManipulatorFns.identity();
    }
    if (isBySegment) {
        finalizerFn = new Function<T, Result<BySegmentResultValue<T>>>() {

            final Function<T, T> baseFinalizer = toolChest.makePostComputeManipulatorFn(query, metricManipulationFn);

            @Override
            public Result<BySegmentResultValue<T>> apply(T input) {
                // noinspection unchecked (input is not actually a T; see class-level javadoc)
                Result<BySegmentResultValueClass<T>> result = (Result<BySegmentResultValueClass<T>>) input;
                if (input == null) {
                    throw new ISE("Cannot have a null result!");
                }
                BySegmentResultValue<T> resultsClass = result.getValue();
                return new Result<>(result.getTimestamp(), new BySegmentResultValueClass<>(Lists.transform(resultsClass.getResults(), baseFinalizer), resultsClass.getSegmentId(), resultsClass.getInterval()));
            }
        };
    } else {
        finalizerFn = toolChest.makePostComputeManipulatorFn(query, metricManipulationFn);
    }
    // noinspection unchecked (Technically unsound, but see class-level javadoc for rationale)
    return (Sequence<T>) Sequences.map(baseRunner.run(queryPlus.withQuery(queryToRun), responseContext), finalizerFn);
}