use of io.prestosql.operator.Operator in project hetu-core by openlookeng.
the class TestLocalExchangeSourceOperator method createOperator.
private LocalExchangeSourceOperator createOperator(LocalExchange.LocalExchangeFactory localExchangeFactory, int totalInputChannels, Session session, int driverId) {
if (pipelineContext == null) {
ScheduledExecutorService scheduler = newScheduledThreadPool(4, daemonThreadsNamed("test-%s"));
ScheduledExecutorService scheduledExecutor = newScheduledThreadPool(2, daemonThreadsNamed("test-scheduledExecutor-%s"));
OperatorFactory operatorFactory = new LocalExchangeSourceOperator.LocalExchangeSourceOperatorFactory(0, new PlanNodeId("test"), localExchangeFactory, totalInputChannels);
pipelineContext = createTaskContext(scheduler, scheduledExecutor, session).addPipelineContext(0, true, true, false);
}
DriverContext driverContext = pipelineContext.addDriverContext(Lifespan.taskWide(), driverId);
OperatorFactory operatorFactory = new LocalExchangeSourceOperator.LocalExchangeSourceOperatorFactory(0, new PlanNodeId("test"), localExchangeFactory, totalInputChannels);
Operator operator = operatorFactory.createOperator(driverContext);
assertEquals(operator.getOperatorContext().getOperatorStats().getSystemMemoryReservation().toBytes(), 0);
return (LocalExchangeSourceOperator) operator;
}
use of io.prestosql.operator.Operator in project hetu-core by openlookeng.
the class TestSpatialJoinOperator method buildIndexSnapshot.
private PagesSpatialIndexFactory buildIndexSnapshot(DriverContext driverContext, SpatialPredicate spatialRelationshipTest, Optional<Integer> radiusChannel, Optional<Integer> partitionChannel, Optional<String> kdbTreeJson, Optional<InternalJoinFilterFunction> filterFunction, RowPagesBuilder buildPages) {
Optional<JoinFilterFunctionCompiler.JoinFilterFunctionFactory> filterFunctionFactory = filterFunction.map(function -> (session, addresses, pages) -> new StandardJoinFilterFunction(function, addresses, pages));
SpatialIndexBuilderOperatorFactory buildOperatorFactory = new SpatialIndexBuilderOperatorFactory(1, new PlanNodeId("test"), buildPages.getTypes(), Ints.asList(1), 0, radiusChannel, partitionChannel, spatialRelationshipTest, kdbTreeJson, filterFunctionFactory, 10_000, new TestingFactory(false));
Operator operator = buildOperatorFactory.createOperator(driverContext);
PagesSpatialIndexFactory pagesSpatialIndexFactory = buildOperatorFactory.getPagesSpatialIndexFactory();
ListenableFuture<PagesSpatialIndex> pagesSpatialIndex = pagesSpatialIndexFactory.createPagesSpatialIndex();
List<Page> buildSideInputs = buildPages.build();
int inputIndex = 0;
boolean restored = false;
Object snapshot = null;
// When build side is not done, keep looping
while (!pagesSpatialIndex.isDone()) {
if (operator.needsInput() && inputIndex < buildSideInputs.size()) {
operator.addInput(buildSideInputs.get(inputIndex));
inputIndex++;
}
// Take snapshot in the middle
if (inputIndex == buildSideInputs.size() / 2 && !restored) {
snapshot = operator.capture(operator.getOperatorContext().getDriverContext().getSerde());
}
// When input pages are used up, restore operator to snapshot and move inputIndex back to when snapshot was taken
if (inputIndex == buildSideInputs.size() && !restored) {
assertTrue(snapshot != null);
operator.restore(snapshot, operator.getOperatorContext().getDriverContext().getSerde());
restored = true;
inputIndex = buildSideInputs.size() / 2;
}
// Used up all provided build side inputs and have done rollback process, finish build operator, this will cause future to be done
if (inputIndex >= buildSideInputs.size() && restored) {
operator.finish();
}
}
return pagesSpatialIndexFactory;
}
use of io.prestosql.operator.Operator in project hetu-core by openlookeng.
the class TestSpatialJoinOperator method testYield.
@Test
public void testYield() {
// create a filter function that yields for every probe match
// verify we will yield #match times totally
TaskContext taskContext = createTaskContext();
DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
// force a yield for every match
AtomicInteger filterFunctionCalls = new AtomicInteger();
InternalJoinFilterFunction filterFunction = new TestInternalJoinFilterFunction(((leftPosition, leftPage, rightPosition, rightPage) -> {
filterFunctionCalls.incrementAndGet();
driverContext.getYieldSignal().forceYieldForTesting();
return true;
}));
RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR)).row(POLYGON_A, "A").pageBreak().row(POLYGON_B, "B");
PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.contains(probe), Optional.empty(), Optional.of(filterFunction), buildPages);
// 10 points in polygon A (x0...x9)
// 10 points in polygons A and B (y0...y9)
// 10 points in polygon B (z0...z9)
// 40 total matches
RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR));
for (int i = 0; i < 10; i++) {
probePages.row(stPoint(1 + 0.1 * i, 1 + 0.1 * i), "x" + i);
}
for (int i = 0; i < 10; i++) {
probePages.row(stPoint(4.5 + 0.01 * i, 4.5 + 0.01 * i), "y" + i);
}
for (int i = 0; i < 10; i++) {
probePages.row(stPoint(6 + 0.1 * i, 6 + 0.1 * i), "z" + i);
}
List<Page> probeInput = probePages.build();
OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.empty(), pagesSpatialIndexFactory);
Operator operator = joinOperatorFactory.createOperator(driverContext);
assertTrue(operator.needsInput());
operator.addInput(probeInput.get(0));
operator.finish();
// we will yield 40 times due to filterFunction
for (int i = 0; i < 40; i++) {
driverContext.getYieldSignal().setWithDelay(5 * SECONDS.toNanos(1), driverContext.getYieldExecutor());
assertNull(operator.getOutput());
assertEquals(filterFunctionCalls.get(), i + 1, "Expected join to stop processing (yield) after calling filter function once");
driverContext.getYieldSignal().reset();
}
// delayed yield is not going to prevent operator from producing a page now (yield won't be forced because filter function won't be called anymore)
driverContext.getYieldSignal().setWithDelay(5 * SECONDS.toNanos(1), driverContext.getYieldExecutor());
Page output = operator.getOutput();
assertNotNull(output);
// make sure we have 40 matches
assertEquals(output.getPositionCount(), 40);
}
use of io.prestosql.operator.Operator in project hetu-core by openlookeng.
the class TestSpatialJoinOperator method testDuplicateProbeFactory.
@Test(dataProvider = "testDuplicateProbeFactoryDataProvider")
public void testDuplicateProbeFactory(boolean createSecondaryOperators) throws Exception {
TaskContext taskContext = createTaskContext();
PipelineContext pipelineContext = taskContext.addPipelineContext(0, true, true, false);
DriverContext probeDriver = pipelineContext.addDriverContext();
DriverContext buildDriver = pipelineContext.addDriverContext();
RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, DOUBLE)).row(stPoint(0, 0), "0_0", 1.5);
PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(buildDriver, (build, probe, r) -> build.distance(probe) <= r.getAsDouble(), Optional.of(2), Optional.empty(), buildPages);
RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR)).row(stPoint(0, 1), "0_1");
OperatorFactory firstFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.empty(), pagesSpatialIndexFactory);
for (int i = 0; i < 3; i++) {
DriverContext secondDriver = pipelineContext.addDriverContext();
OperatorFactory secondFactory = firstFactory.duplicate();
if (createSecondaryOperators) {
try (Operator secondOperator = secondFactory.createOperator(secondDriver)) {
assertEquals(toPages(secondOperator, emptyIterator()), ImmutableList.of());
}
}
secondFactory.noMoreOperators();
}
MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("0_1", "0_0").build();
assertOperatorEquals(firstFactory, probeDriver, probePages.build(), expected);
}
use of io.prestosql.operator.Operator in project hetu-core by openlookeng.
the class TestSpatialJoinOperator method assetResultEqualsSnapshot.
private void assetResultEqualsSnapshot(OperatorFactory operatorFactory, DriverContext driverContext, List<Page> input, MaterializedResult expected) {
Operator operator = operatorFactory.createOperator(driverContext);
ImmutableList.Builder<Page> outputPages = ImmutableList.builder();
int inputIdx = 0;
Object snapshot = null;
boolean restored = false;
for (int loopsSinceLastPage = 0; loopsSinceLastPage < 1_000; loopsSinceLastPage++) {
if (inputIdx < input.size() && operator.needsInput()) {
if (inputIdx == input.size() / 2) {
snapshot = operator.capture(operator.getOperatorContext().getDriverContext().getSerde());
} else if (inputIdx == input.size() - 1 && snapshot != null && !restored) {
operator.restore(snapshot, operator.getOperatorContext().getDriverContext().getSerde());
inputIdx = input.size() / 2;
restored = true;
}
operator.addInput(input.get(inputIdx));
inputIdx++;
loopsSinceLastPage = 0;
// Operator doesn't produce output in the process, but needs getOutput call to clear for next input.
operator.getOutput();
}
}
for (int loopsSinceLastPage = 0; !operator.isFinished() && loopsSinceLastPage < 1_000; loopsSinceLastPage++) {
operator.finish();
Page outputPage = operator.getOutput();
if (outputPage != null && outputPage.getPositionCount() != 0) {
outputPages.add(outputPage);
loopsSinceLastPage = 0;
}
}
assertEquals(operator.isFinished(), true, "Operator did not finish");
assertEquals(operator.needsInput(), false, "Operator still wants input");
assertEquals(operator.isBlocked().isDone(), true, "Operator is blocked");
List<Page> output = outputPages.build();
MaterializedResult actual = toMaterializedResult(driverContext.getSession(), expected.getTypes(), output);
Assert.assertEquals(actual, expected);
}
Aggregations