Merge branch 'metadata-benchmarks' into large-queries-benchmark

Author: madhav-db

src/main/java/com/databricks/jdbc/client/impl/thrift/commons/DatabricksThriftAccessor.java

@@ -27,6 +27,8 @@ public class DatabricksThriftAccessor {
   private static final Logger LOGGER = LoggerFactory.getLogger(DatabricksThriftAccessor.class);
   private final DatabricksConfig databricksConfig;
   private final TCLIService.Client thriftClient;
+  private static final TSparkGetDirectResults DEFAULT_DIRECT_RESULTS =
+      new TSparkGetDirectResults().setMaxRows(DEFAULT_ROW_LIMIT).setMaxBytes(DEFAULT_BYTE_LIMIT);
 
   public DatabricksThriftAccessor(IDatabricksConnectionContext connectionContext) {
     DatabricksHttpTTransport transport =
@@ -117,9 +119,7 @@ public TFetchResultsResp getResultSetResp(
     TFetchResultsReq request =
         new TFetchResultsReq()
             .setOperationHandle(operationHandle)
-            .setIncludeResultSetMetadata(
-                true) // TODO check: even though this is set, we had to make another call for
-            // metadata.
+            .setIncludeResultSetMetadata(true)
             .setFetchType((short) 0) // 0 represents Query output. 1 represents Log
             .setMaxRows(maxRows)
             .setMaxBytes(DEFAULT_BYTE_LIMIT);
@@ -139,26 +139,37 @@ public TFetchResultsResp getResultSetResp(
     }
     verifySuccessStatus(
         response.getStatus().getStatusCode(),
-        String.format("Error while fetching results. TFetchResultsResp {}. "));
+        String.format(
+            "Error while fetching results. TFetchResultsResp {%s}. ", response.toString()));
     return response;
   }
 
   private TFetchResultsResp execute(
       TExecuteStatementReq request, IDatabricksStatement parentStatement)
       throws TException, SQLException {
-    TExecuteStatementResp tExecuteStatementResp = thriftClient.ExecuteStatement(request);
     int maxRows = (parentStatement == null) ? DEFAULT_ROW_LIMIT : parentStatement.getMaxRows();
+    TSparkGetDirectResults directResults =
+        new TSparkGetDirectResults().setMaxBytes(DEFAULT_BYTE_LIMIT).setMaxRows(maxRows);
+    request.setGetDirectResults(directResults);
+    TExecuteStatementResp response = thriftClient.ExecuteStatement(request);
+    if (response.isSetDirectResults()) {
+      return response.getDirectResults().getResultSet();
+    }
     return getResultSetResp(
-        tExecuteStatementResp.getStatus().getStatusCode(),
-        tExecuteStatementResp.getOperationHandle(),
-        tExecuteStatementResp.toString(),
+        response.getStatus().getStatusCode(),
+        response.getOperationHandle(),
+        response.toString(),
         maxRows,
         true);
   }
 
   private TFetchResultsResp listFunctions(TGetFunctionsReq request)
       throws DatabricksHttpException, TException {
     TGetFunctionsResp response = thriftClient.GetFunctions(request);
+    request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
+    if (response.isSetDirectResults()) {
+      return response.getDirectResults().getResultSet();
+    }
     return getResultSetResp(
         response.getStatus().getStatusCode(),
         response.getOperationHandle(),
@@ -169,10 +180,11 @@ private TFetchResultsResp listFunctions(TGetFunctionsReq request)
 
   private TFetchResultsResp listPrimaryKeys(TGetPrimaryKeysReq request)
       throws DatabricksHttpException, TException {
-    request.setGetDirectResults(
-        new TSparkGetDirectResults().setMaxRows(100000).setMaxBytes(1000000));
+    request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
     TGetPrimaryKeysResp response = thriftClient.GetPrimaryKeys(request);
-    System.out.println("Without resultSet is : " + response.toString());
+    if (response.isSetDirectResults()) {
+      return response.getDirectResults().getResultSet();
+    }
     return getResultSetResp(
         response.getStatus().getStatusCode(),
         response.getOperationHandle(),
@@ -183,7 +195,12 @@ private TFetchResultsResp listPrimaryKeys(TGetPrimaryKeysReq request)
 
   private TFetchResultsResp getTables(TGetTablesReq request)
       throws TException, DatabricksHttpException {
+    request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
     TGetTablesResp response = thriftClient.GetTables(request);
+    request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
+    if (response.isSetDirectResults()) {
+      return response.getDirectResults().getResultSet();
+    }
     return getResultSetResp(
         response.getStatus().getStatusCode(),
         response.getOperationHandle(),
@@ -194,7 +211,11 @@ private TFetchResultsResp getTables(TGetTablesReq request)
 
   private TFetchResultsResp getTableTypes(TGetTableTypesReq request)
       throws TException, DatabricksHttpException {
+    request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
     TGetTableTypesResp response = thriftClient.GetTableTypes(request);
+    if (response.isSetDirectResults()) {
+      return response.getDirectResults().getResultSet();
+    }
     return getResultSetResp(
         response.getStatus().getStatusCode(),
         response.getOperationHandle(),
@@ -205,7 +226,11 @@ private TFetchResultsResp getTableTypes(TGetTableTypesReq request)
 
   private TFetchResultsResp getCatalogs(TGetCatalogsReq request)
       throws TException, DatabricksHttpException {
+    request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
     TGetCatalogsResp response = thriftClient.GetCatalogs(request);
+    if (response.isSetDirectResults()) {
+      return response.getDirectResults().getResultSet();
+    }
     return getResultSetResp(
         response.getStatus().getStatusCode(),
         response.getOperationHandle(),
@@ -216,7 +241,11 @@ private TFetchResultsResp getCatalogs(TGetCatalogsReq request)
 
   private TFetchResultsResp listSchemas(TGetSchemasReq request)
       throws TException, DatabricksHttpException {
+    request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
     TGetSchemasResp response = thriftClient.GetSchemas(request);
+    if (response.isSetDirectResults()) {
+      return response.getDirectResults().getResultSet();
+    }
     return getResultSetResp(
         response.getStatus().getStatusCode(),
         response.getOperationHandle(),
@@ -227,7 +256,11 @@ private TFetchResultsResp listSchemas(TGetSchemasReq request)
 
   private TFetchResultsResp getTypeInfo(TGetTypeInfoReq request)
       throws TException, DatabricksHttpException {
+    request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
     TGetTypeInfoResp response = thriftClient.GetTypeInfo(request);
+    if (response.isSetDirectResults()) {
+      return response.getDirectResults().getResultSet();
+    }
     return getResultSetResp(
         response.getStatus().getStatusCode(),
         response.getOperationHandle(),
@@ -238,11 +271,15 @@ private TFetchResultsResp getTypeInfo(TGetTypeInfoReq request)
 
   private TFetchResultsResp listColumns(TGetColumnsReq request)
       throws TException, DatabricksHttpException {
-    TGetColumnsResp tGetColumnsResp = thriftClient.GetColumns(request);
+    request.setGetDirectResults(DEFAULT_DIRECT_RESULTS);
+    TGetColumnsResp response = thriftClient.GetColumns(request);
+    if (response.isSetDirectResults()) {
+      return response.getDirectResults().getResultSet();
+    }
     return getResultSetResp(
-        tGetColumnsResp.getStatus().getStatusCode(),
-        tGetColumnsResp.getOperationHandle(),
-        tGetColumnsResp.toString(),
+        response.getStatus().getStatusCode(),
+        response.getOperationHandle(),
+        response.toString(),
         DEFAULT_ROW_LIMIT,
         false);
   }

src/main/java/com/databricks/jdbc/commons/EnvironmentVariables.java

@@ -12,5 +12,5 @@ public final class EnvironmentVariables {
       false; // By default, we should not process the sql
 
   public static final TProtocolVersion JDBC_THRIFT_VERSION =
-      TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V10;
+      TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V9;
 }

src/test/java/com/databricks/jdbc/client/impl/thrift/DatabricksThriftServiceClientTest.java

@@ -36,7 +36,7 @@ void testCreateSession() throws DatabricksSQLException {
             .setInitialNamespace(getNamespace(CATALOG, SCHEMA))
             .setConfiguration(EMPTY_MAP)
             .setCanUseMultipleCatalogs(true)
-            .setClient_protocol(TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V10);
+            .setClient_protocol(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V9);
     TOpenSessionResp openSessionResp =
         new TOpenSessionResp()
             .setSessionHandle(SESSION_HANDLE)