How to use

Driver

in

com.facebook.presto.operator

@Test
public void testNormalFinish()
{
  List<Type> types = ImmutableList.of(VARCHAR, BIGINT, BIGINT);
  ValuesOperator source = new ValuesOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "values"), rowPagesBuilder(types)
      .addSequencePage(10, 20, 30, 40)
      .build());
  Operator sink = createSinkOperator(types);
  Driver driver = Driver.createDriver(driverContext, source, sink);
  assertSame(driver.getDriverContext(), driverContext);
  assertFalse(driver.isFinished());
  ListenableFuture<?> blocked = driver.processFor(new Duration(1, TimeUnit.SECONDS));
  assertTrue(blocked.isDone());
  assertTrue(driver.isFinished());
  assertTrue(sink.isFinished());
  assertTrue(source.isFinished());
}

private boolean loadNextPage()
{
  Page nextPage = extractNonEmptyPage(pageBuffer);
  while (nextPage == null) {
    if (driver.isFinished()) {
      return false;
    }
    driver.process();
    nextPage = extractNonEmptyPage(pageBuffer);
  }
  currentPage = nextPage;
  return true;
}

@Test
public void testBrokenOperatorCloseWhileProcessing()
    throws Exception
{
  BrokenOperator brokenOperator = new BrokenOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "source"), false);
  final Driver driver = Driver.createDriver(driverContext, brokenOperator, createSinkOperator(ImmutableList.of()));
  assertSame(driver.getDriverContext(), driverContext);
  // block thread in operator processing
  Future<Boolean> driverProcessFor = executor.submit(new Callable<Boolean>()
  {
    @Override
    public Boolean call()
    {
      return driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone();
    }
  });
  brokenOperator.waitForLocked();
  driver.close();
  assertTrue(driver.isFinished());
  try {
    driverProcessFor.get(1, TimeUnit.SECONDS);
    fail("Expected InterruptedException");
  }
  catch (ExecutionException e) {
    assertDriverInterrupted(e.getCause());
  }
}

public static Driver createDriver(DriverContext driverContext, List<Operator> operators)
{
  requireNonNull(driverContext, "driverContext is null");
  requireNonNull(operators, "operators is null");
  Driver driver = new Driver(driverContext, operators);
  driver.initialize();
  return driver;
}

  continue;
