Examples with DataSource org.apache.druid.query.DataSource used on opensource projects

Example 1 with DataSource

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

the class LookupSegmentWrangler method getSegmentsForIntervals.

@Override
public Iterable<Segment> getSegmentsForIntervals(final DataSource dataSource, final Iterable<Interval> intervals) {
    final LookupDataSource lookupDataSource = (LookupDataSource) dataSource;
    final Optional<LookupExtractorFactoryContainer> maybeContainer = lookupProvider.get(lookupDataSource.getLookupName());
    return maybeContainer.map(container -> Collections.<Segment>singletonList(new LookupSegment(lookupDataSource.getLookupName(), container.getLookupExtractorFactory()))).orElse(Collections.emptyList());
}

Example 2 with DataSource

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

the class DataSourceAnalysis method forDataSource.

public static DataSourceAnalysis forDataSource(final DataSource dataSource) {
    // Strip outer queries, retaining querySegmentSpecs as we go down (lowest will become the 'baseQuerySegmentSpec').
    Query<?> baseQuery = null;
    DataSource current = dataSource;
    while (current instanceof QueryDataSource) {
        final Query<?> subQuery = ((QueryDataSource) current).getQuery();
        if (!(subQuery instanceof BaseQuery)) {
            // work properly. All builtin query types are BaseQuery, so we only expect this with funky extension queries.
            throw new IAE("Cannot analyze subquery of class[%s]", subQuery.getClass().getName());
        }
        baseQuery = subQuery;
        current = subQuery.getDataSource();
    }
    if (current instanceof JoinDataSource) {
        final Triple<DataSource, DimFilter, List<PreJoinableClause>> flattened = flattenJoin((JoinDataSource) current);
        return new DataSourceAnalysis(dataSource, flattened.first, baseQuery, flattened.second, flattened.third);
    } else {
        return new DataSourceAnalysis(dataSource, current, baseQuery, null, Collections.emptyList());
    }
}

Example 3 with DataSource

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

the class DruidQuery method toGroupByQuery.

/**
 * Return this query as a GroupBy query, or null if this query is not compatible with GroupBy.
 *
 * @return query or null
 */
@Nullable
private GroupByQuery toGroupByQuery(final QueryFeatureInspector queryFeatureInspector) {
    if (grouping == null) {
        return null;
    }
    if (sorting != null && sorting.getOffsetLimit().hasLimit() && sorting.getOffsetLimit().getLimit() <= 0) {
        // Cannot handle zero or negative limits.
        return null;
    }
    final Pair<DataSource, Filtration> dataSourceFiltrationPair = getFiltration(dataSource, filter, virtualColumnRegistry);
    final DataSource newDataSource = dataSourceFiltrationPair.lhs;
    final Filtration filtration = dataSourceFiltrationPair.rhs;
    final DimFilterHavingSpec havingSpec;
    if (grouping.getHavingFilter() != null) {
        havingSpec = new DimFilterHavingSpec(Filtration.create(grouping.getHavingFilter()).optimizeFilterOnly(grouping.getOutputRowSignature()).getDimFilter(), true);
    } else {
        havingSpec = null;
    }
    final List<PostAggregator> postAggregators = new ArrayList<>(grouping.getPostAggregators());
    if (sorting != null && sorting.getProjection() != null) {
        postAggregators.addAll(sorting.getProjection().getPostAggregators());
    }
    GroupByQuery query = new GroupByQuery(newDataSource, filtration.getQuerySegmentSpec(), getVirtualColumns(true), filtration.getDimFilter(), Granularities.ALL, grouping.getDimensionSpecs(), grouping.getAggregatorFactories(), postAggregators, havingSpec, Optional.ofNullable(sorting).orElse(Sorting.none()).limitSpec(), grouping.getSubtotals().toSubtotalsSpec(grouping.getDimensionSpecs()), ImmutableSortedMap.copyOf(plannerContext.getQueryContext()));
    // We don't apply timestamp computation optimization yet when limit is pushed down. Maybe someday.
    if (query.getLimitSpec() instanceof DefaultLimitSpec && query.isApplyLimitPushDown()) {
        return query;
    }
    Map<String, Object> theContext = new HashMap<>();
    Granularity queryGranularity = null;
    // now, part of the query plan logic is handled in GroupByStrategyV2.
    if (!grouping.getDimensions().isEmpty()) {
        for (DimensionExpression dimensionExpression : grouping.getDimensions()) {
            Granularity granularity = Expressions.toQueryGranularity(dimensionExpression.getDruidExpression(), plannerContext.getExprMacroTable());
            if (granularity == null) {
                continue;
            }
            if (queryGranularity != null) {
                // group by more than one timestamp_floor
                // eg: group by timestamp_floor(__time to DAY),timestamp_floor(__time, to HOUR)
                queryGranularity = null;
                break;
            }
            queryGranularity = granularity;
            int timestampDimensionIndexInDimensions = grouping.getDimensions().indexOf(dimensionExpression);
            // these settings will only affect the most inner query sent to the down streaming compute nodes
            theContext.put(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD, dimensionExpression.getOutputName());
            theContext.put(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX, timestampDimensionIndexInDimensions);
            theContext.put(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY, queryGranularity);
        }
    }
    if (queryGranularity == null) {
        return query;
    }
    return query.withOverriddenContext(theContext);
}

Example 4 with DataSource

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

the class DruidQuery method toTopNQuery.

/**
 * Return this query as a TopN query, or null if this query is not compatible with TopN.
 *
 * @return query or null
 */
