Search in sources :

Example 1 with DriverContext

use of com.facebook.presto.operator.DriverContext in project presto by prestodb.

the class TestOrcPageSourceMemoryTracking method testTableScanOperator.

@Test
public void testTableScanOperator() throws Exception {
    // Numbers used in assertions in this test may change when implementation is modified,
    // feel free to change them if they break in the future
    DriverContext driverContext = testPreparer.newDriverContext();
    SourceOperator operator = testPreparer.newTableScanOperator(driverContext);
    assertEquals(driverContext.getSystemMemoryUsage(), 0);
    long memoryUsage = -1;
    for (int i = 0; i < 20; i++) {
        assertFalse(operator.isFinished());
        Page page = operator.getOutput();
        assertNotNull(page);
        page.getBlock(1);
        if (memoryUsage == -1) {
            memoryUsage = driverContext.getSystemMemoryUsage();
            assertBetweenInclusive(memoryUsage, 460000L, 469999L);
        } else {
            assertEquals(driverContext.getSystemMemoryUsage(), memoryUsage);
        }
    }
    memoryUsage = -1;
    for (int i = 20; i < 40; i++) {
        assertFalse(operator.isFinished());
        Page page = operator.getOutput();
        assertNotNull(page);
        page.getBlock(1);
        if (memoryUsage == -1) {
            memoryUsage = driverContext.getSystemMemoryUsage();
            assertBetweenInclusive(memoryUsage, 460000L, 469999L);
        } else {
            assertEquals(driverContext.getSystemMemoryUsage(), memoryUsage);
        }
    }
    memoryUsage = -1;
    for (int i = 40; i < 50; i++) {
        assertFalse(operator.isFinished());
        Page page = operator.getOutput();
        assertNotNull(page);
        page.getBlock(1);
        if (memoryUsage == -1) {
            memoryUsage = driverContext.getSystemMemoryUsage();
            assertBetweenInclusive(memoryUsage, 360000L, 369999L);
        } else {
            assertEquals(driverContext.getSystemMemoryUsage(), memoryUsage);
        }
    }
    assertFalse(operator.isFinished());
    assertNull(operator.getOutput());
    assertTrue(operator.isFinished());
    assertEquals(driverContext.getSystemMemoryUsage(), 0);
}
Also used : DriverContext(com.facebook.presto.operator.DriverContext) SourceOperator(com.facebook.presto.operator.SourceOperator) Page(com.facebook.presto.spi.Page) Test(org.testng.annotations.Test)

Example 2 with DriverContext

use of com.facebook.presto.operator.DriverContext in project presto by prestodb.

the class LocalQueryRunner method createDrivers.

