Examples with InputFactory io.crate.expression.InputFactory used on opensource projects

Example 21 with InputFactory

use of io.crate.expression.InputFactory in project crate by crate.

the class NodeStatsTest method testNodeStatsIteratorContrat.

@Test
public void testNodeStatsIteratorContrat() throws Exception {
    List<Symbol> toCollect = new ArrayList<>();
    toCollect.add(idRef);
    when(collectPhase.toCollect()).thenReturn(toCollect);
    when(collectPhase.where()).thenReturn(Literal.BOOLEAN_TRUE);
    when(collectPhase.orderBy()).thenReturn(new OrderBy(Collections.singletonList(idRef)));
    List<Object[]> expectedResult = Arrays.asList(new Object[] { "nodeOne" }, new Object[] { "nodeTwo" });
    BatchIteratorTester tester = new BatchIteratorTester(() -> NodeStats.newInstance(transportNodeStatsAction, collectPhase, nodes, txnCtx, new InputFactory(nodeCtx)));
    tester.verifyResultAndEdgeCaseBehaviour(expectedResult);
}

Example 22 with InputFactory

use of io.crate.expression.InputFactory in project crate by crate.

the class NodeStatsTest method testRequestsIfRequired.

@Test
public void testRequestsIfRequired() throws Exception {
    List<Symbol> toCollect = new ArrayList<>();
    toCollect.add(idRef);
    toCollect.add(hostnameRef);
    when(collectPhase.toCollect()).thenReturn(toCollect);
    BatchIterator iterator = NodeStats.newInstance(transportNodeStatsAction, collectPhase, nodes, txnCtx, new InputFactory(nodeCtx));
    iterator.loadNextBatch();
    // Hostnames needs to be collected so requests need to be performed
    verify(transportNodeStatsAction).execute(eq("nodeOne"), any(NodeStatsRequest.class), any(ActionListener.class), eq(TimeValue.timeValueMillis(3000L)));
    verify(transportNodeStatsAction).execute(eq("nodeTwo"), any(NodeStatsRequest.class), any(ActionListener.class), eq(TimeValue.timeValueMillis(3000L)));
    verifyNoMoreInteractions(transportNodeStatsAction);
}

Example 23 with InputFactory

use of io.crate.expression.InputFactory in project crate by crate.

the class GroupByOptimizedIterator method tryOptimizeSingleStringKey.

