Examples with PagesSpatialIndexFactory io.prestosql.operator.PagesSpatialIndexFactory used on opensource projects

Example 1 with PagesSpatialIndexFactory

use of io.prestosql.operator.PagesSpatialIndexFactory 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;
}

Example 2 with PagesSpatialIndexFactory

use of io.prestosql.operator.PagesSpatialIndexFactory in project hetu-core by openlookeng.

the class TestSpatialJoinOperator method assertSpatialJoinSnapshot.

private void assertSpatialJoinSnapshot(TaskContext taskContext, Type joinType, RowPagesBuilder buildPages, RowPagesBuilder probePages, MaterializedResult expected) {
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndexSnapshot(driverContext, (build, probe, r) -> build.contains(probe), Optional.empty(), Optional.empty(), buildPages);
    OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), joinType, probePages.getTypes(), Ints.asList(1), 0, Optional.empty(), pagesSpatialIndexFactory);
    assetResultEqualsSnapshot(joinOperatorFactory, driverContext, probePages.build(), expected);
}

Example 3 with PagesSpatialIndexFactory

use of io.prestosql.operator.PagesSpatialIndexFactory in project hetu-core by openlookeng.

the class TestSpatialJoinOperator method testDistanceQuery.

@Test
public void testDistanceQuery() {
    TaskContext taskContext = createTaskContext();
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, DOUBLE)).row(stPoint(0, 0), "0_0", 1.5).row(null, "null", 1.5).row(stPoint(1, 0), "1_0", 1.5).pageBreak().row(stPoint(3, 0), "3_0", 1.5).pageBreak().row(stPoint(10, 0), "10_0", 1.5);
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (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").row(null, "null").row(stPoint(1, 1), "1_1").pageBreak().row(stPoint(3, 1), "3_1").pageBreak().row(stPoint(10, 1), "10_1");
    OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.empty(), pagesSpatialIndexFactory);
    MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("0_1", "0_0").row("0_1", "1_0").row("1_1", "0_0").row("1_1", "1_0").row("3_1", "3_0").row("10_1", "10_0").build();
    assertOperatorEquals(joinOperatorFactory, driverContext, probePages.build(), expected);
}

Example 4 with PagesSpatialIndexFactory

use of io.prestosql.operator.PagesSpatialIndexFactory in project hetu-core by openlookeng.

the class TestSpatialJoinOperator method buildIndex.

private PagesSpatialIndexFactory buildIndex(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));
    ValuesOperator.ValuesOperatorFactory valuesOperatorFactory = new ValuesOperator.ValuesOperatorFactory(0, new PlanNodeId("test"), buildPages.build());
    SpatialIndexBuilderOperatorFactory buildOperatorFactory = new SpatialIndexBuilderOperatorFactory(1, new PlanNodeId("test"), buildPages.getTypes(), Ints.asList(1), 0, radiusChannel, partitionChannel, spatialRelationshipTest, kdbTreeJson, filterFunctionFactory, 10_000, new TestingFactory(false));
    Driver driver = Driver.createDriver(driverContext, valuesOperatorFactory.createOperator(driverContext), buildOperatorFactory.createOperator(driverContext));
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildOperatorFactory.getPagesSpatialIndexFactory();
    ListenableFuture<PagesSpatialIndex> pagesSpatialIndex = pagesSpatialIndexFactory.createPagesSpatialIndex();
    while (!pagesSpatialIndex.isDone()) {
        driver.process();
    }
    runDriverInThread(executor, driver);
    return pagesSpatialIndexFactory;
}

Example 5 with PagesSpatialIndexFactory

use of io.prestosql.operator.PagesSpatialIndexFactory in project hetu-core by openlookeng.

the class TestSpatialJoinOperator method testDistributedSpatialJoin.

@Test
public void testDistributedSpatialJoin() {
    TaskContext taskContext = createTaskContext();
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, true).addDriverContext();
    RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER)).row(POLYGON_A, "A", 1).row(POLYGON_A, "A", 2).row(null, "null", null).pageBreak().row(POLYGON_B, "B", 0).row(POLYGON_B, "B", 2);
    RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER)).row(POINT_X, "x", 2).row(null, "null", null).row(POINT_Y, "y", 2).pageBreak().row(POINT_Z, "z", 0);
    MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("x", "A").row("y", "A").row("y", "B").row("z", "B").build();
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.contains(probe), Optional.empty(), Optional.of(2), Optional.of(KDB_TREE_JSON), Optional.empty(), buildPages);
    OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.of(2), pagesSpatialIndexFactory);
    assertOperatorEquals(joinOperatorFactory, driverContext, probePages.build(), expected);
}