Beispiel #1
0
  public List<Driver> createDrivers(
      @Language("SQL") String sql, OutputFactory outputFactory, TaskContext taskContext) {
    Statement statement = SqlParser.createStatement(sql);

    if (printPlan) {
      assertFormattedSql(statement);
    }

    PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
    FeaturesConfig featuresConfig = new FeaturesConfig().setExperimentalSyntaxEnabled(true);
    PlanOptimizersFactory planOptimizersFactory =
        new PlanOptimizersFactory(metadata, splitManager, indexManager, featuresConfig);

    QueryExplainer queryExplainer =
        new QueryExplainer(
            session,
            planOptimizersFactory.get(),
            metadata,
            featuresConfig.isExperimentalSyntaxEnabled());
    Analyzer analyzer =
        new Analyzer(
            session,
            metadata,
            Optional.of(queryExplainer),
            featuresConfig.isExperimentalSyntaxEnabled());

    Analysis analysis = analyzer.analyze(statement);

    Plan plan =
        new LogicalPlanner(session, planOptimizersFactory.get(), idAllocator, metadata)
            .plan(analysis);
    if (printPlan) {
      System.out.println(PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata));
    }

    SubPlan subplan =
        new DistributedLogicalPlanner(session, metadata, idAllocator).createSubPlans(plan, true);
    assertTrue(subplan.getChildren().isEmpty(), "Expected subplan to have no children");

    LocalExecutionPlanner executionPlanner =
        new LocalExecutionPlanner(
            new NodeInfo(new NodeConfig().setEnvironment("test").setNodeId("test-node")),
            metadata,
            dataStreamProvider,
            indexManager,
            storageManager,
            recordSinkManager,
            null,
            compiler);

    // plan query
    LocalExecutionPlan localExecutionPlan =
        executionPlanner.plan(
            session, subplan.getFragment().getRoot(), plan.getTypes(), outputFactory);

    // generate sources
    List<TaskSource> sources = new ArrayList<>();
    long sequenceId = 0;
    for (PlanNode sourceNode : subplan.getFragment().getSources()) {
      if (sourceNode instanceof ValuesNode) {
        continue;
      }

      TableScanNode tableScan = (TableScanNode) sourceNode;

      SplitSource splitSource =
          splitManager.getPartitionSplits(tableScan.getTable(), getPartitions(tableScan));

      ImmutableSet.Builder<ScheduledSplit> scheduledSplits = ImmutableSet.builder();
      while (!splitSource.isFinished()) {
        try {
          for (Split split : splitSource.getNextBatch(1000)) {
            scheduledSplits.add(new ScheduledSplit(sequenceId++, split));
          }
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt();
          throw Throwables.propagate(e);
        }
      }

      sources.add(new TaskSource(tableScan.getId(), scheduledSplits.build(), true));
    }

    // create drivers
    List<Driver> drivers = new ArrayList<>();
    Map<PlanNodeId, Driver> driversBySource = new HashMap<>();
    for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) {
      DriverContext driverContext =
          taskContext
              .addPipelineContext(driverFactory.isInputDriver(), driverFactory.isOutputDriver())
              .addDriverContext();
      Driver driver = driverFactory.createDriver(driverContext);
      drivers.add(driver);
      for (PlanNodeId sourceId : driver.getSourceIds()) {
        driversBySource.put(sourceId, driver);
      }
      driverFactory.close();
    }

    // add sources to the drivers
    for (TaskSource source : sources) {
      for (Driver driver : driversBySource.values()) {
        driver.updateSource(source);
      }
    }

    return ImmutableList.copyOf(drivers);
  }
  private SqlTaskExecution(
      TaskStateMachine taskStateMachine,
      TaskContext taskContext,
      SharedBuffer sharedBuffer,
      PlanFragment fragment,
      LocalExecutionPlanner planner,
      TaskExecutor taskExecutor,
      QueryMonitor queryMonitor,
      Executor notificationExecutor) {
    this.taskStateMachine = checkNotNull(taskStateMachine, "taskStateMachine is null");
    this.taskId = taskStateMachine.getTaskId();
    this.taskContext = checkNotNull(taskContext, "taskContext is null");
    this.sharedBuffer = checkNotNull(sharedBuffer, "sharedBuffer is null");

    this.taskExecutor = checkNotNull(taskExecutor, "driverExecutor is null");
    this.notificationExecutor = checkNotNull(notificationExecutor, "notificationExecutor is null");

    this.queryMonitor = checkNotNull(queryMonitor, "queryMonitor is null");

    try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) {
      List<DriverFactory> driverFactories;
      try {
        OutputFactory outputOperatorFactory;
        if (fragment.getOutputPartitioning() == OutputPartitioning.NONE) {
          outputOperatorFactory = new TaskOutputFactory(sharedBuffer);
        } else if (fragment.getOutputPartitioning() == OutputPartitioning.HASH) {
          outputOperatorFactory = new PartitionedOutputFactory(sharedBuffer);
        } else {
          throw new PrestoException(
              NOT_SUPPORTED,
              format("OutputPartitioning %s is not supported", fragment.getOutputPartitioning()));
        }

        LocalExecutionPlan localExecutionPlan =
            planner.plan(
                taskContext.getSession(),
                fragment.getRoot(),
                fragment.getOutputLayout(),
                fragment.getSymbols(),
                fragment.getDistribution(),
                outputOperatorFactory);
        driverFactories = localExecutionPlan.getDriverFactories();
      } catch (Throwable e) {
        // planning failed
        taskStateMachine.failed(e);
        throw Throwables.propagate(e);
      }

      // index driver factories
      DriverSplitRunnerFactory partitionedDriverFactory = null;
      ImmutableList.Builder<DriverSplitRunnerFactory> unpartitionedDriverFactories =
          ImmutableList.builder();
      for (DriverFactory driverFactory : driverFactories) {
        if (driverFactory.getSourceIds().contains(fragment.getPartitionedSource())) {
          checkState(
              partitionedDriverFactory == null, "multiple partitioned sources are not supported");
          partitionedDriverFactory = new DriverSplitRunnerFactory(driverFactory);
        } else {
          unpartitionedDriverFactories.add(new DriverSplitRunnerFactory(driverFactory));
        }
      }
      this.unpartitionedDriverFactories = unpartitionedDriverFactories.build();

      if (fragment.getDistribution() == PlanDistribution.SOURCE) {
        checkArgument(
            partitionedDriverFactory != null,
            "Fragment is partitioned, but no partitioned driver found");
      }
      this.partitionedSourceId = fragment.getPartitionedSource();
      this.partitionedDriverFactory = partitionedDriverFactory;

      // don't register the task if it is already completed (most likely failed during planning
      // above)
      if (!taskStateMachine.getState().isDone()) {
        taskHandle = taskExecutor.addTask(taskId);
        taskStateMachine.addStateChangeListener(
            new RemoveTaskHandleWhenDone(taskExecutor, taskHandle));
      } else {
        taskHandle = null;
      }

      sharedBuffer.addStateChangeListener(
          new CheckTaskCompletionOnBufferFinish(SqlTaskExecution.this));
    }
  }
  private SqlTaskExecution(
      Session session,
      TaskId taskId,
      URI location,
      PlanFragment fragment,
      OutputBuffers outputBuffers,
      LocalExecutionPlanner planner,
      DataSize maxBufferSize,
      TaskExecutor taskExecutor,
      DataSize maxTaskMemoryUsage,
      DataSize operatorPreAllocatedMemory,
      QueryMonitor queryMonitor,
      Executor notificationExecutor,
      boolean cpuTimerEnabled) {
    try (SetThreadName setThreadName = new SetThreadName("Task-%s", taskId)) {
      this.taskId = checkNotNull(taskId, "taskId is null");
      this.location = checkNotNull(location, "location is null");
      this.taskExecutor = checkNotNull(taskExecutor, "driverExecutor is null");
      this.notificationExecutor =
          checkNotNull(notificationExecutor, "notificationExecutor is null");

      this.taskStateMachine = new TaskStateMachine(taskId, notificationExecutor);
      taskStateMachine.addStateChangeListener(
          new StateChangeListener<TaskState>() {
            @Override
            public void stateChanged(TaskState taskState) {
              if (taskState.isDone()) {
                SqlTaskExecution.this.taskExecutor.removeTask(taskHandle);
                // make sure buffers are cleaned up
                if (taskState != TaskState.FAILED) {
                  // don't close buffers for a failed query
                  // closed buffers signal to upstream tasks that everything finished cleanly
                  sharedBuffer.destroy();
                }
              }
            }
          });

      this.taskContext =
          new TaskContext(
              taskStateMachine,
              notificationExecutor,
              session,
              checkNotNull(maxTaskMemoryUsage, "maxTaskMemoryUsage is null"),
              checkNotNull(operatorPreAllocatedMemory, "operatorPreAllocatedMemory is null"),
              cpuTimerEnabled);

      this.sharedBuffer =
          new SharedBuffer(
              taskId,
              notificationExecutor,
              checkNotNull(maxBufferSize, "maxBufferSize is null"),
              outputBuffers);
      sharedBuffer.addStateChangeListener(
          new StateChangeListener<QueueState>() {
            @Override
            public void stateChanged(QueueState taskState) {
              if (taskState == QueueState.FINISHED) {
                checkTaskCompletion();
              }
            }
          });

      this.queryMonitor = checkNotNull(queryMonitor, "queryMonitor is null");

      taskHandle = taskExecutor.addTask(taskId);

      LocalExecutionPlan localExecutionPlan =
          planner.plan(
              session,
              fragment.getRoot(),
              fragment.getSymbols(),
              new TaskOutputFactory(sharedBuffer));
      List<DriverFactory> driverFactories = localExecutionPlan.getDriverFactories();

      // index driver factories
      DriverSplitRunnerFactory partitionedDriverFactory = null;
      ImmutableList.Builder<DriverSplitRunnerFactory> unpartitionedDriverFactories =
          ImmutableList.builder();
      for (DriverFactory driverFactory : driverFactories) {
        if (driverFactory.getSourceIds().contains(fragment.getPartitionedSource())) {
          checkState(
              partitionedDriverFactory == null, "multiple partitioned sources are not supported");
          partitionedDriverFactory = new DriverSplitRunnerFactory(driverFactory);
        } else {
          unpartitionedDriverFactories.add(new DriverSplitRunnerFactory(driverFactory));
        }
      }
      this.unpartitionedDriverFactories = unpartitionedDriverFactories.build();

      if (fragment.getDistribution() == PlanDistribution.SOURCE) {
        checkArgument(
            partitionedDriverFactory != null,
            "Fragment is partitioned, but no partitioned driver found");
      }
      this.partitionedSourceId = fragment.getPartitionedSource();
      this.partitionedDriverFactory = partitionedDriverFactory;
    }
  }
  public List<Driver> createDrivers(
      Session session,
      @Language("SQL") String sql,
      OutputFactory outputFactory,
      TaskContext taskContext) {
    Statement statement = sqlParser.createStatement(sql);

    assertFormattedSql(sqlParser, statement);

    PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
    FeaturesConfig featuresConfig =
        new FeaturesConfig()
            .setExperimentalSyntaxEnabled(true)
            .setDistributedIndexJoinsEnabled(false)
            .setOptimizeHashGeneration(true);
    PlanOptimizersFactory planOptimizersFactory =
        new PlanOptimizersFactory(metadata, sqlParser, indexManager, featuresConfig, true);

    QueryExplainer queryExplainer =
        new QueryExplainer(
            planOptimizersFactory.get(),
            metadata,
            accessControl,
            sqlParser,
            dataDefinitionTask,
            featuresConfig.isExperimentalSyntaxEnabled());
    Analyzer analyzer =
        new Analyzer(
            session,
            metadata,
            sqlParser,
            accessControl,
            Optional.of(queryExplainer),
            featuresConfig.isExperimentalSyntaxEnabled());

    Analysis analysis = analyzer.analyze(statement);
    Plan plan =
        new LogicalPlanner(session, planOptimizersFactory.get(), idAllocator, metadata)
            .plan(analysis);

    if (printPlan) {
      System.out.println(
          PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata, session));
    }

    SubPlan subplan = new PlanFragmenter().createSubPlans(plan);
    if (!subplan.getChildren().isEmpty()) {
      throw new AssertionError("Expected subplan to have no children");
    }

    LocalExecutionPlanner executionPlanner =
        new LocalExecutionPlanner(
            metadata,
            sqlParser,
            pageSourceManager,
            indexManager,
            pageSinkManager,
            null,
            compiler,
            new IndexJoinLookupStats(),
            new CompilerConfig()
                .setInterpreterEnabled(false), // make sure tests fail if compiler breaks
            new TaskManagerConfig().setTaskDefaultConcurrency(4));

    // plan query
    LocalExecutionPlan localExecutionPlan =
        executionPlanner.plan(
            session,
            subplan.getFragment().getRoot(),
            subplan.getFragment().getOutputLayout(),
            plan.getTypes(),
            subplan.getFragment().getDistribution(),
            outputFactory);

    // generate sources
    List<TaskSource> sources = new ArrayList<>();
    long sequenceId = 0;
    for (TableScanNode tableScan : findTableScanNodes(subplan.getFragment().getRoot())) {
      TableLayoutHandle layout = tableScan.getLayout().get();

      SplitSource splitSource = splitManager.getSplits(session, layout);

      ImmutableSet.Builder<ScheduledSplit> scheduledSplits = ImmutableSet.builder();
      while (!splitSource.isFinished()) {
        for (Split split : getFutureValue(splitSource.getNextBatch(1000))) {
          scheduledSplits.add(new ScheduledSplit(sequenceId++, split));
        }
      }

      sources.add(new TaskSource(tableScan.getId(), scheduledSplits.build(), true));
    }

    // create drivers
    List<Driver> drivers = new ArrayList<>();
    Map<PlanNodeId, Driver> driversBySource = new HashMap<>();
    for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) {
      for (int i = 0; i < driverFactory.getDriverInstances(); i++) {
        DriverContext driverContext =
            taskContext
                .addPipelineContext(driverFactory.isInputDriver(), driverFactory.isOutputDriver())
                .addDriverContext();
        Driver driver = driverFactory.createDriver(driverContext);
        drivers.add(driver);
        for (PlanNodeId sourceId : driver.getSourceIds()) {
          driversBySource.put(sourceId, driver);
        }
      }
      driverFactory.close();
    }

    // add sources to the drivers
    for (TaskSource source : sources) {
      for (Driver driver : driversBySource.values()) {
        driver.updateSource(source);
      }
    }

    return ImmutableList.copyOf(drivers);
  }