use of io.crate.expression.reference.doc.lucene.LuceneReferenceResolver in project crate by crate.
the class ReservoirSampler method getSamples.
private Samples getSamples(List<Reference> columns, int maxSamples, DocTableInfo docTable, Random random, Metadata metadata, CoordinatorTxnCtx coordinatorTxnCtx, List<Streamer> streamers, List<Engine.Searcher> searchersToRelease, RamAccounting ramAccounting) {
ramAccounting.addBytes(DataTypes.LONG.fixedSize() * maxSamples);
Reservoir<Long> fetchIdSamples = new Reservoir<>(maxSamples, random);
ArrayList<DocIdToRow> docIdToRowsFunctionPerReader = new ArrayList<>();
long totalNumDocs = 0;
long totalSizeInBytes = 0;
for (String index : docTable.concreteOpenIndices()) {
var indexMetadata = metadata.index(index);
if (indexMetadata == null) {
continue;
}
var indexService = indicesService.indexService(indexMetadata.getIndex());
if (indexService == null) {
continue;
}
var mapperService = indexService.mapperService();
FieldTypeLookup fieldTypeLookup = mapperService::fullName;
var ctx = new DocInputFactory(nodeCtx, new LuceneReferenceResolver(indexService.index().getName(), fieldTypeLookup, docTable.partitionedByColumns())).getCtx(coordinatorTxnCtx);
ctx.add(columns);
List<Input<?>> inputs = ctx.topLevelInputs();
List<? extends LuceneCollectorExpression<?>> expressions = ctx.expressions();
CollectorContext collectorContext = new CollectorContext();
for (LuceneCollectorExpression<?> expression : expressions) {
expression.startCollect(collectorContext);
}
for (IndexShard indexShard : indexService) {
if (!indexShard.routingEntry().primary()) {
continue;
}
try {
Engine.Searcher searcher = indexShard.acquireSearcher("update-table-statistics");
searchersToRelease.add(searcher);
totalNumDocs += searcher.getIndexReader().numDocs();
totalSizeInBytes += indexShard.storeStats().getSizeInBytes();
DocIdToRow docIdToRow = new DocIdToRow(searcher, inputs, expressions);
docIdToRowsFunctionPerReader.add(docIdToRow);
try {
// We do the sampling in 2 phases. First we get the docIds;
// then we retrieve the column values for the sampled docIds.
// we do this in 2 phases because the reservoir sampling might override previously seen
// items and we want to avoid unnecessary disk-lookup
var collector = new ReservoirCollector(fetchIdSamples, searchersToRelease.size() - 1);
searcher.search(new MatchAllDocsQuery(), collector);
} catch (IOException e) {
throw new UncheckedIOException(e);
}
} catch (IllegalIndexShardStateException | AlreadyClosedException ignored) {
}
}
}
var rowAccounting = new RowCellsAccountingWithEstimators(Symbols.typeView(columns), ramAccounting, 0);
ArrayList<Row> records = new ArrayList<>();
for (long fetchId : fetchIdSamples.samples()) {
int readerId = FetchId.decodeReaderId(fetchId);
DocIdToRow docIdToRow = docIdToRowsFunctionPerReader.get(readerId);
Object[] row = docIdToRow.apply(FetchId.decodeDocId(fetchId));
try {
rowAccounting.accountForAndMaybeBreak(row);
} catch (CircuitBreakingException e) {
LOGGER.info("Stopped gathering samples for `ANALYZE` operation because circuit breaker triggered. " + "Generating statistics with {} instead of {} records", records.size(), maxSamples);
break;
}
records.add(new RowN(row));
}
return new Samples(records, streamers, totalNumDocs, totalSizeInBytes);
}
Aggregations