Examples with SQLEngineJob io.cdap.cdap.etl.engine.SQLEngineJob used on opensource projects

Example 1 with SQLEngineJob

use of io.cdap.cdap.etl.engine.SQLEngineJob in project cdap by caskdata.

the class BatchSQLEngineAdapter method pull.

/**
 * Creates a new job to pull a Spark Collection from the SQL engine
 *
 * @param job the job representing the compute stage for the dataset we need to pull.
 * @return Job representing this pull operation.
 */
@SuppressWarnings("unchecked,raw")
public <T> SQLEngineJob<JavaRDD<T>> pull(SQLEngineJob<SQLDataset> job) {
    // If this job already exists, return the existing instance.
    SQLEngineJobKey jobKey = new SQLEngineJobKey(job.getDatasetName(), SQLEngineJobType.PULL);
    if (jobs.containsKey(jobKey)) {
        return (SQLEngineJob<JavaRDD<T>>) jobs.get(jobKey);
    }
    CompletableFuture<JavaRDD<T>> future = new CompletableFuture<>();
    Runnable pullTask = () -> {
        try {
            LOG.debug("Starting pull for dataset '{}'", job.getDatasetName());
            waitForJobAndThrowException(job);
            JavaRDD<T> result = pullInternal(job.waitFor());
            LOG.debug("Completed pull for dataset '{}'", job.getDatasetName());
            future.complete(result);
        } catch (Throwable t) {
            future.completeExceptionally(t);
        }
    };
    executorService.submit(pullTask);
    SQLEngineJob<JavaRDD<T>> pullJob = new SQLEngineJob<>(jobKey, future);
    jobs.put(jobKey, pullJob);
    return pullJob;
}

Example 2 with SQLEngineJob

use of io.cdap.cdap.etl.engine.SQLEngineJob in project cdap by caskdata.

the class BatchSQLEngineAdapter method push.

/**
 * Creates a new job tu push a SparkCollection into the SQL engine.
 *
 * @param datasetName the name of the dataset to push
 * @param schema      the schema for this dataset
 * @param collection  the Spark collection containing the dataset to push
 * @return Job representing this Push operation.
 */
@SuppressWarnings("unchecked,raw")
protected SQLEngineJob<SQLDataset> push(String datasetName, Schema schema, SparkCollection<?> collection) {
    // If this job already exists, return the existing instance.
    SQLEngineJobKey jobKey = new SQLEngineJobKey(datasetName, SQLEngineJobType.PUSH);
    if (jobs.containsKey(jobKey)) {
        return (SQLEngineJob<SQLDataset>) jobs.get(jobKey);
    }
    CompletableFuture<SQLDataset> future = new CompletableFuture<>();
    Runnable pushTask = () -> {
        try {
            LOG.debug("Starting push for dataset '{}'", datasetName);
            SQLDataset result = pushInternal(datasetName, schema, collection);
            LOG.debug("Completed push for dataset '{}'", datasetName);
            future.complete(result);
        } catch (Throwable t) {
            future.completeExceptionally(t);
        }
    };
    executorService.submit(pushTask);
    SQLEngineJob<SQLDataset> job = new SQLEngineJob<>(jobKey, future);
    jobs.put(jobKey, job);
    return job;
}

Example 3 with SQLEngineJob

use of io.cdap.cdap.etl.engine.SQLEngineJob in project cdap by caskdata.

the class BatchSQLEngineAdapter method getDatasetForStage.

/**
 * Function used to fetch the dataset for an input stage.
 *
 * @param stageName
 * @return
 */
private SQLDataset getDatasetForStage(String stageName) {
    // Wait for the previous push or execute jobs to complete
    SQLEngineJobKey pushJobKey = new SQLEngineJobKey(stageName, SQLEngineJobType.PUSH);
    SQLEngineJobKey execJobKey = new SQLEngineJobKey(stageName, SQLEngineJobType.EXECUTE);
    if (jobs.containsKey(pushJobKey)) {
        SQLEngineJob<SQLDataset> job = (SQLEngineJob<SQLDataset>) jobs.get(pushJobKey);
        waitForJobAndThrowException(job);
        return job.waitFor();
    } else if (jobs.containsKey(execJobKey)) {
        SQLEngineJob<SQLDataset> job = (SQLEngineJob<SQLDataset>) jobs.get(execJobKey);
        waitForJobAndThrowException(job);
        return job.waitFor();
    } else {
        throw new IllegalArgumentException("No SQL Engine job exists for stage " + stageName);
    }
}

