[SPARK-58315][SQL] Refactor V2 scan statistics handling

sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsReportStatistics.java

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.connector.read;
 
+import java.util.OptionalLong;
+
 import org.apache.spark.annotation.Evolving;
 
 /**
@@ -36,4 +38,36 @@ public interface SupportsReportStatistics extends Scan {
    * Returns the estimated statistics of this data source scan.
    */
   Statistics estimateStatistics();
+
+  /**
+   * Returns the estimated size in bytes of this scan without computing full statistics.
+   * <p>
+   * When cost-based optimization or plan statistics are disabled, Spark primarily needs the scan's
+   * size in bytes (for example, for broadcast-join thresholding). This method lets connectors serve
+   * that size estimate cheaply and avoid computing the full statistics. The default implementation
+   * delegates to {@link #estimateStatistics()} and returns its {@code sizeInBytes()}, so connectors
+   * that already compute statistics cheaply do not need to override this method.
+   *
+   * @since 4.3.0
+   */
+  default OptionalLong estimateSizeInBytes() {
+    Statistics statistics = estimateStatistics();
+    return statistics != null ? statistics.sizeInBytes() : OptionalLong.empty();
+  }
+
+  /**
+   * Returns whether the statistics reported by this scan already reflect all filters that were
+   * fully pushed down to the data source.
+   * <p>
+   * When {@code true} (the default), the reported statistics describe exactly the data the scan
+   * will produce. When {@code false}, they do <em>not</em> account for the fully pushed filters
+   * (for example, they describe the whole table), so Spark may use those fully pushed filters to
+   * adjust stats. Re-applying those fully pushed filters in Spark should be redundant for query
+   * results because the data source already evaluates them.
+   *
+   * @since 4.3.0
+   */
+  default boolean reflectsFullyPushedDownFilters() {
+    return true;
+  }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala

@@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatisti
 import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, Expression, SortOrder, V2ExpressionUtils}
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, ExposesMetadataColumns, Histogram, HistogramBin, LeafNode, LogicalPlan, Statistics}
+import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils
 import org.apache.spark.sql.catalyst.streaming.{StreamingSourceIdentifyingName, Unassigned}
 import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
 import org.apache.spark.sql.catalyst.util.{removeInternalMetadata, truncatedString, CharVarcharUtils}
@@ -34,6 +35,7 @@ import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReferenc
 import org.apache.spark.sql.connector.read.{Scan, Statistics => V2Statistics, SupportsReportStatistics, SupportsRuntimeV2Filtering}
 import org.apache.spark.sql.connector.read.colstats.{ColumnStatistics, Histogram => V2Histogram, HistogramBin => V2HistogramBin}
 import org.apache.spark.sql.connector.read.streaming.{Offset, SparkDataStream}
