Merge pull request #394 from youngsofun/feat/arrow

feat(jdbc):  arrow format support decimal.

Author: youngsofun

README.md

@@ -249,7 +249,6 @@ public class LoadStreamExample {
         }
     }
 }
-
 ```
 
 ### method `uploadStream` and `downloadStream`
@@ -267,3 +266,34 @@ Download a single file in the stage as `InputStream`
 ```
 InputStream downloadStream(String stageName, String filePathInStage) throws SQLException;
 ```
+
+### Use Arrow Result Format
+
+By default, the driver fetches query results in JSON format. To enable Arrow over HTTP, add
+`query_result_format=arrow` to the JDBC URL:
+
+```java
+String url = "jdbc:databend://localhost:8000/default?query_result_format=arrow";
+Connection conn = DriverManager.getConnection(url, "root", "");
+```
+
+Arrow mode is intended for query result fetching. Internal control queries still use JSON when needed.
+
+Requirements:
+
+1. Databend server must support Arrow result pages.
+2. The JVM must allow Arrow to access `java.nio` internals.
+
+Before starting your application, set:
+
+```shell
+export JAVA_TOOL_OPTIONS='--add-opens=java.base/java.nio=ALL-UNNAMED -Dio.netty.tryReflectionSetAccessible=true'
+```
+
+If you do not want to set `JAVA_TOOL_OPTIONS` globally, pass the same options directly to `java`:
+
+```shell
+java --add-opens=java.base/java.nio=ALL-UNNAMED -Dio.netty.tryReflectionSetAccessible=true -jar your-app.jar
+```
+
+If `query_result_format` is not specified, the driver uses JSON.

databend-jdbc/src/main/java/com/databend/jdbc/internal/query/QueryRequest.java

@@ -2,17 +2,20 @@
 
 import com.databend.jdbc.internal.session.SessionState;
 import com.databend.jdbc.internal.session.PaginationOptions;
+import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.databind.ObjectMapper;
 
+@JsonInclude(JsonInclude.Include.NON_NULL)
 public class QueryRequest {
     private final String sql;
     private final String sessionId;
     private final PaginationOptions paginationOptions;
     private final SessionState session;
     private final StageAttachment stageAttachment;
     private final Integer arrowResultVersionMax;
+    private final ArrowFeatures arrowFeatures;
 
     @JsonCreator
     public QueryRequest(
@@ -21,13 +24,15 @@ public QueryRequest(
             @JsonProperty("pagination") PaginationOptions paginationOptions,
             @JsonProperty("session") SessionState session,
             @JsonProperty("stage_attachment") StageAttachment stageAttachment,
-            @JsonProperty("arrow_result_version_max") Integer arrowResultVersionMax) {
+            @JsonProperty("arrow_result_version_max") Integer arrowResultVersionMax,
+            @JsonProperty("arrow_features") ArrowFeatures arrowFeatures) {
         this.sql = sql;
         this.sessionId = sessionId;
         this.paginationOptions = paginationOptions;
         this.session = session;
         this.stageAttachment = stageAttachment;
         this.arrowResultVersionMax = arrowResultVersionMax;
+        this.arrowFeatures = arrowFeatures;
     }
 
     public static Builder builder() {
@@ -64,6 +69,11 @@ public Integer getArrowResultVersionMax() {
         return arrowResultVersionMax;
     }
 
+    @JsonProperty("arrow_features")
+    public ArrowFeatures getArrowFeatures() {
+        return arrowFeatures;
+    }
+
     @Override
     public String toString() {
         try {
@@ -80,6 +90,7 @@ public static class Builder {
         private SessionState session;
         private StageAttachment stageAttachment;
         private Integer arrowResultVersionMax;
+        private ArrowFeatures arrowFeatures;
 
         public Builder setSql(String sql) {
             this.sql = sql;
@@ -111,8 +122,27 @@ public Builder setArrowResultVersionMax(Integer arrowResultVersionMax) {
             return this;
         }
 
+        public Builder setArrowFeatures(ArrowFeatures arrowFeatures) {
+            this.arrowFeatures = arrowFeatures;
+            return this;
+        }
+
         public QueryRequest build() {
-            return new QueryRequest(sql, sessionId, paginationOptions, session, stageAttachment, arrowResultVersionMax);
+            return new QueryRequest(sql, sessionId, paginationOptions, session, stageAttachment, arrowResultVersionMax, arrowFeatures);
+        }
+    }
+
+    public static class ArrowFeatures {
+        private final Boolean decimal64;
+
+        @JsonCreator
+        public ArrowFeatures(@JsonProperty("decimal64") Boolean decimal64) {
+            this.decimal64 = decimal64;
+        }
+
+        @JsonProperty("decimal64")
+        public Boolean getDecimal64() {
+            return decimal64;
         }
     }
 }

databend-jdbc/src/main/java/com/databend/jdbc/internal/query/RestQueryResultPages.java

@@ -40,6 +40,7 @@
 
 @ThreadSafe
 public class RestQueryResultPages implements QueryResultPages {
+    private static final int ARROW_FEATURE_NEGOTIATION_VERSION = 3;
     public static final MediaType MEDIA_TYPE_JSON = MediaType.parse("application/json; charset=utf-8");
     public static final MediaType MEDIA_TYPE_ARROW = MediaType.parse("application/vnd.apache.arrow.stream");
     public static final JsonCodec<QueryResults> QUERY_RESULTS_CODEC = jsonCodec(QueryResults.class);
@@ -104,7 +105,8 @@ private Request buildQueryRequest(String query, QueryRequestConfig requestConfig
                 .setStageAttachment(requestConfig.getStageAttachment())
                 .setPaginationOptions(requestConfig.getPaginationOptions())
                 .setSql(query)
-                .setArrowResultVersionMax(currentFormat == QueryResultFormat.ARROW ? 2 : null)
+                .setArrowResultVersionMax(currentFormat == QueryResultFormat.ARROW ? ARROW_FEATURE_NEGOTIATION_VERSION : null)
+                .setArrowFeatures(currentFormat == QueryResultFormat.ARROW ? new QueryRequest.ArrowFeatures(false) : null)
                 .build();
         String reqString = req.toString();
         if (reqString == null || reqString.isEmpty()) {

databend-jdbc/src/main/java/com/databend/jdbc/internal/session/DatabendSessionHandle.java

@@ -70,6 +70,7 @@ public class DatabendSessionHandle implements Consumer<SessionState> {
     private static final MediaType MEDIA_TYPE_JSON = MediaType.parse("application/json; charset=utf-8");
     private static final Semver STREAMING_LOAD_MIN_VERSION = new Semver("1.2.781");
     private static final Semver HEARTBEAT_MIN_VERSION = new Semver("1.2.709");
+    private static final int MIN_ARROW_RESULT_VERSION = 3;
     private static final int MAX_STAGE_UPLOAD_RETRY_ATTEMPTS = 20;
     private static final Duration STAGE_UPLOAD_RETRY_TIMEOUT = Duration.ofMinutes(5);
 
@@ -424,10 +425,7 @@ private QueryRequestConfig.Builder makeRequestConfig(String queryId, String host
     }
 
     private boolean supportsArrowTransport() {
-        if (this.serverMaxArrowResultVersion != null) {
-            return this.serverMaxArrowResultVersion > 0;
-        }
-        return this.serverVersion != null && new Capability(this.serverVersion).arrowData();
+        return this.serverMaxArrowResultVersion != null && this.serverMaxArrowResultVersion >= MIN_ARROW_RESULT_VERSION;
     }
 
     private void logout() throws SQLException {

databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDatabaseMetaData.java

@@ -2,7 +2,6 @@
 
 import com.vdurmont.semver4j.Semver;
 import org.testng.Assert;
-import org.testng.SkipException;
 import org.testng.annotations.BeforeTest;
 import org.testng.annotations.Test;
 
@@ -261,9 +260,6 @@ public void testGetColumnTypeWithDecimal() throws Exception {
 
     @Test(groups = {"IT"})
     public void testGetObjectWithDecimal() throws Exception {
-        if ("arrow".equalsIgnoreCase(System.getenv("DATABEND_JDBC_TEST_QUERY_RESULT_FORMAT"))) {
-            throw new SkipException("TODO: re-enable after Arrow Decimal64/Decimal128 compatibility is fixed");
-        }
         try (Connection connection = Utils.createConnection()) {
             connection.createStatement().execute("insert into decimal_test values(1.2)");
             ResultSet rs = connection.createStatement().executeQuery("select * from decimal_test");
@@ -275,9 +271,6 @@ public void testGetObjectWithDecimal() throws Exception {
 
     @Test(groups = {"IT"})
     public void testGetBigDecimal() throws Exception {
-        if ("arrow".equalsIgnoreCase(System.getenv("DATABEND_JDBC_TEST_QUERY_RESULT_FORMAT"))) {
-            throw new SkipException("TODO: re-enable after Arrow Decimal64/Decimal128 compatibility is fixed");
-        }
         String bigDecimalStr = "123.456789012345";
         String scaleBigDecimalStr = "123.46";
         String columnLabel = "a";

databend-jdbc/src/test/java/com/databend/jdbc/TestTypes.java

@@ -5,7 +5,6 @@
 import org.locationtech.jts.io.ParseException;
 import org.locationtech.jts.io.WKBReader;
 import org.testng.Assert;
-import org.testng.SkipException;
 import org.testng.annotations.BeforeSuite;
 import org.testng.annotations.BeforeTest;
 import org.testng.annotations.DataProvider;
@@ -37,11 +36,6 @@ public Object[][] provideFlag() {
     @Test(groups = {"IT"})
     public void testGetDecimalByQueryResultFormat()
             throws SQLException {
-        String queryResultFormat = currentQueryResultFormat();
-        if ("arrow".equals(queryResultFormat)) {
-            throw new SkipException("TODO: re-enable after Arrow Decimal64/Decimal128 compatibility is fixed");
-        }
-
         String sql = "select cast(123.456789012345 as decimal(15, 12)) as a";
         try (Connection connection = Utils.createConnection();
              Statement statement = connection.createStatement()) {
@@ -533,12 +527,4 @@ public void TestBatchInsertWithComplexDataType(boolean presigned,  boolean place
         }
     }
 
-    private static String currentQueryResultFormat() {
-        String queryResultFormat = System.getenv("DATABEND_JDBC_TEST_QUERY_RESULT_FORMAT");
-        if (queryResultFormat == null || queryResultFormat.trim().isEmpty()) {
-            return "json";
-        }
-        return queryResultFormat.trim().toLowerCase();
-    }
-
 }

databend-jdbc/src/test/java/com/databend/jdbc/internal/session/TestDatabendSessionHandle.java

@@ -88,7 +88,7 @@ public void testLoginEnablesArrowTransportWhenServerAdvertisesArrowResultVersion
         HttpServer server = HttpServer.create(new InetSocketAddress(0), 0);
         server.createContext("/v1/session/login", exchange -> {
             try {
-                byte[] response = "{\"version\":\"1.2.100\",\"server_max_arrow_result_version\":2}".getBytes(StandardCharsets.UTF_8);
+                byte[] response = "{\"version\":\"1.2.100\",\"server_max_arrow_result_version\":3}".getBytes(StandardCharsets.UTF_8);
                 exchange.getResponseHeaders().add("Content-Type", "application/json");
                 exchange.sendResponseHeaders(200, response.length);
                 exchange.getResponseBody().write(response);
@@ -134,7 +134,70 @@ public void testLoginEnablesArrowTransportWhenServerAdvertisesArrowResultVersion
             pages.close();
 
             Assert.assertEquals(accept.get(), "application/vnd.apache.arrow.stream");
-            Assert.assertTrue(requestBody.get().contains("\"arrow_result_version_max\":2"));
+            Assert.assertTrue(requestBody.get().contains("\"arrow_result_version_max\":3"));
+            Assert.assertTrue(requestBody.get().contains("\"arrow_features\":{\"decimal64\":false}"));
+        }
+        finally {
+            server.stop(0);
+        }
+    }
+
+    @Test(groups = {"UNIT"})
+    public void testLoginFallsBackToJsonWhenServerArrowResultVersionIsTooLow() throws Exception {
+        AtomicReference<String> accept = new AtomicReference<>();
+        AtomicReference<String> requestBody = new AtomicReference<>();
+
+        HttpServer server = HttpServer.create(new InetSocketAddress(0), 0);
+        server.createContext("/v1/session/login", exchange -> {
+            try {
+                byte[] response = "{\"version\":\"1.2.100\",\"server_max_arrow_result_version\":2}".getBytes(StandardCharsets.UTF_8);
+                exchange.getResponseHeaders().add("Content-Type", "application/json");
+                exchange.sendResponseHeaders(200, response.length);
+                exchange.getResponseBody().write(response);
+            }
+            finally {
+                exchange.close();
+            }
+        });
+        server.createContext("/v1/query", exchange -> {
+            try {
+                accept.set(exchange.getRequestHeaders().getFirst("Accept"));
+                requestBody.set(new String(readAllBytes(exchange), StandardCharsets.UTF_8));
+                byte[] response = "{\"id\":\"qid\",\"node_id\":\"node\",\"session\":{\"database\":\"default\"},\"schema\":[],\"data\":[]}".getBytes(StandardCharsets.UTF_8);
+                exchange.getResponseHeaders().add("Content-Type", "application/json");
+                exchange.sendResponseHeaders(200, response.length);
+                exchange.getResponseBody().write(response);
+            }
+            finally {
+                exchange.close();
+            }
+        });
+        server.start();
+
+        try {
+            DatabendSessionHandle handle = new DatabendSessionHandle(
+                    new OkHttpClient.Builder()
+                            .addInterceptor(userAgentInterceptor(DriverInfo.USER_AGENT_VALUE))
+                            .build(),
+                    SessionHandleConfig.builder()
+                            .setBaseUri(URI.create("http://127.0.0.1:" + server.getAddress().getPort()))
+                            .setInitialSession(SessionState.createDefault())
+                            .setQueryResultFormat(QueryResultFormat.ARROW)
+                            .setQueryTimeoutSecs(30)
+                            .setConnectionTimeoutSecs(30)
+                            .setSocketTimeoutSecs(60)
+                            .setWaitTimeSecs(10)
+                            .setMaxRowsInBuffer(1000)
+                            .setMaxRowsPerPage(1000)
+                            .build(),
+                    null);
+            handle.login();
+            QueryResultPages pages = handle.startQuery("qid", "select 1", null, null);
+            pages.close();
+
+            Assert.assertEquals(accept.get(), "application/json");
+            Assert.assertFalse(requestBody.get().contains("\"arrow_result_version_max\""));
+            Assert.assertFalse(requestBody.get().contains("\"arrow_features\""));
         }
         finally {
             server.stop(0);