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);
}
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);
}
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();
}
};
}
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();
}
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);
}
Aggregations