Examples with SourceAbilitySpec org.apache.flink.table.planner.plan.abilities.source.SourceAbilitySpec used on opensource projects

Example 1 with SourceAbilitySpec

use of org.apache.flink.table.planner.plan.abilities.source.SourceAbilitySpec in project flink by apache.

the class DynamicTableSourceSpec method getTableSource.

private DynamicTableSource getTableSource(FlinkContext flinkContext) {
    if (tableSource == null) {
        final DynamicTableSourceFactory factory = flinkContext.getModuleManager().getFactory(Module::getTableSourceFactory).orElse(null);
        tableSource = FactoryUtil.createDynamicTableSource(factory, contextResolvedTable.getIdentifier(), contextResolvedTable.getResolvedTable(), loadOptionsFromCatalogTable(contextResolvedTable, flinkContext), flinkContext.getTableConfig().getConfiguration(), flinkContext.getClassLoader(), contextResolvedTable.isTemporary());
        if (sourceAbilities != null) {
            RowType newProducedType = (RowType) contextResolvedTable.getResolvedSchema().toSourceRowDataType().getLogicalType();
            for (SourceAbilitySpec spec : sourceAbilities) {
                SourceAbilityContext context = new SourceAbilityContext(flinkContext, newProducedType);
                spec.apply(tableSource, context);
                if (spec.getProducedType().isPresent()) {
                    newProducedType = spec.getProducedType().get();
                }
            }
        }
    }
    return tableSource;
}

Example 2 with SourceAbilitySpec

use of org.apache.flink.table.planner.plan.abilities.source.SourceAbilitySpec in project flink by apache.

the class TemporalTableSourceSpec method getTemporalTable.

@JsonIgnore
public RelOptTable getTemporalTable(FlinkContext flinkContext) {
    if (null != temporalTable) {
        return temporalTable;
    }
    if (null != tableSourceSpec && null != outputType) {
        LookupTableSource lookupTableSource = tableSourceSpec.getLookupTableSource(flinkContext);
        SourceAbilitySpec[] sourceAbilitySpecs = null;
        if (null != tableSourceSpec.getSourceAbilities()) {
            sourceAbilitySpecs = tableSourceSpec.getSourceAbilities().toArray(new SourceAbilitySpec[0]);
        }
        return new TableSourceTable(null, outputType, FlinkStatistic.UNKNOWN(), lookupTableSource, true, tableSourceSpec.getContextResolvedTable(), flinkContext, sourceAbilitySpecs);
    }
    throw new TableException("Can not obtain temporalTable correctly!");
}

Example 3 with SourceAbilitySpec

use of org.apache.flink.table.planner.plan.abilities.source.SourceAbilitySpec in project flink by apache.

the class PushPartitionIntoTableSourceScanRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    Filter filter = call.rel(0);
    LogicalTableScan scan = call.rel(1);
    TableSourceTable tableSourceTable = scan.getTable().unwrap(TableSourceTable.class);
    RelDataType inputFieldTypes = filter.getInput().getRowType();
    List<String> inputFieldNames = inputFieldTypes.getFieldNames();
    List<String> partitionFieldNames = tableSourceTable.contextResolvedTable().<ResolvedCatalogTable>getResolvedTable().getPartitionKeys();
    // extract partition predicates
    RelBuilder relBuilder = call.builder();
    RexBuilder rexBuilder = relBuilder.getRexBuilder();
    Tuple2<Seq<RexNode>, Seq<RexNode>> allPredicates = RexNodeExtractor.extractPartitionPredicateList(filter.getCondition(), FlinkRelOptUtil.getMaxCnfNodeCount(scan), inputFieldNames.toArray(new String[0]), rexBuilder, partitionFieldNames.toArray(new String[0]));
    RexNode partitionPredicate = RexUtil.composeConjunction(rexBuilder, JavaConversions.seqAsJavaList(allPredicates._1));
    if (partitionPredicate.isAlwaysTrue()) {
        return;
    }
    // build pruner
    LogicalType[] partitionFieldTypes = partitionFieldNames.stream().map(name -> {
        int index = inputFieldNames.indexOf(name);
        if (index < 0) {
            throw new TableException(String.format("Partitioned key '%s' isn't found in input columns. " + "Validator should have checked that.", name));
        }
        return inputFieldTypes.getFieldList().get(index).getType();
    }).map(FlinkTypeFactory::toLogicalType).toArray(LogicalType[]::new);
    RexNode finalPartitionPredicate = adjustPartitionPredicate(inputFieldNames, partitionFieldNames, partitionPredicate);
    FlinkContext context = ShortcutUtils.unwrapContext(scan);
    Function<List<Map<String, String>>, List<Map<String, String>>> defaultPruner = partitions -> PartitionPruner.prunePartitions(context.getTableConfig(), partitionFieldNames.toArray(new String[0]), partitionFieldTypes, partitions, finalPartitionPredicate);
    // prune partitions
    List<Map<String, String>> remainingPartitions = readPartitionsAndPrune(rexBuilder, context, tableSourceTable, defaultPruner, allPredicates._1(), inputFieldNames);
    // apply push down
    DynamicTableSource dynamicTableSource = tableSourceTable.tableSource().copy();
    PartitionPushDownSpec partitionPushDownSpec = new PartitionPushDownSpec(remainingPartitions);
    partitionPushDownSpec.apply(dynamicTableSource, SourceAbilityContext.from(scan));
    // build new statistic
    TableStats newTableStat = null;
    if (tableSourceTable.contextResolvedTable().isPermanent()) {
        ObjectIdentifier identifier = tableSourceTable.contextResolvedTable().getIdentifier();
        ObjectPath tablePath = identifier.toObjectPath();
        Catalog catalog = tableSourceTable.contextResolvedTable().getCatalog().get();
        for (Map<String, String> partition : remainingPartitions) {
            Optional<TableStats> partitionStats = getPartitionStats(catalog, tablePath, partition);
            if (!partitionStats.isPresent()) {
                // clear all information before
                newTableStat = null;
                break;
            } else {
                newTableStat = newTableStat == null ? partitionStats.get() : newTableStat.merge(partitionStats.get());
            }
        }
    }
    FlinkStatistic newStatistic = FlinkStatistic.builder().statistic(tableSourceTable.getStatistic()).tableStats(newTableStat).build();
    TableSourceTable newTableSourceTable = tableSourceTable.copy(dynamicTableSource, newStatistic, new SourceAbilitySpec[] { partitionPushDownSpec });
    LogicalTableScan newScan = LogicalTableScan.create(scan.getCluster(), newTableSourceTable, scan.getHints());
    // transform to new node
    RexNode nonPartitionPredicate = RexUtil.composeConjunction(rexBuilder, JavaConversions.seqAsJavaList(allPredicates._2()));
    if (nonPartitionPredicate.isAlwaysTrue()) {
        call.transformTo(newScan);
    } else {
        Filter newFilter = filter.copy(filter.getTraitSet(), newScan, nonPartitionPredicate);
        call.transformTo(newFilter);
    }
}

Example 4 with SourceAbilitySpec

use of org.apache.flink.table.planner.plan.abilities.source.SourceAbilitySpec in project flink by apache.

the class PushWatermarkIntoTableSourceScanRuleBase method getNewScan.

/**
 * It uses the input watermark expression to generate the {@link WatermarkGeneratorSupplier}.
 * After the {@link WatermarkStrategy} is pushed into the scan, it will build a new scan.
 * However, when {@link FlinkLogicalWatermarkAssigner} is the parent of the {@link
 * FlinkLogicalTableSourceScan} it should modify the rowtime type to keep the type of plan is
 * consistent. In other cases, it just keep the data type of the scan as same as before and
 * leave the work when rewriting the projection.
 *
 * <p>NOTES: the row type of the scan is not always as same as the watermark assigner. Because
 * the scan will not add the rowtime column into the row when pushing the watermark assigner
 * into the scan. In some cases, query may have computed columns defined on rowtime column. If
 * modifying the type of the rowtime(with time attribute), it will also influence the type of
 * the computed column. Therefore, if the watermark assigner is not the parent of the scan, set
 * the type of the scan as before and leave the work to projection.
 */
