Search in sources :

Example 6 with AggregationContext

use of io.crate.execution.engine.aggregation.AggregationContext 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());
}
Also used : AggregationContext(io.crate.execution.engine.aggregation.AggregationContext) InputFactory(io.crate.expression.InputFactory) RamAccounting(io.crate.breaker.RamAccounting) AtomicReference(java.util.concurrent.atomic.AtomicReference) Reference(io.crate.metadata.Reference) ShardId(org.elasticsearch.index.shard.ShardId) LuceneQueryBuilder(io.crate.lucene.LuceneQueryBuilder) MappedFieldType(org.elasticsearch.index.mapper.MappedFieldType) InputRow(io.crate.expression.InputRow) QueryShardContext(org.elasticsearch.index.query.QueryShardContext) CollectorContext(io.crate.expression.reference.doc.lucene.CollectorContext) GroupProjection(io.crate.execution.dsl.projection.GroupProjection) SharedShardContext(io.crate.execution.jobs.SharedShardContext) Nullable(javax.annotation.Nullable)

Aggregations

AggregationContext (io.crate.execution.engine.aggregation.AggregationContext)6 CollectExpression (io.crate.execution.engine.collect.CollectExpression)2 Aggregation (io.crate.expression.symbol.Aggregation)2 Function (io.crate.expression.symbol.Function)2 InputColumn (io.crate.expression.symbol.InputColumn)2 Symbol (io.crate.expression.symbol.Symbol)2 CrateDummyClusterServiceUnitTest (io.crate.test.integration.CrateDummyClusterServiceUnitTest)2 BytesRef (org.apache.lucene.util.BytesRef)2 Test (org.junit.Test)2 RamAccounting (io.crate.breaker.RamAccounting)1 Input (io.crate.data.Input)1 Row (io.crate.data.Row)1 RowN (io.crate.data.RowN)1 GroupByOnArrayUnsupportedException (io.crate.exceptions.GroupByOnArrayUnsupportedException)1 GroupProjection (io.crate.execution.dsl.projection.GroupProjection)1 AggregationFunction (io.crate.execution.engine.aggregation.AggregationFunction)1 CountAggregation (io.crate.execution.engine.aggregation.impl.CountAggregation)1 ReaderContext (io.crate.execution.engine.fetch.ReaderContext)1 SharedShardContext (io.crate.execution.jobs.SharedShardContext)1 InputFactory (io.crate.expression.InputFactory)1