respect projection ordering in results (#79)

* respect projection ordering in results * Add all null test case
rockset · Apr 10, 2024 · da1d4a0 · da1d4a0
1 parent 00f48d9
commit da1d4a0
Show file tree

Hide file tree

Showing 5 changed files with 176 additions and 44 deletions.
diff --git a/src/main/java/com/rockset/jdbc/RocksetResultSet.java b/src/main/java/com/rockset/jdbc/RocksetResultSet.java
@@ -44,15 +44,21 @@
 import java.time.temporal.ChronoField;
 import java.util.ArrayList;
 import java.util.Calendar;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.TimeZone;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+
+import org.jetbrains.annotations.Nullable;
 import org.joda.time.DateTimeZone;
 import org.joda.time.LocalTime;
 import org.joda.time.format.DateTimeFormat;
@@ -1567,55 +1573,49 @@ private static void log(String msg) {
 
   private static List<Column> getColumns(QueryResponse response) throws SQLException {
     List<Column> out = new ArrayList<Column>();
-    ObjectMapper mapper = new ObjectMapper();
+
 
     try {
-      if (response.getResults().size() > 0) {
-        Set<String> fieldNames = new HashSet<>();
-        // Loop through all the rows to get the fields and (their first
-        // non-null) types.
+      if (!response.getResults().isEmpty()) {
+        List<String> orderedFields = Collections.emptyList();
+        // column fields will be null if a wildcard is used in the query
+        // otherwise, we want to return results in the order set by the projection
+        if(response.getColumnFields() != null) {
+          orderedFields = response.getColumnFields()
+                  .stream()
+                  .map(QueryFieldType::getName)
+                  .collect(Collectors.toList());
+        }
+        // Loop through all the rows to get the fields and types
+        Map<String, Column.ColumnTypes> fieldTypes = new HashMap<>();
         for (int i = 0; i < response.getResults().size(); ++i) {
           log("Extracting column information from record " + i + " in resultset");
           Object onedoc = response.getResults().get(i);
-          JsonNode docRootNode = mapper.readTree(mapper.writeValueAsString(onedoc));
-
-          Iterator<Map.Entry<String, JsonNode>> fields = docRootNode.fields();
-          while (fields.hasNext()) {
-            Map.Entry<String, JsonNode> field = fields.next();
-            String fieldName = field.getKey();
-            if (fieldNames.contains(fieldName)) {
-              // This fieldname was already found to have a non-null type
-              // in a previous record.
-              continue;
-            }
-            JsonNode value = field.getValue();
-            Column.ColumnTypes type = Column.ColumnTypes.fromValue(value.getNodeType().toString());
-            // Skip over the fields with null type unless all values for that
-            // field are null
-            if (type.equals(Column.ColumnTypes.NULL) && i != response.getResults().size() - 1) {
-              continue;
-            }
-            if (type.equals(Column.ColumnTypes.STRING)) {
-              java.time.format.DateTimeFormatter format = TIMESTAMP_PARSE_FORMAT;
-              try {
-                LocalDateTime.parse(value.asText(), format);
-                type = Column.ColumnTypes.TIMESTAMP;
-              } catch (DateTimeParseException e) {
-                // ignore
-              }
+          JsonNode docRootNode = OBJECT_MAPPER.readTree(OBJECT_MAPPER.writeValueAsString(onedoc));
+
+            for (Iterator<Map.Entry<String, JsonNode>> it = docRootNode.fields(); it.hasNext(); ) {
+                Map.Entry<String, JsonNode> field = it.next();
+                String fieldName = field.getKey();
+                if (fieldTypes.containsKey(fieldName) && !fieldTypes.get(fieldName).equals(Column.ColumnTypes.NULL)) {
+                  // This fieldname was already found to have a non-null type
+                  // in a previous record.
+                  continue;
+                }
+                JsonNode value = field.getValue();
+                Column.ColumnTypes type = valueToColumnType(value);
+                log("getColumns::column name " + fieldName + " type: " + type.toString());
+                fieldTypes.put(fieldName, type);
             }
-            if (type.equals(Column.ColumnTypes.OBJECT)) {
-              if (value.get("__rockset_type") != null) {
-                type = Column.ColumnTypes.fromValue(value.get("__rockset_type").asText());
-              }
-            }
-            log("getColumns::column name " + fieldName + " type: " + type.toString());
-            Column c = new Column(fieldName, type);
-            out.add(c);
-            fieldNames.add(fieldName);
-          }
         }
-      } else if (response.getColumnFields() != null && response.getColumnFields().size() > 0) {
+
+        // If we know the desired field ordering through explicit projection, use that ordering.
+        // Otherwise, just iterate over in arbitrary ordering
+        Collection<String> fields = orderedFields.isEmpty() ? fieldTypes.keySet() : orderedFields;
+        for(String field : fields){
+          out.add(new Column(field, fieldTypes.get(field)));
+        }
+
+      } else if (response.getColumnFields() != null && !response.getColumnFields().isEmpty()) {
         // If this is not a select star query, and has returned 0 rows.
         // Extrapolate the fields from query response's getColumnFields
         log("Extracting column information from explicit fields");
@@ -1637,6 +1637,26 @@ private static List<Column> getColumns(QueryResponse response) throws SQLExcepti
     }
   }
 
+  @Nullable
+  private static Column.ColumnTypes valueToColumnType(JsonNode value) {
+    Column.ColumnTypes type = Column.ColumnTypes.fromValue(value.getNodeType().toString());
+    if (type.equals(Column.ColumnTypes.STRING)) {
+      java.time.format.DateTimeFormatter format = TIMESTAMP_PARSE_FORMAT;
+      try {
+        LocalDateTime.parse(value.asText(), format);
+        type = Column.ColumnTypes.TIMESTAMP;
+      } catch (DateTimeParseException e) {
+        // ignore
+      }
+    }
+    if (type.equals(Column.ColumnTypes.OBJECT)) {
+      if (value.get("__rockset_type") != null) {
+        type = Column.ColumnTypes.fromValue(value.get("__rockset_type").asText());
+      }
+    }
+    return type;
+  }
+
   private static Map<String, Integer> getFieldMap(List<Column> columns) {
     Map<String, Integer> map = new HashMap<>();
     for (int i = 0; i < columns.size(); i++) {

diff --git a/src/test/java/com/rockset/client/TestWorkspace.java b/src/test/java/com/rockset/client/TestWorkspace.java
@@ -60,7 +60,8 @@ public void testDeleteWorkspace() throws Exception {
 
     // wait for collection to go away
     Awaitility.await("Waiting for collection to be cleaned up ")
-        .atMost(60, TimeUnit.SECONDS)
+        .atMost(3, TimeUnit.MINUTES)
+            .pollInterval(1, TimeUnit.SECONDS)
         .until(
             (Callable<Boolean>)
                 () -> {
@@ -70,7 +71,6 @@ public void testDeleteWorkspace() throws Exception {
                   } catch (Exception e) {
                     return true; // collection deleted
                   }
-                  Thread.sleep(1000);
                   return false;
                 });
 

diff --git a/src/test/java/com/rockset/jdbc/TestTable.java b/src/test/java/com/rockset/jdbc/TestTable.java
@@ -6,6 +6,8 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.Collections2;
+import com.google.common.collect.Lists;
 import com.rockset.client.RocksetClient;
 import com.rockset.client.model.Collection;
 import com.rockset.client.model.CreateCollectionRequest;
@@ -20,15 +22,23 @@
 import java.sql.DatabaseMetaData;
 import java.sql.Date;
 import java.sql.DriverManager;
+import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.sql.Time;
 import java.sql.Types;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Properties;
+import java.util.Set;
+import java.util.stream.Collectors;
+
 import okhttp3.MediaType;
 import okhttp3.MultipartBody;
 import okhttp3.OkHttpClient;
@@ -38,6 +48,7 @@
 import org.apache.commons.lang3.RandomStringUtils;
 import org.testng.Assert;
 import org.testng.annotations.BeforeSuite;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 //
@@ -360,6 +371,68 @@ public void testGetColumnsAllTypes() throws Exception {
     }
   }
 
+  @Test(dataProvider = "columnOrderings")
+  public void testColumnOrdering(Map<String, String> projectionTypes, String queryFormatString, String filePath) throws Exception {
+    List<String> collections = generateCollectionNames(/*numCollections*/ 1);
+    Connection conn = null;
+    try {
+      createCollections(collections);
+      waitCollections(collections);
+
+      String collection = collections.get(0);
+      uploadFile(collection, filePath, null);
+      waitNumberDocs(collection, 2);
+
+      conn = DriverManager.getConnection(DB_URL, property);
+
+
+
+      // When there is a wildcard in the projection, no ordering is guaranteed
+      String query = String.format("select * EXCEPT(_id, _meta, _event_time) from %s", collection);
+      try (PreparedStatement stmt = conn.prepareStatement(query);
+           ResultSet rs = stmt.executeQuery()) {
+
+        ResultSetMetaData rsmd = rs.getMetaData();
+        int cc = rsmd.getColumnCount();
+        while (rs.next()) {
+          Set<String> colNames = new HashSet<>();
+          for (int i = 1; i <= cc; i++) {
+            colNames.add(rsmd.getColumnName(i));
+          }
+          Assert.assertEquals(colNames, projectionTypes.keySet());
+        }
+      }
+
+      java.util.Collection<List<String>> projectionOrderings = Collections2.permutations(new ArrayList<>(projectionTypes.keySet()));
+
+      for(List<String> projections : projectionOrderings) {
+        List<String> expectedTypes = projections.stream().map(projectionTypes::get).collect(Collectors.toList());
+        List<String> queryStringParams = new ArrayList<>(projections);
+        queryStringParams.add(collection);
+        query = String.format(queryFormatString, queryStringParams.toArray());
+        try (PreparedStatement stmt = conn.prepareStatement(query);
+             ResultSet rs = stmt.executeQuery()) {
+
+          ResultSetMetaData rsmd = rs.getMetaData();
+          int cc = rsmd.getColumnCount();
+          List<String> colNames = new ArrayList<>();
+          List<String> colTypes = new ArrayList<>();
+          for (int i = 1; i <= cc; i++) {
+            colNames.add(rsmd.getColumnName(i));
+            colTypes.add(rsmd.getColumnTypeName(i));
+          }
+          Assert.assertEquals(colNames, Lists.newArrayList(projections));
+          Assert.assertEquals(colTypes, expectedTypes);
+        }
+      }
+
+
+    } finally {
+      cleanup(collections, null, conn);
+    }
+  }
+
+
   private void assertNextEquals(ResultSet rs, String expectedColumnName, int expectedType)
       throws SQLException {
     int columnNameIndex = rs.findColumn("COLUMN_NAME");
@@ -527,4 +600,39 @@ private void waitNumberDocs(String collectionName, int expectedDocs) throws Exce
       Thread.sleep(1000);
     }
   }
+
+  @DataProvider(name="columnOrderings")
+  Object[][] columnOrderings(){
+
+    Map<String, String> nullProjectionTypes = new HashMap<>();
+    nullProjectionTypes.put("a", "null");
+    nullProjectionTypes.put("b", "null");
+
+    String nullProjectionQueryFormat = "select %s, %s from %s ";
+    String nullProjectionFile = "src/test/resources/all_nulls.json";
+
+    // id is non null in all records
+    // name is null in one of the records
+    // age is null in all the records
+    // mixed is a string or array, however we query such that string is the first non-null type in the result set
+    // the type in the column will be string
+    // When there is a single non-null field, the type should be of the first non-null field
+    Map<String, String> mixedProjectionTypes = new HashMap<>();
+    mixedProjectionTypes.put("age", "null");
+    mixedProjectionTypes.put("id", "bigint");
+    mixedProjectionTypes.put("name", "varchar");
+    mixedProjectionTypes.put("mixed", "varchar");
+
+    String mixedProjectionQueryFormat = "select %s, %s, %s, %s from %s ORDER BY id ASC";
+    String mixedProjectionFile = "src/test/resources/mixed_nulls.json";
+
+
+
+    return new Object[][]{{
+      nullProjectionTypes, nullProjectionQueryFormat, nullProjectionFile
+    }, {
+      mixedProjectionTypes, mixedProjectionQueryFormat, mixedProjectionFile
+    }};
+
+  }
 }
diff --git a/src/test/resources/all_nulls.json b/src/test/resources/all_nulls.json
@@ -0,0 +1,2 @@
+{"a":  null, "b":  null}
+{"a":  null, "b":  null}
diff --git a/src/test/resources/mixed_nulls.json b/src/test/resources/mixed_nulls.json
@@ -0,0 +1,2 @@
+{"id":  1, "name": null, "age": null, "mixed": "hello"  }
+{"id":  2, "name": "bob", "age": null, "mixed": ["a", "b"] }