Examples with CollectorContext io.crate.expression.reference.doc.lucene.CollectorContext used on opensource projects

Example 1 with CollectorContext

use of io.crate.expression.reference.doc.lucene.CollectorContext in project crate by crate.

the class LuceneBatchIteratorBenchmark method createLuceneBatchIterator.

@Setup
public void createLuceneBatchIterator() throws Exception {
    IndexWriter iw = new IndexWriter(new ByteBuffersDirectory(), new IndexWriterConfig(new StandardAnalyzer()));
    String columnName = "x";
    for (int i = 0; i < 10_000_000; i++) {
        Document doc = new Document();
        doc.add(new NumericDocValuesField(columnName, i));
        iw.addDocument(doc);
    }
    iw.commit();
    iw.forceMerge(1, true);
    indexSearcher = new IndexSearcher(DirectoryReader.open(iw));
    IntegerColumnReference columnReference = new IntegerColumnReference(columnName);
    columnRefs = Collections.singletonList(columnReference);
    collectorContext = new CollectorContext();
}

Example 2 with CollectorContext

use of io.crate.expression.reference.doc.lucene.CollectorContext in project crate by crate.

the class LuceneShardCollectorProvider method getOrderedCollector.

@Override
public OrderedDocCollector getOrderedCollector(RoutedCollectPhase phase, SharedShardContext sharedShardContext, CollectTask collectTask, boolean requiresRepeat) {
    RoutedCollectPhase collectPhase = phase.normalize(shardNormalizer, collectTask.txnCtx());
    CollectorContext collectorContext;
    InputFactory.Context<? extends LuceneCollectorExpression<?>> ctx;
    var searcher = sharedShardContext.acquireSearcher("ordered-collector: " + formatSource(phase));
    collectTask.addSearcher(sharedShardContext.readerId(), searcher);
    IndexService indexService = sharedShardContext.indexService();
    QueryShardContext queryShardContext = indexService.newQueryShardContext();
    final var queryContext = luceneQueryBuilder.convert(collectPhase.where(), collectTask.txnCtx(), indexService.mapperService(), indexShard.shardId().getIndexName(), queryShardContext, table, indexService.cache());
    ctx = docInputFactory.extractImplementations(collectTask.txnCtx(), collectPhase);
    collectorContext = new CollectorContext(sharedShardContext.readerId());
    int batchSize = collectPhase.shardQueueSize(localNodeId.get());
    if (LOGGER.isTraceEnabled()) {
        LOGGER.trace("[{}][{}] creating LuceneOrderedDocCollector. Expected number of rows to be collected: {}", sharedShardContext.indexShard().routingEntry().currentNodeId(), sharedShardContext.indexShard().shardId(), batchSize);
    }
    OptimizeQueryForSearchAfter optimizeQueryForSearchAfter = new OptimizeQueryForSearchAfter(collectPhase.orderBy(), queryContext.queryShardContext(), fieldTypeLookup);
    return new LuceneOrderedDocCollector(indexShard.shardId(), searcher.item(), queryContext.query(), queryContext.minScore(), Symbols.containsColumn(collectPhase.toCollect(), DocSysColumns.SCORE), batchSize, collectTask.getRamAccounting(), collectorContext, optimizeQueryForSearchAfter, LuceneSortGenerator.generateLuceneSort(collectTask.txnCtx(), collectorContext, collectPhase.orderBy(), docInputFactory, fieldTypeLookup), ctx.topLevelInputs(), ctx.expressions());
}

Example 3 with CollectorContext

use of io.crate.expression.reference.doc.lucene.CollectorContext in project crate by crate.

the class SortSymbolVisitor method customSortField.

