fields support NLPchina#146

eliranmoyal · eliranmoyal · commit 63cb84a5a03b · 2016-02-27T18:16:37.000+02:00
diff --git a/src/main/java/org/elasticsearch/plugin/nlpcn/executors/CSVResultsExtractor.java b/src/main/java/org/elasticsearch/plugin/nlpcn/executors/CSVResultsExtractor.java
@@ -4,6 +4,7 @@
 import org.elasticsearch.cluster.routing.allocation.decider.Decision;
 import org.elasticsearch.common.base.Joiner;
 import org.elasticsearch.search.SearchHit;
+import org.elasticsearch.search.SearchHitField;
 import org.elasticsearch.search.SearchHits;
 import org.elasticsearch.search.aggregations.Aggregation;
 import org.elasticsearch.search.aggregations.AggregationBuilder;
@@ -234,7 +235,11 @@ private List<String> createHeadersAndFillDocsMap(boolean flat, SearchHit[] hits,
         Set<String> csvHeaders = new HashSet<>();
         for(SearchHit hit : hits){
             Map<String, Object> doc = hit.sourceAsMap();
-            mergeHeaders(csvHeaders,doc,flat);
+            Map<String, SearchHitField> fields = hit.getFields();
+            for(SearchHitField searchHitField : fields.values()){
+                doc.put(searchHitField.getName(),searchHitField.value());
+            }
+            mergeHeaders(csvHeaders, doc, flat);
             if(this.includeScore){
                 doc.put("_score",hit.score());
             }
diff --git a/src/test/java/org/nlpcn/es4sql/CSVResultsExtractorTests.java b/src/test/java/org/nlpcn/es4sql/CSVResultsExtractorTests.java
@@ -383,12 +383,27 @@ public void includeScoreAndType() throws SqlParseException, SQLFeatureNotSupport
         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
     * geo
      */
 
+/* test for scripted fields
+    @Test
+    public void scriptedField() throws SqlParseException, SQLFeatureNotSupportedException, Exception {
+        String query = String.format("select age+1 as agePlusOne ,age , firstname from %s/account where age =  31 limit 1", TEST_INDEX);
+        CSVResult csvResult = getCsvResult(false, query,false,false);
+        List<String> headers = csvResult.getHeaders();
+        Assert.assertEquals(3,headers.size());
+        Assert.assertTrue(headers.contains("agePlusOne"));
+        Assert.assertTrue(headers.contains("age"));
+        Assert.assertTrue(headers.contains("firstname"));
+        List<String> lines = csvResult.getLines();
+        Assert.assertTrue(lines.get(0).contains("32,31") || lines.get(0).contains("32.0,31.0"));
+    }
+*/
 
     private CSVResult getCsvResult(boolean flat, String query) throws SqlParseException, SQLFeatureNotSupportedException, Exception, CsvExtractorException {
         return getCsvResult(flat,query,false,false);