public List<Driver> createDrivers(Session session, @Language("SQL") String sql, OutputFactory outputFactory, TaskContext taskContext) {
    Plan plan = createPlan(session, sql);
    if (printPlan) {
        System.out.println(PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata, session));
    }
    SubPlan subplan = PlanFragmenter.createSubPlans(session, metadata, plan);
    if (!subplan.getChildren().isEmpty()) {
        throw new AssertionError("Expected subplan to have no children");
    }
    LocalExecutionPlanner executionPlanner = new LocalExecutionPlanner(metadata, sqlParser, Optional.empty(), pageSourceManager, indexManager, nodePartitioningManager, pageSinkManager, null, expressionCompiler, joinFilterFunctionCompiler, new IndexJoinLookupStats(), // make sure tests fail if compiler breaks
    new CompilerConfig().setInterpreterEnabled(false), new TaskManagerConfig().setTaskConcurrency(4), spillerFactory, blockEncodingSerde, new PagesIndex.TestingFactory(), new JoinCompiler(), new LookupJoinOperators(new JoinProbeCompiler()));
    // plan query
    LocalExecutionPlan localExecutionPlan = executionPlanner.plan(session, subplan.getFragment().getRoot(), subplan.getFragment().getPartitioningScheme().getOutputLayout(), plan.getTypes(), outputFactory);
    // generate sources
    List<TaskSource> sources = new ArrayList<>();
    long sequenceId = 0;
    for (TableScanNode tableScan : findTableScanNodes(subplan.getFragment().getRoot())) {
        TableLayoutHandle layout = tableScan.getLayout().get();
        SplitSource splitSource = splitManager.getSplits(session, layout);
        ImmutableSet.Builder<ScheduledSplit> scheduledSplits = ImmutableSet.builder();
        while (!splitSource.isFinished()) {
            for (Split split : getFutureValue(splitSource.getNextBatch(1000))) {
                scheduledSplits.add(new ScheduledSplit(sequenceId++, tableScan.getId(), split));
            }
        }
        sources.add(new TaskSource(tableScan.getId(), scheduledSplits.build(), true));
    }
    // create drivers
    List<Driver> drivers = new ArrayList<>();
    Map<PlanNodeId, DriverFactory> driverFactoriesBySource = new HashMap<>();
    for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) {
        for (int i = 0; i < driverFactory.getDriverInstances().orElse(1); i++) {
            if (driverFactory.getSourceId().isPresent()) {
                checkState(driverFactoriesBySource.put(driverFactory.getSourceId().get(), driverFactory) == null);
            } else {
                DriverContext driverContext = taskContext.addPipelineContext(driverFactory.getPipelineId(), driverFactory.isInputDriver(), driverFactory.isOutputDriver()).addDriverContext();
                Driver driver = driverFactory.createDriver(driverContext);
                drivers.add(driver);
            }
        }
    }
    // add sources to the drivers
    for (TaskSource source : sources) {
        DriverFactory driverFactory = driverFactoriesBySource.get(source.getPlanNodeId());
        checkState(driverFactory != null);
        for (ScheduledSplit split : source.getSplits()) {
            DriverContext driverContext = taskContext.addPipelineContext(driverFactory.getPipelineId(), driverFactory.isInputDriver(), driverFactory.isOutputDriver()).addDriverContext();
            Driver driver = driverFactory.createDriver(driverContext);
            driver.updateSource(new TaskSource(split.getPlanNodeId(), ImmutableSet.of(split), true));
            drivers.add(driver);
        }
    }
    for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) {
        driverFactory.close();
    }
    return ImmutableList.copyOf(drivers);
}
Also used : DriverContext(com.facebook.presto.operator.DriverContext) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) Driver(com.facebook.presto.operator.Driver) TaskManagerConfig(com.facebook.presto.execution.TaskManagerConfig) PagesIndex(com.facebook.presto.operator.PagesIndex) PlanNodeId(com.facebook.presto.sql.planner.plan.PlanNodeId) JoinProbeCompiler(com.facebook.presto.sql.gen.JoinProbeCompiler) ImmutableSet(com.google.common.collect.ImmutableSet) DriverFactory(com.facebook.presto.operator.DriverFactory) LookupJoinOperators(com.facebook.presto.operator.LookupJoinOperators) JoinCompiler(com.facebook.presto.sql.gen.JoinCompiler) ScheduledSplit(com.facebook.presto.ScheduledSplit) LocalExecutionPlanner(com.facebook.presto.sql.planner.LocalExecutionPlanner) IndexJoinLookupStats(com.facebook.presto.operator.index.IndexJoinLookupStats) LocalExecutionPlan(com.facebook.presto.sql.planner.LocalExecutionPlanner.LocalExecutionPlan) Plan(com.facebook.presto.sql.planner.Plan) SubPlan(com.facebook.presto.sql.planner.SubPlan) CompilerConfig(com.facebook.presto.sql.planner.CompilerConfig) TableLayoutHandle(com.facebook.presto.metadata.TableLayoutHandle) Constraint(com.facebook.presto.spi.Constraint) LocalExecutionPlan(com.facebook.presto.sql.planner.LocalExecutionPlanner.LocalExecutionPlan) TableScanNode(com.facebook.presto.sql.planner.plan.TableScanNode) SplitSource(com.facebook.presto.split.SplitSource) ScheduledSplit(com.facebook.presto.ScheduledSplit) Split(com.facebook.presto.metadata.Split) SubPlan(com.facebook.presto.sql.planner.SubPlan) TaskSource(com.facebook.presto.TaskSource)

Example 3 with DriverContext

use of com.facebook.presto.operator.DriverContext in project presto by prestodb.

the class LocalQueryRunner method createTableScanOperator.

