Examples with CreatableRelationProvider org.apache.spark.sql.sources.CreatableRelationProvider used on opensource projects

Example 1 with CreatableRelationProvider

use of org.apache.spark.sql.sources.CreatableRelationProvider in project OpenLineage by OpenLineage.

the class BigQueryNodeVisitor method bigQuerySupplier.

private Optional<Supplier<BigQueryRelation>> bigQuerySupplier(LogicalPlan plan) {
    // SaveIntoDataSourceCommand is a special case because it references a CreatableRelationProvider
    // Every other write instance references a LogicalRelation(BigQueryRelation, _, _, _)
    SQLContext sqlContext = context.getSparkSession().get().sqlContext();
    if (plan instanceof SaveIntoDataSourceCommand) {
        SaveIntoDataSourceCommand saveCommand = (SaveIntoDataSourceCommand) plan;
        CreatableRelationProvider relationProvider = saveCommand.dataSource();
        if (relationProvider instanceof BigQueryRelationProvider) {
            return Optional.of(() -> (BigQueryRelation) ((BigQueryRelationProvider) relationProvider).createRelation(sqlContext, saveCommand.options(), saveCommand.schema()));
        }
    } else {
        if (plan instanceof LogicalRelation && ((LogicalRelation) plan).relation() instanceof BigQueryRelation) {
            return Optional.of(() -> (BigQueryRelation) ((LogicalRelation) plan).relation());
        }
    }
    return Optional.empty();
}