@Nullable
private TopNQuery toTopNQuery(final QueryFeatureInspector queryFeatureInspector) {
    // Must be allowed by the QueryMaker.
    if (!queryFeatureInspector.feature(QueryFeature.CAN_RUN_TOPN)) {
        return null;
    }
    // Must have GROUP BY one column, no GROUPING SETS, ORDER BY ≤ 1 column, LIMIT > 0 and ≤ maxTopNLimit,
    // no OFFSET, no HAVING.
    final boolean topNOk = grouping != null && grouping.getDimensions().size() == 1 && !grouping.getSubtotals().hasEffect(grouping.getDimensionSpecs()) && sorting != null && (sorting.getOrderBys().size() <= 1 && sorting.getOffsetLimit().hasLimit() && sorting.getOffsetLimit().getLimit() > 0 && sorting.getOffsetLimit().getLimit() <= plannerContext.getPlannerConfig().getMaxTopNLimit() && !sorting.getOffsetLimit().hasOffset()) && grouping.getHavingFilter() == null;
    if (!topNOk) {
        return null;
    }
    final DimensionSpec dimensionSpec = Iterables.getOnlyElement(grouping.getDimensions()).toDimensionSpec();
    // grouping col cannot be type array
    if (dimensionSpec.getOutputType().isArray()) {
        return null;
    }
    final OrderByColumnSpec limitColumn;
    if (sorting.getOrderBys().isEmpty()) {
        limitColumn = new OrderByColumnSpec(dimensionSpec.getOutputName(), OrderByColumnSpec.Direction.ASCENDING, Calcites.getStringComparatorForValueType(dimensionSpec.getOutputType()));
    } else {
        limitColumn = Iterables.getOnlyElement(sorting.getOrderBys());
    }
    final TopNMetricSpec topNMetricSpec;
    if (limitColumn.getDimension().equals(dimensionSpec.getOutputName())) {
        // DimensionTopNMetricSpec is exact; always return it even if allowApproximate is false.
        final DimensionTopNMetricSpec baseMetricSpec = new DimensionTopNMetricSpec(null, limitColumn.getDimensionComparator());
        topNMetricSpec = limitColumn.getDirection() == OrderByColumnSpec.Direction.ASCENDING ? baseMetricSpec : new InvertedTopNMetricSpec(baseMetricSpec);
    } else if (plannerContext.getPlannerConfig().isUseApproximateTopN()) {
        // ORDER BY metric
        final NumericTopNMetricSpec baseMetricSpec = new NumericTopNMetricSpec(limitColumn.getDimension());
        topNMetricSpec = limitColumn.getDirection() == OrderByColumnSpec.Direction.ASCENDING ? new InvertedTopNMetricSpec(baseMetricSpec) : baseMetricSpec;
    } else {
        return null;
    }
    final Pair<DataSource, Filtration> dataSourceFiltrationPair = getFiltration(dataSource, filter, virtualColumnRegistry);
    final DataSource newDataSource = dataSourceFiltrationPair.lhs;
    final Filtration filtration = dataSourceFiltrationPair.rhs;
    final List<PostAggregator> postAggregators = new ArrayList<>(grouping.getPostAggregators());
    if (sorting.getProjection() != null) {
        postAggregators.addAll(sorting.getProjection().getPostAggregators());
    }
    return new TopNQuery(newDataSource, getVirtualColumns(true), dimensionSpec, topNMetricSpec, Ints.checkedCast(sorting.getOffsetLimit().getLimit()), filtration.getQuerySegmentSpec(), filtration.getDimFilter(), Granularities.ALL, grouping.getAggregatorFactories(), postAggregators, ImmutableSortedMap.copyOf(plannerContext.getQueryContext()));
}

Example 5 with DataSource

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

the class DruidUnionDataSourceRel method toDruidQuery.

@Override
public DruidQuery toDruidQuery(final boolean finalizeAggregations) {
    final List<TableDataSource> dataSources = new ArrayList<>();
    RowSignature signature = null;
    for (final RelNode relNode : unionRel.getInputs()) {
        final DruidRel<?> druidRel = (DruidRel<?>) relNode;
        if (!DruidRels.isScanOrMapping(druidRel, false)) {
            getPlannerContext().setPlanningError("SQL requires union between inputs that are not simple table scans " + "and involve a filter or aliasing");
            throw new CannotBuildQueryException(druidRel);
        }
        final DruidQuery query = druidRel.toDruidQuery(false);
        final DataSource dataSource = query.getDataSource();
        if (!(dataSource instanceof TableDataSource)) {
            getPlannerContext().setPlanningError("SQL requires union with input of '%s' type that is not supported." + " Union operation is only supported between regular tables. ", dataSource.getClass().getSimpleName());
            throw new CannotBuildQueryException(druidRel);
        }
        if (signature == null) {
            signature = query.getOutputRowSignature();
        }
        if (signature.getColumnNames().equals(query.getOutputRowSignature().getColumnNames())) {
            dataSources.add((TableDataSource) dataSource);
        } else {
            getPlannerContext().setPlanningError("There is a mismatch between the output row signature of input tables and the row signature of union output.");
            throw new CannotBuildQueryException(druidRel);
        }
    }
    if (signature == null) {
        // No inputs.
        throw new CannotBuildQueryException(unionRel);
    }
    // creation time.
    if (!signature.getColumnNames().equals(unionColumnNames)) {
        throw new CannotBuildQueryException(unionRel);
    }
    return partialQuery.build(new UnionDataSource(dataSources), signature, getPlannerContext(), getCluster().getRexBuilder(), finalizeAggregations);
}