Examples with RowConsumer io.crate.data.RowConsumer used on opensource projects

Example 16 with RowConsumer

use of io.crate.data.RowConsumer in project crate by crate.

the class ProjectingRowConsumerTest method testConsumerRequiresScrollAndProjectorsDontSupportScrolling.

@Test
public void testConsumerRequiresScrollAndProjectorsDontSupportScrolling() {
    List<Symbol> arguments = Arrays.asList(Literal.of(2), new InputColumn(1, DataTypes.INTEGER));
    EqOperator op = (EqOperator) nodeCtx.functions().get(null, EqOperator.NAME, arguments, SearchPath.pathWithPGCatalogAndDoc());
    Function function = new Function(op.signature(), arguments, EqOperator.RETURN_TYPE);
    FilterProjection filterProjection = new FilterProjection(function, Arrays.asList(new InputColumn(0), new InputColumn(1)));
    RowConsumer delegateConsumerRequiresScroll = new DummyRowConsumer(true);
    RowConsumer projectingConsumer = ProjectingRowConsumer.create(delegateConsumerRequiresScroll, Collections.singletonList(filterProjection), UUID.randomUUID(), txnCtx, RamAccounting.NO_ACCOUNTING, memoryManager, projectorFactory);
    assertThat(projectingConsumer.requiresScroll(), is(true));
}

Example 17 with RowConsumer

use of io.crate.data.RowConsumer in project crate by crate.

the class BatchPortalTest method testEachStatementReceivesCorrectParams.

@Test
public void testEachStatementReceivesCorrectParams() throws Throwable {
    SQLExecutor sqlExecutor = SQLExecutor.builder(clusterService).addTable("create table t1 (x int)").build();
    Plan insertPlan = new Plan() {

        @Override
        public StatementType type() {
            return StatementType.INSERT;
        }

        @Override
        public void executeOrFail(DependencyCarrier executor, PlannerContext plannerContext, RowConsumer consumer, Row params, SubQueryResults subQueryResults) {
            consumer.accept(InMemoryBatchIterator.of(params, null), null);
        }
    };
    Planner planner = new Planner(Settings.EMPTY, clusterService, sqlExecutor.nodeCtx, new TableStats(), null, null, sqlExecutor.schemas(), new StubUserManager(), mock(SessionSettingRegistry.class)) {

        @Override
        public Plan plan(AnalyzedStatement analyzedStatement, PlannerContext plannerContext) {
            return insertPlan;
        }
    };
    DependencyCarrier executor = mock(DependencyCarrier.class, Answers.RETURNS_MOCKS);
    Session session = new Session(sqlExecutor.nodeCtx, sqlExecutor.analyzer, planner, new JobsLogs(() -> false), false, executor, AccessControl.DISABLED, SessionContext.systemSessionContext());
    session.parse("S_1", "insert into t1(x) values(1)", Collections.emptyList());
    session.bind("Portal", "S_1", Collections.emptyList(), null);
    final ArrayList<Object[]> s1Rows = new ArrayList<>();
    session.execute("Portal", 0, new BaseResultReceiver() {

        @Override
        public void setNextRow(Row row) {
            s1Rows.add(row.materialize());
        }
    });
    session.parse("S_2", "insert into t1(x) values(?)", Collections.emptyList());
    session.bind("Portal", "S_2", Collections.singletonList(2), null);
    final ArrayList<Object[]> s2Rows = new ArrayList<>();
    session.execute("Portal", 0, new BaseResultReceiver() {

        @Override
        public void setNextRow(Row row) {
            s2Rows.add(row.materialize());
        }
    });
    session.sync().get(5, TimeUnit.SECONDS);
    assertThat(s1Rows, contains(emptyArray()));
    assertThat(s2Rows, contains(arrayContaining(is(2))));
}

Example 18 with RowConsumer

use of io.crate.data.RowConsumer in project crate by crate.

the class JobSetup method prepareOnHandler.

public List<CompletableFuture<StreamBucket>> prepareOnHandler(SessionSettings sessionInfo, Collection<? extends NodeOperation> nodeOperations, RootTask.Builder taskBuilder, List<Tuple<ExecutionPhase, RowConsumer>> handlerPhases, SharedShardContexts sharedShardContexts) {
    Context context = new Context(clusterService.localNode().getId(), sessionInfo, taskBuilder, LOGGER, distributingConsumerFactory, nodeOperations, sharedShardContexts);
    for (Tuple<ExecutionPhase, RowConsumer> handlerPhase : handlerPhases) {
        context.registerLeaf(handlerPhase.v1(), handlerPhase.v2());
    }
    registerContextPhases(nodeOperations, context);
    LOGGER.trace("prepareOnHandler: nodeOperations={}, handlerPhases={}, targetSourceMap={}", nodeOperations, handlerPhases, context.opCtx.targetToSourceMap);
    IntHashSet leafs = new IntHashSet();
    for (Tuple<ExecutionPhase, RowConsumer> handlerPhase : handlerPhases) {
        ExecutionPhase phase = handlerPhase.v1();
        createContexts(phase, context);
        leafs.add(phase.phaseId());
    }
    leafs.addAll(context.opCtx.findLeafs());
    for (IntCursor cursor : leafs) {
        prepareSourceOperations(cursor.value, context);
    }
    assert context.opCtx.allNodeOperationContextsBuilt() : "some nodeOperations haven't been processed";
    return context.directResponseFutures;
}

