Fix dynamic pruning regression in GpuFileSourceScanExec (#11191)

Signed-off-by: Jason Lowe <[email protected]>
NVIDIA · Jul 16, 2024 · 2b09372 · 2b09372
1 parent a41a616
commit 2b09372
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/GpuFileSourceScanExec.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/GpuFileSourceScanExec.scala
@@ -558,11 +558,12 @@ case class GpuFileSourceScanExec(
     val partitions = StaticPartitionShims.getStaticPartitions(fsRelation).getOrElse {
       val openCostInBytes = fsRelation.sparkSession.sessionState.conf.filesOpenCostInBytes
       val maxSplitBytes =
-        FilePartition.maxSplitBytes(fsRelation.sparkSession, selectedPartitions)
+        FilePartition.maxSplitBytes(fsRelation.sparkSession, dynamicallySelectedPartitions)
       logInfo(s"Planning scan with bin packing, max size: $maxSplitBytes bytes, " +
         s"open cost is considered as scanning $openCostInBytes bytes.")
 
-      val splitFiles = FilePartitionShims.splitFiles(selectedPartitions, relation, maxSplitBytes)
+      val splitFiles = FilePartitionShims.splitFiles(dynamicallySelectedPartitions, relation,
+        maxSplitBytes)
 
       FilePartition.getFilePartitions(relation.sparkSession, splitFiles, maxSplitBytes)
     }