public OperatorFactory createTableScanOperator(Session session, int operatorId, PlanNodeId planNodeId, String tableName, String... columnNames) {
    checkArgument(session.getCatalog().isPresent(), "catalog not set");
    checkArgument(session.getSchema().isPresent(), "schema not set");
    // look up the table
    QualifiedObjectName qualifiedTableName = new QualifiedObjectName(session.getCatalog().get(), session.getSchema().get(), tableName);
    TableHandle tableHandle = metadata.getTableHandle(session, qualifiedTableName).orElse(null);
    checkArgument(tableHandle != null, "Table %s does not exist", qualifiedTableName);
    // lookup the columns
    Map<String, ColumnHandle> allColumnHandles = metadata.getColumnHandles(session, tableHandle);
    ImmutableList.Builder<ColumnHandle> columnHandlesBuilder = ImmutableList.builder();
    ImmutableList.Builder<Type> columnTypesBuilder = ImmutableList.builder();
    for (String columnName : columnNames) {
        ColumnHandle columnHandle = allColumnHandles.get(columnName);
        checkArgument(columnHandle != null, "Table %s does not have a column %s", tableName, columnName);
        columnHandlesBuilder.add(columnHandle);
        ColumnMetadata columnMetadata = metadata.getColumnMetadata(session, tableHandle, columnHandle);
        columnTypesBuilder.add(columnMetadata.getType());
    }
    List<ColumnHandle> columnHandles = columnHandlesBuilder.build();
    List<Type> columnTypes = columnTypesBuilder.build();
    // get the split for this table
    List<TableLayoutResult> layouts = metadata.getLayouts(session, tableHandle, Constraint.alwaysTrue(), Optional.empty());
    Split split = getLocalQuerySplit(session, layouts.get(0).getLayout().getHandle());
    return new OperatorFactory() {

        @Override
        public List<Type> getTypes() {
            return columnTypes;
        }

        @Override
        public Operator createOperator(DriverContext driverContext) {
            OperatorContext operatorContext = driverContext.addOperatorContext(operatorId, planNodeId, "BenchmarkSource");
            ConnectorPageSource pageSource = pageSourceManager.createPageSource(session, split, columnHandles);
            return new PageSourceOperator(pageSource, columnTypes, operatorContext);
        }

        @Override
        public void close() {
        }

        @Override
        public OperatorFactory duplicate() {
            throw new UnsupportedOperationException();
        }
    };
}
Also used : ColumnHandle(com.facebook.presto.spi.ColumnHandle) DriverContext(com.facebook.presto.operator.DriverContext) ColumnMetadata(com.facebook.presto.spi.ColumnMetadata) ImmutableList(com.google.common.collect.ImmutableList) TableLayoutResult(com.facebook.presto.metadata.TableLayoutResult) ConnectorPageSource(com.facebook.presto.spi.ConnectorPageSource) QualifiedObjectName(com.facebook.presto.metadata.QualifiedObjectName) PageSourceOperator(com.facebook.presto.operator.PageSourceOperator) Type(com.facebook.presto.spi.type.Type) OperatorFactory(com.facebook.presto.operator.OperatorFactory) OperatorContext(com.facebook.presto.operator.OperatorContext) TableHandle(com.facebook.presto.metadata.TableHandle) ScheduledSplit(com.facebook.presto.ScheduledSplit) Split(com.facebook.presto.metadata.Split)

Example 4 with DriverContext

use of com.facebook.presto.operator.DriverContext in project presto by prestodb.

the class TestMemoryPools method setupConsumeRevocableMemory.

