MaxPassageFinder 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!"); } List<Long> docIds; List untyped = queryParams.getList("working"); if (untyped.get(0) instanceof String) { List<String> working = (List<String>) untyped; docIds = retrieval.getDocumentIds(working); } else if (untyped.get(0) instanceof Long) { docIds = (List<Long>) untyped; } else { throw new IllegalArgumentException("Bad form with your annotateMe list"); } Collections.sort(docIds); int requested = docIds.size(); int passageSize = (int) queryParams.getLong("passageSize"); int passageShift = (int) queryParams.getLong("passageShift"); boolean annotate = queryParams.get("annotate", false); 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(); ArrayList<ScoredPassage> perDocument = new ArrayList<ScoredPassage>(requested); for (Long docId : docIds) { if (docId < 0) { continue; } iterator.syncTo(docId); context.document = docId; documentLengths.syncTo(docId); int length = documentLengths.length(context); context.begin = 0; context.end = Math.min(passageSize, length); ScoredPassage documentBest = null; boolean lastIteration = false; boolean bFoundMatch = false; while (context.begin < length && !lastIteration) { if (context.end >= length) { lastIteration = true; } if (iterator.hasMatch(context)) { bFoundMatch = true; } double score = iterator.score(context); if (bFoundMatch && (documentBest == null || documentBest.score < score)) { documentBest = new ScoredPassage(docId, score, context.begin, context.end); if (annotate) { documentBest.annotation = iterator.getAnnotatedNode(context); } } context.begin += passageShift; context.end = Math.max(context.begin, Math.min(passageSize + context.begin, length)); } if (documentBest != null) { perDocument.add(documentBest); } iterator.movePast(docId); } if (perDocument.size() == 0) { return perDocument.toArray(new ScoredPassage[perDocument.size()]); } Collections.sort(perDocument, new ScoredPassage.ScoredPassageComparator()); Collections.reverse(perDocument); return perDocument.toArray(new ScoredPassage[perDocument.size()]); } MaxPassageFinder(LocalRetrieval lr); @Override ScoredDocument[] execute(Node queryTree, Parameters queryParams);  }
@Test(expected = IllegalArgumentException.class) public void testMissingPassageQuery() throws Exception { Parameters badParms = Parameters.create(); badParms.set("requested", 100); badParms.set("passageSize", 10); badParms.set("passageShift", 5); badParms.set("working", Arrays.asList(new Long[]{2l})); LocalRetrieval ret = new LocalRetrieval(index.getAbsolutePath()); Node query = StructuredQuery.parse("test"); query = ret.transformQuery(query, badParms); MaxPassageFinder model = new MaxPassageFinder(ret); model.execute(query, badParms); }
@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(); queryParams.set("passageQuery", true); queryParams.set("passageSize", 10); queryParams.set("passageShift", 5); List docs = new ArrayList<Long>(); for (int i = 0; i < 100; i++) { docs.add(new Long(i)); } queryParams.set("working", docs); Node query = StructuredQuery.parse("#combine( test text 0 1 )"); query = ret.transformQuery(query, queryParams); MaxPassageFinder model = new MaxPassageFinder(ret); ScoredPassage[] results = (ScoredPassage[]) model.execute(query, queryParams); assertEquals(results.length, 100); for (int i = 0; i < 100; i++) { assertEquals(results[i].document, i); assertEquals(results[i].begin, 0); assertEquals(results[i].end, 10); if (i > 0) { assert (CmpUtil.compare(results[i].score, results[i - 1].score) == 0); } } query = StructuredQuery.parse("michael zarozinski"); query = ret.transformQuery(query, queryParams); results = (ScoredPassage[]) model.execute(query, queryParams); assertEquals(0, results.length); queryParams.set("working", Arrays.asList(new String[]{"d-1", "d-98", "d-99"})); query = StructuredQuery.parse("#combine( test text 108 )"); query = ret.transformQuery(query, queryParams); results = (ScoredPassage[]) model.execute(query, queryParams); assertEquals(results.length, 3); assertEquals(results[0].document, 99); }