Examples with SupportsProjectionPushDown org.apache.flink.table.connector.source.abilities.SupportsProjectionPushDown used on opensource projects

Example 1 with SupportsProjectionPushDown

use of org.apache.flink.table.connector.source.abilities.SupportsProjectionPushDown in project flink by apache.

the class PushProjectIntoTableSourceScanRule method performPushDown.

private RowType performPushDown(TableSourceTable source, NestedSchema projectedSchema, RowType producedType, List<SourceAbilitySpec> abilitySpecs) {
    final int numPhysicalColumns;
    final List<NestedColumn> projectedMetadataColumns;
    if (supportsMetadata(source.tableSource())) {
        final List<String> declaredMetadataKeys = createRequiredMetadataKeys(source.contextResolvedTable().getResolvedSchema(), source.tableSource());
        numPhysicalColumns = producedType.getFieldCount() - declaredMetadataKeys.size();
        projectedMetadataColumns = IntStream.range(0, declaredMetadataKeys.size()).mapToObj(i -> producedType.getFieldNames().get(numPhysicalColumns + i)).map(fieldName -> projectedSchema.columns().get(fieldName)).filter(Objects::nonNull).collect(Collectors.toList());
    } else {
        numPhysicalColumns = producedType.getFieldCount();
        projectedMetadataColumns = Collections.emptyList();
    }
    final int[][] physicalProjections;
    if (supportsProjectionPushDown(source.tableSource())) {
        projectedMetadataColumns.forEach(metaColumn -> projectedSchema.columns().remove(metaColumn.name()));
        physicalProjections = NestedProjectionUtil.convertToIndexArray(projectedSchema);
        projectedMetadataColumns.forEach(metaColumn -> projectedSchema.columns().put(metaColumn.name(), metaColumn));
    } else {
        physicalProjections = IntStream.range(0, numPhysicalColumns).mapToObj(columnIndex -> new int[] { columnIndex }).toArray(int[][]::new);
    }
    final int[][] projectedFields = Stream.concat(Stream.of(physicalProjections), projectedMetadataColumns.stream().map(NestedColumn::indexInOriginSchema).map(columnIndex -> new int[] { columnIndex })).toArray(int[][]::new);
    int newIndex = physicalProjections.length;
    for (NestedColumn metaColumn : projectedMetadataColumns) {
        metaColumn.setIndexOfLeafInNewSchema(newIndex++);
    }
    if (supportsProjectionPushDown(source.tableSource())) {
        final RowType projectedPhysicalType = (RowType) Projection.of(physicalProjections).project(producedType);
        abilitySpecs.add(new ProjectPushDownSpec(physicalProjections, projectedPhysicalType));
    }
    final RowType newProducedType = (RowType) Projection.of(projectedFields).project(producedType);
    if (supportsMetadata(source.tableSource())) {
        final List<String> projectedMetadataKeys = projectedMetadataColumns.stream().map(NestedColumn::name).collect(Collectors.toList());
        abilitySpecs.add(new ReadingMetadataSpec(projectedMetadataKeys, newProducedType));
    }
    return newProducedType;
}