@Nullable
static BatchIterator<Row> tryOptimizeSingleStringKey(IndexShard indexShard, DocTableInfo table, LuceneQueryBuilder luceneQueryBuilder, FieldTypeLookup fieldTypeLookup, BigArrays bigArrays, InputFactory inputFactory, DocInputFactory docInputFactory, RoutedCollectPhase collectPhase, CollectTask collectTask) {
    Collection<? extends Projection> shardProjections = shardProjections(collectPhase.projections());
    GroupProjection groupProjection = getSingleStringKeyGroupProjection(shardProjections);
    if (groupProjection == null) {
        return null;
    }
    assert groupProjection.keys().size() == 1 : "Must have 1 key if getSingleStringKeyGroupProjection returned a projection";
    Reference keyRef = getKeyRef(collectPhase.toCollect(), groupProjection.keys().get(0));
    if (keyRef == null) {
        // group by on non-reference
        return null;
    }
    keyRef = (Reference) DocReferences.inverseSourceLookup(keyRef);
    MappedFieldType keyFieldType = fieldTypeLookup.get(keyRef.column().fqn());
    if (keyFieldType == null || !keyFieldType.hasDocValues()) {
        return null;
    }
    if (Symbols.containsColumn(collectPhase.toCollect(), DocSysColumns.SCORE) || Symbols.containsColumn(collectPhase.where(), DocSysColumns.SCORE)) {
        // to keep the optimized implementation a bit simpler
        return null;
    }
    if (hasHighCardinalityRatio(() -> indexShard.acquireSearcher("group-by-cardinality-check"), keyFieldType.name())) {
        return null;
    }
    ShardId shardId = indexShard.shardId();
    SharedShardContext sharedShardContext = collectTask.sharedShardContexts().getOrCreateContext(shardId);
    var searcher = sharedShardContext.acquireSearcher("group-by-ordinals:" + formatSource(collectPhase));
    collectTask.addSearcher(sharedShardContext.readerId(), searcher);
    final QueryShardContext queryShardContext = sharedShardContext.indexService().newQueryShardContext();
    InputFactory.Context<? extends LuceneCollectorExpression<?>> docCtx = docInputFactory.getCtx(collectTask.txnCtx());
    docCtx.add(collectPhase.toCollect().stream()::iterator);
    InputFactory.Context<CollectExpression<Row, ?>> ctxForAggregations = inputFactory.ctxForAggregations(collectTask.txnCtx());
    ctxForAggregations.add(groupProjection.values());
    final List<CollectExpression<Row, ?>> aggExpressions = ctxForAggregations.expressions();
    List<AggregationContext> aggregations = ctxForAggregations.aggregations();
    List<? extends LuceneCollectorExpression<?>> expressions = docCtx.expressions();
    RamAccounting ramAccounting = collectTask.getRamAccounting();
    CollectorContext collectorContext = new CollectorContext(sharedShardContext.readerId());
    InputRow inputRow = new InputRow(docCtx.topLevelInputs());
    LuceneQueryBuilder.Context queryContext = luceneQueryBuilder.convert(collectPhase.where(), collectTask.txnCtx(), indexShard.mapperService(), indexShard.shardId().getIndexName(), queryShardContext, table, sharedShardContext.indexService().cache());
    return getIterator(bigArrays, searcher.item(), keyRef.column().fqn(), aggregations, expressions, aggExpressions, ramAccounting, collectTask.memoryManager(), collectTask.minNodeVersion(), inputRow, queryContext.query(), collectorContext, groupProjection.mode());
}

Example 24 with InputFactory

use of io.crate.expression.InputFactory in project crate by crate.

the class ProjectingRowConsumerTest method prepare.

@Before
public void prepare() {
    nodeCtx = createNodeContext();
    memoryManager = new OnHeapMemoryManager(usedBytes -> {
    });
    projectorFactory = new ProjectionToProjectorVisitor(clusterService, new NodeLimits(new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS)), new NoneCircuitBreakerService(), nodeCtx, THREAD_POOL, Settings.EMPTY, mock(TransportActionProvider.class, Answers.RETURNS_DEEP_STUBS), new InputFactory(nodeCtx), new EvaluatingNormalizer(nodeCtx, RowGranularity.SHARD, r -> Literal.ofUnchecked(r.valueType(), r.valueType().implicitCast("1")), null), t -> null, t -> null, Version.CURRENT, new ShardId("dummy", UUID.randomUUID().toString(), 0), Map.of(LocalFsFileOutputFactory.NAME, new LocalFsFileOutputFactory()));
}

Example 25 with InputFactory

use of io.crate.expression.InputFactory in project crate by crate.

the class CheckConstraintsTest method setUpExecutor.

@Before
public void setUpExecutor() throws Exception {
    SQLExecutor sqlExecutor = SQLExecutor.builder(clusterService).addTable("CREATE TABLE t (" + "    id int," + "    qty int," + "    sentinel boolean CONSTRAINT sentinel CHECK(sentinel)," + "    CONSTRAINT id_is_even CHECK(id % 2 = 0))").build();
    DocTableInfo docTableInfo = sqlExecutor.resolveTableInfo("t");
    TransactionContext txnCtx = CoordinatorTxnCtx.systemTransactionContext();
    checkConstraints = new CheckConstraints(txnCtx, new InputFactory(sqlExecutor.nodeCtx), FromSourceRefResolver.WITHOUT_PARTITIONED_BY_REFS, docTableInfo);
}