Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -1215,18 +1215,11 @@ public Response search(@QueryParam("q") String queryTerm) {
HashSet<String> userAndRoles = new HashSet<>();
userAndRoles.add(principal);
userAndRoles.addAll(roles);
List<Map<String, String>> notesFound = noteSearchService.query(queryTerm);
for (int i = 0; i < notesFound.size(); i++) {
String[] ids = notesFound.get(i).get("id").split("/", 2);
String noteId = ids[0];
if (!authorizationService.isOwner(noteId, userAndRoles) &&
!authorizationService.isReader(noteId, userAndRoles) &&
!authorizationService.isWriter(noteId, userAndRoles) &&
!authorizationService.isRunner(noteId, userAndRoles)) {
notesFound.remove(i);
i--;
}
}
// isReader() already covers owners, writers and runners. Handing the check to the
// search service keeps it in front of the cutoff, so a caller with access to few notes
// is still served a full result set of the notes it may read.
Comment on lines +1218 to +1220

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
// isReader() already covers owners, writers and runners. Handing the check to the
// search service keeps it in front of the cutoff, so a caller with access to few notes
// is still served a full result set of the notes it may read.

List<Map<String, String>> notesFound = noteSearchService.query(queryTerm,
noteId -> authorizationService.isReader(noteId, userAndRoles));
LOGGER.info("{} notes found", notesFound.size());
return new JsonResponse<>(Status.OK, notesFound).build();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
Expand Down Expand Up @@ -509,15 +510,11 @@ private String buildParagraphText(String noteName, Paragraph p) {
// ---- SearchService implementation ----

@Override
// TODO(ZEPPELIN-6414): Accept user/roles (or a readability Predicate) and apply the auth
// filter before Phase-1 table collection and before the top-K cutoff. Currently the REST
// layer filters after truncation, which can hide results the caller is authorized for and
// lets inaccessible notes contaminate the table-boost ranking. Requires a SearchService
// interface change that also affects LuceneSearch.
public List<Map<String, String>> query(String queryStr) {
public List<Map<String, String>> query(String queryStr, Predicate<String> readable) {
if (StringUtils.isBlank(queryStr) || index.isEmpty()) {
return Collections.emptyList();
}
Map<String, Boolean> readableNotes = new HashMap<>();

float[] queryEmbedding = embed(queryStr);
String queryLower = queryStr.toLowerCase(Locale.ROOT);
Expand All @@ -527,6 +524,12 @@ public List<Map<String, String>> query(String queryStr) {
indexLock.readLock().lock();
try {
for (Map.Entry<String, IndexEntry> entry : index.entrySet()) {
// Dropping the entries here keeps them out of the table weights below and out of
// the cutoff, so the caller is served its own top results and not what is left of
// everyone's top results.
if (!readableNotes.computeIfAbsent(noteIdOf(entry.getKey()), readable::test)) {
continue;
}
float sim = cosineSimilarity(queryEmbedding, entry.getValue().embedding);
IndexEntry ie = entry.getValue();
if (ie.text != null && ie.text.toLowerCase(Locale.ROOT).contains(queryLower)) {
Expand Down Expand Up @@ -612,6 +615,16 @@ public List<Map<String, String>> query(String queryStr) {
return results;
}

/**
* The key of an indexed entry is either a noteId or a noteId followed by the paragraph.
*
* @see #formatId(String, Paragraph)
*/
private static String noteIdOf(String docId) {
int separator = docId.indexOf('/');
return separator < 0 ? docId : docId.substring(0, separator);
}

@Override
public void addNoteIndex(String noteId) {
try {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,13 @@
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Predicate;
import javax.annotation.PreDestroy;
import jakarta.inject.Inject;

Expand Down Expand Up @@ -76,6 +79,12 @@ public class LuceneSearch extends SearchService {
private static final String SEARCH_FIELD_TITLE = "header";
private static final String PARAGRAPH = "paragraph";
private static final String ID_FIELD = "id";
/** Number of results a query returns at most. */
private static final int MAX_RESULTS = 20;
/** Number of hits pulled from the index per round while looking for readable ones. */
private static final int HIT_BATCH_SIZE = 20;
/** Only the id is needed to tell whether the caller may read a hit. */
private static final Set<String> ID_FIELD_ONLY = Collections.singleton(ID_FIELD);

private final Directory indexDirectory;
private final IndexWriter indexWriter;
Expand Down Expand Up @@ -113,7 +122,7 @@ public LuceneSearch(ZeppelinConfiguration zConf, Notebook notebook) throws IOExc
* @see org.apache.zeppelin.search.Search#query(java.lang.String)
*/
@Override
public List<Map<String, String>> query(String queryStr) {
public List<Map<String, String>> query(String queryStr, Predicate<String> readable) {
if (null == indexDirectory) {
throw new IllegalStateException(
"Something went wrong on instance creation time, index dir is null");
Expand All @@ -134,7 +143,7 @@ public List<Map<String, String>> query(String queryStr) {
SimpleHTMLFormatter htmlFormatter = new SimpleHTMLFormatter();
Highlighter highlighter = new Highlighter(htmlFormatter, new QueryScorer(query));

result = doSearch(indexSearcher, query, analyzer, highlighter);
result = doSearch(indexSearcher, query, analyzer, highlighter, readable);
} catch (IOException e) {
LOGGER.error("Failed to open index dir {}, make sure indexing finished OK", indexDirectory, e);
} catch (ParseException e) {
Expand All @@ -144,64 +153,40 @@ public List<Map<String, String>> query(String queryStr) {
}

private List<Map<String, String>> doSearch(
IndexSearcher searcher, Query query, Analyzer analyzer, Highlighter highlighter) {
IndexSearcher searcher, Query query, Analyzer analyzer, Highlighter highlighter,
Predicate<String> readable) {
List<Map<String, String>> matchingParagraphs = new ArrayList<>();
ScoreDoc[] hits;
Map<String, Boolean> readableNotes = new HashMap<>();
try {
hits = searcher.search(query, 20).scoreDocs;
for (int i = 0; i < hits.length; i++) {
LOGGER.debug("doc={} score={}", hits[i].doc, hits[i].score);

int id = hits[i].doc;
Document doc = searcher.doc(id);
String path = doc.get(ID_FIELD);
if (path != null) {
LOGGER.debug( "{}. {}", (i + 1), path);
String title = doc.get("title");
if (title != null) {
LOGGER.debug(" Title: {}", doc.get("title"));
// Walk the hits in score order and keep the ones the caller may read until the result
// set is full or the hits run out. Reading is checked here and not on the result set,
// because a cut that runs first would hide results the caller is allowed to see.
ScoreDoc lastHit = null;
while (matchingParagraphs.size() < MAX_RESULTS) {
ScoreDoc[] hits = (lastHit == null
? searcher.search(query, HIT_BATCH_SIZE)
: searcher.searchAfter(lastHit, query, HIT_BATCH_SIZE)).scoreDocs;
if (hits.length == 0) {
break;
}
lastHit = hits[hits.length - 1];
Comment on lines +165 to +172

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about growing the batch size as the walk goes deeper?

HIT_BATCH_SIZE is fixed at 20, so walking all N hits calls searchAfter N/20 times. searchAfter re-runs the query from the start on each round (the score floor set by after does prune, but the posting lists are still traversed again), which comes to 500 re-executions for 10,000 hits.

Doubling the batch each round takes that from O(N/20) to O(log N), about 9 rounds under the same conditions:

int batchSize = HIT_BATCH_SIZE;
while (matchingParagraphs.size() < MAX_RESULTS) {
  ...
  batchSize = Math.min(batchSize * 2, MAX_HIT_BATCH_SIZE);
}

And when hits.length < batchSize the index is exhausted, so could the loop stop right there? It looks like it currently runs one more search just to get an empty result back.

For what it is worth, the walk only goes deep on installations locked down with zeppelin.notebook.public=false. The default is true, and then the readers set is empty so isMember() returns true (AuthorizationService.java:271) and the first batch ends it. Worth noting anyway, since that setting is exactly the population this bug affects.

for (ScoreDoc hit : hits) {
if (matchingParagraphs.size() >= MAX_RESULTS) {
break;
}

String text = doc.get(SEARCH_FIELD_TEXT);
String header = doc.get(SEARCH_FIELD_TITLE);
String fragment = "";

if (text != null) {
TokenStream tokenStream =
TokenSources.getTokenStream(
searcher.getIndexReader(), id, SEARCH_FIELD_TEXT, analyzer);
TextFragment[] frags = highlighter.getBestTextFragments(tokenStream, text, true, 3);
LOGGER.debug(" {} fragments found for query '{}'", frags.length, query);
for (TextFragment frag : frags) {
if ((frag != null) && (frag.getScore() > 0)) {
LOGGER.debug(" Fragment: {}", frag);
}
}
fragment = (frags != null && frags.length > 0) ? frags[0].toString() : "";
LOGGER.debug("doc={} score={}", hit.doc, hit.score);
// Read the id alone to decide on a hit. The rest of the document is only worth
// loading for the hits that end up in the result set.
String path = searcher.doc(hit.doc, ID_FIELD_ONLY).get(ID_FIELD);
if (path == null) {
LOGGER.info("No {} for this document", ID_FIELD);
continue;
}

if (header != null) {
TokenStream tokenTitle =
TokenSources.getTokenStream(
searcher.getIndexReader(), id, SEARCH_FIELD_TITLE, analyzer);
TextFragment[] frgTitle = highlighter.getBestTextFragments(tokenTitle, header, true, 3);
header = (frgTitle != null && frgTitle.length > 0) ? frgTitle[0].toString() : "";
} else {
header = "";
if (!readableNotes.computeIfAbsent(noteIdOf(path), readable::test)) {
continue;
}
matchingParagraphs.add(
ImmutableMap.<String, String>builder()
.put("id", path)
.put("name", title)
.put("snippet", fragment)
.put("text", text)
.put("header", header)
.put("title", header)
.put("tables", "")
.put("output", "")
.build());
} else {
LOGGER.info("{}. No {} for this document", i + 1, ID_FIELD);
toMatch(searcher, query, analyzer, highlighter, hit.doc, searcher.doc(hit.doc)));
}
}
} catch (IOException | InvalidTokenOffsetsException e) {
Expand All @@ -210,6 +195,55 @@ private List<Map<String, String>> doSearch(
return matchingParagraphs;
}

private Map<String, String> toMatch(IndexSearcher searcher, Query query, Analyzer analyzer,
Highlighter highlighter, int docId, Document doc)
throws IOException, InvalidTokenOffsetsException {
String path = doc.get(ID_FIELD);
String title = doc.get("title");
String text = doc.get(SEARCH_FIELD_TEXT);
String header = doc.get(SEARCH_FIELD_TITLE);
String fragment = "";

if (text != null) {
TokenStream tokenStream =
TokenSources.getTokenStream(
searcher.getIndexReader(), docId, SEARCH_FIELD_TEXT, analyzer);
TextFragment[] frags = highlighter.getBestTextFragments(tokenStream, text, true, 3);
LOGGER.debug(" {} fragments found for query '{}'", frags.length, query);
fragment = (frags != null && frags.length > 0) ? frags[0].toString() : "";
}

if (header != null) {
TokenStream tokenTitle =
TokenSources.getTokenStream(
searcher.getIndexReader(), docId, SEARCH_FIELD_TITLE, analyzer);
TextFragment[] frgTitle = highlighter.getBestTextFragments(tokenTitle, header, true, 3);
header = (frgTitle != null && frgTitle.length > 0) ? frgTitle[0].toString() : "";
} else {
header = "";
}
return ImmutableMap.<String, String>builder()
.put("id", path)
.put("name", title)
.put("snippet", fragment)
.put("text", text)
.put("header", header)
.put("title", header)
.put("tables", "")
.put("output", "")
.build();
}

/**
* The id of an indexed document is either a noteId or a noteId followed by the paragraph.
*
* @see #formatId(String, Paragraph)
*/
private static String noteIdOf(String documentId) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could noteIdOf move up to SearchService?

There are four copies of the same method right now: LuceneSearch, EmbeddingSearch, and both test classes.

The tests already pull formatId, which plays the same role, from the production class as a static import:

// LuceneSearchTest.java:18
import static org.apache.zeppelin.search.LuceneSearch.formatId;

Putting noteIdOf on SearchService as a package-private static and importing it the same way would bring the four copies down to one.

int separator = documentId.indexOf('/');
return separator < 0 ? documentId : documentId.substring(0, separator);
}

/* (non-Javadoc)
* @see org.apache.zeppelin.search.Search#updateIndexDoc(org.apache.zeppelin.notebook.Note)
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;

public class NoSearchService extends SearchService {

Expand All @@ -30,7 +31,7 @@ public NoSearchService() {
}

@Override
public List<Map<String, String>> query(String queryStr) {
public List<Map<String, String>> query(String queryStr, Predicate<String> readable) {
return Collections.emptyList();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;
import org.apache.zeppelin.notebook.NoteEventAsyncListener;
import javax.annotation.PreDestroy;

Expand All @@ -39,9 +40,12 @@ protected SearchService(String name) {
* Full-text search in all the notes
*
* @param queryStr a query
* @param readable tells for a noteId whether the caller may read it. Entries the caller
* cannot read are dropped before the result set is cut down, so that the
* caller is not served fewer results than it is allowed to see.
* @return A list of matching paragraphs (id, text, snippet w/ highlight)
*/
public abstract List<Map<String, String>> query(String queryStr);
public abstract List<Map<String, String>> query(String queryStr, Predicate<String> readable);

/**
* Updates note index for the given note, only update index of note meta info,
Expand Down
Loading
Loading