mkhludnev commented on code in PR #4749: URL: https://github.com/apache/solr/pull/4749#discussion_r3825859089
########## solr/core/src/java/org/apache/solr/search/join/aijoin/AIJoinUtil.java: ########## @@ -0,0 +1,580 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.solr.search.join.aijoin; + +import java.io.IOException; +import java.util.Arrays; +import java.util.HashSet; +import java.util.Set; +import java.util.regex.Pattern; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.FieldInfosFormat; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.DocValues; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.FilterCodecReader; +import org.apache.lucene.index.FilterLeafReader; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.ParallelCompositeReader; +import org.apache.lucene.index.ParallelLeafReader; +import org.apache.lucene.index.SegmentCommitInfo; +import org.apache.lucene.index.SegmentReader; +import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.index.SortedSetDocValues; +import org.apache.lucene.index.TermsEnum; +import org.apache.lucene.search.BulkScorer; +import org.apache.lucene.search.DocIdSet; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.DocIdStream; +import org.apache.lucene.search.LeafCollector; +import org.apache.lucene.search.Scorable; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.store.FilterDirectory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.util.Accountable; +import org.apache.lucene.util.BitDocIdSet; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.FixedBitSet; +import org.apache.lucene.util.RamUsageEstimator; +import org.apache.lucene.util.RoaringDocIdSet; +import org.apache.lucene.util.StringHelper; +import org.slf4j.Logger; +import org.slf4j.event.Level; + +/** + * Column-building and addressing helpers for the auxiliary join index managed by {@link + * AIJoinIndex}: for every (from-segment, to-segment) pair it produces a SORTED_NUMERIC column named + * {@link #pairFieldName}, whose doc number is the from-side doc id and whose value is the to-side + * doc id whose {@code toField} term equals the from doc's {@code fromField} term, plus two + * companion edges columns persisting the pair's {min, max} from-doc and to-doc bounds. + */ +final class AIJoinUtil { + + /** Suffix of the always-written column persisting a pair's {min, max} from-doc edges. */ + static final String FROM_EDGES_PREFIX = "fromDoc_edges_"; // TODO reduce to the singe letter + + /** Suffix of the always-written column persisting a pair's {min, max} to-doc edges. */ + static final String TO_EDGES_PREFIX = "toDoc_edges_"; + + /** main join colums for join index to_doc_num[from_docnum] */ + static final String TO_DOC_VAL_BY_FROM_DOCNUM = "join_toDoc_"; + + static final String TO_COUNT_PREFIX = "num_toDoc_"; + + private AIJoinUtil() {} + + /** + * Configurable level for the {@code AIJOIN evt=...} diagnostic logs; defaults to {@code INFO}, + * override with the {@code solr.aijoin.log.level} system property (or {@code + * SOLR_AIJOIN_LOG_LEVEL} env var). + */ + static final Level AIJOIN_LOG_LEVEL = Level.TRACE; + + /** Whether the AIJOIN diagnostic logs would emit at the configured level. */ + static boolean diagnosticsEnabled(Logger log) { + return log.isEnabledForLevel(AIJOIN_LOG_LEVEL); + } + + /** Emits an AIJOIN diagnostic line at the configured level. */ + static void logDiagnostic(Logger log, String message, Object... args) { + log.atLevel(AIJOIN_LOG_LEVEL).log(message, args); + } + + /** + * A pair's {min, max} from-doc and to-doc bounds and match count, common to both a pair freshly + * built on demand ({@link JoinColumnModel#edges()}) and one already persisted in the join index + * ({@link Edges}, loaded through {@link #loadEdges}), so code walking matches doesn't need to + * care which one backs it. + */ + interface DocEdges { + int[] fromDocEdges(); + + int[] toDocEdges(); + + /** this is rather doubtful */ + int toCount(); + } + + /** A self-contained {@link DocEdges} value, with no addressing information of its own. */ + record Edges(int[] fromDocEdges, int[] toDocEdges, int toCount) implements DocEdges {} + + /** + * The from-doc-to-to-doc map produced by {@link #computeDocMapping}, paired with its resolved + * {@link #edges()}. {@link #toDocByFromDoc()} mirrors the on-disk column's read API, so freshly + * built pairs (not yet flushed to the join index) and pairs loaded from the join index can be + * walked by the same code. + */ + static final class JoinColumnModel { + private final int[] toDocByFromDoc; + private final DocEdges edges; + + JoinColumnModel(int[] toDocByFromDoc, DocEdges edges) { + this.toDocByFromDoc = toDocByFromDoc; + this.edges = edges; + } + + /** + * Returns a fresh single-valued cursor over the from-doc -> to-doc map, positioned before doc + * 0. + */ + SortedNumericDocValues toDocByFromDoc() { + return new ArrayBackedSortedNumericDocValues(toDocByFromDoc); + } + + DocEdges edges() { + return edges; + } + } + + /** + * Adapts an int-array from-doc -> to-doc map (as produced by {@link #computeDocMapping}, {@code + * -1} meaning no value) to the {@link SortedNumericDocValues} read API, so it can be consumed the + * same way as the on-disk join column. Always single-valued until M:N pairs are supported. + */ + private static final class ArrayBackedSortedNumericDocValues extends SortedNumericDocValues { + private final int[] toDocByFromDoc; + private int doc = -1; + + ArrayBackedSortedNumericDocValues(int[] toDocByFromDoc) { + this.toDocByFromDoc = toDocByFromDoc; + } + + @Override + public long nextValue() { + return toDocByFromDoc[doc]; + } + + @Override + public int docValueCount() { + return 1; + } + + @Override + public boolean advanceExact(int target) { + doc = target; + return target < toDocByFromDoc.length && toDocByFromDoc[target] >= 0; + } + + @Override + public int docID() { + return doc; + } + + @Override + public int nextDoc() { + return advance(doc + 1); + } + + @Override + public int advance(int target) { + while (target < toDocByFromDoc.length && toDocByFromDoc[target] < 0) { + target++; + } + doc = target < toDocByFromDoc.length ? target : NO_MORE_DOCS; + return doc; + } + + @Override + public long cost() { + return toDocByFromDoc.length; + } + } + + /** + * Builds the join column for one (from-segment, to-segment) pair: resolves every from-side doc to + * its matching to-side doc id, along with the pair's from-doc and to-doc bounds. From-side terms + * are hashed by {@link ForeignKeyColumn}; each to-side term is looked up in that hash to map + * from-side ords to to-side ords. + * + * <p>Docs already deleted at build time are skipped, purely to avoid persisting entries nobody + * can ever match -- deletes are otherwise re-checked live at query time (from-side in {@code + * ToLeafJoinContext}, to-side by the searcher's own {@code acceptDocs}), since a pair's cached + * mapping outlives whatever gets deleted after it was built. + */ + static JoinColumnModel computeDocMapping( + LeafReaderContext toContext, String toField, ForeignKeyColumn fromSideData) + throws IOException { + assert fromSideData != null; + + long[] toOrdByFromOrd = new long[fromSideData.getFromValuesCount()]; + Arrays.fill(toOrdByFromOrd, -1L); + SortedSetDocValues toDV = DocValues.getSortedSet(toContext.reader(), toField); + Bits toLiveDocs = toContext.reader().getLiveDocs(); + TermsEnum toTerms = toDV.termsEnum(); + // resolve from-side ords to to-side ords: look each to-side term up in the from-side hash. + boolean termsAreDisjoint = true; + for (BytesRef term = toTerms.next(); term != null; term = toTerms.next()) { + int fromOrd = fromSideData.getFromTermOrdOrDashOne(term); + if (fromOrd != -1) { + toOrdByFromOrd[fromOrd] = (int) toTerms.ord(); + termsAreDisjoint = false; + } + } + // TODO: this degrades M:N joins to M:1. Both toDocByToOrd and toDocByFromDoc keep a single + // to-side doc per slot, so when several to docs share a term (non-unique toField) or a + // fromSideData + // doc is multi-valued with several matching terms, later assignments overwrite earlier ones + // and only the last match survives. The read side (AIJoinQuery) already consumes all + // docValueCount() values per doc, so only this writer needs to learn to emit multiple + // to docs per fromSideData doc. + if (!termsAreDisjoint) { + int[] toDocByToOrd = new int[Math.toIntExact(toDV.getValueCount())]; + Arrays.fill(toDocByToOrd, -1); + for (int toDoc = toDV.nextDoc(); + toDoc != DocIdSetIterator.NO_MORE_DOCS; + toDoc = toDV.nextDoc()) { + if (toLiveDocs != null && !toLiveDocs.get(toDoc)) { + continue; + } + for (int i = 0; i < toDV.docValueCount(); i++) { + long toOrd = toDV.nextOrd(); + toDocByToOrd[(int) toOrd] = toDoc; Review Comment: let's limit to primitive 1:M -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
