Examples with DimensionExpression org.apache.druid.sql.calcite.aggregation.DimensionExpression used on opensource projects

Example 6 with DimensionExpression

use of org.apache.druid.sql.calcite.aggregation.DimensionExpression in project druid by druid-io.

the class DruidQuery method getVirtualColumns.

private VirtualColumns getVirtualColumns(final boolean includeDimensions) {
    // 'sourceRowSignature' could provide a list of all defined virtual columns while constructing a query, but we
    // still want to collect the set of VirtualColumns this way to ensure we only add what is still being used after
    // the various transforms and optimizations
    Set<VirtualColumn> virtualColumns = new HashSet<>();
    // rewrite any "specialized" virtual column expressions as top level virtual columns so that their native
    // implementation can be used instead of being composed as part of some expression tree in an expresson virtual
    // column
    Set<String> specialized = new HashSet<>();
    virtualColumnRegistry.visitAllSubExpressions((expression) -> {
        switch(expression.getType()) {
            case SPECIALIZED:
                // add the expression to the top level of the registry as a standalone virtual column
                final String name = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(expression, expression.getDruidType());
                specialized.add(name);
                // replace with an identifier expression of the new virtual column name
                return DruidExpression.ofColumn(expression.getDruidType(), name);
            default:
                // do nothing
                return expression;
        }
    });
    // we always want to add any virtual columns used by the query level DimFilter
    if (filter != null) {
        for (String columnName : filter.getRequiredColumns()) {
            if (virtualColumnRegistry.isVirtualColumnDefined(columnName)) {
                virtualColumns.add(virtualColumnRegistry.getVirtualColumn(columnName));
            }
        }
    }
    if (selectProjection != null) {
        for (String columnName : selectProjection.getVirtualColumns()) {
            if (virtualColumnRegistry.isVirtualColumnDefined(columnName)) {
                virtualColumns.add(virtualColumnRegistry.getVirtualColumn(columnName));
            }
        }
    }
    if (grouping != null) {
        if (includeDimensions) {
            for (DimensionExpression expression : grouping.getDimensions()) {
                if (virtualColumnRegistry.isVirtualColumnDefined(expression.getVirtualColumn())) {
                    virtualColumns.add(virtualColumnRegistry.getVirtualColumn(expression.getVirtualColumn()));
                }
            }
        }
        for (Aggregation aggregation : grouping.getAggregations()) {
            virtualColumns.addAll(virtualColumnRegistry.getAllVirtualColumns(aggregation.getRequiredColumns()));
        }
    }
    if (sorting != null && sorting.getProjection() != null && grouping == null) {
        for (String columnName : sorting.getProjection().getVirtualColumns()) {
            if (virtualColumnRegistry.isVirtualColumnDefined(columnName)) {
                virtualColumns.add(virtualColumnRegistry.getVirtualColumn(columnName));
            }
        }
    }
    if (dataSource instanceof JoinDataSource) {
        for (String expression : ((JoinDataSource) dataSource).getVirtualColumnCandidates()) {
            if (virtualColumnRegistry.isVirtualColumnDefined(expression)) {
                virtualColumns.add(virtualColumnRegistry.getVirtualColumn(expression));
            }
        }
    }
    for (String columnName : specialized) {
        if (virtualColumnRegistry.isVirtualColumnDefined(columnName)) {
            virtualColumns.add(virtualColumnRegistry.getVirtualColumn(columnName));
        }
    }
    // sort for predictable output
    List<VirtualColumn> columns = new ArrayList<>(virtualColumns);
    columns.sort(Comparator.comparing(VirtualColumn::getOutputName));
    return VirtualColumns.create(columns);
}