[GitHub] [spark] somani commented on a change in pull request #32583: [SPARK-35437][SQL] Use expressions to filter Hive partitions at client side

[GitBox <gi...@apache.org>]

somani commented on a change in pull request #32583:
URL: https://github.com/apache/spark/pull/32583#discussion_r683466898



##########
File path: sql/hive/src/main/scala/org/apache/spark/sql/hive/client/HiveShim.scala
##########
@@ -900,13 +904,80 @@ private[client] class Shim_v0_13 extends Shim_v0_12 {
               s"${SQLConf.HIVE_METASTORE_PARTITION_PRUNING_FALLBACK_ON_EXCEPTION.key} " +
               " to false and let the query fail instead.", ex)
             // HiveShim clients are expected to handle a superset of the requested partitions
-            getAllPartitionsMethod.invoke(hive, table).asInstanceOf[JSet[Partition]]
+            prunePartitionsEvalClientSide(hive, table, catalogTable, predicates)
           case ex: InvocationTargetException if ex.getCause.isInstanceOf[MetaException] =>
             throw QueryExecutionErrors.getPartitionMetadataByFilterError(ex)
         }
       }
 
-    partitions.asScala.toSeq
+    partitions.toSeq
+  }
+
+  def prunePartitionsEvalClientSide(hive: Hive,
+      table: Table,
+      catalogTable: CatalogTable,
+      predicates: Seq[Expression]): Seq[Partition] = {
+
+    val timeZoneId = SQLConf.get.sessionLocalTimeZone
+
+    // Because there is no way to know whether the partition properties has timeZone,
+    // client-side filtering cannot be used with TimeZoneAwareExpression.
+    def hasTimeZoneAwareExpression(e: Expression): Boolean = {
+      e.collectFirst {
+        case t: TimeZoneAwareExpression => t
+      }.isDefined
+    }
+
+    if (!SQLConf.get.metastorePartitionPruningEvalClientSide ||
+      predicates.isEmpty ||
+      predicates.exists(hasTimeZoneAwareExpression)) {
+      getAllPartitionsMethod.invoke(hive, table).asInstanceOf[JSet[Partition]].asScala.toSeq
+    } else {
+      try {
+        val partitionSchema = CharVarcharUtils.replaceCharVarcharWithStringInSchema(
+          catalogTable.partitionSchema)
+        val partitionColumnNames = catalogTable.partitionColumnNames.toSet
+
+        val nonPartitionPruningPredicates = predicates.filterNot {
+          _.references.map(_.name).toSet.subsetOf(partitionColumnNames)
+        }
+        if (nonPartitionPruningPredicates.nonEmpty) {
+          throw QueryCompilationErrors.nonPartitionPruningPredicatesNotExpectedError(
+            nonPartitionPruningPredicates)
+        }
+
+        val boundPredicate =
+          Predicate.createInterpreted(predicates.reduce(And).transform {
+            case att: AttributeReference =>
+              val index = partitionSchema.indexWhere(_.name == att.name)
+              BoundReference(index, partitionSchema(index).dataType, nullable = true)
+          })
+
+        def toRow(spec: TablePartitionSpec): InternalRow = {
+          InternalRow.fromSeq(partitionSchema.map { field =>
+            val partValue = if (spec(field.name) == ExternalCatalogUtils.DEFAULT_PARTITION_NAME) {
+              null
+            } else {
+              spec(field.name)
+            }
+            Cast(Literal(partValue), field.dataType, Option(timeZoneId)).eval()
+          })
+        }
+
+        val allPartitionNames = hive.getPartitionNames(
+          table.getDbName, table.getTableName, -1).asScala
+        val partSpec = allPartitionNames.filter { p =>

Review comment:
       I think we should change the variable name to partNames, as there is another getPartitionsByNames() API which takes partSpec as an input, and is more expensive as that one does an extra hive metastore API call.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org