csv results type and score support NLPchina#150

eliranmoyal · eliranmoyal · commit 4de3f626cee1 · 2016-02-26T16:06:39.000+02:00
diff --git a/src/main/java/org/elasticsearch/plugin/nlpcn/executors/CSVResultRestExecutor.java b/src/main/java/org/elasticsearch/plugin/nlpcn/executors/CSVResultRestExecutor.java
@@ -26,7 +26,9 @@ public void execute(Client client, Map<String, String> params, QueryAction query
         if(params.containsKey("separator")){
          separator = params.get("separator");
         }
-        CSVResult result  = new CSVResultsExtractor().extractResults(queryResult,flat,separator);
+        boolean includeScore = Boolean.getBoolean(params.getOrDefault("_score", "false"));
+        boolean includeType = Boolean.getBoolean(params.getOrDefault("_type", "false"));
+        CSVResult result  = new CSVResultsExtractor(includeScore,includeType).extractResults(queryResult,flat,separator);
         String newLine = "\n";
         if(params.containsKey("newLine")){
          newLine = params.get("newLine");
diff --git a/src/main/java/org/elasticsearch/plugin/nlpcn/executors/CSVResultsExtractor.java b/src/main/java/org/elasticsearch/plugin/nlpcn/executors/CSVResultsExtractor.java
@@ -26,9 +26,13 @@
  * Created by Eliran on 27/12/2015.
  */
 public class CSVResultsExtractor {
+    private final boolean includeType;
+    private final boolean includeScore;
     private int currentLineIndex;
 
-    public CSVResultsExtractor() {
+    public CSVResultsExtractor(boolean includeScore, boolean includeType) {
+        this.includeScore = includeScore;
+        this.includeType = includeType;
         this.currentLineIndex = 0;
     }
 
@@ -231,8 +235,20 @@ private List<String> createHeadersAndFillDocsMap(boolean flat, SearchHit[] hits,
         for(SearchHit hit : hits){
             Map<String, Object> doc = hit.sourceAsMap();
             mergeHeaders(csvHeaders,doc,flat);
+            if(this.includeScore){
+                doc.put("_score",hit.score());
+            }
+            if(this.includeType){
+                doc.put("_type",hit.type());
+            }
             docsAsMap.add(doc);
         }
+        if (this.includeScore){
+            csvHeaders.add("_score");
+        }
+        if (this.includeType){
+            csvHeaders.add("_type");
+        }
         return new ArrayList<>(csvHeaders);
     }
 
diff --git a/src/test/java/org/nlpcn/es4sql/CSVResultsExtractorTests.java b/src/test/java/org/nlpcn/es4sql/CSVResultsExtractorTests.java
@@ -341,8 +341,48 @@ public void percentileAggregationTest() throws SqlParseException, SQLFeatureNotS
         Assert.assertEquals("32.0,32.0,34.0,36.0,38.0,40.0,40.0",lines.get(0));
     }
 
+    @Test
+    public void includeTypeAndNotScore() throws SqlParseException, SQLFeatureNotSupportedException, Exception {
+        String query = String.format("select age , firstname from %s/account where age > 31 limit 2", TEST_INDEX);
+        CSVResult csvResult = getCsvResult(false, query,false,true);
+        List<String> headers = csvResult.getHeaders();
+        Assert.assertEquals(3,headers.size());
+        Assert.assertTrue(headers.contains("age"));
+        Assert.assertTrue(headers.contains("firstname"));
+        Assert.assertTrue(headers.contains("_type"));
+        List<String> lines = csvResult.getLines();
+        Assert.assertTrue(lines.get(0).contains(",account"));
+        Assert.assertTrue(lines.get(1).contains(",account"));
+    }
 
+    @Test
+    public void includeScoreAndNotType() throws SqlParseException, SQLFeatureNotSupportedException, Exception {
+        String query = String.format("select age , firstname from %s/account where age > 31 limit 2", TEST_INDEX);
+        CSVResult csvResult = getCsvResult(false, query,true,false);
+        List<String> headers = csvResult.getHeaders();
+        Assert.assertEquals(3,headers.size());
+        Assert.assertTrue(headers.contains("age"));
+        Assert.assertTrue(headers.contains("firstname"));
+        Assert.assertTrue(headers.contains("_score"));
+        List<String> lines = csvResult.getLines();
+        Assert.assertTrue(lines.get(0).contains(",1.0"));
+        Assert.assertTrue(lines.get(1).contains(",1.0"));
+    }
 
+    @Test
+    public void includeScoreAndType() throws SqlParseException, SQLFeatureNotSupportedException, Exception {
+        String query = String.format("select age , firstname from %s/account where age > 31 limit 2", TEST_INDEX);
+        CSVResult csvResult = getCsvResult(false, query,true,true);
+        List<String> headers = csvResult.getHeaders();
+        Assert.assertEquals(4,headers.size());
+        Assert.assertTrue(headers.contains("age"));
+        Assert.assertTrue(headers.contains("firstname"));
+        Assert.assertTrue(headers.contains("_score"));
+        Assert.assertTrue(headers.contains("_type"));
+        List<String> lines = csvResult.getLines();
+        Assert.assertTrue(lines.get(0).contains(",account,1.0"));
+        Assert.assertTrue(lines.get(1).contains(",account,1.0"));
+    }
     /* todo: more tests:
     * filter/nested and than metric
     * histogram
@@ -351,12 +391,14 @@ public void percentileAggregationTest() throws SqlParseException, SQLFeatureNotS
 
 
     private CSVResult getCsvResult(boolean flat, String query) throws SqlParseException, SQLFeatureNotSupportedException, Exception, CsvExtractorException {
+        return getCsvResult(flat,query,false,false);
+    }
+
+    private CSVResult getCsvResult(boolean flat, String query,boolean includeScore , boolean includeType) throws SqlParseException, SQLFeatureNotSupportedException, Exception, CsvExtractorException {
         SearchDao searchDao = MainTestSuite.getSearchDao();
         QueryAction queryAction = searchDao.explain(query);
         Object execution =  QueryActionElasticExecutor.executeAnyAction(searchDao.getClient(), queryAction);
-        return new CSVResultsExtractor().extractResults(execution, flat, ",");
+        return new CSVResultsExtractor(includeScore,includeType).extractResults(execution, flat, ",");
     }
 
-
-
 }
