use of io.trino.split.SplitSource in project trino by trinodb.
the class ExtractSpatialJoins method loadKdbTree.
private static KdbTree loadKdbTree(String tableName, Session session, Metadata metadata, SplitManager splitManager, PageSourceManager pageSourceManager) {
QualifiedObjectName name = toQualifiedObjectName(tableName, session.getCatalog().get(), session.getSchema().get());
TableHandle tableHandle = metadata.getTableHandle(session, name).orElseThrow(() -> new TrinoException(INVALID_SPATIAL_PARTITIONING, format("Table not found: %s", name)));
Map<String, ColumnHandle> columnHandles = metadata.getColumnHandles(session, tableHandle);
List<ColumnHandle> visibleColumnHandles = columnHandles.values().stream().filter(handle -> !metadata.getColumnMetadata(session, tableHandle, handle).isHidden()).collect(toImmutableList());
checkSpatialPartitioningTable(visibleColumnHandles.size() == 1, "Expected single column for table %s, but found %s columns", name, columnHandles.size());
ColumnHandle kdbTreeColumn = Iterables.getOnlyElement(visibleColumnHandles);
Optional<KdbTree> kdbTree = Optional.empty();
try (SplitSource splitSource = splitManager.getSplits(session, tableHandle, UNGROUPED_SCHEDULING, EMPTY, alwaysTrue())) {
while (!Thread.currentThread().isInterrupted()) {
SplitBatch splitBatch = getFutureValue(splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), 1000));
List<Split> splits = splitBatch.getSplits();
for (Split split : splits) {
try (ConnectorPageSource pageSource = pageSourceManager.createPageSource(session, split, tableHandle, ImmutableList.of(kdbTreeColumn), DynamicFilter.EMPTY)) {
do {
getFutureValue(pageSource.isBlocked());
Page page = pageSource.getNextPage();
if (page != null && page.getPositionCount() > 0) {
checkSpatialPartitioningTable(kdbTree.isEmpty(), "Expected exactly one row for table %s, but found more", name);
checkSpatialPartitioningTable(page.getPositionCount() == 1, "Expected exactly one row for table %s, but found %s rows", name, page.getPositionCount());
String kdbTreeJson = VARCHAR.getSlice(page.getBlock(0), 0).toStringUtf8();
try {
kdbTree = Optional.of(KdbTreeUtils.fromJson(kdbTreeJson));
} catch (IllegalArgumentException e) {
checkSpatialPartitioningTable(false, "Invalid JSON string for KDB tree: %s", e.getMessage());
}
}
} while (!pageSource.isFinished());
} catch (IOException e) {
throw new UncheckedIOException(e);
}
}
if (splitBatch.isLastBatch()) {
break;
}
}
}
checkSpatialPartitioningTable(kdbTree.isPresent(), "Expected exactly one row for table %s, but got none", name);
return kdbTree.get();
}
use of io.trino.split.SplitSource in project trino by trinodb.
the class TestKuduIntegrationDynamicFilter method testIncompleteDynamicFilterTimeout.
@Test(timeOut = 30_000)
public void testIncompleteDynamicFilterTimeout() throws Exception {
QueryRunner runner = getQueryRunner();
TransactionManager transactionManager = runner.getTransactionManager();
TransactionId transactionId = transactionManager.beginTransaction(false);
Session session = Session.builder(getSession()).setCatalogSessionProperty("kudu", "dynamic_filtering_wait_timeout", "1s").build().beginTransactionId(transactionId, transactionManager, new AllowAllAccessControl());
QualifiedObjectName tableName = new QualifiedObjectName("kudu", "tpch", "orders");
Optional<TableHandle> tableHandle = runner.getMetadata().getTableHandle(session, tableName);
assertTrue(tableHandle.isPresent());
SplitSource splitSource = runner.getSplitManager().getSplits(session, tableHandle.get(), UNGROUPED_SCHEDULING, new IncompleteDynamicFilter(), alwaysTrue());
List<Split> splits = new ArrayList<>();
while (!splitSource.isFinished()) {
splits.addAll(splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), 1000).get().getSplits());
}
splitSource.close();
assertFalse(splits.isEmpty());
}
use of io.trino.split.SplitSource in project trino by trinodb.
the class AbstractOperatorBenchmark method getLocalQuerySplit.
private Split getLocalQuerySplit(Session session, TableHandle handle) {
SplitSource splitSource = localQueryRunner.getSplitManager().getSplits(session, handle, UNGROUPED_SCHEDULING, EMPTY, alwaysTrue());
List<Split> splits = new ArrayList<>();
while (!splitSource.isFinished()) {
splits.addAll(getNextBatch(splitSource));
}
checkArgument(splits.size() == 1, "Expected only one split for a local query, but got %s splits", splits.size());
return splits.get(0);
}
use of io.trino.split.SplitSource in project trino by trinodb.
the class LocalQueryRunner method createDrivers.
private List<Driver> createDrivers(Session session, Plan plan, OutputFactory outputFactory, TaskContext taskContext) {
if (printPlan) {
System.out.println(PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), plannerContext.getMetadata(), plannerContext.getFunctionManager(), plan.getStatsAndCosts(), session, 0, false));
}
SubPlan subplan = createSubPlans(session, plan, true);
if (!subplan.getChildren().isEmpty()) {
throw new AssertionError("Expected subplan to have no children");
}
TableExecuteContextManager tableExecuteContextManager = new TableExecuteContextManager();
tableExecuteContextManager.registerTableExecuteContextForQuery(taskContext.getQueryContext().getQueryId());
LocalExecutionPlanner executionPlanner = new LocalExecutionPlanner(plannerContext, new TypeAnalyzer(plannerContext, statementAnalyzerFactory), Optional.empty(), pageSourceManager, indexManager, nodePartitioningManager, pageSinkManager, null, expressionCompiler, pageFunctionCompiler, joinFilterFunctionCompiler, new IndexJoinLookupStats(), this.taskManagerConfig, spillerFactory, singleStreamSpillerFactory, partitioningSpillerFactory, new PagesIndex.TestingFactory(false), joinCompiler, operatorFactories, new OrderingCompiler(plannerContext.getTypeOperators()), new DynamicFilterConfig(), blockTypeOperators, tableExecuteContextManager, exchangeManagerRegistry);
// plan query
StageExecutionDescriptor stageExecutionDescriptor = subplan.getFragment().getStageExecutionDescriptor();
LocalExecutionPlan localExecutionPlan = executionPlanner.plan(taskContext, stageExecutionDescriptor, subplan.getFragment().getRoot(), subplan.getFragment().getPartitioningScheme().getOutputLayout(), plan.getTypes(), subplan.getFragment().getPartitionedSources(), outputFactory);
// generate splitAssignments
List<SplitAssignment> splitAssignments = new ArrayList<>();
long sequenceId = 0;
for (TableScanNode tableScan : findTableScanNodes(subplan.getFragment().getRoot())) {
TableHandle table = tableScan.getTable();
SplitSource splitSource = splitManager.getSplits(session, table, stageExecutionDescriptor.isScanGroupedExecution(tableScan.getId()) ? GROUPED_SCHEDULING : UNGROUPED_SCHEDULING, EMPTY, alwaysTrue());
ImmutableSet.Builder<ScheduledSplit> scheduledSplits = ImmutableSet.builder();
while (!splitSource.isFinished()) {
for (Split split : getNextBatch(splitSource)) {
scheduledSplits.add(new ScheduledSplit(sequenceId++, tableScan.getId(), split));
}
}
splitAssignments.add(new SplitAssignment(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(), false).addDriverContext();
Driver driver = driverFactory.createDriver(driverContext);
drivers.add(driver);
}
}
}
// add split assignments to the drivers
ImmutableSet<PlanNodeId> partitionedSources = ImmutableSet.copyOf(subplan.getFragment().getPartitionedSources());
for (SplitAssignment splitAssignment : splitAssignments) {
DriverFactory driverFactory = driverFactoriesBySource.get(splitAssignment.getPlanNodeId());
checkState(driverFactory != null);
boolean partitioned = partitionedSources.contains(driverFactory.getSourceId().get());
for (ScheduledSplit split : splitAssignment.getSplits()) {
DriverContext driverContext = taskContext.addPipelineContext(driverFactory.getPipelineId(), driverFactory.isInputDriver(), driverFactory.isOutputDriver(), partitioned).addDriverContext();
Driver driver = driverFactory.createDriver(driverContext);
driver.updateSplitAssignment(new SplitAssignment(split.getPlanNodeId(), ImmutableSet.of(split), true));
drivers.add(driver);
}
}
for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) {
driverFactory.noMoreDrivers();
}
return ImmutableList.copyOf(drivers);
}
Aggregations