Spark: Add compute stats to scan builder also (apache#5136)

InvisibleProgrammer · Jun 27, 2022 · 313f497 · 313f497
1 parent 35b8558
commit 313f497
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 3 deletions.
diff --git a/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java b/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java
@@ -39,14 +39,18 @@
 import org.apache.spark.sql.SparkSession;
 import org.apache.spark.sql.connector.read.Scan;
 import org.apache.spark.sql.connector.read.ScanBuilder;
+import org.apache.spark.sql.connector.read.Statistics;
 import org.apache.spark.sql.connector.read.SupportsPushDownFilters;
 import org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns;
+import org.apache.spark.sql.connector.read.SupportsReportStatistics;
 import org.apache.spark.sql.sources.Filter;
 import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
 import org.apache.spark.sql.util.CaseInsensitiveStringMap;
 
-public class SparkScanBuilder implements ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns {
+public class SparkScanBuilder implements ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns,
+    SupportsReportStatistics {
+
   private static final Filter[] NO_FILTERS = new Filter[0];
 
   private final SparkSession spark;
@@ -167,4 +171,14 @@ public Scan buildMergeScan() {
         spark, table, readConf, caseSensitive, ignoreResiduals,
         schemaWithMetadataColumns(), filterExpressions, options);
   }
+
+  @Override
+  public Statistics estimateStatistics() {
+    return ((SparkBatchScan) build()).estimateStatistics();
+  }
+
+  @Override
+  public StructType readSchema() {
+    return build().readSchema();
+  }
 }
diff --git a/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java b/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java
@@ -39,14 +39,18 @@
 import org.apache.spark.sql.SparkSession;
 import org.apache.spark.sql.connector.read.Scan;
 import org.apache.spark.sql.connector.read.ScanBuilder;
+import org.apache.spark.sql.connector.read.Statistics;
 import org.apache.spark.sql.connector.read.SupportsPushDownFilters;
 import org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns;
+import org.apache.spark.sql.connector.read.SupportsReportStatistics;
 import org.apache.spark.sql.sources.Filter;
 import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
 import org.apache.spark.sql.util.CaseInsensitiveStringMap;
 
-public class SparkScanBuilder implements ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns {
+public class SparkScanBuilder implements ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns,
+    SupportsReportStatistics {
+
   private static final Filter[] NO_FILTERS = new Filter[0];
 
   private final SparkSession spark;
@@ -167,4 +171,14 @@ public Scan buildMergeScan() {
         spark, table, readConf, caseSensitive, ignoreResiduals,
         schemaWithMetadataColumns(), filterExpressions, options);
   }
+
+  @Override
+  public Statistics estimateStatistics() {
+    return ((SparkBatchScan) build()).estimateStatistics();
+  }
+
+  @Override
+  public StructType readSchema() {
+    return build().readSchema();
+  }
 }
diff --git a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java
@@ -44,14 +44,18 @@
 import org.apache.spark.sql.SparkSession;
 import org.apache.spark.sql.connector.read.Scan;
 import org.apache.spark.sql.connector.read.ScanBuilder;
+import org.apache.spark.sql.connector.read.Statistics;
 import org.apache.spark.sql.connector.read.SupportsPushDownFilters;
 import org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns;
+import org.apache.spark.sql.connector.read.SupportsReportStatistics;
 import org.apache.spark.sql.sources.Filter;
 import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
 import org.apache.spark.sql.util.CaseInsensitiveStringMap;
 
-public class SparkScanBuilder implements ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns {
+public class SparkScanBuilder implements ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns,
+    SupportsReportStatistics {
+
   private static final Filter[] NO_FILTERS = new Filter[0];
 
   private final SparkSession spark;
@@ -278,4 +282,14 @@ private TableScan configureSplitPlanning(TableScan scan) {
 
     return configuredScan;
   }
+
+  @Override
+  public Statistics estimateStatistics() {
+    return ((SparkScan) build()).estimateStatistics();
+  }
+
+  @Override
+  public StructType readSchema() {
+    return build().readSchema();
+  }
 }