GenerateWorkingSetQueries extends AppFunction { @Override public void run(Parameters parameters, PrintStream output) throws Exception { List<Parameters> queries = JSONQueryFormat.collectQueries(parameters); if (!parameters.containsKey("index") || queries.isEmpty()) { output.println(getHelpString()); return; } Retrieval retrieval = RetrievalFactory.create(parameters); List<Parameters> wsQueries = new ArrayList<Parameters>(); int topK = (int) parameters.get("topK", 10000); boolean numbers = parameters.get("wsNumbers", false); if (numbers) { try { LocalRetrieval lr = (LocalRetrieval) retrieval; } catch (Exception e) { output.println("--wsNumbers=true requires a LocalRetrieval.\n" + "Got: " + retrieval.getClass() + "\n" + "Required: " + LocalRetrieval.class + "\n" + getHelpString()); return; } } QuerySetJudgments qrels = null; if (parameters.isString("qrels")) { boolean binaryJudgments = parameters.get("binary", false); boolean positiveJudgments = parameters.get("postive", true); qrels = new QuerySetJudgments(parameters.getString("qrels"), binaryJudgments, positiveJudgments); } for (Parameters queryParams : queries) { Parameters wsQ = queryParams.clone(); wsQueries.add(wsQ); HashSet<String> workingSetNames = new HashSet<String>(); HashSet<Long> workingSetNumbers = new HashSet<Long>(); queryParams.set("requested", topK); Node query = StructuredQuery.parse(queryParams.getString("text")); query = retrieval.transformQuery(query, queryParams); List<ScoredDocument> results = retrieval.executeQuery(query, queryParams).scoredDocuments; for (ScoredDocument r : results) { workingSetNames.add(r.documentName); if (numbers) { workingSetNumbers.add(r.document); } } if (qrels != null) { QueryJudgments qrel = qrels.get(queryParams.getString("number")); if (qrel != null) { for (String docName : qrel.getDocumentSet()) { if (qrel.isRelevant(docName)) { try { long docId = ((LocalRetrieval) retrieval).getDocumentId(docName); workingSetNames.add(docName); if (numbers) { workingSetNumbers.add(docId); } } catch (IOException e) { } } } } } List<String> wsNames = new ArrayList<String>(workingSetNames); Collections.sort(wsNames); wsQ.set("working", wsNames); if (numbers) { List<Long> wsNumbers = new ArrayList<Long>(workingSetNumbers); Collections.sort(wsNumbers); wsQ.set("workingIds", wsNumbers); } } Parameters wsParameters = Parameters.create(); wsParameters.set("queries", wsQueries); if (parameters.isString("output")) { File o = new File(parameters.getString("output")); FSUtil.makeParentDirectories(o); StreamUtil.copyStringToFile(wsParameters.toPrettyString(), o); } else { output.println(wsParameters.toPrettyString()); } }  @Override String getName(); @Override String getHelpString(); @Override void run(Parameters parameters, PrintStream output);  }
@Test public void testSomeMethod() throws Exception { File out = new File(dataFolder, "out1"); Parameters p = Parameters.create(); p.set("index", index.getAbsolutePath()); p.set("topK", 10); p.set("output", out.getAbsolutePath()); p.set("wsNumders", true); ArrayList<Parameters> qs = new ArrayList<Parameters>(); qs.add(Parameters.parseString("{\"number\" : \"q1\", \"text\" : \"#combine(t1 t2)\"}")); qs.add(Parameters.parseString("{\"number\" : \"q2\", \"text\" : \"#combine(t3 t4)\"}")); qs.add(Parameters.parseString("{\"number\" : \"q3\", \"text\" : \"#combine(t5 t6)\"}")); p.set("queries", qs); GenerateWorkingSetQueries generator = new GenerateWorkingSetQueries(); generator.run(p, System.err); Parameters outputParams = Parameters.parseFile(out); List<Parameters> outQueries = outputParams.getList("queries", Parameters.class); assertEquals( outQueries.get(0).getList("working").size(), 10 ); assertEquals( outQueries.get(1).getList("working").size(), 10 ); assertEquals( outQueries.get(2).getList("working").size(), 10 ); }