Examples with LocalStateStoreProvider io.prestosql.statestore.LocalStateStoreProvider used on opensource projects

Example 11 with LocalStateStoreProvider

use of io.prestosql.statestore.LocalStateStoreProvider in project hetu-core by openlookeng.

the class TestHiveIntegrationSmokeTest method testRuseExchangeGroupSplitsMatchingBetweenProducerConsumer.

@Test
public void testRuseExchangeGroupSplitsMatchingBetweenProducerConsumer() {
    setUpNodes();
    NodeTaskMap nodeTasks = new NodeTaskMap(new FinalizerService());
    StageId stageId = new StageId(new QueryId("query"), 0);
    UUID uuid = UUID.randomUUID();
    PlanFragment testFragmentProducer = createTableScanPlanFragment("build", ReuseExchangeOperator.STRATEGY.REUSE_STRATEGY_PRODUCER, uuid, 1);
    PlanNodeId tableScanNodeId = new PlanNodeId("plan_id");
    StageExecutionPlan producerStageExecutionPlan = new StageExecutionPlan(testFragmentProducer, ImmutableMap.of(tableScanNodeId, new ConnectorAwareSplitSource(CONNECTOR_ID, createFixedSplitSource(0, TestingSplit::createRemoteSplit))), ImmutableList.of(), ImmutableMap.of(tableScanNodeId, new TableInfo(new QualifiedObjectName("test", TEST_SCHEMA, "test"), TupleDomain.all())));
    SqlStageExecution producerStage = createSqlStageExecution(stageId, new TestSqlTaskManager.MockLocationFactory().createStageLocation(stageId), producerStageExecutionPlan.getFragment(), producerStageExecutionPlan.getTables(), new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor), TEST_SESSION_REUSE, true, nodeTasks, remoteTaskExecutor, new NoOpFailureDetector(), new SplitSchedulerStats(), new DynamicFilterService(new LocalStateStoreProvider(new SeedStoreManager(new FileSystemClientManager()))), new QuerySnapshotManager(stageId.getQueryId(), NOOP_SNAPSHOT_UTILS, TEST_SESSION));
    Set<Split> splits = createAndGetSplits(10);
    Multimap<InternalNode, Split> producerAssignment = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values()), Optional.of(producerStage)).getAssignments();
    PlanFragment testFragmentConsumer = createTableScanPlanFragment("build", ReuseExchangeOperator.STRATEGY.REUSE_STRATEGY_CONSUMER, uuid, 1);
    StageExecutionPlan consumerStageExecutionPlan = new StageExecutionPlan(testFragmentConsumer, ImmutableMap.of(tableScanNodeId, new ConnectorAwareSplitSource(CONNECTOR_ID, createFixedSplitSource(0, TestingSplit::createRemoteSplit))), ImmutableList.of(), ImmutableMap.of(tableScanNodeId, new TableInfo(new QualifiedObjectName("test", TEST_SCHEMA, "test"), TupleDomain.all())));
    SqlStageExecution stage = createSqlStageExecution(stageId, new TestSqlTaskManager.MockLocationFactory().createStageLocation(stageId), consumerStageExecutionPlan.getFragment(), consumerStageExecutionPlan.getTables(), new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor), TEST_SESSION_REUSE, true, nodeTasks, remoteTaskExecutor, new NoOpFailureDetector(), new SplitSchedulerStats(), new DynamicFilterService(new LocalStateStoreProvider(new SeedStoreManager(new FileSystemClientManager()))), new QuerySnapshotManager(stageId.getQueryId(), NOOP_SNAPSHOT_UTILS, TEST_SESSION));
    Multimap<InternalNode, Split> consumerAssignment = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values()), Optional.of(stage)).getAssignments();
    assertEquals(consumerAssignment.size(), consumerAssignment.size());
    for (InternalNode node : consumerAssignment.keySet()) {
        List<Split> splitList = new ArrayList<>();
        List<Split> splitList2 = new ArrayList<>();
        boolean b = producerAssignment.containsEntry(node, consumerAssignment.get(node));
        Collection<Split> producerSplits = producerAssignment.get(node);
        Collection<Split> consumerSplits = producerAssignment.get(node);
        producerSplits.forEach(s -> splitList.add(s));
        List<Split> splitList1 = splitList.get(0).getSplits();
        consumerSplits.forEach(s -> splitList2.add(s));
        int i = 0;
        for (Split split3 : splitList1) {
            SplitKey splitKey1 = new SplitKey(split3, TEST_CATALOG, TEST_SCHEMA, TEST_TABLE);
            SplitKey splitKey2 = new SplitKey(splitList1.get(i), TEST_CATALOG, TEST_SCHEMA, TEST_TABLE);
            boolean f = splitKey1.equals(splitKey2);
            assertEquals(true, f);
            i++;
        }
    }
}

