Remove operator summaries from stored query completion event

trinodb · May 29, 2024 · f7367b7 · f7367b7
1 parent 7f1cb1d
commit f7367b7
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 2 deletions.
diff --git a/benchto-driver/pom.xml b/benchto-driver/pom.xml
@@ -103,7 +103,10 @@
             <groupId>org.apache.httpcomponents</groupId>
             <artifactId>httpclient</artifactId>
         </dependency>
-
+        <dependency>
+            <groupId>com.google.code.gson</groupId>
+            <artifactId>gson</artifactId>
+        </dependency>
         <!-- Utils -->
         <dependency>
             <groupId>com.google.guava</groupId>

diff --git a/...driver/src/main/java/io/trino/benchto/driver/presto/PrestoQueryCompletionEventLoader.java b/...driver/src/main/java/io/trino/benchto/driver/presto/PrestoQueryCompletionEventLoader.java
@@ -13,13 +13,22 @@
  */
 package io.trino.benchto.driver.presto;
 
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParser;
+import com.google.gson.internal.Streams;
+import com.google.gson.stream.JsonReader;
+import com.google.gson.stream.JsonWriter;
 import io.trino.benchto.driver.Measurable;
 import io.trino.benchto.driver.execution.QueryExecutionResult;
 import io.trino.benchto.driver.listeners.queryinfo.QueryCompletionEventProvider;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.stereotype.Component;
 
+import java.io.IOException;
+import java.io.StringReader;
+import java.io.StringWriter;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
@@ -38,9 +47,28 @@ public CompletableFuture<Optional<String>> loadQueryCompletionEvent(Measurable m
     {
         if (measurable instanceof QueryExecutionResult executionResult) {
             if (executionResult.getPrestoQueryId().isPresent()) {
-                return completedFuture(prestoClient.getQueryCompletionEvent(executionResult.getPrestoQueryId().get()));
+                Optional<String> queryCompletionEvent = prestoClient.getQueryCompletionEvent(executionResult.getPrestoQueryId().get());
+                queryCompletionEvent = queryCompletionEvent.map(event -> {
+                    JsonElement root = JsonParser.parseReader(new JsonReader(new StringReader(event)));
+                    JsonObject rootObject = root.getAsJsonObject();
+                    JsonObject statistics = rootObject.getAsJsonObject("statistics");
+                    // remove huge operatorSummaries entry as it is already part of query_info
+                    statistics.remove("operatorSummaries");
+                    StringWriter stringWriter = new StringWriter();
+                    JsonWriter jsonWriter = new JsonWriter(stringWriter);
+                    jsonWriter.setLenient(true);
+                    try {
+                        Streams.write(root, jsonWriter);
+                    }
+                    catch (IOException e) {
+                        throw new RuntimeException(e);
+                    }
+                    return stringWriter.toString();
+                });
+                return completedFuture(queryCompletionEvent);
             }
         }
         return completedFuture(Optional.empty());
     }
 }
+