From 4563dde4e8033dfa6718a04d86e0a3797336ec52 Mon Sep 17 00:00:00 2001
From: Yuya Ebihara <ebyhry@gmail.com>
Date: Wed, 23 Mar 2022 12:32:02 +0900
Subject: [PATCH] Add processed input bytes and rows to event listener
 QueryStatistics

---
 .../main/java/io/trino/event/QueryMonitor.java |  4 ++++
 .../spi/eventlistener/QueryStatistics.java     | 18 ++++++++++++++++++
 .../httpquery/TestHttpEventListener.java       |  2 ++
 .../execution/TestEventListenerWithSplits.java |  2 ++
 4 files changed, 26 insertions(+)

diff --git a/core/trino-main/src/main/java/io/trino/event/QueryMonitor.java b/core/trino-main/src/main/java/io/trino/event/QueryMonitor.java
index 24398c03883c..33649fbd7500 100644
--- a/core/trino-main/src/main/java/io/trino/event/QueryMonitor.java
+++ b/core/trino-main/src/main/java/io/trino/event/QueryMonitor.java
@@ -195,6 +195,8 @@ public void queryImmediateFailureEvent(BasicQueryInfo queryInfo, ExecutionFailur
                         0,
                         0,
                         0,
+                        0,
+                        0,
                         ImmutableList.of(),
                         0,
                         true,
@@ -273,6 +275,8 @@ private QueryStatistics createQueryStatistics(QueryInfo queryInfo)
                 queryStats.getPeakTaskTotalMemory().toBytes(),
                 queryStats.getPhysicalInputDataSize().toBytes(),
                 queryStats.getPhysicalInputPositions(),
+                queryStats.getProcessedInputDataSize().toBytes(),
+                queryStats.getProcessedInputPositions(),
                 queryStats.getInternalNetworkInputDataSize().toBytes(),
                 queryStats.getInternalNetworkInputPositions(),
                 queryStats.getRawInputDataSize().toBytes(),
diff --git a/core/trino-spi/src/main/java/io/trino/spi/eventlistener/QueryStatistics.java b/core/trino-spi/src/main/java/io/trino/spi/eventlistener/QueryStatistics.java
index 11ae32c181be..ff99f57ce456 100644
--- a/core/trino-spi/src/main/java/io/trino/spi/eventlistener/QueryStatistics.java
+++ b/core/trino-spi/src/main/java/io/trino/spi/eventlistener/QueryStatistics.java
@@ -43,6 +43,8 @@ public class QueryStatistics
     private final long peakTaskTotalMemory;
     private final long physicalInputBytes;
     private final long physicalInputRows;
+    private final long processedInputBytes;
+    private final long processedInputRows;
     private final long internalNetworkBytes;
     private final long internalNetworkRows;
     private final long totalBytes;
@@ -90,6 +92,8 @@ public QueryStatistics(
             long peakTaskTotalMemory,
             long physicalInputBytes,
             long physicalInputRows,
+            long processedInputBytes,
+            long processedInputRows,
             long internalNetworkBytes,
             long internalNetworkRows,
             long totalBytes,
@@ -122,6 +126,8 @@ public QueryStatistics(
         this.peakTaskTotalMemory = peakTaskTotalMemory;
         this.physicalInputBytes = physicalInputBytes;
         this.physicalInputRows = physicalInputRows;
+        this.processedInputBytes = processedInputBytes;
+        this.processedInputRows = processedInputRows;
         this.internalNetworkBytes = internalNetworkBytes;
         this.internalNetworkRows = internalNetworkRows;
         this.totalBytes = totalBytes;
@@ -230,6 +236,18 @@ public long getPhysicalInputRows()
         return physicalInputRows;
     }
 
+    @JsonProperty
+    public long getProcessedInputBytes()
+    {
+        return processedInputBytes;
+    }
+
+    @JsonProperty
+    public long getProcessedInputRows()
+    {
+        return processedInputRows;
+    }
+
     @JsonProperty
     public long getInternalNetworkBytes()
     {
diff --git a/plugin/trino-http-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpEventListener.java b/plugin/trino-http-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpEventListener.java
index b3493568ae2e..984b444e640b 100644
--- a/plugin/trino-http-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpEventListener.java
+++ b/plugin/trino-http-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpEventListener.java
@@ -164,6 +164,8 @@ public class TestHttpEventListener
                 0L,
                 0L,
                 0L,
+                0L,
+                0L,
                 0.0f,
                 Collections.emptyList(),
                 0,
diff --git a/testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerWithSplits.java b/testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerWithSplits.java
index 1a2dbce355a2..294b4f2346c2 100644
--- a/testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerWithSplits.java
+++ b/testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerWithSplits.java
@@ -164,6 +164,8 @@ public void testSplitsForNormalQuery()
         // Deterministic statistics
         assertEquals(statistics.getPhysicalInputBytes(), 0);
         assertEquals(statistics.getPhysicalInputRows(), expectedCompletedPositions);
+        assertEquals(statistics.getProcessedInputBytes(), 0);
+        assertEquals(statistics.getProcessedInputRows(), expectedCompletedPositions);
         assertEquals(statistics.getInternalNetworkBytes(), 381);
         assertEquals(statistics.getInternalNetworkRows(), 3);
         assertEquals(statistics.getTotalBytes(), 0);