Example 4 with SQLEngineJob

use of io.cdap.cdap.etl.engine.SQLEngineJob in project cdap by caskdata.

the class BatchSQLEngineAdapter method tryRelationalTransform.

/**
 * This method is called when engine is present and is willing to try performing a relational transform.
 *
 * @param stageSpec stage specification
 * @param transform transform plugin
 * @param input     input collections
 * @return resulting collection or empty optional if tranform can't be done with this engine
 */
public Optional<SQLEngineJob<SQLDataset>> tryRelationalTransform(StageSpec stageSpec, RelationalTransform transform, Map<String, SparkCollection<Object>> input) {
    String stageName = stageSpec.getName();
    Map<String, Relation> inputRelations = input.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> sqlEngine.getRelation(new SQLRelationDefinition(e.getKey(), stageSpec.getInputSchemas().get(e.getKey())))));
    BasicRelationalTransformContext pluginContext = new BasicRelationalTransformContext(getSQLRelationalEngine(), inputRelations, stageSpec.getInputSchemas(), stageSpec.getOutputSchema());
    if (!transform.transform(pluginContext)) {
        // Plugin was not able to do relational tranform with this engine
        return Optional.empty();
    }
    if (pluginContext.getOutputRelation() == null) {
        // Plugin said that tranformation was success but failed to set output
        throw new IllegalStateException("Plugin " + transform + " did not produce a relational output");
    }
    if (!pluginContext.getOutputRelation().isValid()) {
        // An output is set to invalid relation, probably some of transforms are not supported by an engine
        return Optional.empty();
    }
    // Ensure input and output schemas for this stage are supported by the engine
    if (stageSpec.getInputSchemas().values().stream().anyMatch(s -> !sqlEngine.supportsInputSchema(s))) {
        return Optional.empty();
    }
    if (!sqlEngine.supportsOutputSchema(stageSpec.getOutputSchema())) {
        return Optional.empty();
    }
    // Validate transformation definition with engine
    SQLTransformDefinition transformDefinition = new SQLTransformDefinition(stageName, pluginContext.getOutputRelation(), stageSpec.getOutputSchema(), Collections.emptyMap(), Collections.emptyMap());
    if (!sqlEngine.canTransform(transformDefinition)) {
        return Optional.empty();
    }
    return Optional.of(runJob(stageSpec.getName(), SQLEngineJobType.EXECUTE, () -> {
        // Push all stages that need to be pushed to execute this aggregation
        input.forEach((name, collection) -> {
            if (!exists(name)) {
                push(name, stageSpec.getInputSchemas().get(name), collection);
            }
        });
        // Initialize metrics collector
        DefaultStageMetrics stageMetrics = new DefaultStageMetrics(metrics, stageName);
        StageStatisticsCollector statisticsCollector = statsCollectors.get(stageName);
        // Collect input datasets and execute transformation
        Map<String, SQLDataset> inputDatasets = input.keySet().stream().collect(Collectors.toMap(Function.identity(), this::getDatasetForStage));
        // Count input records
        for (SQLDataset inputDataset : inputDatasets.values()) {
            countRecordsIn(inputDataset, statisticsCollector, stageMetrics);
        }
        // Execute transform
        SQLTransformRequest sqlContext = new SQLTransformRequest(inputDatasets, stageSpec.getName(), pluginContext.getOutputRelation(), stageSpec.getOutputSchema());
        SQLDataset transformed = sqlEngine.transform(sqlContext);
        // Count output records
        countRecordsOut(transformed, statisticsCollector, stageMetrics);
        return transformed;
    }));
}