Optional<PlanNodeId> sourceId = driver.getSourceId();
if (!sourceId.isPresent()) {
  continue;
  continue;
driver.updateSource(sourceUpdate);

@Test
public void testNormalFinish()
{
  List<Type> types = ImmutableList.<Type>of(VARCHAR, BIGINT, BIGINT);
  ValuesOperator source = new ValuesOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "values"), types, rowPagesBuilder(types)
      .addSequencePage(10, 20, 30, 40)
      .build());
  MaterializingOperator sink = createSinkOperator(source);
  Driver driver = new Driver(driverContext, source, sink);
  assertSame(driver.getDriverContext(), driverContext);
  assertFalse(driver.isFinished());
  ListenableFuture<?> blocked = driver.processFor(new Duration(1, TimeUnit.SECONDS));
  assertTrue(blocked.isDone());
  assertTrue(driver.isFinished());
  assertTrue(sink.isFinished());
  assertTrue(source.isFinished());
}

  @Test(dataProvider = "hashEnabledValues", expectedExceptions = ExceededMemoryLimitException.class, expectedExceptionsMessageRegExp = "Query exceeded local memory limit of.*")
  public void testMemoryLimit(boolean hashEnabled)
      throws Exception
  {
    DriverContext driverContext = createTaskContext(executor, TEST_SESSION, new DataSize(100, BYTE))
        .addPipelineContext(true, true)
        .addDriverContext();

    OperatorContext operatorContext = driverContext.addOperatorContext(0, new PlanNodeId("test"), ValuesOperator.class.getSimpleName());
    List<Type> buildTypes = ImmutableList.<Type>of(BIGINT);
    RowPagesBuilder rowPagesBuilder = rowPagesBuilder(hashEnabled, Ints.asList(0), buildTypes);
    Operator buildOperator = new ValuesOperator(operatorContext, buildTypes, rowPagesBuilder
        .addSequencePage(10000, 20)
        .build());
    SetBuilderOperatorFactory setBuilderOperatorFactory = new SetBuilderOperatorFactory(1, new PlanNodeId("test"), buildOperator.getTypes().get(0), 0, rowPagesBuilder.getHashChannel(), 10);
    Operator setBuilderOperator = setBuilderOperatorFactory.createOperator(driverContext);

    Driver driver = new Driver(driverContext, buildOperator, setBuilderOperator);
    while (!driver.isFinished()) {
      driver.process();
    }
  }
}

  @Test(dataProvider = "hashEnabledValues", expectedExceptions = ExceededMemoryLimitException.class, expectedExceptionsMessageRegExp = "Query exceeded per-node user memory limit of.*")
  public void testMemoryLimit(boolean hashEnabled)
  {
    DriverContext driverContext = createTaskContext(executor, scheduledExecutor, TEST_SESSION, new DataSize(100, BYTE))
        .addPipelineContext(0, true, true, false)
        .addDriverContext();

    OperatorContext operatorContext = driverContext.addOperatorContext(0, new PlanNodeId("test"), ValuesOperator.class.getSimpleName());
    List<Type> buildTypes = ImmutableList.of(BIGINT);
    RowPagesBuilder rowPagesBuilder = rowPagesBuilder(hashEnabled, Ints.asList(0), buildTypes);
    Operator buildOperator = new ValuesOperator(operatorContext, rowPagesBuilder
        .addSequencePage(10000, 20)
        .build());
    SetBuilderOperatorFactory setBuilderOperatorFactory = new SetBuilderOperatorFactory(
        1,
        new PlanNodeId("test"),
        buildTypes.get(0),
        0,
        rowPagesBuilder.getHashChannel(),
        10,
        new JoinCompiler(createTestMetadataManager(), new FeaturesConfig()));
    Operator setBuilderOperator = setBuilderOperatorFactory.createOperator(driverContext);

    Driver driver = Driver.createDriver(driverContext, buildOperator, setBuilderOperator);
    while (!driver.isFinished()) {
      driver.process();
    }
  }
}

/**
 * Runs Driver in another thread until it is finished
 */
private static void runDriverInThread(ExecutorService executor, Driver driver)
{
  executor.execute(() -> {
    if (!driver.isFinished()) {
      try {
        driver.process();
      }
      catch (PrestoException e) {
        driver.getDriverContext().failed(e);
        throw e;
      }
      runDriverInThread(executor, driver);
    }
  });
}

PlanNodeId sourcePlanNodeId = driverFactory.getSourceId().get();
ScheduledSplit split = new ScheduledSplit(0, sourcePlanNodeId, new Split(INDEX_CONNECTOR_ID, new ConnectorTransactionHandle() {}, new IndexSplit(recordSetForLookupSource)));
driver.updateSource(new TaskSource(sourcePlanNodeId, ImmutableSet.of(split), true));
while (!driver.isFinished()) {
  ListenableFuture<?> process = driver.process();
  checkState(process.isDone(), "Driver should never block");

@Test
public void testBrokenOperatorCloseWhileProcessing()
    throws Exception
{
  BrokenOperator brokenOperator = new BrokenOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "source"), false);
  final Driver driver = new Driver(driverContext, brokenOperator, createSinkOperator(brokenOperator));
  assertSame(driver.getDriverContext(), driverContext);
  // block thread in operator processing
  Future<Boolean> driverProcessFor = executor.submit(new Callable<Boolean>()
  {
    @Override
    public Boolean call()
        throws Exception
    {
      return driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone();
    }
  });
  brokenOperator.waitForLocked();
  driver.close();
  assertTrue(driver.isFinished());
  try {
    driverProcessFor.get(1, TimeUnit.SECONDS);
    fail("Expected InterruptedException");
  }
  catch (ExecutionException e) {
    checkArgument(getRootCause(e) instanceof InterruptedException, "Expected root cause exception to be an instance of InterruptedException");
  }
}

