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

Example 1 with PagesSpatialIndex

use of io.prestosql.operator.PagesSpatialIndex 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 PagesSpatialIndex

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