Skip to content

Commit 79ddff4

Browse files
samikshya-dbsamikshya-db
authored
Improve perf - by enabling setMetadata (#548)
1 parent 1c8f0f9 commit 79ddff4

File tree

3 files changed

+29
-30
lines changed

3 files changed

+29
-30
lines changed

‎src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -141,18 +141,17 @@ private TFetchResultsResp getResultSetResp(
141141
TFetchResultsReq request =
142142
new TFetchResultsReq()
143143
.setOperationHandle(operationHandle)
144-
.setIncludeResultSetMetadata(true)
145144
.setFetchType((short) 0) // 0 represents Query output. 1 represents Log
146145
.setMaxRows(maxRows)
147146
.setMaxBytes(DEFAULT_BYTE_LIMIT);
147+
if (fetchMetadata) {
148+
request.setIncludeResultSetMetadata(true);
149+
}
148150
TFetchResultsResp response = null;
149151
DatabricksHttpTTransport transport =
150152
(DatabricksHttpTTransport) getThriftClient().getInputProtocol().getTransport();
151153
try {
152154
response = getThriftClient().FetchResults(request);
153-
if (fetchMetadata) {
154-
response.setResultSetMetadata(getResultSetMetadata(operationHandle));
155-
}
156155
} catch (TException e) {
157156
String errorMessage =
158157
String.format(
@@ -415,7 +414,7 @@ private TFetchResultsResp getTypeInfo(TGetTypeInfoReq request)
415414
response.getOperationHandle(),
416415
response.toString(),
417416
DEFAULT_ROW_LIMIT,
418-
true);
417+
false);
419418
} finally {
420419
transport.close();
421420
}

‎src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,7 @@ public DatabricksResultSet listTypeInfo(IDatabricksSession session)
198198
TGetTypeInfoReq request =
199199
new TGetTypeInfoReq().setSessionHandle(session.getSessionInfo().sessionHandle());
200200
TFetchResultsResp response = (TFetchResultsResp) thriftAccessor.getThriftResponse(request);
201-
return getTypeInfoResult(extractValues(response.getResults().getColumns()));
201+
return getTypeInfoResult(extractValuesColumnar(response.getResults().getColumns()));
202202
}
203203

204204
@Override

‎src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java‎

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -38,16 +38,6 @@ public class DatabricksThriftAccessorTest {
3838
static THandleIdentifier handleIdentifier = new THandleIdentifier().setGuid(TEST_BYTES);
3939
private static final TOperationHandle tOperationHandle =
4040
new TOperationHandle().setOperationId(handleIdentifier).setHasResultSet(false);
41-
private static final TFetchResultsReq fetchResultsReq =
42-
new TFetchResultsReq()
43-
.setOperationHandle(tOperationHandle)
44-
.setIncludeResultSetMetadata(true)
45-
.setFetchType((short) 0)
46-
.setMaxRows(DEFAULT_ROW_LIMIT)
47-
.setMaxBytes(DEFAULT_BYTE_LIMIT);
48-
private static final TGetResultSetMetadataReq resultSetMetadataReq =
49-
new TGetResultSetMetadataReq().setOperationHandle(tOperationHandle);
50-
5141
private static final TRowSet rowSet = new TRowSet().setResultLinks(new ArrayList<>(2));
5242

5343
private static final TFetchResultsResp response =
@@ -64,8 +54,6 @@ public class DatabricksThriftAccessorTest {
6454
new TGetResultSetMetadataResp()
6555
.setResultFormat(TSparkRowSetType.COLUMN_BASED_SET)
6656
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS)));
67-
private static final TGetResultSetMetadataResp metadataResp =
68-
new TGetResultSetMetadataResp().setResultFormat(TSparkRowSetType.COLUMN_BASED_SET);
6957

7058
private static final String NEW_ACCESS_TOKEN = "new-access-token";
7159

