pan3793 commented on code in PR #58756:
URL: https://github.com/apache/spark/pull/58756#discussion_r4036688997
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedDeleteFilterSuite.scala:
##########
@@ -357,6 +386,35 @@ class DataSourceV2EnhancedDeleteFilterSuite extends
SharedSparkSession {
}
}
+ /**
+ * Asserts that the group-based plan's row-level scan was pruned by one
PartitionPredicate with
+ * the given references, and that only the partitions with the given `dep`
values, the first
+ * partition field, were replaced.
+ */
+ private def assertRowLevelScanPrunedByPartitionPredicate(
+ plan: SparkPlan,
+ expectedOrdinals: Array[Int],
+ expectedPartitionFieldNames: Array[String],
+ expectedReplacedDeps: Set[String]): Unit = {
+ assert(plan.isInstanceOf[ReplaceDataExec],
+ s"Expected ReplaceDataExec but got: ${plan.getClass.getSimpleName}")
+ val scans = collect(plan) { case s: BatchScanExec => s }
+ val scan = scans.map(_.scan).collectFirst {
+ case s:
InMemoryPartitionPredicateDeleteTable#PartitionPredicateRowLevelBatchScan => s
+ }.getOrElse(fail("Expected the row-level scan of the in-memory table"))
+ assertPartitionFieldReferences(
+ scan.pushedPartitionPredicates.toArray, Seq(expectedOrdinals),
expectedPartitionFieldNames)
+
+ val table = scans.map(_.table).collectFirst {
+ case RowLevelOperationTable(t: InMemoryPartitionPredicateDeleteTable, _)
=> t
+ }.getOrElse(fail("Expected the row-level operation table"))
+ val replacedDeps = table.replacedPartitions.map(_.head.toString)
+ assert(
+ replacedDeps.toSet === expectedReplacedDeps &&
Review Comment:
Done in 0abbdb2.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryPartitionPredicateDeleteTable.scala:
##########
@@ -107,6 +113,72 @@ class InMemoryPartitionPredicateDeleteTable(
}
}
+ /**
+ * Row-level scans push V2 predicates iteratively, so a group-based
operation receives a
+ * second-pass [[PartitionPredicate]] the same way a metadata-only DELETE
does. Only partition
+ * predicates prune, by partition key; a data predicate is always returned
since the scan
+ * cannot filter rows.
+ */
+ override protected def newRowLevelScanBuilder(
+ options: CaseInsensitiveStringMap)(
+ onBuild: BatchScanBaseClass => Unit): ScanBuilder = {
+ new PartitionPredicateRowLevelScanBuilder(onBuild)
+ }
+
+ class PartitionPredicateRowLevelScanBuilder(onBuild: BatchScanBaseClass =>
Unit)
+ extends ScanBuilder with SupportsPushDownV2Filters with
SupportsPushDownRequiredColumns {
+
+ private var readSchema: StructType = schema
+ private val pushed = ArrayBuffer.empty[Predicate]
+
+ override def supportsIterativePushdown(): Boolean = true
+
+ override def pushPredicates(predicates: Array[Predicate]):
Array[Predicate] = {
+ val (accepted, returned) = predicates.partition {
+ case _: PartitionPredicate => acceptPartitionPredicates
+ case p => refsOnlyPartCols(p) &&
InMemoryTableWithV2Filter.supportsPredicates(Array(p))
+ }
+ pushed ++= accepted
+ returned
+ }
+
+ override def pushedPredicates(): Array[Predicate] = pushed.toArray
+
+ override def pruneColumns(requiredSchema: StructType): Unit = {
+ val metadataNames = metadataColumns.map(_.name).toSet
+ val schemaNames = schema.map(_.name).toSet
+ readSchema = StructType(requiredSchema.filter {
+ case MetadataStructFieldWithLogicalName(_, name) =>
metadataNames.contains(name)
+ case f => schemaNames.contains(f.name)
+ })
+ }
+
+ override def build(): Scan = {
+ val (partPreds, stdPreds) =
pushed.toArray.partition(_.isInstanceOf[PartitionPredicate])
+ val partitionPredicates =
partPreds.map(_.asInstanceOf[PartitionPredicate])
+ val keys = InMemoryTableWithV2Filter.filtersToKeys(
+ data.map(_.key).toImmutableArraySeq,
Review Comment:
Done in 0abbdb2; `build()` reads it once.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
