Merge pull request #224 from datafuselabs/fix/handle-nan-data

fix: handle nan data in response
databendlabs · Jun 1, 2024 · a079e3c · a079e3c
2 parents 56eb163 + 2c7fdae
commit a079e3c
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 4 deletions.
diff --git a/databend-client/src/main/java/com/databend/client/DatabendClientV1.java b/databend-client/src/main/java/com/databend/client/DatabendClientV1.java
@@ -159,7 +159,8 @@ private boolean executeInternal(Request request, OptionalLong materializedJsonSi
                 response = JsonResponse.execute(QUERY_RESULTS_CODEC, httpClient, request, materializedJsonSizeLimit);
             } catch (RuntimeException e) {
                 cause = e;
-                continue;
+//                continue;
+                throw new RuntimeException("Query failed: " + e.getMessage());
             }
 
             if ((response.getStatusCode() == HTTP_OK) && response.hasValue() && (response.getValue().getError() == null)) {

diff --git a/databend-client/src/main/java/com/databend/client/JsonResponse.java b/databend-client/src/main/java/com/databend/client/JsonResponse.java
@@ -68,7 +68,7 @@ private JsonResponse(int statusCode, String statusMessage, Headers headers, @Nul
         this.hasValue = (exception == null);
     }
 
-    public static <T> JsonResponse<T> execute(JsonCodec<T> codec, OkHttpClient client, Request request, OptionalLong materializedJsonSizeLimit)
+    public static <T> JsonResponse<T> execute(JsonCodec<T> codec, OkHttpClient client, Request request, OptionalLong materializedJsonSizeLimit) throws RuntimeException
     {
         try (Response response = client.newCall(request).execute()) {
             // TODO: fix in OkHttp: https://github.com/square/okhttp/issues/3111
@@ -106,6 +106,7 @@ public static <T> JsonResponse<T> execute(JsonCodec<T> codec, OkHttpClient clien
                         message = format("Unable to create %s from JSON response", codec.getType());
                     }
                     exception = new IllegalArgumentException(message, e);
+                    throw exception;
                 }
                 return new JsonResponse<>(response.code(), response.message(), response.headers(), body, value, exception);
             }

diff --git a/databend-client/src/main/java/com/databend/client/data/Int8Handler.java b/databend-client/src/main/java/com/databend/client/data/Int8Handler.java
@@ -465,6 +465,9 @@ private Float parseNullableValue(Object value) {
         if (value == null || value.equals("NULL")) {
             return null;
         }
+        if (value.equals("NaN") || value.equals("nan")) {
+            return Float.NaN;
+        }
         if (value instanceof String) {
             return Float.parseFloat((String) value);
         }
@@ -478,6 +481,9 @@ private float parseNonNullValue(Object value) {
         if (value == null || value.equals("NULL")) {
             throw new IllegalArgumentException("Float32 type is not nullable");
         }
+        if (value.equals("NaN") || value.equals("nan")) {
+            return Float.NaN;
+        }
         if (value instanceof String) {
             return Float.parseFloat((String) value);
         }
@@ -518,6 +524,9 @@ private Double parseNullableValue(Object value) {
         if (value == null || value.equals("NULL")) {
             return null;
         }
+        if (value.equals("NaN") || value.equals("nan")) {
+            return Double.NaN;
+        }
         if (value instanceof String) {
             return Double.parseDouble((String) value);
         }
@@ -531,6 +540,9 @@ private double parseNonNullValue(Object value) {
         if (value == null || value.equals("NULL")) {
             throw new IllegalArgumentException("Float64 type is not nullable");
         }
+        if (value.equals("NaN") || value.equals("nan")) {
+            return Double.NaN;
+        }
         if (value instanceof String) {
             return Double.parseDouble((String) value);
         }
@@ -631,7 +643,7 @@ private BigDecimal parseNullableValue(Object value) {
         if (value == null || value.equals("NULL")) {
             return null;
         }
-        if(value instanceof Integer){
+        if (value instanceof Integer) {
             return BigDecimal.valueOf((int) value);
         }
         if (value instanceof String) {
@@ -647,7 +659,7 @@ private BigDecimal parseNonNullValue(Object value) {
         if (value == null || value.equals("NULL")) {
             throw new IllegalArgumentException("decimal type is not nullable");
         }
-        if(value instanceof Integer){
+        if (value instanceof Integer) {
             return BigDecimal.valueOf((int) value);
         }
         if (value instanceof String) {

diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java
@@ -525,4 +525,31 @@ public void shouldBuildStageAttachmentWithFileFormatOptions() throws SQLExceptio
         Assertions.assertEquals("true", stageAttachment.getCopyOptions().get("PURGE"));
         Assertions.assertEquals("\\N", stageAttachment.getCopyOptions().get("NULL_DISPLAY"));
     }
+
+    @Test
+    public void testSelectWithClusterKey() throws SQLException {
+        Connection conn = createConnection();
+        conn.createStatement().execute("drop table if exists default.test_clusterkey");
+        conn.createStatement().execute("create table default.test_clusterkey (a int, b string)");
+        String insertSql = "insert into default.test_clusterkey values (?,?)";
+        try (PreparedStatement statement = conn.prepareStatement(insertSql)) {
+            statement.setInt(1, 1);
+            statement.setString(2, "b");
+            statement.addBatch();
+            statement.setInt(1, 2);
+            statement.setString(2, "c");
+            statement.addBatch();
+            int[] result = statement.executeBatch();
+            System.out.println(result);
+            Assertions.assertEquals(2, result.length);
+        }
+        conn.createStatement().execute("alter table default.test_clusterkey cluster by (a)");
+        String selectSQL = "select * from clustering_information('default','test_clusterkey')";
+        try (PreparedStatement statement = conn.prepareStatement(selectSQL)) {
+            ResultSet rs = statement.executeQuery();
+            while (rs.next()) {
+                Assertions.assertEquals("NaN", rs.getString(5));
+            }
+        }
+    }
 }