private RevocableMemoryOperator setupConsumeRevocableMemory(DataSize reservedPerPage, long numberOfPages) {
    AtomicReference<RevocableMemoryOperator> createOperator = new AtomicReference<>();
    setUp(() -> {
        DriverContext driverContext = taskContext.addPipelineContext(0, false, false, false).addDriverContext();
        OperatorContext revokableOperatorContext = driverContext.addOperatorContext(Integer.MAX_VALUE, new PlanNodeId("revokable_operator"), TableScanOperator.class.getSimpleName());
        OutputFactory outputFactory = new PageConsumerOutputFactory(types -> (page -> {
        }));
        Operator outputOperator = outputFactory.createOutputOperator(2, new PlanNodeId("output"), ImmutableList.of(), Function.identity(), Optional.empty(), new TestingPagesSerdeFactory()).createOperator(driverContext);
        RevocableMemoryOperator revocableMemoryOperator = new RevocableMemoryOperator(revokableOperatorContext, reservedPerPage, numberOfPages);
        createOperator.set(revocableMemoryOperator);
        Driver driver = Driver.createDriver(driverContext, revocableMemoryOperator, outputOperator);
        return ImmutableList.of(driver);
    });
    return createOperator.get();
}
Also used : PlanNodeId(com.facebook.presto.spi.plan.PlanNodeId) Page(com.facebook.presto.common.Page) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) Assert.assertNull(org.testng.Assert.assertNull) MEGABYTE(io.airlift.units.DataSize.Unit.MEGABYTE) Supplier(com.google.common.base.Supplier) Assert.assertEquals(org.testng.Assert.assertEquals) Test(org.testng.annotations.Test) AfterMethod(org.testng.annotations.AfterMethod) AtomicReference(java.util.concurrent.atomic.AtomicReference) Function(java.util.function.Function) JsonCodec.listJsonCodec(com.facebook.airlift.json.JsonCodec.listJsonCodec) SpillSpaceTracker(com.facebook.presto.spiller.SpillSpaceTracker) GIGABYTE(io.airlift.units.DataSize.Unit.GIGABYTE) ImmutableList(com.google.common.collect.ImmutableList) Map(java.util.Map) Operator(com.facebook.presto.operator.Operator) LocalQueryRunner.queryRunnerWithInitialTransaction(com.facebook.presto.testing.LocalQueryRunner.queryRunnerWithInitialTransaction) Assert.assertFalse(org.testng.Assert.assertFalse) LocalQueryRunner(com.facebook.presto.testing.LocalQueryRunner) PageConsumerOutputFactory(com.facebook.presto.testing.PageConsumerOperator.PageConsumerOutputFactory) TaskContext(com.facebook.presto.operator.TaskContext) TestingTaskContext.createTaskContext(com.facebook.presto.testing.TestingTaskContext.createTaskContext) ImmutableMap(com.google.common.collect.ImmutableMap) Session(com.facebook.presto.Session) Assert.fail(org.testng.Assert.fail) LocalMemoryContext(com.facebook.presto.memory.context.LocalMemoryContext) TpchConnectorFactory(com.facebook.presto.tpch.TpchConnectorFactory) OutputFactory(com.facebook.presto.operator.OutputFactory) TestingGcMonitor(com.facebook.airlift.stats.TestingGcMonitor) Driver(com.facebook.presto.operator.Driver) TestingSession.testSessionBuilder(com.facebook.presto.testing.TestingSession.testSessionBuilder) TestingPagesSerdeFactory(com.facebook.presto.execution.buffer.TestingPagesSerdeFactory) Preconditions.checkState(com.google.common.base.Preconditions.checkState) MemoryPoolId(com.facebook.presto.spi.memory.MemoryPoolId) AtomicLong(java.util.concurrent.atomic.AtomicLong) Futures(com.google.common.util.concurrent.Futures) DataSize(io.airlift.units.DataSize) List(java.util.List) Predicate(com.google.common.base.Predicate) QueryId(com.facebook.presto.spi.QueryId) Optional(java.util.Optional) Assert.assertTrue(org.testng.Assert.assertTrue) DriverContext(com.facebook.presto.operator.DriverContext) TableScanOperator(com.facebook.presto.operator.TableScanOperator) TaskMemoryReservationSummary(com.facebook.presto.operator.TaskMemoryReservationSummary) OperatorContext(com.facebook.presto.operator.OperatorContext) BYTE(io.airlift.units.DataSize.Unit.BYTE) Operator(com.facebook.presto.operator.Operator) TableScanOperator(com.facebook.presto.operator.TableScanOperator) DriverContext(com.facebook.presto.operator.DriverContext) TableScanOperator(com.facebook.presto.operator.TableScanOperator) PageConsumerOutputFactory(com.facebook.presto.testing.PageConsumerOperator.PageConsumerOutputFactory) TestingPagesSerdeFactory(com.facebook.presto.execution.buffer.TestingPagesSerdeFactory) Driver(com.facebook.presto.operator.Driver) AtomicReference(java.util.concurrent.atomic.AtomicReference) PlanNodeId(com.facebook.presto.spi.plan.PlanNodeId) OperatorContext(com.facebook.presto.operator.OperatorContext) PageConsumerOutputFactory(com.facebook.presto.testing.PageConsumerOperator.PageConsumerOutputFactory) OutputFactory(com.facebook.presto.operator.OutputFactory)