private SortField customSortField(String name, final Symbol symbol, final SortSymbolContext context) {
    InputFactory.Context<? extends LuceneCollectorExpression<?>> inputContext = docInputFactory.getCtx(context.txnCtx);
    final Input<?> input = inputContext.add(symbol);
    final List<? extends LuceneCollectorExpression<?>> expressions = inputContext.expressions();
    final CollectorContext collectorContext = context.context;
    final boolean nullFirst = context.nullFirst;
    return new SortField(name, new FieldComparatorSource() {

        @Override
        public FieldComparator<?> newComparator(String fieldName, int numHits, int sortPos, boolean reversed) {
            for (int i = 0; i < expressions.size(); i++) {
                expressions.get(i).startCollect(collectorContext);
            }
            @SuppressWarnings("unchecked") DataType<Object> dataType = (DataType<Object>) symbol.valueType();
            Object nullSentinel = NullSentinelValues.nullSentinel(dataType, NullValueOrder.fromFlag(nullFirst), reversed);
            return new InputFieldComparator(numHits, expressions, input, // for `null` sentinels we need to have a comparator that can deal with that
            nullSentinel == null ? nullFirst ^ reversed ? Comparator.nullsFirst(dataType) : Comparator.nullsLast(dataType) : dataType, nullSentinel);
        }
    }, context.reverseFlag);
}

Example 4 with CollectorContext

use of io.crate.expression.reference.doc.lucene.CollectorContext in project crate by crate.

the class OrderedLuceneBatchIteratorFactoryTest method createOrderedCollector.

private LuceneOrderedDocCollector createOrderedCollector(IndexSearcher searcher, int shardId) {
    CollectorContext collectorContext = new CollectorContext();
    List<LuceneCollectorExpression<?>> expressions = Collections.singletonList(new OrderByCollectorExpression(reference, orderBy, o -> o));
    return new LuceneOrderedDocCollector(new ShardId("dummy", UUIDs.randomBase64UUID(), shardId), searcher, new MatchAllDocsQuery(), null, false, // batchSize < 10 to have at least one searchMore call.
    5, RamAccounting.NO_ACCOUNTING, collectorContext, f -> null, new Sort(new SortedNumericSortField(columnName, SortField.Type.LONG, reverseFlags[0])), expressions, expressions);
}

Example 5 with CollectorContext

use of io.crate.expression.reference.doc.lucene.CollectorContext in project crate by crate.

the class DocValuesGroupByOptimizedIteratorTest method test_group_by_doc_values_optimized_iterator_for_many_keys.

@Test
public void test_group_by_doc_values_optimized_iterator_for_many_keys() throws Exception {
    SumAggregation<?> sumAggregation = (SumAggregation<?>) functions.getQualified(Signature.aggregate(SumAggregation.NAME, DataTypes.LONG.getTypeSignature(), DataTypes.LONG.getTypeSignature()), List.of(DataTypes.LONG), DataTypes.LONG);
    var sumDocValuesAggregator = sumAggregation.getDocValueAggregator(List.of(new Reference(new ReferenceIdent(RelationName.fromIndexName("test"), "z"), RowGranularity.DOC, DataTypes.LONG, ColumnPolicy.DYNAMIC, IndexType.PLAIN, true, true, 0, null)), mock(DocTableInfo.class), List.of());
    var keyExpressions = List.of(new BytesRefColumnReference("x"), new LongColumnReference("y"));
    var keyRefs = List.of(new Reference(new ReferenceIdent(RelationName.fromIndexName("test"), "x"), RowGranularity.DOC, DataTypes.STRING, ColumnPolicy.DYNAMIC, IndexType.PLAIN, true, true, 1, null), new Reference(new ReferenceIdent(RelationName.fromIndexName("test"), "y"), RowGranularity.DOC, DataTypes.LONG, ColumnPolicy.DYNAMIC, IndexType.PLAIN, true, true, 2, null));
    var it = DocValuesGroupByOptimizedIterator.GroupByIterator.forManyKeys(List.of(sumDocValuesAggregator), indexSearcher, keyRefs, keyExpressions, RamAccounting.NO_ACCOUNTING, null, null, new MatchAllDocsQuery(), new CollectorContext());
    var rowConsumer = new TestingRowConsumer();
    rowConsumer.accept(it, null);
    assertThat(rowConsumer.getResult(), containsInAnyOrder(new Object[] { "0", 0L, 6L }, new Object[] { "1", 1L, 4L }));
}