@@ -104,8 +92,7 @@ void testExecute() throws TException, SQLException {
10492
new TExecuteStatementResp()
10593
.setOperationHandle(tOperationHandle)
10694
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
107-
when(thriftClient.GetResultSetMetadata(resultSetMetadataReq)).thenReturn(metadataResp);
108-
when(thriftClient.FetchResults(fetchResultsReq)).thenReturn(response);
95+
when(thriftClient.FetchResults(getFetchResultsRequest(true))).thenReturn(response);
10996
when(thriftClient.ExecuteStatement(request)).thenReturn(tExecuteStatementResp);
11097
TGetOperationStatusReq operationStatusReq =
11198
new TGetOperationStatusReq()
@@ -203,7 +190,7 @@ void testListPrimaryKeys() throws TException, DatabricksSQLException {
203190
new TGetPrimaryKeysResp()
204191
.setOperationHandle(tOperationHandle)
205192
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
206-
when(thriftClient.FetchResults(fetchResultsReq)).thenReturn(response);
193+
when(thriftClient.FetchResults(getFetchResultsRequest(false))).thenReturn(response);
207194
when(thriftClient.GetPrimaryKeys(request)).thenReturn(tGetPrimaryKeysResp);
208195
TFetchResultsResp actualResponse = (TFetchResultsResp) accessor.getThriftResponse(request);
209196
assertEquals(actualResponse, response);
@@ -231,7 +218,7 @@ void testListFunctions() throws TException, DatabricksSQLException {
231218
new TGetFunctionsResp()
232219
.setOperationHandle(tOperationHandle)
233220
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
234-
when(thriftClient.FetchResults(fetchResultsReq)).thenReturn(response);
221+
when(thriftClient.FetchResults(getFetchResultsRequest(false))).thenReturn(response);
235222
when(thriftClient.GetFunctions(request)).thenReturn(tGetFunctionsResp);
236223
TFetchResultsResp actualResponse = (TFetchResultsResp) accessor.getThriftResponse(request);
237224
assertEquals(actualResponse, response);
@@ -259,7 +246,7 @@ void testListSchemas() throws TException, DatabricksSQLException {
259246
new TGetSchemasResp()
260247
.setOperationHandle(tOperationHandle)
261248
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
262-
when(thriftClient.FetchResults(fetchResultsReq)).thenReturn(response);
249+
when(thriftClient.FetchResults(getFetchResultsRequest(false))).thenReturn(response);
263250
when(thriftClient.GetSchemas(request)).thenReturn(tGetSchemasResp);
264251
TFetchResultsResp actualResponse = (TFetchResultsResp) accessor.getThriftResponse(request);
265252
assertEquals(actualResponse, response);
@@ -287,7 +274,7 @@ void testListColumns() throws TException, DatabricksSQLException {
287274
new TGetColumnsResp()
288275
.setOperationHandle(tOperationHandle)
289276
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
290-
when(thriftClient.FetchResults(fetchResultsReq)).thenReturn(response);
277+
when(thriftClient.FetchResults(getFetchResultsRequest(false))).thenReturn(response);
291278
when(thriftClient.GetColumns(request)).thenReturn(tGetColumnsResp);
292279
TFetchResultsResp actualResponse = (TFetchResultsResp) accessor.getThriftResponse(request);
293280
assertEquals(actualResponse, response);
@@ -315,7 +302,7 @@ void testListCatalogs() throws TException, DatabricksSQLException {
315302
new TGetCatalogsResp()
316303
.setOperationHandle(tOperationHandle)
317304
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
318-
when(thriftClient.FetchResults(fetchResultsReq)).thenReturn(response);
305+
when(thriftClient.FetchResults(getFetchResultsRequest(false))).thenReturn(response);
319306
when(thriftClient.GetCatalogs(request)).thenReturn(tGetCatalogsResp);
320307
TFetchResultsResp actualResponse = (TFetchResultsResp) accessor.getThriftResponse(request);
321308
assertEquals(actualResponse, response);
@@ -343,7 +330,7 @@ void testListTables() throws TException, DatabricksSQLException {
343330
new TGetTablesResp()
344331
.setOperationHandle(tOperationHandle)
345332
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
346-
when(thriftClient.FetchResults(fetchResultsReq)).thenReturn(response);
333+
when(thriftClient.FetchResults(getFetchResultsRequest(false))).thenReturn(response);
347334
when(thriftClient.GetTables(request)).thenReturn(tGetTablesResp);
348335
TFetchResultsResp actualResponse = (TFetchResultsResp) accessor.getThriftResponse(request);
349336
assertEquals(actualResponse, response);
@@ -371,7 +358,7 @@ void testListTableTypes() throws TException, DatabricksSQLException {
371358
new TGetTableTypesResp()
372359
.setOperationHandle(tOperationHandle)
373360
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
374-
when(thriftClient.FetchResults(fetchResultsReq)).thenReturn(response);
361+
when(thriftClient.FetchResults(getFetchResultsRequest(false))).thenReturn(response);
375362
when(thriftClient.GetTableTypes(request)).thenReturn(tGetTableTypesResp);
376363
TFetchResultsResp actualResponse = (TFetchResultsResp) accessor.getThriftResponse(request);
377364
assertEquals(actualResponse, response);
@@ -399,7 +386,7 @@ void testTypeInfo() throws TException, DatabricksSQLException {
399386
new TGetTypeInfoResp()
400387
.setOperationHandle(tOperationHandle)
401388
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
402-
when(thriftClient.FetchResults(fetchResultsReq)).thenReturn(response);
389+
when(thriftClient.FetchResults(getFetchResultsRequest(false))).thenReturn(response);
403390
when(thriftClient.GetTypeInfo(request)).thenReturn(tGetTypeInfoResp);
404391
TFetchResultsResp actualResponse = (TFetchResultsResp) accessor.getThriftResponse(request);
405392
assertEquals(actualResponse, response);
@@ -428,7 +415,7 @@ void testAccessorWhenFetchResultsThrowsError() throws TException {
428415
.setOperationHandle(tOperationHandle)
429416
.setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS));
430417
when(thriftClient.GetTables(request)).thenReturn(tGetTablesResp);
431-
when(thriftClient.FetchResults(fetchResultsReq)).thenThrow(new TException());
418+
when(thriftClient.FetchResults(getFetchResultsRequest(false))).thenThrow(new TException());
432419
assertThrows(DatabricksSQLException.class, () -> accessor.getThriftResponse(request));
433420
}
434421

@@ -441,9 +428,22 @@ void testAccessorDuringThriftError() throws TException {
441428
}
442429

443430
@Test
444-
void testResetAccessToken() throws Exception {
431+
void testResetAccessToken() {
445432
accessor = new DatabricksThriftAccessor(thriftClient, config, connectionContext);
446433
accessor.resetAccessToken(NEW_ACCESS_TOKEN);
447434
verify(config).setToken(NEW_ACCESS_TOKEN);
448435
}
436+
437+
private TFetchResultsReq getFetchResultsRequest(boolean includeMetadata) {
438+
TFetchResultsReq request =
439+
new TFetchResultsReq()
440+
.setOperationHandle(tOperationHandle)
441+
.setFetchType((short) 0)
442+
.setMaxRows(DEFAULT_ROW_LIMIT)
443+
.setMaxBytes(DEFAULT_BYTE_LIMIT);
444+
if (includeMetadata) {
445+
request.setIncludeResultSetMetadata(true);
446+
}
447+
return request;
448+
}
449449
}

0 commit comments

Comments
 (0)