use of io.trino.connector.MockConnectorColumnHandle 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.MockConnectorColumnHandle 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.MockConnectorColumnHandle in project trino by trinodb.
the class TestPushDistinctLimitIntoTableScan method testNoEffect.
@Test
public void testNoEffect() {
AtomicInteger applyCallCounter = new AtomicInteger();
testApplyAggregation = (session, handle, aggregates, assignments, groupingSets) -> {
applyCallCounter.incrementAndGet();
return Optional.empty();
};
tester().assertThat(rule).on(p -> {
Symbol regionkey = p.symbol("regionkey");
return p.distinctLimit(10, List.of(regionkey), p.tableScan(tableHandle, ImmutableList.of(regionkey), ImmutableMap.of(regionkey, new MockConnectorColumnHandle("regionkey", BIGINT))));
}).doesNotFire();
assertThat(applyCallCounter).as("applyCallCounter").hasValue(1);
}
use of io.trino.connector.MockConnectorColumnHandle in project trino by trinodb.
the class TestPruneTableScanColumns method testPushColumnPruningProjection.
@Test
public void testPushColumnPruningProjection() {
try (RuleTester ruleTester = defaultRuleTester()) {
String mockCatalog = "mock_catalog";
String testSchema = "test_schema";
String testTable = "test_table";
SchemaTableName testSchemaTable = new SchemaTableName(testSchema, testTable);
ColumnHandle columnHandleA = new MockConnectorColumnHandle("cola", DATE);
ColumnHandle columnHandleB = new MockConnectorColumnHandle("colb", DOUBLE);
Map<String, ColumnHandle> assignments = ImmutableMap.of("cola", columnHandleA, "colb", columnHandleB);
// Create catalog with applyProjection
MockConnectorFactory factory = MockConnectorFactory.builder().withListSchemaNames(connectorSession -> ImmutableList.of(testSchema)).withListTables((connectorSession, schema) -> testSchema.equals(schema) ? ImmutableList.of(testSchemaTable) : ImmutableList.of()).withGetColumns(schemaTableName -> assignments.entrySet().stream().map(entry -> new ColumnMetadata(entry.getKey(), ((MockConnectorColumnHandle) entry.getValue()).getType())).collect(toImmutableList())).withApplyProjection(this::mockApplyProjection).build();
ruleTester.getQueryRunner().createCatalog(mockCatalog, factory, ImmutableMap.of());
ruleTester.assertThat(new PruneTableScanColumns(ruleTester.getMetadata())).on(p -> {
Symbol symbolA = p.symbol("cola", DATE);
Symbol symbolB = p.symbol("colb", DOUBLE);
return p.project(Assignments.of(p.symbol("x"), symbolB.toSymbolReference()), p.tableScan(new TableHandle(new CatalogName(mockCatalog), new MockConnectorTableHandle(testSchemaTable), MockConnectorTransactionHandle.INSTANCE), ImmutableList.of(symbolA, symbolB), ImmutableMap.of(symbolA, columnHandleA, symbolB, columnHandleB)));
}).withSession(testSessionBuilder().setCatalog(mockCatalog).setSchema(testSchema).build()).matches(strictProject(ImmutableMap.of("expr", PlanMatchPattern.expression("COLB")), tableScan(new MockConnectorTableHandle(testSchemaTable, TupleDomain.all(), Optional.of(ImmutableList.of(columnHandleB)))::equals, TupleDomain.all(), ImmutableMap.of("COLB", columnHandleB::equals))));
}
}
use of io.trino.connector.MockConnectorColumnHandle in project trino by trinodb.
the class TestPruneTableScanColumns method mockApplyProjection.
private Optional<ProjectionApplicationResult<ConnectorTableHandle>> mockApplyProjection(ConnectorSession session, ConnectorTableHandle tableHandle, List<ConnectorExpression> projections, Map<String, ColumnHandle> assignments) {
MockConnectorTableHandle handle = (MockConnectorTableHandle) tableHandle;
List<Variable> variables = projections.stream().map(Variable.class::cast).collect(toImmutableList());
List<ColumnHandle> newColumns = variables.stream().map(variable -> assignments.get(variable.getName())).collect(toImmutableList());
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)), projections, variables.stream().map(variable -> new Assignment(variable.getName(), assignments.get(variable.getName()), ((MockConnectorColumnHandle) assignments.get(variable.getName())).getType())).collect(toImmutableList()), false));
}
Aggregations