Example 19 with RowConsumer

use of io.crate.data.RowConsumer in project crate by crate.

the class RestoreSnapshotPlan method executeOrFail.

@Override
public void executeOrFail(DependencyCarrier dependencies, PlannerContext plannerContext, RowConsumer consumer, Row parameters, SubQueryResults subQueryResults) {
    BoundRestoreSnapshot stmt = bind(restoreSnapshot, plannerContext.transactionContext(), dependencies.nodeContext(), parameters, subQueryResults, dependencies.schemas());
    var settings = stmt.settings();
    boolean ignoreUnavailable = IGNORE_UNAVAILABLE.get(settings);
    var transportActionProvider = dependencies.transportActionProvider();
    resolveIndexNames(restoreSnapshot.repository(), stmt.restoreTables(), ignoreUnavailable, transportActionProvider.transportGetSnapshotsAction()).whenComplete((ResolveIndicesAndTemplatesContext ctx, Throwable t) -> {
        if (t == null) {
            String[] indexNames = ctx.resolvedIndices().toArray(new String[0]);
            String[] templateNames = stmt.includeTables() && stmt.restoreTables().isEmpty() ? new String[] { ALL_TEMPLATES } : ctx.resolvedTemplates().toArray(new String[0]);
            // ignore_unavailable as set by statement
            IndicesOptions indicesOptions = IndicesOptions.fromOptions(ignoreUnavailable, true, true, false, IndicesOptions.lenientExpandOpen());
            RestoreSnapshotRequest request = new RestoreSnapshotRequest(restoreSnapshot.repository(), restoreSnapshot.snapshot()).indices(indexNames).templates(templateNames).indicesOptions(indicesOptions).settings(settings).waitForCompletion(WAIT_FOR_COMPLETION.get(settings)).includeIndices(stmt.includeTables()).includeAliases(stmt.includeTables()).includeCustomMetadata(stmt.includeCustomMetadata()).customMetadataTypes(stmt.customMetadataTypes()).includeGlobalSettings(stmt.includeGlobalSettings()).globalSettings(stmt.globalSettings());
            transportActionProvider.transportRestoreSnapshotAction().execute(request, new OneRowActionListener<>(consumer, r -> new Row1(r == null ? -1L : 1L)));
        }
    });
}

Example 20 with RowConsumer

use of io.crate.data.RowConsumer in project crate by crate.

the class ExplainPlan method executeOrFail.

@Override
public void executeOrFail(DependencyCarrier dependencies, PlannerContext plannerContext, RowConsumer consumer, Row params, SubQueryResults subQueryResults) {
    if (context != null) {
        assert subPlan instanceof LogicalPlan : "subPlan must be a LogicalPlan";
        LogicalPlan plan = (LogicalPlan) subPlan;
        /**
         * EXPLAIN ANALYZE does not support analyzing {@link io.crate.planner.MultiPhasePlan}s
         */
        if (plan.dependencies().isEmpty()) {
            UUID jobId = plannerContext.jobId();
            BaseResultReceiver resultReceiver = new BaseResultReceiver();
            RowConsumer noopRowConsumer = new RowConsumerToResultReceiver(resultReceiver, 0, t -> {
            });
            Timer timer = context.createTimer(Phase.Execute.name());
            timer.start();
            NodeOperationTree operationTree = LogicalPlanner.getNodeOperationTree(plan, dependencies, plannerContext, params, subQueryResults);
            resultReceiver.completionFuture().whenComplete(createResultConsumer(dependencies, consumer, jobId, timer, operationTree));
            LogicalPlanner.executeNodeOpTree(dependencies, plannerContext.transactionContext(), jobId, noopRowConsumer, true, operationTree);
        } else {
            consumer.accept(null, new UnsupportedOperationException("EXPLAIN ANALYZE does not support profiling multi-phase plans, " + "such as queries with scalar subselects."));
        }
    } else {
        if (subPlan instanceof LogicalPlan) {
            PrintContext printContext = new PrintContext();
            ((LogicalPlan) subPlan).print(printContext);
            consumer.accept(InMemoryBatchIterator.of(new Row1(printContext.toString()), SENTINEL), null);
        } else if (subPlan instanceof CopyFromPlan) {
            ExecutionPlan executionPlan = CopyFromPlan.planCopyFromExecution(((CopyFromPlan) subPlan).copyFrom(), dependencies.clusterService().state().nodes(), plannerContext, params, subQueryResults);
            String planAsJson = DataTypes.STRING.implicitCast(PlanPrinter.objectMap(executionPlan));
            consumer.accept(InMemoryBatchIterator.of(new Row1(planAsJson), SENTINEL), null);
        } else {
            consumer.accept(InMemoryBatchIterator.of(new Row1("EXPLAIN not supported for " + subPlan.getClass().getSimpleName()), SENTINEL), null);
        }
    }
}