Example 5 with DriverContext

use of com.facebook.presto.operator.DriverContext in project presto by prestodb.

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));
    addGeometryPartitionRows(buildPages, POLYGON_A, "A");
    buildPages.row(null, "null", null);
    buildPages.pageBreak();
    addGeometryPartitionRows(buildPages, POLYGON_B, "B");
    addGeometryPartitionRows(buildPages, POLYGON_C, "C");
    addGeometryPartitionRows(buildPages, POLYGON_D, "D");
    RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER));
    addGeometryPartitionRows(probePages, POINT_X, "x");
    probePages.row(null, "null", null);
    addGeometryPartitionRows(probePages, POINT_Y, "y");
    probePages.pageBreak();
    addGeometryPartitionRows(probePages, POINT_Z, "z");
    addGeometryPartitionRows(probePages, POINT_W, "w");
    addGeometryPartitionRows(probePages, POINT_V, "v");
    addGeometryPartitionRows(probePages, MULTIPOINT_U, "u");
    probePages.pageBreak();
    addGeometryPartitionRows(probePages, MULTIPOINT_T, "t");
    addGeometryPartitionRows(probePages, POINT_S, "s");
    addGeometryPartitionRows(probePages, MULTIPOINT_R, "r");
    addGeometryPartitionRows(probePages, POINT_Q, "q");
    MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("x", "A").row("y", "A").row("y", "B").row("z", "B").row("v", "C").row("u", "C").row("t", "C").row("s", "D").row("r", "C").row("r", "D").build();
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.intersects(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);
    assertOperatorEqualsIgnoreOrder(joinOperatorFactory, driverContext, probePages.build(), expected);
}
Also used : SpatialJoinOperatorFactory(com.facebook.presto.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) PlanNodeId(com.facebook.presto.spi.plan.PlanNodeId) DriverContext(com.facebook.presto.operator.DriverContext) TestingTaskContext(com.facebook.presto.testing.TestingTaskContext) TaskContext(com.facebook.presto.operator.TaskContext) RowPagesBuilder(com.facebook.presto.RowPagesBuilder) SpatialJoinOperatorFactory(com.facebook.presto.operator.SpatialJoinOperator.SpatialJoinOperatorFactory) OperatorFactory(com.facebook.presto.operator.OperatorFactory) SpatialIndexBuilderOperatorFactory(com.facebook.presto.operator.SpatialIndexBuilderOperator.SpatialIndexBuilderOperatorFactory) MaterializedResult(com.facebook.presto.testing.MaterializedResult) PagesSpatialIndexFactory(com.facebook.presto.operator.PagesSpatialIndexFactory) Test(org.testng.annotations.Test)

Aggregations

DriverContext (com.facebook.presto.operator.DriverContext)24 PlanNodeId (com.facebook.presto.spi.plan.PlanNodeId)13 Test (org.testng.annotations.Test)12 TaskContext (com.facebook.presto.operator.TaskContext)10 OperatorFactory (com.facebook.presto.operator.OperatorFactory)8 Driver (com.facebook.presto.operator.Driver)6 OperatorContext (com.facebook.presto.operator.OperatorContext)6 PagesSpatialIndexFactory (com.facebook.presto.operator.PagesSpatialIndexFactory)5 SpatialIndexBuilderOperatorFactory (com.facebook.presto.operator.SpatialIndexBuilderOperator.SpatialIndexBuilderOperatorFactory)5 SpatialJoinOperatorFactory (com.facebook.presto.operator.SpatialJoinOperator.SpatialJoinOperatorFactory)5 ListenableFuture (com.google.common.util.concurrent.ListenableFuture)5 RowPagesBuilder (com.facebook.presto.RowPagesBuilder)4 Page (com.facebook.presto.common.Page)4 LocalMemoryContext (com.facebook.presto.memory.context.LocalMemoryContext)4 Split (com.facebook.presto.metadata.Split)4 DriverFactory (com.facebook.presto.operator.DriverFactory)4 PipelineContext (com.facebook.presto.operator.PipelineContext)4 QueryId (com.facebook.presto.spi.QueryId)4 MaterializedResult (com.facebook.presto.testing.MaterializedResult)4 TestingTaskContext (com.facebook.presto.testing.TestingTaskContext)4