use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.
the class TestTableScanRedirectionWithPushdown method testRedirectionAfterPredicatePushdownIntoTableScan.
@Test
public void testRedirectionAfterPredicatePushdownIntoTableScan() {
// the connector can detect a filter
try (LocalQueryRunner queryRunner = createLocalQueryRunner(getMockApplyRedirectAfterPredicatePushdown(REDIRECTION_MAPPING_A, Optional.empty()), Optional.empty(), Optional.of(getMockApplyFilter(ImmutableSet.of(SOURCE_COLUMN_HANDLE_A, DESTINATION_COLUMN_HANDLE_A))))) {
assertPlan(queryRunner, "SELECT source_col_a FROM test_table WHERE source_col_a = 1", output(ImmutableList.of("DEST_COL"), tableScan(new MockConnectorTableHandle(DESTINATION_TABLE, TupleDomain.withColumnDomains(ImmutableMap.of(DESTINATION_COLUMN_HANDLE_A, singleValue(INTEGER, 1L))), Optional.empty())::equals, TupleDomain.withColumnDomains(ImmutableMap.of(DESTINATION_COLUMN_HANDLE_A::equals, singleValue(INTEGER, 1L))), ImmutableMap.of("DEST_COL", DESTINATION_COLUMN_HANDLE_A::equals))));
assertPlan(queryRunner, "SELECT source_col_a FROM test_table", output(ImmutableList.of("SOURCE_COL"), tableScan(TEST_TABLE, ImmutableMap.of("SOURCE_COL", SOURCE_COLUMN_NAME_A))));
}
}
use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.
the class TestTableScanRedirectionWithPushdown method createLocalQueryRunner.
private LocalQueryRunner createLocalQueryRunner(ApplyTableScanRedirect applyTableScanRedirect, Optional<ApplyProjection> applyProjection, Optional<ApplyFilter> applyFilter) {
LocalQueryRunner queryRunner = LocalQueryRunner.create(MOCK_SESSION);
MockConnectorFactory.Builder builder = MockConnectorFactory.builder().withGetTableHandle((session, schemaTableName) -> new MockConnectorTableHandle(schemaTableName)).withGetColumns(name -> {
if (name.equals(SOURCE_TABLE)) {
return ImmutableList.of(new ColumnMetadata(SOURCE_COLUMN_NAME_A, INTEGER), new ColumnMetadata(SOURCE_COLUMN_NAME_B, INTEGER), new ColumnMetadata(SOURCE_COLUMN_NAME_C, VARCHAR), new ColumnMetadata(SOURCE_COLUMN_NAME_D, ROW_TYPE));
} else if (name.equals(DESTINATION_TABLE)) {
return ImmutableList.of(new ColumnMetadata(DESTINATION_COLUMN_NAME_A, INTEGER), new ColumnMetadata(DESTINATION_COLUMN_NAME_B, INTEGER), new ColumnMetadata(DESTINATION_COLUMN_NAME_C, ROW_TYPE), new ColumnMetadata(DESTINATION_COLUMN_NAME_D, BOGUS));
}
throw new IllegalArgumentException();
}).withApplyTableScanRedirect(applyTableScanRedirect);
applyProjection.ifPresent(builder::withApplyProjection);
applyFilter.ifPresent(builder::withApplyFilter);
queryRunner.createCatalog(MOCK_CATALOG, builder.build(), ImmutableMap.of());
return queryRunner;
}
use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.
the class TestTableScanRedirectionWithPushdown method testRedirectionBeforeDeferencePushdown.
@Test
public void testRedirectionBeforeDeferencePushdown() {
// the connector can detect that source_col_a and source_col_d is projected
try (LocalQueryRunner queryRunner = createLocalQueryRunner(mockApplyRedirectAfterProjectionPushdown(ROW_TYPE_REDIRECTION_MAPPING_AD, Optional.of(ImmutableSet.of(SOURCE_COLUMN_HANDLE_A, SOURCE_COLUMN_HANDLE_D))), Optional.of(this::mockApplyProjection), Optional.empty())) {
// Pushdown of dereference for source_col_d.a into table scan results in a new column handle
// Table scan redirection would not take place if dereference pushdown has already taken place before redirection
ColumnHandle destinationColumnHandleC0 = new MockConnectorColumnHandle(DESTINATION_COLUMN_NAME_C + "#0", BIGINT);
assertPlan(queryRunner, "SELECT source_col_a, source_col_d.a FROM test_table", output(ImmutableList.of("DEST_COL_A", "DEST_COL_C#0"), tableScan(new MockConnectorTableHandle(DESTINATION_TABLE, TupleDomain.all(), Optional.of(ImmutableList.of(DESTINATION_COLUMN_HANDLE_A, destinationColumnHandleC0)))::equals, TupleDomain.all(), ImmutableMap.of("DEST_COL_A", DESTINATION_COLUMN_HANDLE_A::equals, "DEST_COL_C#0", destinationColumnHandleC0::equals))));
}
}
use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.
the class TestTableScanRedirectionWithPushdown method mockApplyProjection.
private Optional<ProjectionApplicationResult<ConnectorTableHandle>> mockApplyProjection(ConnectorSession session, ConnectorTableHandle tableHandle, List<ConnectorExpression> projections, Map<String, ColumnHandle> assignments) {
MockConnectorTableHandle handle = (MockConnectorTableHandle) tableHandle;
ImmutableList.Builder<ColumnHandle> newColumnsBuilder = ImmutableList.builder();
ImmutableList.Builder<ConnectorExpression> outputExpressions = ImmutableList.builder();
ImmutableList.Builder<Assignment> outputAssignments = ImmutableList.builder();
for (ConnectorExpression projection : projections) {
String newVariableName;
ColumnHandle newColumnHandle;
if (projection instanceof Variable) {
newVariableName = ((Variable) projection).getName();
newColumnHandle = assignments.get(newVariableName);
} else if (projection instanceof FieldDereference) {
FieldDereference dereference = (FieldDereference) projection;
if (!(dereference.getTarget() instanceof Variable)) {
throw new UnsupportedOperationException();
}
String dereferenceTargetName = ((Variable) dereference.getTarget()).getName();
newVariableName = ((MockConnectorColumnHandle) assignments.get(dereferenceTargetName)).getName() + "#" + dereference.getField();
newColumnHandle = new MockConnectorColumnHandle(newVariableName, projection.getType());
} else {
throw new UnsupportedOperationException();
}
Variable newVariable = new Variable(newVariableName, projection.getType());
newColumnsBuilder.add(newColumnHandle);
outputExpressions.add(newVariable);
outputAssignments.add(new Assignment(newVariableName, newColumnHandle, projection.getType()));
}
List<ColumnHandle> newColumns = newColumnsBuilder.build();
if (handle.getColumns().isPresent() && newColumns.equals(handle.getColumns().get())) {
return Optional.empty();
}
return Optional.of(new ProjectionApplicationResult<>(new MockConnectorTableHandle(handle.getTableName(), handle.getConstraint(), Optional.of(newColumns)), outputExpressions.build(), outputAssignments.build(), false));
}
use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.
the class TestPushProjectionIntoTableScan method mockApplyProjection.
private Optional<ProjectionApplicationResult<ConnectorTableHandle>> mockApplyProjection(ConnectorSession session, ConnectorTableHandle tableHandle, List<ConnectorExpression> projections, Map<String, ColumnHandle> assignments) {
// Prepare new table handle
SchemaTableName inputSchemaTableName = ((MockConnectorTableHandle) tableHandle).getTableName();
SchemaTableName projectedTableName = new SchemaTableName(inputSchemaTableName.getSchemaName(), "projected_" + inputSchemaTableName.getTableName());
// Prepare new column handles
ImmutableList.Builder<ConnectorExpression> outputExpressions = ImmutableList.builder();
ImmutableList.Builder<Assignment> outputAssignments = ImmutableList.builder();
ImmutableList.Builder<ColumnHandle> projectedColumnsBuilder = ImmutableList.builder();
for (ConnectorExpression projection : projections) {
String variablePrefix;
if (projection instanceof Variable) {
variablePrefix = "projected_variable_";
} else if (projection instanceof FieldDereference) {
variablePrefix = "projected_dereference_";
} else if (projection instanceof Constant) {
variablePrefix = "projected_constant_";
} else if (projection instanceof Call) {
variablePrefix = "projected_call_";
} else {
throw new UnsupportedOperationException();
}
String newVariableName = variablePrefix + projection.toString();
Variable newVariable = new Variable(newVariableName, projection.getType());
ColumnHandle newColumnHandle = new TpchColumnHandle(newVariableName, projection.getType());
outputExpressions.add(newVariable);
outputAssignments.add(new Assignment(newVariableName, newColumnHandle, projection.getType()));
projectedColumnsBuilder.add(newColumnHandle);
}
return Optional.of(new ProjectionApplicationResult<>(new MockConnectorTableHandle(projectedTableName, TupleDomain.all(), Optional.of(projectedColumnsBuilder.build())), outputExpressions.build(), outputAssignments.build(), false));
}
Aggregations