Search in sources :

Example 1 with GEOMETRY

use of com.facebook.presto.plugin.geospatial.GeometryType.GEOMETRY in project presto by prestodb.

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);
}
Also used : Page(com.facebook.presto.common.Page) TestingFactory(com.facebook.presto.operator.PagesIndex.TestingFactory) MaterializedResult.resultBuilder(com.facebook.presto.testing.MaterializedResult.resultBuilder) Test(org.testng.annotations.Test) AfterMethod(org.testng.annotations.AfterMethod) GeoFunctions.spatialPartitions(com.facebook.presto.plugin.geospatial.GeoFunctions.spatialPartitions) Type(com.facebook.presto.sql.planner.plan.SpatialJoinNode.Type) RowPagesBuilder(com.facebook.presto.RowPagesBuilder) SpatialJoinOperatorFactory(com.facebook.presto.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) Executors.newScheduledThreadPool(java.util.concurrent.Executors.newScheduledThreadPool) GeoFunctions.stGeometryFromText(com.facebook.presto.plugin.geospatial.GeoFunctions.stGeometryFromText) Slices(io.airlift.slice.Slices) StandardJoinFilterFunction(com.facebook.presto.operator.StandardJoinFilterFunction) Operator(com.facebook.presto.operator.Operator) SynchronousQueue(java.util.concurrent.SynchronousQueue) DOUBLE(com.facebook.presto.common.type.DoubleType.DOUBLE) BeforeMethod(org.testng.annotations.BeforeMethod) LEFT(com.facebook.presto.sql.planner.plan.SpatialJoinNode.Type.LEFT) Assert.assertNotNull(org.testng.Assert.assertNotNull) Rectangle(com.facebook.presto.geospatial.Rectangle) Threads.daemonThreadsNamed(com.facebook.airlift.concurrent.Threads.daemonThreadsNamed) OperatorAssertion.assertOperatorEqualsIgnoreOrder(com.facebook.presto.operator.OperatorAssertion.assertOperatorEqualsIgnoreOrder) PagesSpatialIndex(com.facebook.presto.operator.PagesSpatialIndex) List(java.util.List) KdbTreeUtils(com.facebook.presto.geospatial.KdbTreeUtils) InternalJoinFilterFunction(com.facebook.presto.operator.InternalJoinFilterFunction) INTEGER(com.facebook.presto.common.type.IntegerType.INTEGER) Optional(java.util.Optional) SpatialPredicate(com.facebook.presto.operator.SpatialIndexBuilderOperator.SpatialPredicate) RowPagesBuilder.rowPagesBuilder(com.facebook.presto.RowPagesBuilder.rowPagesBuilder) PlanNodeId(com.facebook.presto.spi.plan.PlanNodeId) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) Slice(io.airlift.slice.Slice) Assert.assertNull(org.testng.Assert.assertNull) ThreadPoolExecutor(java.util.concurrent.ThreadPoolExecutor) VARCHAR(com.facebook.presto.common.type.VarcharType.VARCHAR) Assert.assertEquals(org.testng.Assert.assertEquals) ValuesOperator(com.facebook.presto.operator.ValuesOperator) PrestoException(com.facebook.presto.spi.PrestoException) OperatorFactory(com.facebook.presto.operator.OperatorFactory) TEST_SESSION(com.facebook.presto.SessionTestUtils.TEST_SESSION) ImmutableList(com.google.common.collect.ImmutableList) PagesSpatialIndexFactory(com.facebook.presto.operator.PagesSpatialIndexFactory) SpatialIndexBuilderOperatorFactory(com.facebook.presto.operator.SpatialIndexBuilderOperator.SpatialIndexBuilderOperatorFactory) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) KdbTree(com.facebook.presto.geospatial.KdbTree) TestingTaskContext(com.facebook.presto.testing.TestingTaskContext) ExecutorService(java.util.concurrent.ExecutorService) GEOMETRY(com.facebook.presto.plugin.geospatial.GeometryType.GEOMETRY) TaskContext(com.facebook.presto.operator.TaskContext) GeoFunctions.stPoint(com.facebook.presto.plugin.geospatial.GeoFunctions.stPoint) JoinFilterFunctionCompiler(com.facebook.presto.sql.gen.JoinFilterFunctionCompiler) Driver(com.facebook.presto.operator.Driver) Ints(com.google.common.primitives.Ints) MaterializedResult(com.facebook.presto.testing.MaterializedResult) Assert.assertTrue(org.testng.Assert.assertTrue) DriverContext(com.facebook.presto.operator.DriverContext) Block(com.facebook.presto.common.block.Block) INNER(com.facebook.presto.sql.planner.plan.SpatialJoinNode.Type.INNER) SECONDS(java.util.concurrent.TimeUnit.SECONDS) Operator(com.facebook.presto.operator.Operator) ValuesOperator(com.facebook.presto.operator.ValuesOperator) DriverContext(com.facebook.presto.operator.DriverContext) TestingTaskContext(com.facebook.presto.testing.TestingTaskContext) TaskContext(com.facebook.presto.operator.TaskContext) RowPagesBuilder(com.facebook.presto.RowPagesBuilder) Page(com.facebook.presto.common.Page) InternalJoinFilterFunction(com.facebook.presto.operator.InternalJoinFilterFunction) GeoFunctions.stPoint(com.facebook.presto.plugin.geospatial.GeoFunctions.stPoint) SpatialJoinOperatorFactory(com.facebook.presto.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) PlanNodeId(com.facebook.presto.spi.plan.PlanNodeId) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) SpatialJoinOperatorFactory(com.facebook.presto.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) OperatorFactory(com.facebook.presto.operator.OperatorFactory) SpatialIndexBuilderOperatorFactory(com.facebook.presto.operator.SpatialIndexBuilderOperator.SpatialIndexBuilderOperatorFactory) PagesSpatialIndexFactory(com.facebook.presto.operator.PagesSpatialIndexFactory) Test(org.testng.annotations.Test)