+import org.apache.spark.sql.internal.connector.V2StatisticsUtils
 import org.apache.spark.sql.types.{DataType, StructType}
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 import org.apache.spark.util.ArrayImplicits._
@@ -95,7 +97,7 @@ abstract class DataSourceV2RelationBase(
       table.asReadable.newScanBuilder(options).build() match {
         case r: SupportsReportStatistics =>
           val statistics = r.estimateStatistics()
-          DataSourceV2Relation.transformV2Stats(statistics, None, conf.defaultSizeInBytes, output)
+          DataSourceV2Relation.transformV2Stats(statistics, conf.defaultSizeInBytes, output)
         case _ =>
           Statistics(sizeInBytes = conf.defaultSizeInBytes)
       }
@@ -199,15 +201,32 @@ case class DataSourceV2ScanRelation(
   }
 
   override def computeStats(): Statistics = {
-    scan match {
-      case r: SupportsReportStatistics =>
-        val statistics = r.estimateStatistics()
-        DataSourceV2Relation.transformV2Stats(statistics, None, conf.defaultSizeInBytes, output)
-      case _ =>
-        Statistics(sizeInBytes = conf.defaultSizeInBytes)
+    if (conf.cboEnabled || conf.planStatsEnabled) {
+      computeFullStats()
+    } else {
+      computeSizeOnlyStats()
+    }
+  }
+
+  private def computeFullStats(): Statistics = {
+    V2StatisticsUtils.computeStats(scan) match {
+      case Some(v2Stats) =>
+        DataSourceV2Relation.transformV2Stats(v2Stats, conf.defaultSizeInBytes, output)
+      case _ => defaultSizeOnlyStats
+    }
+  }
+
+  private def computeSizeOnlyStats(): Statistics = {
+    V2StatisticsUtils.computeSizeInBytes(scan, EstimationUtils.getSizePerRow(output)) match {
+      case Some(sizeInBytes) => Statistics(sizeInBytes = sizeInBytes)
+      case _ => defaultSizeOnlyStats
     }
   }
 
+  private def defaultSizeOnlyStats: Statistics = {
+    Statistics(sizeInBytes = conf.defaultSizeInBytes)
+  }
+
   override def doCanonicalize(): DataSourceV2ScanRelation = {
     this.copy(
       relation = this.relation.copy(
@@ -276,7 +295,7 @@ case class StreamingDataSourceV2ScanRelation(
   override def computeStats(): Statistics = scan match {
     case r: SupportsReportStatistics =>
       val statistics = r.estimateStatistics()
-      DataSourceV2Relation.transformV2Stats(statistics, None, conf.defaultSizeInBytes, output)
+      DataSourceV2Relation.transformV2Stats(statistics, conf.defaultSizeInBytes, output)
     case _ =>
       Statistics(sizeInBytes = conf.defaultSizeInBytes)
   }
@@ -419,13 +438,12 @@ object DataSourceV2Relation {
    */
   def transformV2Stats(
       v2Statistics: V2Statistics,
-      defaultRowCount: Option[BigInt],
       defaultSizeInBytes: Long,
       output: Seq[Attribute] = Seq.empty): Statistics = {
     val numRows: Option[BigInt] = if (v2Statistics.numRows().isPresent) {
       Some(v2Statistics.numRows().getAsLong)
     } else {
-      defaultRowCount
+      None
     }
 
     var colStats: Seq[(Attribute, ColumnStat)] = Seq.empty[(Attribute, ColumnStat)]
@@ -463,9 +481,20 @@ object DataSourceV2Relation {
         })
       })
     }
+    val attributeStats = AttributeMap(colStats)
+    // Prefer the source-reported size. Otherwise infer a projection-aware size from the row count
+    // (numRows * outputRowSize via getOutputSize). Fall back to the default size when neither is
+    // available.
+    val sizeInBytes = if (v2Statistics.sizeInBytes().isPresent) {
+      BigInt(v2Statistics.sizeInBytes().getAsLong)
+    } else if (numRows.isDefined) {
+      EstimationUtils.getOutputSize(output, numRows.get, attributeStats)
+    } else {
+      BigInt(defaultSizeInBytes)
+    }
     Statistics(
-      sizeInBytes = v2Statistics.sizeInBytes().orElse(defaultSizeInBytes),
+      sizeInBytes = sizeInBytes,
       rowCount = numRows,
-      attributeStats = AttributeMap(colStats))
+      attributeStats = attributeStats)
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/V2StatisticsUtils.scala

@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.internal.connector
+
+import java.util.OptionalLong
+
+import org.apache.spark.sql.connector.read.{Scan, Statistics, SupportsReportStatistics}
+
+object V2StatisticsUtils {
+
+  def isNotEmpty(stats: Statistics): Boolean = {
+    stats != null && hasAnyValue(stats)
+  }
+
+  private def hasAnyValue(stats: Statistics): Boolean = {
+    stats.sizeInBytes().isPresent ||
+      stats.numRows().isPresent ||
+      (stats.columnStats() != null && !stats.columnStats().isEmpty)
+  }
+
+  def computeStats(scan: Scan): Option[Statistics] = scan match {
+    case s: SupportsReportStatistics => Some(s.estimateStatistics()).filter(isNotEmpty)
+    case _ => None
+  }
+
+  def computeSizeInBytes(
+      scan: Scan,
+      avgRowSize: => BigInt): Option[BigInt] = {
+    extractSizeInBytes(scan).orElse(extractRowCount(scan).map(_ * avgRowSize))
+  }
+
+  private def extractSizeInBytes(scan: Scan): Option[BigInt] = scan match {
+    case s: SupportsReportStatistics => toBigInt(s.estimateSizeInBytes())
+    case _ => None
+  }
+
+  private def extractRowCount(scan: Scan): Option[BigInt] = scan match {
+    case s: SupportsReportStatistics =>
+      Option(s.estimateStatistics()).flatMap(stats => toBigInt(stats.numRows()))
+    case _ =>
+      None
+  }
+
+  private def toBigInt(value: OptionalLong): Option[BigInt] = {
+    if (value.isPresent) Some(BigInt(value.getAsLong)) else None
+  }
+}