[spark] Avoid unnecessary get splits in spark scan (apache#3645)

tongcheng-elong · Jul 1, 2024 · 1fde4a4 · 1fde4a4
1 parent 37fb8f5
commit 1fde4a4
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonBaseScan.scala
@@ -19,6 +19,7 @@
 package org.apache.paimon.spark
 
 import org.apache.paimon.{stats, CoreOptions}
+import org.apache.paimon.annotation.VisibleForTesting
 import org.apache.paimon.predicate.{Predicate, PredicateBuilder}
 import org.apache.paimon.spark.metric.SparkMetricRegistry
 import org.apache.paimon.spark.schema.PaimonMetadataColumn
@@ -92,6 +93,7 @@ abstract class PaimonBaseScan(
     _readBuilder
   }
 
+  @VisibleForTesting
   def getOriginSplits: Array[Split] = {
     readBuilder
       .newScan()

diff --git a/...ark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonInputPartition.scala b/...ark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonInputPartition.scala
@@ -22,7 +22,11 @@ import org.apache.paimon.table.source.Split
 
 import org.apache.spark.sql.connector.read.InputPartition
 
-case class PaimonInputPartition(splits: Seq[Split]) extends InputPartition {}
+case class PaimonInputPartition(splits: Seq[Split]) extends InputPartition {
+  def rowCount(): Long = {
+    splits.map(_.rowCount()).sum
+  }
+}
 
 object PaimonInputPartition {
   def apply(split: Split): PaimonInputPartition = {

diff --git a/...n-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonStatistics.scala b/...n-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonStatistics.scala
@@ -34,7 +34,7 @@ import scala.collection.JavaConverters._
 
 case class PaimonStatistics[T <: PaimonBaseScan](scan: T) extends Statistics {
 
-  private lazy val rowCount: Long = scan.getOriginSplits.map(_.rowCount).sum
+  private lazy val rowCount: Long = scan.getInputPartitions.map(_.rowCount()).sum
 
   private lazy val scannedTotalSize: Long = rowCount * scan.readSchema().defaultSize