Example 2 with GEOMETRY

use of com.facebook.presto.plugin.geospatial.GeometryType.GEOMETRY in project presto by prestodb.

the class AbstractTestGeoAggregationFunctions method assertAggregatedGeometries.

protected void assertAggregatedGeometries(String testDescription, String expectedWkt, String... wkts) {
    List<Slice> geometrySlices = Arrays.stream(wkts).map(text -> text == null ? null : OGCGeometry.fromText(text)).map(input -> input == null ? null : EsriGeometrySerde.serialize(input)).collect(Collectors.toList());
    // Add a custom equality assertion because the resulting geometry may have
    // its constituent points in a different order
    BiFunction<Object, Object, Boolean> equalityFunction = (left, right) -> {
        if (left == null && right == null) {
            return true;
        }
        if (left == null || right == null) {
            return false;
        }
        OGCGeometry leftGeometry = OGCGeometry.fromText(left.toString());
        OGCGeometry rightGeometry = OGCGeometry.fromText(right.toString());
        // Check for equality by getting the difference
        return leftGeometry.difference(rightGeometry).isEmpty() && rightGeometry.difference(leftGeometry).isEmpty();
    };
    // Test in forward and reverse order to verify that ordering doesn't affect the output
    assertAggregation(function, equalityFunction, testDescription, new Page(BlockAssertions.createSlicesBlock(geometrySlices)), expectedWkt);
    Collections.reverse(geometrySlices);
    assertAggregation(function, equalityFunction, testDescription, new Page(BlockAssertions.createSlicesBlock(geometrySlices)), expectedWkt);
}
Also used : FunctionAndTypeManager(com.facebook.presto.metadata.FunctionAndTypeManager) GEOMETRY(com.facebook.presto.plugin.geospatial.GeometryType.GEOMETRY) Page(com.facebook.presto.common.Page) Arrays(java.util.Arrays) Slice(io.airlift.slice.Slice) AbstractTestFunctions(com.facebook.presto.operator.scalar.AbstractTestFunctions) OGCGeometry(com.esri.core.geometry.ogc.OGCGeometry) BeforeClass(org.testng.annotations.BeforeClass) BiFunction(java.util.function.BiFunction) InternalAggregationFunction(com.facebook.presto.operator.aggregation.InternalAggregationFunction) Collectors(java.util.stream.Collectors) TypeSignatureProvider.fromTypes(com.facebook.presto.sql.analyzer.TypeSignatureProvider.fromTypes) EsriGeometrySerde(com.facebook.presto.geospatial.serde.EsriGeometrySerde) FunctionExtractor.extractFunctions(com.facebook.presto.metadata.FunctionExtractor.extractFunctions) BlockAssertions(com.facebook.presto.block.BlockAssertions) List(java.util.List) GeoPlugin(com.facebook.presto.plugin.geospatial.GeoPlugin) Collections(java.util.Collections) Type(com.facebook.presto.common.type.Type) AggregationTestUtils.assertAggregation(com.facebook.presto.operator.aggregation.AggregationTestUtils.assertAggregation) OGCGeometry(com.esri.core.geometry.ogc.OGCGeometry) Slice(io.airlift.slice.Slice) Page(com.facebook.presto.common.Page)

Aggregations

Page (com.facebook.presto.common.Page)2 GEOMETRY (com.facebook.presto.plugin.geospatial.GeometryType.GEOMETRY)2 OGCGeometry (com.esri.core.geometry.ogc.OGCGeometry)1 Threads.daemonThreadsNamed (com.facebook.airlift.concurrent.Threads.daemonThreadsNamed)1 RowPagesBuilder (com.facebook.presto.RowPagesBuilder)1 RowPagesBuilder.rowPagesBuilder (com.facebook.presto.RowPagesBuilder.rowPagesBuilder)1 TEST_SESSION (com.facebook.presto.SessionTestUtils.TEST_SESSION)1 BlockAssertions (com.facebook.presto.block.BlockAssertions)1 Block (com.facebook.presto.common.block.Block)1 DOUBLE (com.facebook.presto.common.type.DoubleType.DOUBLE)1 INTEGER (com.facebook.presto.common.type.IntegerType.INTEGER)1 Type (com.facebook.presto.common.type.Type)1 VARCHAR (com.facebook.presto.common.type.VarcharType.VARCHAR)1 KdbTree (com.facebook.presto.geospatial.KdbTree)1 KdbTreeUtils (com.facebook.presto.geospatial.KdbTreeUtils)1 Rectangle (com.facebook.presto.geospatial.Rectangle)1 EsriGeometrySerde (com.facebook.presto.geospatial.serde.EsriGeometrySerde)1 FunctionAndTypeManager (com.facebook.presto.metadata.FunctionAndTypeManager)1 FunctionExtractor.extractFunctions (com.facebook.presto.metadata.FunctionExtractor.extractFunctions)1 Driver (com.facebook.presto.operator.Driver)1