Example 12 with LocalStateStoreProvider

use of io.prestosql.statestore.LocalStateStoreProvider in project hetu-core by openlookeng.

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(), metadata, plan.getStatsAndCosts(), session, 0, false));
    }
    SubPlan subplan = planFragmenter.createSubPlans(session, plan, true, WarningCollector.NOOP);
    if (!subplan.getChildren().isEmpty()) {
        throw new AssertionError("Expected subplan to have no children");
    }
    NodeInfo nodeInfo = new NodeInfo("test");
    FileSystemClientManager fileSystemClientManager = new FileSystemClientManager();
    SeedStoreManager seedStoreManager = new SeedStoreManager(fileSystemClientManager);
    StateStoreProvider stateStoreProvider = new LocalStateStoreProvider(seedStoreManager);
    LocalExecutionPlanner executionPlanner = new LocalExecutionPlanner(metadata, new TypeAnalyzer(sqlParser, metadata), Optional.empty(), pageSourceManager, indexManager, nodePartitioningManager, pageSinkManager, null, expressionCompiler, pageFunctionCompiler, joinFilterFunctionCompiler, new IndexJoinLookupStats(), this.taskManagerConfig, spillerFactory, singleStreamSpillerFactory, partitioningSpillerFactory, new PagesIndex.TestingFactory(false), joinCompiler, new LookupJoinOperators(), new OrderingCompiler(), nodeInfo, stateStoreProvider, new StateStoreListenerManager(stateStoreProvider), new DynamicFilterCacheManager(), heuristicIndexerManager, cubeManager);
    // 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(), null, outputFactory, Optional.empty(), Optional.empty(), null);
    // generate sources
    List<TaskSource> sources = 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, null, Optional.empty(), Collections.emptyMap(), ImmutableSet.of(), tableScan.getStrategy() != ReuseExchangeOperator.STRATEGY.REUSE_STRATEGY_DEFAULT, tableScan.getId());
        ImmutableSet.Builder<ScheduledSplit> scheduledSplits = ImmutableSet.builder();
        while (!splitSource.isFinished()) {
            for (Split split : getNextBatch(splitSource)) {
                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(), false).addDriverContext();
                Driver driver = driverFactory.createDriver(driverContext);
                drivers.add(driver);
            }
        }
    }
    // add sources to the drivers
    ImmutableSet<PlanNodeId> partitionedSources = ImmutableSet.copyOf(subplan.getFragment().getPartitionedSources());
    for (TaskSource source : sources) {
        DriverFactory driverFactory = driverFactoriesBySource.get(source.getPlanNodeId());
        checkState(driverFactory != null);
        boolean partitioned = partitionedSources.contains(driverFactory.getSourceId().get());
        for (ScheduledSplit split : source.getSplits()) {
            DriverContext driverContext = taskContext.addPipelineContext(driverFactory.getPipelineId(), driverFactory.isInputDriver(), driverFactory.isOutputDriver(), partitioned).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.noMoreDrivers();
    }
    return ImmutableList.copyOf(drivers);
}

Example 13 with LocalStateStoreProvider

use of io.prestosql.statestore.LocalStateStoreProvider in project hetu-core by openlookeng.

the class TestSqlStageExecution method testFinalStageInfoInternal.

private void testFinalStageInfoInternal() throws Exception {
    NodeTaskMap nodeTaskMap = new NodeTaskMap(new FinalizerService());
    StageId stageId = new StageId(new QueryId("query"), 0);
    SqlStageExecution stage = createSqlStageExecution(stageId, new MockLocationFactory().createStageLocation(stageId), createExchangePlanFragment(), ImmutableMap.of(), new MockRemoteTaskFactory(executor, scheduledExecutor), TEST_SESSION, true, nodeTaskMap, executor, new NoOpFailureDetector(), new SplitSchedulerStats(), new DynamicFilterService(new LocalStateStoreProvider(new SeedStoreManager(new FileSystemClientManager()))), new QuerySnapshotManager(stageId.getQueryId(), NOOP_SNAPSHOT_UTILS, TEST_SESSION));
    stage.setOutputBuffers(createInitialEmptyOutputBuffers(ARBITRARY));
    // add listener that fetches stage info when the final status is available
    SettableFuture<StageInfo> finalStageInfo = SettableFuture.create();
    stage.addFinalStageInfoListener(finalStageInfo::set);
    // in a background thread add a ton of tasks
    CountDownLatch latch = new CountDownLatch(1000);
    Future<?> addTasksTask = executor.submit(() -> {
        try {
            for (int i = 0; i < 1_000_000; i++) {
                if (Thread.interrupted()) {
                    return;
                }
                InternalNode node = new InternalNode("source" + i, URI.create("http://10.0.0." + (i / 10_000) + ":" + (i % 10_000)), NodeVersion.UNKNOWN, false);
                stage.scheduleTask(node, i, OptionalInt.empty());
                latch.countDown();
            }
        } finally {
            while (latch.getCount() > 0) {
                latch.countDown();
            }
        }
    });
    // wait for some tasks to be created, and then abort the query
    latch.await(1, MINUTES);
    assertFalse(stage.getStageInfo().getTasks().isEmpty());
    stage.abort();
    // once the final stage info is available, verify that it is complete
    StageInfo stageInfo = finalStageInfo.get(1, MINUTES);
    assertFalse(stageInfo.getTasks().isEmpty());
    assertTrue(stageInfo.isCompleteInfo());
    assertSame(stage.getStageInfo(), stageInfo);
    // cancel the background thread adding tasks
    addTasksTask.cancel(true);
}

