Examples with ReaderContext io.crate.execution.engine.fetch.ReaderContext used on opensource projects

Example 1 with ReaderContext

use of io.crate.execution.engine.fetch.ReaderContext in project crate by crate.

the class GroupByOptimizedIterator method applyAggregatesGroupedByKey.

private static Map<BytesRef, Object[]> applyAggregatesGroupedByKey(BigArrays bigArrays, IndexSearcher indexSearcher, String keyColumnName, List<AggregationContext> aggregations, List<? extends LuceneCollectorExpression<?>> expressions, List<CollectExpression<Row, ?>> aggExpressions, RamAccounting ramAccounting, MemoryManager memoryManager, Version minNodeVersion, InputRow inputRow, Query query, AtomicReference<Throwable> killed) throws IOException {
    final HashMap<BytesRef, Object[]> statesByKey = new HashMap<>();
    final Weight weight = indexSearcher.createWeight(indexSearcher.rewrite(query), ScoreMode.COMPLETE_NO_SCORES, 1f);
    final List<LeafReaderContext> leaves = indexSearcher.getTopReaderContext().leaves();
    Object[] nullStates = null;
    for (LeafReaderContext leaf : leaves) {
        raiseIfClosedOrKilled(killed);
        Scorer scorer = weight.scorer(leaf);
        if (scorer == null) {
            continue;
        }
        var readerContext = new ReaderContext(leaf);
        for (int i = 0, expressionsSize = expressions.size(); i < expressionsSize; i++) {
            expressions.get(i).setNextReader(readerContext);
        }
        SortedSetDocValues values = DocValues.getSortedSet(leaf.reader(), keyColumnName);
        try (ObjectArray<Object[]> statesByOrd = bigArrays.newObjectArray(values.getValueCount())) {
            DocIdSetIterator docs = scorer.iterator();
            Bits liveDocs = leaf.reader().getLiveDocs();
            for (int doc = docs.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = docs.nextDoc()) {
                raiseIfClosedOrKilled(killed);
                if (docDeleted(liveDocs, doc)) {
                    continue;
                }
                for (int i = 0, expressionsSize = expressions.size(); i < expressionsSize; i++) {
                    expressions.get(i).setNextDocId(doc);
                }
                for (int i = 0, expressionsSize = aggExpressions.size(); i < expressionsSize; i++) {
                    aggExpressions.get(i).setNextRow(inputRow);
                }
                if (values.advanceExact(doc)) {
                    long ord = values.nextOrd();
                    Object[] states = statesByOrd.get(ord);
                    if (states == null) {
                        statesByOrd.set(ord, initStates(aggregations, ramAccounting, memoryManager, minNodeVersion));
                    } else {
                        aggregateValues(aggregations, ramAccounting, memoryManager, states);
                    }
                    if (values.nextOrd() != SortedSetDocValues.NO_MORE_ORDS) {
                        throw new GroupByOnArrayUnsupportedException(keyColumnName);
                    }
                } else {
                    if (nullStates == null) {
                        nullStates = initStates(aggregations, ramAccounting, memoryManager, minNodeVersion);
                    } else {
                        aggregateValues(aggregations, ramAccounting, memoryManager, nullStates);
                    }
                }
            }
            for (long ord = 0; ord < statesByOrd.size(); ord++) {
                raiseIfClosedOrKilled(killed);
                Object[] states = statesByOrd.get(ord);
                if (states == null) {
                    continue;
                }
                BytesRef sharedKey = values.lookupOrd(ord);
                Object[] prevStates = statesByKey.get(sharedKey);
                if (prevStates == null) {
                    ramAccounting.addBytes(StringSizeEstimator.estimateSize(sharedKey) + HASH_MAP_ENTRY_OVERHEAD);
                    statesByKey.put(BytesRef.deepCopyOf(sharedKey), states);
                } else {
                    for (int i = 0; i < aggregations.size(); i++) {
                        AggregationContext aggregation = aggregations.get(i);
                        // noinspection unchecked
                        prevStates[i] = aggregation.function().reduce(ramAccounting, prevStates[i], states[i]);
                    }
                }
            }
        }
    }
    if (nullStates != null) {
        statesByKey.put(null, nullStates);
    }
    return statesByKey;
}

Example 2 with ReaderContext

use of io.crate.execution.engine.fetch.ReaderContext in project crate by crate.

the class ScoreDocRowFunction method apply.

@Nullable
@Override
public Row apply(@Nullable ScoreDoc input) {
    onScoreDoc.run();
    if (input == null) {
        return null;
    }
    FieldDoc fieldDoc = (FieldDoc) input;
    scorer.score(fieldDoc.score);
    for (int i = 0; i < orderByCollectorExpressions.size(); i++) {
        orderByCollectorExpressions.get(i).setNextFieldDoc(fieldDoc);
    }
    List<LeafReaderContext> leaves = indexReader.leaves();
    int readerIndex = ReaderUtil.subIndex(fieldDoc.doc, leaves);
    LeafReaderContext subReaderContext = leaves.get(readerIndex);
    int subDoc = fieldDoc.doc - subReaderContext.docBase;
    var readerContext = new ReaderContext(subReaderContext);
    for (LuceneCollectorExpression<?> expression : expressions) {
        try {
            expression.setNextReader(readerContext);
            expression.setNextDocId(subDoc);
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }
    return inputRow;
}

Example 3 with ReaderContext

use of io.crate.execution.engine.fetch.ReaderContext in project crate by crate.

the class BooleanColumnReferenceTest method testBooleanExpression.

@Test
public void testBooleanExpression() throws Exception {
    BooleanColumnReference booleanColumn = new BooleanColumnReference(column);
    booleanColumn.startCollect(ctx);
    booleanColumn.setNextReader(new ReaderContext(readerContext));
    IndexSearcher searcher = new IndexSearcher(readerContext.reader());
    TopDocs topDocs = searcher.search(new MatchAllDocsQuery(), 20);
    int i = 0;
    for (ScoreDoc doc : topDocs.scoreDocs) {
        booleanColumn.setNextDocId(doc.doc);
        assertThat(booleanColumn.value(), is(i % 2 == 0));
        i++;
    }
}

Example 4 with ReaderContext

use of io.crate.execution.engine.fetch.ReaderContext in project crate by crate.

the class FloatColumnReferenceTest method testFieldCacheExpression.

@Test
public void testFieldCacheExpression() throws Exception {
    FloatColumnReference floatColumn = new FloatColumnReference(column);
    floatColumn.startCollect(ctx);
    floatColumn.setNextReader(new ReaderContext(readerContext));
    IndexSearcher searcher = new IndexSearcher(readerContext.reader());
    TopDocs topDocs = searcher.search(new MatchAllDocsQuery(), 10);
    float f = -0.5f;
    for (ScoreDoc doc : topDocs.scoreDocs) {
        floatColumn.setNextDocId(doc.doc);
        assertThat(floatColumn.value(), is(f));
        f++;
    }
}

Example 5 with ReaderContext

use of io.crate.execution.engine.fetch.ReaderContext in project crate by crate.

the class IpColumnReferenceTest method testNullDocValuesDoNotResultInNPE.

@Test
public void testNullDocValuesDoNotResultInNPE() throws IOException {
    IpColumnReference ref = new IpColumnReference("missing_column");
    ref.startCollect(ctx);
    ref.setNextReader(new ReaderContext(readerContext));
    ref.setNextDocId(0);
    assertThat(ref.value(), Matchers.nullValue());
}