Examples with OffsetLimit org.apache.druid.sql.calcite.planner.OffsetLimit used on opensource projects

Example 1 with OffsetLimit

use of org.apache.druid.sql.calcite.planner.OffsetLimit in project druid by druid-io.

the class DruidQuery method computeSorting.

@Nonnull
private static Sorting computeSorting(final PartialDruidQuery partialQuery, final PlannerContext plannerContext, final RowSignature rowSignature, @Nullable final VirtualColumnRegistry virtualColumnRegistry) {
    final Sort sort = Preconditions.checkNotNull(partialQuery.getSort(), "sort");
    final Project sortProject = partialQuery.getSortProject();
    // Extract limit and offset.
    final OffsetLimit offsetLimit = OffsetLimit.fromSort(sort);
    // Extract orderBy column specs.
    final List<OrderByColumnSpec> orderBys = new ArrayList<>(sort.getChildExps().size());
    for (int sortKey = 0; sortKey < sort.getChildExps().size(); sortKey++) {
        final RexNode sortExpression = sort.getChildExps().get(sortKey);
        final RelFieldCollation collation = sort.getCollation().getFieldCollations().get(sortKey);
        final OrderByColumnSpec.Direction direction;
        final StringComparator comparator;
        if (collation.getDirection() == RelFieldCollation.Direction.ASCENDING) {
            direction = OrderByColumnSpec.Direction.ASCENDING;
        } else if (collation.getDirection() == RelFieldCollation.Direction.DESCENDING) {
            direction = OrderByColumnSpec.Direction.DESCENDING;
        } else {
            throw new ISE("Don't know what to do with direction[%s]", collation.getDirection());
        }
        final SqlTypeName sortExpressionType = sortExpression.getType().getSqlTypeName();
        if (SqlTypeName.NUMERIC_TYPES.contains(sortExpressionType) || SqlTypeName.TIMESTAMP == sortExpressionType || SqlTypeName.DATE == sortExpressionType) {
            comparator = StringComparators.NUMERIC;
        } else {
            comparator = StringComparators.LEXICOGRAPHIC;
        }
        if (sortExpression.isA(SqlKind.INPUT_REF)) {
            final RexInputRef ref = (RexInputRef) sortExpression;
            final String fieldName = rowSignature.getColumnName(ref.getIndex());
            orderBys.add(new OrderByColumnSpec(fieldName, direction, comparator));
        } else {
            // We don't support sorting by anything other than refs which actually appear in the query result.
            throw new CannotBuildQueryException(sort, sortExpression);
        }
    }
    // Extract any post-sort Projection.
    final Projection projection;
    if (sortProject == null) {
        projection = null;
    } else if (partialQuery.getAggregate() == null) {
        if (virtualColumnRegistry == null) {
            throw new ISE("Must provide 'virtualColumnRegistry' for pre-aggregation Projection!");
        }
        projection = Projection.preAggregation(sortProject, plannerContext, rowSignature, virtualColumnRegistry);
    } else {
        projection = Projection.postAggregation(sortProject, plannerContext, rowSignature, "s");
    }
    return Sorting.create(orderBys, offsetLimit, projection);
}

Example 2 with OffsetLimit

use of org.apache.druid.sql.calcite.planner.OffsetLimit in project druid by druid-io.

the class SortCollapseRule method onMatch.

@Override
public void onMatch(final RelOptRuleCall call) {
    final Sort outerSort = call.rel(0);
    final Sort innerSort = call.rel(1);
    if (outerSort.collation.getFieldCollations().isEmpty() || outerSort.collation.getFieldCollations().equals(innerSort.collation.getFieldCollations())) {
        final OffsetLimit offsetLimit = OffsetLimit.fromSort(innerSort).andThen(OffsetLimit.fromSort(outerSort));
        final Sort combined = innerSort.copy(innerSort.getTraitSet(), innerSort.getInput(), innerSort.getCollation(), offsetLimit.getOffsetAsRexNode(call.builder().getRexBuilder()), offsetLimit.getLimitAsRexNode(call.builder().getRexBuilder()));
        call.transformTo(combined);
        call.getPlanner().setImportance(outerSort, 0.0);
    }
}