private static void processRow(final Driver joinDriver, final TaskStateMachine taskStateMachine)
{
  joinDriver.getDriverContext().getYieldSignal().setWithDelay(TimeUnit.SECONDS.toNanos(1), joinDriver.getDriverContext().getYieldExecutor());
  joinDriver.process();
  joinDriver.getDriverContext().getYieldSignal().reset();
  checkErrors(taskStateMachine);
}

@Override
public synchronized boolean isFinished()
{
  if (closed) {
    return true;
  }
  return driver != null && driver.isFinished();
}

private void assertDriversProgress(Predicate<OperatorContext> reason)
{
  do {
    assertFalse(isOperatorBlocked(drivers, reason));
    boolean progress = false;
    for (Driver driver : drivers) {
      ListenableFuture<?> blocked = driver.process();
      progress = progress | blocked.isDone();
    }
    // query should not block
    assertTrue(progress);
  }
  while (!drivers.stream().allMatch(Driver::isFinished));
}

@Test(invocationCount = 1_000, timeOut = 10_000)
public void testConcurrentClose()
{
  List<Type> types = ImmutableList.of(VARCHAR, BIGINT, BIGINT);
  OperatorContext operatorContext = driverContext.addOperatorContext(0, new PlanNodeId("test"), "values");
  ValuesOperator source = new ValuesOperator(operatorContext, rowPagesBuilder(types)
      .addSequencePage(10, 20, 30, 40)
      .build());
  Operator sink = createSinkOperator(types);
  Driver driver = Driver.createDriver(driverContext, source, sink);
  // let these threads race
  scheduledExecutor.submit(() -> driver.processFor(new Duration(1, TimeUnit.NANOSECONDS))); // don't want to call isFinishedInternal in processFor
  scheduledExecutor.submit(() -> driver.close());
  while (!driverContext.isDone()) {
    Uninterruptibles.sleepUninterruptibly(1, TimeUnit.MILLISECONDS);
  }
}

Driver driver = Driver.createDriver(driverContext,
    valuesOperator,
    nestedLoopBuildOperator);
  driver.process();

  @Override
  public void close()
  {
    driver.close();
    currentPage = null;
  }
}

@Override
public ListenableFuture<?> processFor(Duration duration)
{
  Driver driver;
  synchronized (this) {
    // if close() was called before we get here, there's not point in even creating the driver
    if (closed) {
      return Futures.immediateFuture(null);
    }
    if (this.driver == null) {
      this.driver = driverSplitRunnerFactory.createDriver(driverContext, partitionedSplit);
    }
    driver = this.driver;
  }
  return driver.processFor(duration);
}

public IndexedData streamIndexDataForSingleKey(UpdateRequest updateRequest)
{
  Page indexKeyTuple = updateRequest.getPage().getRegion(0, 1);
  PageBuffer pageBuffer = new PageBuffer(100);
  DriverFactory driverFactory = indexBuildDriverFactoryProvider.createStreaming(pageBuffer, indexKeyTuple);
  Driver driver = driverFactory.createDriver(pipelineContext.addDriverContext());
  PageRecordSet pageRecordSet = new PageRecordSet(keyTypes, indexKeyTuple);
  PlanNodeId planNodeId = driverFactory.getSourceId().get();
  ScheduledSplit split = new ScheduledSplit(0, planNodeId, new Split(INDEX_CONNECTOR_ID, new ConnectorTransactionHandle() {}, new IndexSplit(pageRecordSet)));
  driver.updateSource(new TaskSource(planNodeId, ImmutableSet.of(split), true));
  return new StreamingIndexedData(outputTypes, keyTypes, indexKeyTuple, pageBuffer, driver);
}

public synchronized DriverContext getDriverContext()
{
  if (driver == null) {
    return null;
  }
  return driver.getDriverContext();
}