Search in sources :

Example 1 with SpatialJoinOperatorFactory

use of io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory in project trino by trinodb.

the class TestSpatialJoinOperator method assertSpatialJoin.

private void assertSpatialJoin(TaskContext taskContext, Type joinType, RowPagesBuilder buildPages, RowPagesBuilder probePages, MaterializedResult expected) {
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(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);
    assertOperatorEquals(joinOperatorFactory, driverContext, probePages.build(), expected);
}
Also used : SpatialJoinOperatorFactory(io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) PlanNodeId(io.trino.sql.planner.plan.PlanNodeId) DriverContext(io.trino.operator.DriverContext) OperatorFactory(io.trino.operator.OperatorFactory) SpatialIndexBuilderOperatorFactory(io.trino.operator.SpatialIndexBuilderOperator.SpatialIndexBuilderOperatorFactory) SpatialJoinOperatorFactory(io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) PagesSpatialIndexFactory(io.trino.operator.PagesSpatialIndexFactory)

Example 2 with SpatialJoinOperatorFactory

use of io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory in project trino by trinodb.

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);
}
Also used : SpatialJoinOperatorFactory(io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) PlanNodeId(io.trino.sql.planner.plan.PlanNodeId) Operator(io.trino.operator.Operator) ValuesOperator(io.trino.operator.ValuesOperator) DriverContext(io.trino.operator.DriverContext) TestingTaskContext(io.trino.testing.TestingTaskContext) TaskContext(io.trino.operator.TaskContext) RowPagesBuilder(io.trino.RowPagesBuilder) PipelineContext(io.trino.operator.PipelineContext) OperatorFactory(io.trino.operator.OperatorFactory) SpatialIndexBuilderOperatorFactory(io.trino.operator.SpatialIndexBuilderOperator.SpatialIndexBuilderOperatorFactory) SpatialJoinOperatorFactory(io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) PagesSpatialIndexFactory(io.trino.operator.PagesSpatialIndexFactory) MaterializedResult(io.trino.testing.MaterializedResult) GeoFunctions.stPoint(io.trino.plugin.geospatial.GeoFunctions.stPoint) Test(org.testng.annotations.Test)

Example 3 with SpatialJoinOperatorFactory

use of io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory in project trino by trinodb.

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);
}
Also used : SpatialJoinOperatorFactory(io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) PlanNodeId(io.trino.sql.planner.plan.PlanNodeId) DriverContext(io.trino.operator.DriverContext) TestingTaskContext(io.trino.testing.TestingTaskContext) TaskContext(io.trino.operator.TaskContext) RowPagesBuilder(io.trino.RowPagesBuilder) OperatorFactory(io.trino.operator.OperatorFactory) SpatialIndexBuilderOperatorFactory(io.trino.operator.SpatialIndexBuilderOperator.SpatialIndexBuilderOperatorFactory) SpatialJoinOperatorFactory(io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) PagesSpatialIndexFactory(io.trino.operator.PagesSpatialIndexFactory) MaterializedResult(io.trino.testing.MaterializedResult) Test(org.testng.annotations.Test)

Example 4 with SpatialJoinOperatorFactory

use of io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory in project trino by trinodb.

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

Example 5 with SpatialJoinOperatorFactory

use of io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory in project trino by trinodb.

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);
}
Also used : SpatialJoinOperatorFactory(io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) PlanNodeId(io.trino.sql.planner.plan.PlanNodeId) DriverContext(io.trino.operator.DriverContext) TestingTaskContext(io.trino.testing.TestingTaskContext) TaskContext(io.trino.operator.TaskContext) RowPagesBuilder(io.trino.RowPagesBuilder) OperatorFactory(io.trino.operator.OperatorFactory) SpatialIndexBuilderOperatorFactory(io.trino.operator.SpatialIndexBuilderOperator.SpatialIndexBuilderOperatorFactory) SpatialJoinOperatorFactory(io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) MaterializedResult(io.trino.testing.MaterializedResult) PagesSpatialIndexFactory(io.trino.operator.PagesSpatialIndexFactory) Test(org.testng.annotations.Test)

Aggregations

DriverContext (io.trino.operator.DriverContext)6 OperatorFactory (io.trino.operator.OperatorFactory)6 PagesSpatialIndexFactory (io.trino.operator.PagesSpatialIndexFactory)6 SpatialIndexBuilderOperatorFactory (io.trino.operator.SpatialIndexBuilderOperator.SpatialIndexBuilderOperatorFactory)6 SpatialJoinOperatorFactory (io.trino.operator.SpatialJoinOperator.SpatialJoinOperatorFactory)6 PlanNodeId (io.trino.sql.planner.plan.PlanNodeId)6 RowPagesBuilder (io.trino.RowPagesBuilder)5 TaskContext (io.trino.operator.TaskContext)5 MaterializedResult (io.trino.testing.MaterializedResult)5 TestingTaskContext (io.trino.testing.TestingTaskContext)5 Test (org.testng.annotations.Test)5 Operator (io.trino.operator.Operator)2 PipelineContext (io.trino.operator.PipelineContext)2 ValuesOperator (io.trino.operator.ValuesOperator)2 GeoFunctions.stPoint (io.trino.plugin.geospatial.GeoFunctions.stPoint)2 ImmutableList (com.google.common.collect.ImmutableList)1 Ints (com.google.common.primitives.Ints)1 ListenableFuture (com.google.common.util.concurrent.ListenableFuture)1 Threads.daemonThreadsNamed (io.airlift.concurrent.Threads.daemonThreadsNamed)1 Slice (io.airlift.slice.Slice)1