From 28091e6e3dbaa653d8228c1db4795a49959197a3 Mon Sep 17 00:00:00 2001
From: guojunfei <970763131@qq.com>
Date: Wed, 20 Apr 2022 19:55:28 +0800
Subject: [PATCH] add API to get statistics from pageSource

---
 .../ScanFilterAndProjectOperator.java         |  7 ++++-
 .../prestosql/operator/TableScanOperator.java | 31 +++++++++++++++----
 .../spi/connector/ConnectorPageSource.java    | 11 +++++++
 3 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/presto-main/src/main/java/io/prestosql/operator/ScanFilterAndProjectOperator.java b/presto-main/src/main/java/io/prestosql/operator/ScanFilterAndProjectOperator.java
index 5046517cb..417b3cc6a 100644
--- a/presto-main/src/main/java/io/prestosql/operator/ScanFilterAndProjectOperator.java
+++ b/presto-main/src/main/java/io/prestosql/operator/ScanFilterAndProjectOperator.java
@@ -493,7 +493,12 @@ public class ScanFilterAndProjectOperator
             page = recordMaterializedBytes(page, sizeInBytes -> processedBytes += sizeInBytes);
 
             // update operator stats
-            processedPositions += page.getPositionCount();
+            if (pageSource.getCompletedPositionCount().isPresent()) {
+                processedPositions = pageSource.getCompletedPositionCount().getAsLong();
+            }
+            else {
+                processedPositions += page.getPositionCount();
+            }
             physicalBytes = pageSource.getCompletedBytes();
             readTimeNanos = pageSource.getReadTimeNanos();
 
diff --git a/presto-main/src/main/java/io/prestosql/operator/TableScanOperator.java b/presto-main/src/main/java/io/prestosql/operator/TableScanOperator.java
index ccf28b635..6dcd2654b 100644
--- a/presto-main/src/main/java/io/prestosql/operator/TableScanOperator.java
+++ b/presto-main/src/main/java/io/prestosql/operator/TableScanOperator.java
@@ -267,6 +267,8 @@ public class TableScanOperator
 
     private boolean finished;
 
+    // completedPositionCount is used only if connectorPageSource.getCompletedPositionCount is present.
+    private long completedPositionCount;
     private long completedBytes;
     private long readTimeNanos;
     Optional<TableScanNode> tableScanNodeOptional;
@@ -707,17 +709,34 @@ public class TableScanOperator
         }
 
         Page page = source.getNextPage();
+
+        // if pageSource.getCompletedPositionCount is present, get operator statistics from pageSource
+        if (source.getCompletedPositionCount().isPresent()) {
+            long endCompletedPositionCount = source.getCompletedPositionCount().getAsLong();
+            long endCompletedBytes = source.getCompletedBytes();
+            long endReadTimeNanos = source.getReadTimeNanos();
+            long currentPositionCount = endCompletedPositionCount - completedPositionCount;
+            long currentCompletedBytes = endCompletedBytes - completedBytes;
+            operatorContext.recordPhysicalInputWithTiming(currentCompletedBytes, currentPositionCount, endReadTimeNanos - readTimeNanos);
+            operatorContext.recordProcessedInput(currentCompletedBytes, currentPositionCount);
+            completedPositionCount = endCompletedPositionCount;
+            completedBytes = endCompletedBytes;
+            readTimeNanos = endReadTimeNanos;
+        }
+
         if (page != null) {
             // assure the page is in memory before handing to another operator
             page = page.getLoadedPage();
 
             // update operator stats
-            long endCompletedBytes = source.getCompletedBytes();
-            long endReadTimeNanos = source.getReadTimeNanos();
-            operatorContext.recordPhysicalInputWithTiming(endCompletedBytes - completedBytes, page.getPositionCount(), endReadTimeNanos - readTimeNanos);
-            operatorContext.recordProcessedInput(page.getSizeInBytes(), page.getPositionCount());
-            completedBytes = endCompletedBytes;
-            readTimeNanos = endReadTimeNanos;
+            if (!source.getCompletedPositionCount().isPresent()) {
+                long endCompletedBytes = source.getCompletedBytes();
+                long endReadTimeNanos = source.getReadTimeNanos();
+                operatorContext.recordPhysicalInputWithTiming(endCompletedBytes - completedBytes, page.getPositionCount(), endReadTimeNanos - readTimeNanos);
+                operatorContext.recordProcessedInput(page.getSizeInBytes(), page.getPositionCount());
+                completedBytes = endCompletedBytes;
+                readTimeNanos = endReadTimeNanos;
+            }
 
             // pull bloomFilter from stateStore and filter page
             if (existsCrossFilter) {
diff --git a/presto-spi/src/main/java/io/prestosql/spi/connector/ConnectorPageSource.java b/presto-spi/src/main/java/io/prestosql/spi/connector/ConnectorPageSource.java
index f0db11772..0049cc63f 100644
--- a/presto-spi/src/main/java/io/prestosql/spi/connector/ConnectorPageSource.java
+++ b/presto-spi/src/main/java/io/prestosql/spi/connector/ConnectorPageSource.java
@@ -17,6 +17,7 @@ import io.prestosql.spi.Page;
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.util.OptionalLong;
 import java.util.concurrent.CompletableFuture;
 
 public interface ConnectorPageSource
@@ -75,4 +76,14 @@ public interface ConnectorPageSource
     {
         return false;
     }
+
+    /**
+     * Some components can push down some operators to other component for calculation.
+     * This API is used to return the number of position count before push down for statistics.
+     * If this API is not used, just return empty.
+     */
+    default OptionalLong getCompletedPositionCount()
+    {
+        return OptionalLong.empty();
+    }
 }