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

Example 1 with DruidUnionDataSourceRel

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

the class DruidUnionDataSourceRuleTest method test_getColumnNamesIfTableOrUnion_unionProject.

@Test
public void test_getColumnNamesIfTableOrUnion_unionProject() {
    final DruidUnionDataSourceRel druidRel = DruidRelsTest.mockDruidRel(DruidUnionDataSourceRel.class, rel -> EasyMock.expect(rel.getUnionColumnNames()).andReturn(fooDruidTable.getRowSignature().getColumnNames()), PartialDruidQuery.Stage.SELECT_PROJECT, null, DruidRelsTest.mockNonMappingProject(), null);
    Assert.assertEquals(Optional.of(ImmutableList.of("__time", "col1", "col2")), DruidUnionDataSourceRule.getColumnNamesIfTableOrUnion(druidRel, null));
}

Example 2 with DruidUnionDataSourceRel

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

the class DruidUnionDataSourceRule method onMatch.

@Override
public void onMatch(final RelOptRuleCall call) {
    final Union unionRel = call.rel(0);
    final DruidRel<?> firstDruidRel = call.rel(1);
    final DruidQueryRel secondDruidRel = call.rel(2);
    if (firstDruidRel instanceof DruidUnionDataSourceRel) {
        // Unwrap and flatten the inputs to the Union.
        final RelNode newUnionRel = call.builder().pushAll(firstDruidRel.getInputs()).push(secondDruidRel).union(true, firstDruidRel.getInputs().size() + 1).build();
        call.transformTo(DruidUnionDataSourceRel.create((Union) newUnionRel, getColumnNamesIfTableOrUnion(firstDruidRel, plannerContext).get(), firstDruidRel.getPlannerContext()));
    } else {
        // Sanity check.
        if (!(firstDruidRel instanceof DruidQueryRel)) {
            throw new ISE("Expected first rel to be a DruidQueryRel, but it was %s", firstDruidRel.getClass().getName());
        }
        call.transformTo(DruidUnionDataSourceRel.create(unionRel, getColumnNamesIfTableOrUnion(firstDruidRel, plannerContext).get(), firstDruidRel.getPlannerContext()));
    }
}

Example 3 with DruidUnionDataSourceRel

use of org.apache.druid.sql.calcite.rel.DruidUnionDataSourceRel 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();
    }
}

Example 4 with DruidUnionDataSourceRel

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

the class DruidUnionDataSourceRuleTest method test_getColumnNamesIfTableOrUnion_unionScan.

@Test
public void test_getColumnNamesIfTableOrUnion_unionScan() {
    final DruidUnionDataSourceRel druidRel = DruidRelsTest.mockDruidRel(DruidUnionDataSourceRel.class, rel -> EasyMock.expect(rel.getUnionColumnNames()).andReturn(fooDruidTable.getRowSignature().getColumnNames()), PartialDruidQuery.Stage.SCAN, null, null, null);
    Assert.assertEquals(Optional.of(ImmutableList.of("__time", "col1", "col2")), DruidUnionDataSourceRule.getColumnNamesIfTableOrUnion(druidRel, null));
}

Example 5 with DruidUnionDataSourceRel

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

the class DruidUnionDataSourceRuleTest method test_getColumnNamesIfTableOrUnion_unionMapping.

@Test
public void test_getColumnNamesIfTableOrUnion_unionMapping() {
    final Project project = DruidRelsTest.mockMappingProject(ImmutableList.of(2, 1), 3);
    final Mappings.TargetMapping mapping = project.getMapping();
    final String[] mappedColumnNames = new String[mapping.getTargetCount()];
    final List<String> columnNames = fooDruidTable.getRowSignature().getColumnNames();
    for (int i = 0; i < columnNames.size(); i++) {
        mappedColumnNames[mapping.getTargetOpt(i)] = columnNames.get(i);
    }
    final DruidUnionDataSourceRel druidRel = DruidRelsTest.mockDruidRel(DruidUnionDataSourceRel.class, rel -> EasyMock.expect(rel.getUnionColumnNames()).andReturn(Arrays.asList(mappedColumnNames)), PartialDruidQuery.Stage.SELECT_PROJECT, null, project, null);
    Assert.assertEquals(Optional.of(ImmutableList.of("col2", "col1")), DruidUnionDataSourceRule.getColumnNamesIfTableOrUnion(druidRel, null));
}