Skip to content

Commit fc70a1f

Browse files
madhav-dbmadhav-db
authored
Merge pull request #158 from databricks/madhav-db/large-queries-benchmark
Large Queries Benchmark
2 parents 32c7d09 + 4b3c0ba commit fc70a1f

File tree

3 files changed

+192
-0
lines changed

3 files changed

+192
-0
lines changed

‎pom.xml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,7 @@
206206
<excludes>
207207
<exclude>**/integration/**/*.java</exclude>
208208
<exclude>**/*MetadataBenchmarkingTest.java</exclude>
209+
<exclude>**/*LargeQueriesBenchmarkingTest.java</exclude>
209210
</excludes>
210211
<argLine>
211212
@{argLine}

‎src/test/java/com/databricks/jdbc/integration/IntegrationTestUtil.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,11 @@ public static String getDatabricksBenchfoodHost() {
4343
return System.getenv("DATABRICKS_BENCHFOOD_HOST");
4444
}
4545

46+
public static String getDatabricksBenchmarkingHost() {
47+
// includes port
48+
return System.getenv("DATABRICKS_BENCHMARKING_HOST");
49+
}
50+
4651
public static String getDatabricksToken() {
4752
return System.getenv("DATABRICKS_TOKEN");
4853
}
@@ -51,6 +56,10 @@ public static String getDatabricksBenchfoodToken() {
5156
return System.getenv("DATABRICKS_BENCHFOOD_TOKEN");
5257
}
5358

59+
public static String getDatabricksBenchmarkingToken() {
60+
return System.getenv("DATABRICKS_BENCHMARKING_TOKEN");
61+
}
62+
5463
public static String getDatabricksHTTPPath() {
5564
return System.getenv("DATABRICKS_HTTP_PATH");
5665
}
@@ -59,6 +68,10 @@ public static String getDatabricksBenchfoodHTTPPath() {
5968
return System.getenv("DATABRICKS_BENCHFOOD_HTTP_PATH");
6069
}
6170

71+
public static String getDatabricksBenchmarkingHTTPPath() {
72+
return System.getenv("DATABRICKS_BENCHMARKING_HTTP_PATH");
73+
}
74+
6275
public static String getDatabricksCatalog() {
6376
return System.getenv("DATABRICKS_CATALOG");
6477
}
@@ -82,6 +95,12 @@ public static Connection getBenchfoodJDBCConnection() throws SQLException {
8295
getBenchfoodJDBCUrl(), getDatabricksUser(), getDatabricksBenchfoodToken());
8396
}
8497

98+
public static Connection getBenchmarkingJDBCConnection() throws SQLException {
99+
// add support for properties
100+
return DriverManager.getConnection(
101+
getBenchmarkingJDBCUrl(), getDatabricksUser(), getDatabricksBenchmarkingToken());
102+
}
103+
85104
public static void resetJDBCConnection() {
86105
JDBCConnection = null;
87106
}
@@ -107,6 +126,15 @@ public static String getBenchfoodJDBCUrl() {
107126
return String.format(template, host, httpPath);
108127
}
109128

