RankedPassageModel extends ProcessingModel { @Override public ScoredDocument[] execute(Node queryTree, Parameters queryParams) throws Exception { PassageScoringContext context = new PassageScoringContext(); context.cachable = false; if (!queryParams.get("passageQuery", false)) { throw new IllegalArgumentException("passageQuery must be true for passage retrieval to work!"); } int requested = (int) queryParams.get("requested", 1000); int passageSize = (int) queryParams.getLong("passageSize"); int passageShift = (int) queryParams.getLong("passageShift"); if (passageSize <= 0 || passageShift <= 0) { throw new IllegalArgumentException("passageSize/passageShift must be specified as positive integers."); } ScoreIterator iterator = (ScoreIterator) retrieval.createIterator(queryParams, queryTree); LengthsIterator documentLengths = retrieval.getDocumentLengthsIterator(); FixedSizeMinHeap<ScoredPassage> queue = new FixedSizeMinHeap(ScoredPassage.class, requested, new ScoredPassage.ScoredPassageComparator()); while (!iterator.isDone()) { long document = iterator.currentCandidate(); context.document = document; documentLengths.syncTo(document); int length = documentLengths.length(context); context.begin = 0; context.end = Math.min(passageSize, length); iterator.syncTo(document); boolean lastIteration = false; if (!iterator.hasMatch(context)) { lastIteration = true; } while (context.begin < length && !lastIteration) { if (context.end >= length) { lastIteration = true; } double score = iterator.score(context); if (requested < 0 || queue.size() < requested || queue.peek().score < score) { ScoredPassage scored = new ScoredPassage(document, score, context.begin, context.end); queue.offer(scored); } context.begin += passageShift; context.end = Math.max(context.begin, Math.min(passageSize + context.begin, length)); } iterator.movePast(document); } return toReversedArray(queue); } RankedPassageModel(LocalRetrieval lr); @Override ScoredDocument[] execute(Node queryTree, Parameters queryParams);  }
@Test public void testEntireCollection() throws Exception { Parameters globals = Parameters.create(); globals.set("passageQuery", true); LocalRetrieval ret = new LocalRetrieval(index.getAbsolutePath(), globals); Parameters queryParams = Parameters.create(); int req = 10; queryParams.set("requested", req); queryParams.set("passageQuery", true); queryParams.set("passageSize", 10); queryParams.set("passageShift", 5); Node query = StructuredQuery.parse("#combine( test text 0 1 )"); query = ret.transformQuery(query, queryParams); RankedPassageModel model = new RankedPassageModel(ret); ScoredPassage[] results = (ScoredPassage[]) model.execute(query, queryParams); assertEquals(results.length, req); for (int i = 0; i < req; i++) { assertEquals(results[i].document, i); assertEquals(results[i].begin, 0); assertEquals(results[i].end, 10); assertEquals(results[i].rank, i + 1); if (i > 0) { assert (CmpUtil.compare(results[i].score, results[i - 1].score) == 0); } } query = StructuredQuery.parse("#combine( test text 99 )"); query = ret.transformQuery(query, queryParams); results = (ScoredPassage[]) model.execute(query, queryParams); assertEquals(results.length, req); assertEquals(results[0].document, 94); assertEquals(results[0].begin, 100); assertEquals(results[0].end, 106); assertEquals(results[0].rank, 1); assertEquals(results[0].score, -4.776027, 0.000001); assertEquals(results[1].document, 90); assertEquals(results[1].begin, 95); assertEquals(results[1].end, 102); assertEquals(results[1].rank, 2); assertEquals(results[1].score, -4.776691, 0.000001); assertEquals(results[2].document, 95); assertEquals(results[2].begin, 100); assertEquals(results[2].end, 107); assertEquals(results[2].rank, 3); assertEquals(results[2].score, -4.776691, 0.000001); }
@Test public void testWhiteList() throws Exception { Parameters globals = Parameters.create(); globals.set("passageQuery", true); LocalRetrieval ret = new LocalRetrieval(index.getAbsolutePath(), globals); Parameters queryParams = Parameters.create(); queryParams.set("requested", 100); queryParams.set("passageQuery", true); queryParams.set("passageSize", 10); queryParams.set("passageShift", 5); Node query = StructuredQuery.parse("#combine( test text 0 1 )"); query = ret.transformQuery(query, queryParams); WorkingSetPassageModel model = new WorkingSetPassageModel(ret); queryParams.set("working", Arrays.asList(new Long[]{2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L, 10L, 11L})); ScoredPassage[] results = (ScoredPassage[]) model.execute(query, queryParams); assertEquals(31, results.length); for (int i = 0; i < 10; i++) { assertEquals(results[i].document, i + 2); assertEquals(results[i].begin, 0); assertEquals(results[i].end, 10); assertEquals(results[i].rank, i + 1); assertEquals(results[i].score, -4.085500, 0.000001); } queryParams.set("working", Arrays.asList(new String[]{"d-2", "d-3", "d-4", "d-5", "d-6", "d-7", "d-8", "d-9", "d-10", "d-11"})); results = (ScoredPassage[]) model.execute(query, queryParams); assertEquals(31, results.length); for (int i = 0; i < 10; i++) { assertEquals(results[i].document, i + 2); assertEquals(results[i].begin, 0); assertEquals(results[i].end, 10); assertEquals(results[i].rank, i + 1); assertEquals(results[i].score, -4.085500, 0.000001); } query = StructuredQuery.parse("#combine( test text 80 )"); query = ret.transformQuery(query, queryParams); queryParams.set("working", Arrays.asList(new Long[]{0L, 1L, 2L, 3L, 4L, 89L, 90L, 91L, 92L, 93L})); results = (ScoredPassage[]) model.execute(query, queryParams); assertEquals(results.length, 100); assertEquals(results[0].document, 89); assertEquals(results[0].begin, 75); assertEquals(results[0].end, 85); assertEquals(results[0].rank, 1); assertEquals(results[0].score, -4.49422735, 0.000001); assertEquals(results[1].document, 89); assertEquals(results[1].begin, 80); assertEquals(results[1].end, 90); assertEquals(results[1].rank, 2); assertEquals(results[1].score, -4.49422735, 0.000001); assertEquals(results[10].document, 0); assertEquals(results[10].begin, 0); assertEquals(results[10].end, 10); assertEquals(results[10].rank, 11); assertEquals(results[10].score, -4.51151864, 0.000001); assertEquals(results[15].document, 89); assertEquals(results[15].begin, 0); assertEquals(results[15].end, 10); assertEquals(results[15].rank, 16); assertEquals(results[15].score, -4.51151864, 0.000001); }