protected FlinkLogicalTableSourceScan getNewScan(FlinkLogicalWatermarkAssigner watermarkAssigner, RexNode watermarkExpr, FlinkLogicalTableSourceScan scan, TableConfig tableConfig, boolean useWatermarkAssignerRowType) {
    final TableSourceTable tableSourceTable = scan.getTable().unwrap(TableSourceTable.class);
    final DynamicTableSource newDynamicTableSource = tableSourceTable.tableSource().copy();
    final boolean isSourceWatermark = newDynamicTableSource instanceof SupportsSourceWatermark && hasSourceWatermarkDeclaration(watermarkExpr);
    final RelDataType newType;
    if (useWatermarkAssignerRowType) {
        // project is trivial and set rowtime type in scan
        newType = watermarkAssigner.getRowType();
    } else {
        // project add/delete columns and set the rowtime column type in project
        newType = scan.getRowType();
    }
    final RowType producedType = (RowType) FlinkTypeFactory.toLogicalType(newType);
    final SourceAbilityContext abilityContext = SourceAbilityContext.from(scan);
    final SourceAbilitySpec abilitySpec;
    if (isSourceWatermark) {
        final SourceWatermarkSpec sourceWatermarkSpec = new SourceWatermarkSpec(true, producedType);
        sourceWatermarkSpec.apply(newDynamicTableSource, abilityContext);
        abilitySpec = sourceWatermarkSpec;
    } else {
        final Duration idleTimeout = tableConfig.getConfiguration().get(ExecutionConfigOptions.TABLE_EXEC_SOURCE_IDLE_TIMEOUT);
        final long idleTimeoutMillis;
        if (!idleTimeout.isZero() && !idleTimeout.isNegative()) {
            idleTimeoutMillis = idleTimeout.toMillis();
        } else {
            idleTimeoutMillis = -1L;
        }
        final WatermarkPushDownSpec watermarkPushDownSpec = new WatermarkPushDownSpec(watermarkExpr, idleTimeoutMillis, producedType);
        watermarkPushDownSpec.apply(newDynamicTableSource, abilityContext);
        abilitySpec = watermarkPushDownSpec;
    }
    TableSourceTable newTableSourceTable = tableSourceTable.copy(newDynamicTableSource, newType, new SourceAbilitySpec[] { abilitySpec });
    return FlinkLogicalTableSourceScan.create(scan.getCluster(), scan.getHints(), newTableSourceTable);
}

Example 5 with SourceAbilitySpec

use of org.apache.flink.table.planner.plan.abilities.source.SourceAbilitySpec in project flink by apache.

the class PushLocalAggIntoScanRuleBase method pushLocalAggregateIntoScan.

protected void pushLocalAggregateIntoScan(RelOptRuleCall call, BatchPhysicalGroupAggregateBase localAgg, BatchPhysicalTableSourceScan oldScan, int[] calcRefFields) {
    RowType inputType = FlinkTypeFactory.toLogicalRowType(oldScan.getRowType());
    List<int[]> groupingSets = Collections.singletonList(ArrayUtils.addAll(localAgg.grouping(), localAgg.auxGrouping()));
    List<AggregateCall> aggCallList = JavaScalaConversionUtil.toJava(localAgg.getAggCallList());
    // map arg index in aggregate to field index in scan through referred fields by calc.
    if (calcRefFields != null) {
        groupingSets = translateGroupingArgIndex(groupingSets, calcRefFields);
        aggCallList = translateAggCallArgIndex(aggCallList, calcRefFields);
    }
    RowType producedType = FlinkTypeFactory.toLogicalRowType(localAgg.getRowType());
    TableSourceTable oldTableSourceTable = oldScan.tableSourceTable();
    DynamicTableSource newTableSource = oldScan.tableSource().copy();
    boolean isPushDownSuccess = AggregatePushDownSpec.apply(inputType, groupingSets, aggCallList, producedType, newTableSource, SourceAbilityContext.from(oldScan));
    if (!isPushDownSuccess) {
        // aggregate push down failed, just return without changing any nodes.
        return;
    }
    // create new source table with new spec and statistic.
    AggregatePushDownSpec aggregatePushDownSpec = new AggregatePushDownSpec(inputType, groupingSets, aggCallList, producedType);
    TableSourceTable newTableSourceTable = oldTableSourceTable.copy(newTableSource, localAgg.getRowType(), new SourceAbilitySpec[] { aggregatePushDownSpec }).copy(FlinkStatistic.UNKNOWN());
    // transform to new nodes.
    BatchPhysicalTableSourceScan newScan = oldScan.copy(oldScan.getTraitSet(), newTableSourceTable);
    BatchPhysicalExchange oldExchange = call.rel(0);
    BatchPhysicalExchange newExchange = oldExchange.copy(oldExchange.getTraitSet(), newScan, oldExchange.getDistribution());
    call.transformTo(newExchange);
}