[SPARK-32284][SQL] Avoid expanding too many CNF predicates in partition pruning

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/predicates.scala

@@ -267,20 +267,13 @@ trait PredicateHelper extends Logging {
 
   /**
    * Convert an expression to conjunctive normal form for predicate pushdown and partition pruning.
-   * When expanding predicates, this method groups expressions by their references for reducing
-   * the size of pushed down predicates and corresponding codegen. In partition pruning strategies,
-   * we split filters by [[splitConjunctivePredicates]] and partition filters by judging if it's
-   * references is subset of partCols, if we combine expressions group by reference when expand
-   * predicate of [[Or]], it won't impact final predicate pruning result since
-   * [[splitConjunctivePredicates]] won't split [[Or]] expression.
    *
    * @param condition condition need to be converted
    * @return the CNF result as sequence of disjunctive expressions. If the number of expressions
    *         exceeds threshold on converting `Or`, `Seq.empty` is returned.
    */
-  def CNFWithGroupExpressionsByReference(condition: Expression): Seq[Expression] = {
-    conjunctiveNormalForm(condition, (expressions: Seq[Expression]) =>
-        expressions.groupBy(e => AttributeSet(e.references)).map(_._2.reduceLeft(And)).toSeq)
+  def CNFConversion(condition: Expression): Seq[Expression] = {
+    conjunctiveNormalForm(condition, (expressions: Seq[Expression]) => expressions)
   }
 
   /**

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PruneFileSourcePartitions.scala

@@ -53,11 +53,17 @@ private[sql] object PruneFileSourcePartitions
     val partitionColumns =
       relation.resolve(partitionSchema, sparkSession.sessionState.analyzer.resolver)
     val partitionSet = AttributeSet(partitionColumns)
-    val (partitionFilters, dataFilters) = normalizedFilters.partition(f =>
+    val (partitionFilters, remainingFilters) = normalizedFilters.partition(f =>
       f.references.subsetOf(partitionSet)
     )
 
-    (ExpressionSet(partitionFilters), dataFilters)
+    // Try extracting more convertible partition filters from the remaining filters by converting
+    // them into CNF.
+    val remainingFilterInCnf = remainingFilters.flatMap(CNFConversion)
+    val extraPartitionFilters =
+      remainingFilterInCnf.filter(f => f.references.subsetOf(partitionSet))
+
+    (ExpressionSet(partitionFilters ++ extraPartitionFilters), remainingFilters)
   }
 
   private def rebuildPhysicalOperation(
@@ -88,12 +94,9 @@ private[sql] object PruneFileSourcePartitions
             _,
             _))
         if filters.nonEmpty && fsRelation.partitionSchemaOption.isDefined =>
-      val predicates = CNFWithGroupExpressionsByReference(filters.reduceLeft(And))
-      val finalPredicates = if (predicates.nonEmpty) predicates else filters
       val (partitionKeyFilters, _) = getPartitionKeyFiltersAndDataFilters(
-        fsRelation.sparkSession, logicalRelation, partitionSchema, finalPredicates,
+        fsRelation.sparkSession, logicalRelation, partitionSchema, filters,
         logicalRelation.output)
-
       if (partitionKeyFilters.nonEmpty) {
         val prunedFileIndex = catalogFileIndex.filterPartitions(partitionKeyFilters.toSeq)
         val prunedFsRelation =

sql/hive/src/main/scala/org/apache/spark/sql/hive/execution/PruneHiveTablePartitions.scala

@@ -27,6 +27,7 @@ import org.apache.spark.sql.catalyst.planning.PhysicalOperation
 import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, Project}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.execution.datasources.DataSourceStrategy
+import org.apache.spark.sql.execution.datasources.PruneFileSourcePartitions.CNFConversion
 import org.apache.spark.sql.internal.SQLConf
 
 /**
@@ -54,9 +55,15 @@ private[sql] class PruneHiveTablePartitions(session: SparkSession)
     val normalizedFilters = DataSourceStrategy.normalizeExprs(
       filters.filter(f => f.deterministic && !SubqueryExpression.hasSubquery(f)), relation.output)
     val partitionColumnSet = AttributeSet(relation.partitionCols)
-    ExpressionSet(normalizedFilters.filter { f =>
+    val (partitionFilters, remainingFilters) = normalizedFilters.partition { f =>
       !f.references.isEmpty && f.references.subsetOf(partitionColumnSet)
-    })
+    }
+    // Try extracting more convertible partition filters from the remaining filters by converting
+    // them into CNF.
+    val remainingFilterInCnf = remainingFilters.flatMap(CNFConversion)
+    val extraPartitionFilters = remainingFilterInCnf.filter(f =>
+      !f.references.isEmpty && f.references.subsetOf(partitionColumnSet))
+    ExpressionSet(partitionFilters ++ extraPartitionFilters)
   }
 
   /**
@@ -103,7 +110,7 @@ private[sql] class PruneHiveTablePartitions(session: SparkSession)
   override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
     case op @ PhysicalOperation(projections, filters, relation: HiveTableRelation)
       if filters.nonEmpty && relation.isPartitioned && relation.prunedPartitions.isEmpty =>
-      val predicates = CNFWithGroupExpressionsByReference(filters.reduceLeft(And))
+      val predicates = CNFConversion(filters.reduceLeft(And))
       val finalPredicates = if (predicates.nonEmpty) predicates else filters
       val partitionKeyFilters = getPartitionKeyFilters(finalPredicates, relation)
       if (partitionKeyFilters.nonEmpty) {

sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/PrunePartitionSuiteBase.scala

@@ -67,6 +67,31 @@ abstract class PrunePartitionSuiteBase extends QueryTest with SQLTestUtils with
     }
   }
 
+  test("SPARK-32284: Avoid pushing down too many predicates in partition pruning") {
+    withTempView("temp") {
+      withTable("t") {
+        sql(
+          s"""
+             |CREATE TABLE t(i INT, p0 INT, p1 INT)
+             |USING $format
+             |PARTITIONED BY (p0, p1)""".stripMargin)
+
+        spark.range(0, 10, 1).selectExpr("id as col")
+          .createOrReplaceTempView("temp")
+
+        for (part <- (0 to 25)) {
+          sql(
+            s"""
+               |INSERT OVERWRITE TABLE t PARTITION (p0='$part', p1='$part')
+               |SELECT col FROM temp""".stripMargin)
+        }
+        val scale = 20
+        val predicate = (1 to scale).map(i => s"(p0 = '$i' AND p1 = '$i')").mkString(" OR ")
+        assertPrunedPartitions(s"SELECT * FROM t WHERE $predicate", scale)
+      }
+    }
+  }
+
   protected def assertPrunedPartitions(query: String, expected: Long): Unit = {
     val plan = sql(query).queryExecution.sparkPlan
     assert(getScanExecPartitionSize(plan) == expected)