129+
public static String getBenchmarkingJDBCUrl() {
130+
String template =
131+
"jdbc:databricks://%s/default;transportMode=http;ssl=1;AuthMech=3;httpPath=%s";
132+
String host = getDatabricksBenchmarkingHost();
133+
String httpPath = getDatabricksBenchmarkingHTTPPath();
134+
135+
return String.format(template, host, httpPath);
136+
}
137+
110138
public static boolean executeSQL(String sql) {
111139
try {
112140
if (JDBCConnection == null) JDBCConnection = getValidJDBCConnection();

‎src/test/java/com/databricks/jdbc/integration/benchmarking/LargeQueriesBenchmarkingTest.java‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
package com.databricks.jdbc.integration.benchmarking;
2+
3+
import static com.databricks.jdbc.integration.IntegrationTestUtil.*;
4+
5+
import java.sql.*;
6+
import java.time.LocalDateTime;
7+
import java.util.Arrays;
8+
import java.util.Enumeration;
9+
import java.util.Random;
10+
import org.junit.jupiter.api.AfterEach;
11+
import org.junit.jupiter.api.BeforeEach;
12+
import org.junit.jupiter.api.Test;
13+
14+
public class LargeQueriesBenchmarkingTest {
15+
private Connection connection;
16+
17+
private String SCHEMA_NAME = "main";
18+
private String TABLE_NAME = "tpcds_sf100_delta.catalog_sales";
19+
20+
private static final String RESULTS_TABLE =
21+
"main.jdbc_large_queries_benchmarking_schema.benchmarking_results";
22+
private int ATTEMPTS = 10;
23+
24+
private int ROWS = 1000000;
25+
26+
long timesForOSSDriver[] = new long[ATTEMPTS];
27+
long timesForDatabricksDriver[] = new long[ATTEMPTS];
28+
29+
long totalTimeForOSSDriver = 0;
30+
long totalTimeForDatabricksDriver = 0;
31+
32+
@BeforeEach
33+
void setUp() throws SQLException {
34+
connection = getBenchmarkingJDBCConnection();
35+
}
36+
37+
@AfterEach
38+
void tearDown() throws SQLException {
39+
DriverManager.registerDriver(new com.databricks.jdbc.driver.DatabricksDriver());
40+
insertBenchmarkingDataIntoBenchfood();
41+
connection.close();
42+
}
43+
44+
@Test
45+
void testLargeQueries() throws SQLException {
46+
// Currently connection is held by OSS driver
47+
long startTime = System.currentTimeMillis();
48+
measureLargeQueriesPerformance(1);
49+
long endTime = System.currentTimeMillis();
50+
System.out.println(
51+
"Time taken to execute large queries by OSS Driver: "
52+
+ (endTime - startTime)
53+
+ "ms "
54+
+ "for "
55+
+ ROWS
56+
+ " rows and "
57+
+ ATTEMPTS
58+
+ " attempts");
59+
60+
System.out.println(
61+
"Driver used : "
62+
+ connection.getMetaData().getDriverVersion()
63+
+ " "
64+
+ connection.getMetaData().getDriverName());
65+
66+
connection.close();
67+
68+
Enumeration<Driver> drivers = DriverManager.getDrivers();
69+
70+
while (drivers.hasMoreElements()) {
71+
Driver driver = drivers.nextElement();
72+
if (driver.getClass().getName().contains("DatabricksDriver")) {
73+
DriverManager.deregisterDriver(driver);
74+
}
75+
}
76+
77+
connection =
78+
DriverManager.getConnection(
79+
getBenchmarkingJDBCUrl(), "token", getDatabricksBenchmarkingToken());
80+
81+
startTime = System.currentTimeMillis();
82+
measureLargeQueriesPerformance(2);
83+
endTime = System.currentTimeMillis();
84+
System.out.println(
85+
"Time taken to execute large queries by Databricks Driver: "
86+
+ (endTime - startTime)
87+
+ "ms "
88+
+ "for "
89+
+ ROWS
90+
+ " rows and "
91+
+ ATTEMPTS
92+
+ " attempts");
93+
94+
System.out.println(
95+
"Driver used : "
96+
+ connection.getMetaData().getDriverVersion()
97+
+ " "
98+
+ connection.getMetaData().getDriverName());
99+
100+
connection.close();
101+
}
102+
103+
void measureLargeQueriesPerformance(int recording) {
104+
Random random = new Random();
105+
for (int i = 0; i < ATTEMPTS; i++) {
106+
System.out.println("Attempt: " + i);
107+
int offset =
108+
i * 1000000 + random.nextInt(1000000); // Randomization to avoid possible query caching
109+
try (Statement statement = connection.createStatement()) {
110+
long startTime = System.currentTimeMillis();
111+
ResultSet rs =
112+
statement.executeQuery(
113+
"SELECT * FROM "
114+
+ SCHEMA_NAME
115+
+ "."
116+
+ TABLE_NAME
117+
+ " LIMIT "
118+
+ ROWS
119+
+ " OFFSET "
120+
+ offset);
121+
while (rs.next()) {}
122+
long endTime = System.currentTimeMillis();
123+
if (recording == 1) {
124+
timesForOSSDriver[i] = endTime - startTime;
125+
} else {
126+
timesForDatabricksDriver[i] = endTime - startTime;
127+
}
128+
} catch (SQLException e) {
129+
e.printStackTrace();
130+
}
131+
}
132+
}
133+
134+
private void insertBenchmarkingDataIntoBenchfood() throws SQLException {
135+
136+
Arrays.sort(timesForOSSDriver);
137+
Arrays.sort(timesForDatabricksDriver);
138+
139+
// Removing min and max times for better average calculation
140+
for (int i = 1; i < ATTEMPTS - 1; i++) {
141+
totalTimeForOSSDriver += timesForOSSDriver[i];
142+
totalTimeForDatabricksDriver += timesForDatabricksDriver[i];
143+
}
144+
145+
connection = getBenchfoodJDBCConnection();
146+
147+
String sql =
148+
"INSERT INTO "
149+
+ RESULTS_TABLE
150+
+ "(DateTime, OSS_AVG, DATABRICKS_AVG, OSS_TOTAL, DATABRICKS_TOTAL) VALUES (?, ?, ?, ?, ?)";
151+
152+
try (PreparedStatement stmt = connection.prepareStatement(sql)) {
153+
stmt.setTimestamp(1, java.sql.Timestamp.valueOf(LocalDateTime.now()));
154+
stmt.setDouble(2, ((totalTimeForOSSDriver * 1.0) / (ATTEMPTS - 2)));
155+
stmt.setDouble(3, ((totalTimeForDatabricksDriver * 1.0) / (ATTEMPTS - 2)));
156+
stmt.setLong(4, totalTimeForOSSDriver);
157+
stmt.setLong(5, totalTimeForDatabricksDriver);
158+
stmt.executeUpdate();
159+
} catch (SQLException e) {
160+
e.printStackTrace();
161+
}
162+
}
163+
}

0 commit comments

Comments
 (0)