From 10a76c685c81725ca74f6150aa595a7e5022671a Mon Sep 17 00:00:00 2001
From: Dai <daichen@amazon.com>
Date: Tue, 19 May 2020 20:38:25 -0700
Subject: [PATCH] Refactor

---
 .../sql/protocol/response/QueryResponse.java     | 16 ++++++++++------
 ...java => SimpleJsonResponseFormatterTest.java} |  2 +-
 2 files changed, 11 insertions(+), 7 deletions(-)
 rename protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/{JsonResponseFormatterTest.java => SimpleJsonResponseFormatterTest.java} (98%)

diff --git a/protocol/src/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResponse.java b/protocol/src/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResponse.java
index 3517cbdc2d..2a41cea00a 100644
--- a/protocol/src/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResponse.java
+++ b/protocol/src/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResponse.java
@@ -27,7 +27,7 @@
 import java.util.Map;
 
 /**
- * Query response that encapsulates query results and isolate expression related class from formatter implementation.
+ * Query response that encapsulates query results and isolate {@link ExprValue} related from formatter implementation.
  */
 @RequiredArgsConstructor
 public class QueryResponse implements Iterable<Object[]> {
@@ -61,16 +61,16 @@ public Iterator<Object[]> iterator() {
     }
 
     private Map<String, ExprValue> getFirstTupleValue() {
-        // Assume expression is always tuple on first level and columns of all tuples are exactly same
+        // Assume expression is always tuple on first level
+        //  and columns (keys) of all tuple values are exactly same
         ExprValue firstValue = exprValues.iterator().next();
         return ExprValueUtils.getTupleValue(firstValue);
     }
 
     private Map<String, String> populateColumnNameAndTypes(Map<String, ExprValue> tupleValue) {
-        Map<String, String> colNameTypes = new LinkedHashMap<>(); // Maintain original order in tuple expression
-        tupleValue.forEach((name, expr) -> {
-            colNameTypes.put(name, expr.type().name().toLowerCase());
-        });
+        // Use linked hashmap to maintain original order in tuple expression
+        Map<String, String> colNameTypes = new LinkedHashMap<>();
+        tupleValue.forEach((name, expr) -> colNameTypes.put(name, getTypeString(expr)));
         return colNameTypes;
     }
 
@@ -80,4 +80,8 @@ private Object[] convertExprValuesToValues(Collection<ExprValue> exprValues) {
                           toArray(Object[]::new);
     }
 
+    private String getTypeString(ExprValue exprValue) {
+        return exprValue.type().name().toLowerCase();
+    }
+
 }
diff --git a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/JsonResponseFormatterTest.java b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatterTest.java
similarity index 98%
rename from protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/JsonResponseFormatterTest.java
rename to protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatterTest.java
index 057978cad8..d8b1b2f9c0 100644
--- a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/JsonResponseFormatterTest.java
+++ b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatterTest.java
@@ -27,7 +27,7 @@
 import static com.amazon.opendistroforelasticsearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
-class JsonResponseFormatterTest {
+class SimpleJsonResponseFormatterTest {
 
     @Test
     void formatResponse() {