Examples with PartialDruidQuery org.apache.druid.sql.calcite.rel.PartialDruidQuery used on opensource projects

Example 1 with PartialDruidQuery

use of org.apache.druid.sql.calcite.rel.PartialDruidQuery in project druid by druid-io.

the class DruidUnionDataSourceRule method getColumnNamesIfTableOrUnion.

static Optional<List<String>> getColumnNamesIfTableOrUnion(final DruidRel<?> druidRel, @Nullable PlannerContext plannerContext) {
    final PartialDruidQuery partialQuery = druidRel.getPartialDruidQuery();
    final Optional<DruidTable> druidTable = DruidRels.druidTableIfLeafRel(druidRel).filter(table -> table.getDataSource() instanceof TableDataSource);
    if (druidTable.isPresent() && DruidRels.isScanOrMapping(druidRel, false)) {
        if (partialQuery.stage() == PartialDruidQuery.Stage.SCAN) {
            return Optional.of(druidTable.get().getRowSignature().getColumnNames());
        } else {
            // Sanity check. Expected to be true due to the "scan or mapping" check.
            if (partialQuery.stage() != PartialDruidQuery.Stage.SELECT_PROJECT) {
                throw new ISE("Expected stage %s but got %s", PartialDruidQuery.Stage.SELECT_PROJECT, partialQuery.stage());
            }
            // Apply the mapping (with additional sanity checks).
            final RowSignature tableSignature = druidTable.get().getRowSignature();
            final Mappings.TargetMapping mapping = partialQuery.getSelectProject().getMapping();
            if (mapping.getSourceCount() != tableSignature.size()) {
                throw new ISE("Expected mapping with %d columns but got %d columns", tableSignature.size(), mapping.getSourceCount());
            }
            final List<String> retVal = new ArrayList<>();
            for (int i = 0; i < mapping.getTargetCount(); i++) {
                final int sourceField = mapping.getSourceOpt(i);
                retVal.add(tableSignature.getColumnName(sourceField));
            }
            return Optional.of(retVal);
        }
    } else if (!druidTable.isPresent() && druidRel instanceof DruidUnionDataSourceRel) {
        return Optional.of(((DruidUnionDataSourceRel) druidRel).getUnionColumnNames());
    } else if (druidTable.isPresent()) {
        if (null != plannerContext) {
            plannerContext.setPlanningError("SQL requires union between inputs that are not simple table scans " + "and involve a filter or aliasing. Or column types of tables being unioned are not of same type.");
        }
        return Optional.empty();
    } else {
        if (null != plannerContext) {
            plannerContext.setPlanningError("SQL requires union with input of a datasource type that is not supported." + " Union operation is only supported between regular tables. ");
        }
        return Optional.empty();
    }
}