Example 14 with LocalStateStoreProvider

use of io.prestosql.statestore.LocalStateStoreProvider in project hetu-core by openlookeng.

the class TestUtil method getTestStage.

public static SqlStageExecution getTestStage(RowExpression expression) {
    StageId stageId = new StageId(new QueryId("query"), 0);
    SqlStageExecution stage = createSqlStageExecution(stageId, new TestSqlTaskManager.MockLocationFactory().createStageLocation(stageId), createExchangePlanFragment(expression), new HashMap<>(), new MockRemoteTaskFactory(executor, scheduledExecutor), TEST_SESSION, true, nodeTaskMap, executor, new NoOpFailureDetector(), new SplitSchedulerStats(), new DynamicFilterService(new LocalStateStoreProvider(new SeedStoreManager(new FileSystemClientManager()))), new QuerySnapshotManager(stageId.getQueryId(), TestingSnapshotUtils.NOOP_SNAPSHOT_UTILS, TEST_SESSION));
    stage.setOutputBuffers(createInitialEmptyOutputBuffers(ARBITRARY));
    return stage;
}

Example 15 with LocalStateStoreProvider

use of io.prestosql.statestore.LocalStateStoreProvider in project hetu-core by openlookeng.

the class DistributedQueryRunnerWithStateStore method setupStateStore.

public void setupStateStore() throws Exception {
    int port = SslSocketUtil.getAvailablePort();
    for (TestingPrestoServer server : servers) {
        server.installPlugin(new StateStoreManagerPlugin());
        // State Store
        StateStoreLauncher launcher = server.getInstance(Key.get(StateStoreLauncher.class));
        if (launcher instanceof EmbeddedStateStoreLauncher) {
            Set<String> ips = Sets.newHashSet(Arrays.asList("127.0.0.1"));
            Map<String, String> stateStoreProperties = new HashMap<>();
            stateStoreProperties.put(DISCOVERY_PORT_CONFIG_NAME, port + "");
            stateStoreProperties.putIfAbsent(HazelcastConstants.DISCOVERY_MODE_CONFIG_NAME, HazelcastConstants.DISCOVERY_MODE_TCPIP);
            this.stateStores.add(((EmbeddedStateStoreLauncher) launcher).launchStateStore(ips, stateStoreProperties));
        }
        launcher.launchStateStore();
        StateStoreProvider provider = server.getInstance(Key.get(StateStoreProvider.class));
        Seed seed = new FileBasedSeed("127.0.0.1:" + port, 0);
        SeedStore seedStore = new SeedStore() {

            @Override
            public Collection<Seed> add(Collection<Seed> seeds) throws IOException {
                return null;
            }

            @Override
            public Collection<Seed> get() throws IOException {
                return new ArrayList<>((Arrays.asList(seed)));
            }

            @Override
            public Collection<Seed> remove(Collection<Seed> seeds) throws IOException {
                return null;
            }

            @Override
            public Seed create(Map<String, String> properties) {
                return null;
            }

            @Override
            public String getName() {
                return null;
            }

            @Override
            public void setName(String name) {
            }
        };
        server.getInstance(Key.get(SeedStoreManager.class)).setSeedStore(SeedStoreSubType.HAZELCAST, seedStore);
        if (provider instanceof LocalStateStoreProvider) {
            Map<String, String> stateStoreProperties = new HashMap<>();
            stateStoreProperties.putIfAbsent(HazelcastConstants.DISCOVERY_MODE_CONFIG_NAME, HazelcastConstants.DISCOVERY_MODE_TCPIP);
            stateStoreProperties.put(DISCOVERY_PORT_CONFIG_NAME, port + "");
            ((LocalStateStoreProvider) provider).setStateStore("hazelcast", stateStoreProperties);
            ((LocalStateStoreProvider) provider).createStateCollections();
        }
        provider.loadStateStore();
        this.providers.add(provider);
    }
}