[GLUTEN-11692][CORE] Driver side should avoid submitting subqueries while validating file formats

gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala

@@ -144,7 +144,7 @@ trait LeafTransformSupport extends TransformSupport with LeafExecNode {
   /** Returns the partitions generated by this data source scan and tied with ReadFileFormat. */
   def getPartitionWithReadFileFormats: Seq[(Partition, ReadFileFormat)]
 
-  /** Returns distinct partition's ReadFileFormats. */
+  /** Returns distinct ReadFileFormats of all partitions, used to perform validation. */
   def getDistinctPartitionReadFileFormats: Set[ReadFileFormat] = Set.empty
 }
 

gluten-substrait/src/main/scala/org/apache/spark/sql/hive/HiveTableScanExecTransformer.scala

@@ -77,7 +77,16 @@ case class HiveTableScanExecTransformer(
     partitionWithReadFileFormats
 
   override def getDistinctPartitionReadFileFormats: Set[ReadFileFormat] =
-    distinctReadFileFormats
+    if (
+      relation.isPartitioned &&
+      basePrunedPartitions.exists(_.getInputFormatClass != tableDesc.getInputFileFormatClass)
+    ) {
+      basePrunedPartitions.map {
+        partition => getReadFileFormat(HiveClientImpl.fromHivePartition(partition).storage)
+      }.toSet
+    } else {
+      Set(fileFormat)
+    }
 
   override def getPartitionSchema: StructType = relation.tableMeta.partitionSchema
 
@@ -120,9 +129,6 @@ case class HiveTableScanExecTransformer(
 
   @transient private lazy val partitions: Seq[Partition] = partitionWithReadFileFormats.unzip._1
 
-  @transient private lazy val distinctReadFileFormats: Set[ReadFileFormat] =
-    partitionWithReadFileFormats.iterator.map(_._2).toSet
-
   @transient override lazy val fileFormat: ReadFileFormat =
     getReadFileFormat(relation.tableMeta.storage)
 

shims/spark33/src/main/scala/org/apache/spark/sql/hive/execution/AbstractHiveTableScanExec.scala

@@ -173,26 +173,30 @@ abstract private[hive] class AbstractHiveTableScanExec(
     }
   }
 
-  @transient lazy val prunedPartitions: Seq[HivePartition] = {
+  // This is used on the driver side, so it is important to avoid executing subqueries
+  @transient lazy val basePrunedPartitions: Seq[HivePartition] = {
+    if (relation.prunedPartitions.nonEmpty) {
+      relation.prunedPartitions.get.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
+    } else {
+      rawPartitions
+    }
+  }
+
+  @transient lazy val prunedPartitions: Seq[HivePartition] =
     if (relation.prunedPartitions.nonEmpty) {
-      val hivePartitions =
-        relation.prunedPartitions.get.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
       if (partitionPruningPred.forall(!ExecSubqueryExpression.hasSubquery(_))) {
-        hivePartitions
+        basePrunedPartitions
       } else {
-        prunePartitions(hivePartitions)
+        prunePartitions(basePrunedPartitions)
       }
+    } else if (
+      sparkSession.sessionState.conf.metastorePartitionPruning &&
+      partitionPruningPred.nonEmpty
+    ) {
+      basePrunedPartitions
     } else {
-      if (
-        sparkSession.sessionState.conf.metastorePartitionPruning &&
-        partitionPruningPred.nonEmpty
-      ) {
-        rawPartitions
-      } else {
-        prunePartitions(rawPartitions)
-      }
+      prunePartitions(basePrunedPartitions)
     }
-  }
 
   // exposed for tests
   @transient lazy val rawPartitions: Seq[HivePartition] = {

shims/spark34/src/main/scala/org/apache/spark/sql/hive/execution/AbstractHiveTableScanExec.scala

@@ -175,26 +175,30 @@ abstract private[hive] class AbstractHiveTableScanExec(
     }
   }
 
-  @transient lazy val prunedPartitions: Seq[HivePartition] = {
+  // This is used on the driver side, so it is important to avoid executing subqueries
+  @transient lazy val basePrunedPartitions: Seq[HivePartition] = {
+    if (relation.prunedPartitions.nonEmpty) {
+      relation.prunedPartitions.get.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
+    } else {
+      rawPartitions
+    }
+  }
+
+  @transient lazy val prunedPartitions: Seq[HivePartition] =
     if (relation.prunedPartitions.nonEmpty) {
-      val hivePartitions =
-        relation.prunedPartitions.get.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
       if (partitionPruningPred.forall(!ExecSubqueryExpression.hasSubquery(_))) {
-        hivePartitions
+        basePrunedPartitions
       } else {
-        prunePartitions(hivePartitions)
+        prunePartitions(basePrunedPartitions)
       }
+    } else if (
+      sparkSession.sessionState.conf.metastorePartitionPruning &&
+      partitionPruningPred.nonEmpty
+    ) {
+      basePrunedPartitions
     } else {
-      if (
-        sparkSession.sessionState.conf.metastorePartitionPruning &&
-        partitionPruningPred.nonEmpty
-      ) {
-        rawPartitions
-      } else {
-        prunePartitions(rawPartitions)
-      }
+      prunePartitions(basePrunedPartitions)
     }
-  }
 
   // exposed for tests
   @transient lazy val rawPartitions: Seq[HivePartition] = {

shims/spark35/src/main/scala/org/apache/spark/sql/hive/execution/AbstractHiveTableScanExec.scala

@@ -175,26 +175,30 @@ abstract private[hive] class AbstractHiveTableScanExec(
     }
   }
 
-  @transient lazy val prunedPartitions: Seq[HivePartition] = {
+  // This is used on the driver side, so it is important to avoid executing subqueries
+  @transient lazy val basePrunedPartitions: Seq[HivePartition] = {
+    if (relation.prunedPartitions.nonEmpty) {
+      relation.prunedPartitions.get.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
+    } else {
+      rawPartitions
+    }
+  }
+
+  @transient lazy val prunedPartitions: Seq[HivePartition] =
     if (relation.prunedPartitions.nonEmpty) {
-      val hivePartitions =
-        relation.prunedPartitions.get.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
       if (partitionPruningPred.forall(!ExecSubqueryExpression.hasSubquery(_))) {
-        hivePartitions
+        basePrunedPartitions
       } else {
-        prunePartitions(hivePartitions)
+        prunePartitions(basePrunedPartitions)
       }
+    } else if (
+      sparkSession.sessionState.conf.metastorePartitionPruning &&
+      partitionPruningPred.nonEmpty
+    ) {
+      basePrunedPartitions
     } else {
-      if (
-        sparkSession.sessionState.conf.metastorePartitionPruning &&
-        partitionPruningPred.nonEmpty
-      ) {
-        rawPartitions
-      } else {
-        prunePartitions(rawPartitions)
-      }
+      prunePartitions(basePrunedPartitions)
     }
-  }
 
   // exposed for tests
   @transient lazy val rawPartitions: Seq[HivePartition] = {

shims/spark40/src/main/scala/org/apache/spark/sql/hive/execution/AbstractHiveTableScanExec.scala

@@ -175,26 +175,30 @@ abstract private[hive] class AbstractHiveTableScanExec(
     }
   }
 
-  @transient lazy val prunedPartitions: Seq[HivePartition] = {
+  // This is used on the driver side, so it is important to avoid executing subqueries
+  @transient lazy val basePrunedPartitions: Seq[HivePartition] = {
+    if (relation.prunedPartitions.nonEmpty) {
+      relation.prunedPartitions.get.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
+    } else {
+      rawPartitions
+    }
+  }
+
+  @transient lazy val prunedPartitions: Seq[HivePartition] =
     if (relation.prunedPartitions.nonEmpty) {
-      val hivePartitions =
-        relation.prunedPartitions.get.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
       if (partitionPruningPred.forall(!ExecSubqueryExpression.hasSubquery(_))) {
-        hivePartitions
+        basePrunedPartitions
       } else {
-        prunePartitions(hivePartitions)
+        prunePartitions(basePrunedPartitions)
       }
+    } else if (
+      sparkSession.sessionState.conf.metastorePartitionPruning &&
+      partitionPruningPred.nonEmpty
+    ) {
+      basePrunedPartitions
     } else {
-      if (
-        sparkSession.sessionState.conf.metastorePartitionPruning &&
-        partitionPruningPred.nonEmpty
-      ) {
-        rawPartitions
-      } else {
-        prunePartitions(rawPartitions)
-      }
+      prunePartitions(basePrunedPartitions)
     }
-  }
 
   // exposed for tests
   @transient lazy val rawPartitions: Seq[HivePartition] = {

shims/spark41/src/main/scala/org/apache/spark/sql/hive/execution/AbstractHiveTableScanExec.scala

@@ -175,26 +175,30 @@ abstract private[hive] class AbstractHiveTableScanExec(
     }
   }
 
-  @transient lazy val prunedPartitions: Seq[HivePartition] = {
+  // This is used on the driver side, so it is important to avoid executing subqueries
+  @transient lazy val basePrunedPartitions: Seq[HivePartition] = {
+    if (relation.prunedPartitions.nonEmpty) {
+      relation.prunedPartitions.get.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
+    } else {
+      rawPartitions
+    }
+  }
+
+  @transient lazy val prunedPartitions: Seq[HivePartition] =
     if (relation.prunedPartitions.nonEmpty) {
-      val hivePartitions =
-        relation.prunedPartitions.get.map(HiveClientImpl.toHivePartition(_, hiveQlTable))
       if (partitionPruningPred.forall(!ExecSubqueryExpression.hasSubquery(_))) {
-        hivePartitions
+        basePrunedPartitions
       } else {
-        prunePartitions(hivePartitions)
+        prunePartitions(basePrunedPartitions)
       }
+    } else if (
+      sparkSession.sessionState.conf.metastorePartitionPruning &&
+      partitionPruningPred.nonEmpty
+    ) {
+      basePrunedPartitions
     } else {
-      if (
-        sparkSession.sessionState.conf.metastorePartitionPruning &&
-        partitionPruningPred.nonEmpty
-      ) {
-        rawPartitions
-      } else {
-        prunePartitions(rawPartitions)
-      }
+      prunePartitions(basePrunedPartitions)
     }
-  }
 
   // exposed for tests
   @transient lazy val rawPartitions: Seq[HivePartition] = {