diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala index f21f53a28300d..168f30623ee4c 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala @@ -295,7 +295,12 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper // build a plan to write the row delta to the table val writeRelation = relation.copy(table = operationTable) val projections = buildWriteDeltaProjections(mergeRowsPlan, rowAttrs, rowIdAttrs, metadataAttrs) - WriteDelta(writeRelation, cond, mergeRowsPlan, relation, projections) + val groupFilterCond = if (notMatchedBySourceActions.isEmpty && groupFilterEnabled) { + Some(toGroupFilterCondition(relation, source, cond)) + } else { + None + } + WriteDelta(writeRelation, cond, mergeRowsPlan, relation, projections, groupFilterCond) } private def chooseWriteDeltaJoinType( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala index 3c41b6bfa5683..f235374bd5d6f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala @@ -174,7 +174,8 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { // build a plan to write the row delta to the table val writeRelation = relation.copy(table = operationTable) val projections = buildWriteDeltaProjections(rowDeltaPlan, rowAttrs, rowIdAttrs, metadataAttrs) - WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections) + val groupFilterCond = if (groupFilterEnabled) Some(cond) else None + WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond) } // this method assumes the assignments have been already aligned before diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceNullWithFalseInPredicate.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceNullWithFalseInPredicate.scala index 9a676571d1071..3d14bffb2fb92 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceNullWithFalseInPredicate.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceNullWithFalseInPredicate.scala @@ -60,7 +60,10 @@ object ReplaceNullWithFalseInPredicate extends Rule[LogicalPlan] { val newCond = replaceNullWithFalse(cond) val newGroupFilterCond = groupFilterCond.map(replaceNullWithFalse) rd.copy(condition = newCond, groupFilterCondition = newGroupFilterCond) - case wd @ WriteDelta(_, cond, _, _, _, _) => wd.copy(condition = replaceNullWithFalse(cond)) + case wd @ WriteDelta(_, cond, _, _, _, groupFilterCond, _) => + val newCond = replaceNullWithFalse(cond) + val newGroupFilterCond = groupFilterCond.map(replaceNullWithFalse) + wd.copy(condition = newCond, groupFilterCondition = newGroupFilterCond) case d @ DeleteFromTable(_, cond) => d.copy(condition = replaceNullWithFalse(cond)) case u @ UpdateTable(_, _, Some(cond)) => u.copy(condition = Some(replaceNullWithFalse(cond))) case m: MergeIntoTable => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala index b95c4624b8c50..c4af18fc8739a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala @@ -432,7 +432,7 @@ object ExtractSingleColumnNullAwareAntiJoin extends JoinSelectionHelper with Pre * - the read relation that can be either [[DataSourceV2Relation]] or [[DataSourceV2ScanRelation]] * depending on whether the planning has already happened; */ -object GroupBasedRowLevelOperation { +object GroupBasedRowLevelOperation extends RowLevelOperationExtractor { type ReturnType = (ReplaceData, Expression, Option[Expression], LogicalPlan) def unapply(plan: LogicalPlan): Option[ReturnType] = plan match { @@ -445,8 +445,34 @@ object GroupBasedRowLevelOperation { case _ => None } +} + +/** + * An extractor for row-level commands such as DELETE, UPDATE, MERGE that were rewritten using plans + * that operate on individual rows (row deltas). + * + * This class extracts the following entities: + * - the delta-based rewrite plan; + * - the condition that defines matching rows; + * - the group filter condition; + * - the read relation that can be either [[DataSourceV2Relation]] or [[DataSourceV2ScanRelation]] + * depending on whether the planning has already happened; + */ +object DeltaBasedRowLevelOperation extends RowLevelOperationExtractor { + type ReturnType = (WriteDelta, Expression, Option[Expression], LogicalPlan) + + def unapply(plan: LogicalPlan): Option[ReturnType] = plan match { + case wd @ WriteDelta(ExtractV2Table(table), cond, query, _, _, groupFilterCond, _) => + val readRelation = findReadRelation(table, query, allowMultipleReads = false) + readRelation.map((wd, cond, groupFilterCond, _)) + + case _ => + None + } +} - private def findReadRelation( +trait RowLevelOperationExtractor { + protected def findReadRelation( table: Table, plan: LogicalPlan, allowMultipleReads: Boolean): Option[LogicalPlan] = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala index 0eded2d9dbdf9..fa1ed6f0df6d8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala @@ -425,6 +425,7 @@ case class ReplaceData( * @param query a query with a delta of records that should written * @param originalTable a plan for the original table for which the row-level command was triggered * @param projections projections for row ID, row, metadata attributes + * @param groupFilterCondition a condition that can be used to filter groups at runtime * @param write a logical write, if already constructed */ case class WriteDelta( @@ -433,6 +434,7 @@ case class WriteDelta( query: LogicalPlan, originalTable: NamedRelation, projections: WriteDeltaProjections, + groupFilterCondition: Option[Expression] = None, write: Option[DeltaWrite] = None) extends RowLevelWrite { override val isByName: Boolean = false diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 99a5339950a58..bd0be0f75fe91 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -746,15 +746,15 @@ object SQLConf { val RUNTIME_ROW_LEVEL_OPERATION_GROUP_FILTER_ENABLED = buildConf("spark.sql.optimizer.runtime.rowLevelOperationGroupFilter.enabled") - .doc("Enables runtime group filtering for group-based row-level operations. " + - "Data sources that replace groups of data (e.g. files, partitions) may prune entire " + - "groups using provided data source filters when planning a row-level operation scan. " + - "However, such filtering is limited as not all expressions can be converted into data " + - "source filters and some expressions can only be evaluated by Spark (e.g. subqueries). " + - "Since rewriting groups is expensive, Spark can execute a query at runtime to find what " + - "records match the condition of the row-level operation. The information about matching " + - "records will be passed back to the row-level operation scan, allowing data sources to " + - "discard groups that don't have to be rewritten.") + .doc("Enables runtime filtering for group-based and delta-based row-level operations. " + + "Data sources may prune entire file groups at runtime when planning a row-level " + + "operation scan. Planning-time filter pushdown is limited as not all expressions can " + + "be converted into data source filters and some expressions can only be evaluated by " + + "Spark (e.g. subqueries). Since rewriting groups or scanning unnecessary files is " + + "expensive, Spark can execute a lightweight query at runtime to find what records match " + + "the condition of the row-level operation. The information about matching records will " + + "be passed back to the row-level operation scan, allowing data sources to skip files " + + "that don't have to be processed.") .version("3.4.0") .booleanConf .createWithDefault(true) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala index e03928867e24d..a5bc6ad4bdcf5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala @@ -491,7 +491,7 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat r.name) :: Nil case wd @ WriteDelta(_: DataSourceV2Relation, _, query, r: DataSourceV2Relation, projections, - Some(write)) => + _, Some(write)) => WriteDeltaExec( planLater(query), refreshCache(r), // use the original relation to refresh the cache diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala index 29bbfacaea004..c02fdec1fab60 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala @@ -114,7 +114,7 @@ object OptimizeMetadataOnlyDeleteFromTable extends Rule[LogicalPlan] with Predic val command = rd.operation.command Some(rd, command, cond, originalTable) - case wd @ WriteDelta(_, cond, _, originalTable, _, _) => + case wd @ WriteDelta(_, cond, _, originalTable, _, _, _) => val command = wd.operation.command Some(wd, command, cond, originalTable) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala index d8e871bcf4824..44c4613510928 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala @@ -112,7 +112,7 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper { val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog) rd.copy(write = Some(write), query = newQuery) - case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, None) => + case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _, None) => val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap) val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions, projections) val deltaWrite = deltaWriteBuilder.build() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala index 41971e60f5737..9f8409efa360e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala @@ -21,11 +21,10 @@ import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, DynamicPruningExpression, Expression, InSubquery, ListQuery, PredicateHelper, V2ExpressionUtils} import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.optimizer.RewritePredicateSubquery -import org.apache.spark.sql.catalyst.planning.GroupBasedRowLevelOperation -import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan} +import org.apache.spark.sql.catalyst.planning.{DeltaBasedRowLevelOperation, GroupBasedRowLevelOperation} +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan, RowLevelWrite} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering -import org.apache.spark.sql.connector.write.RowLevelOperation.Command import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Implicits, DataSourceV2Relation, DataSourceV2ScanRelation, ExtractV2Scan} import org.apache.spark.util.ArrayImplicits._ @@ -34,13 +33,16 @@ import org.apache.spark.util.ArrayImplicits._ * A rule that assigns a subquery to filter groups in row-level operations at runtime. * * Data skipping during job planning for row-level operations is limited to expressions that can be - * converted to data source filters. Since not all expressions can be pushed down that way and - * rewriting groups is expensive, Spark allows data sources to filter group at runtime. - * If the primary scan in a group-based row-level operation supports runtime filtering, this rule - * will inject a subquery to find all rows that match the condition so that data sources know - * exactly which groups must be rewritten. + * converted to data source filters. Since not all expressions can be pushed down that way, Spark + * allows data sources to filter groups at runtime. If the primary scan in a row-level operation + * supports runtime filtering, this rule will inject a subquery to find all rows that match the + * condition so that data sources know exactly which groups have changes. * - * Note this rule only applies to group-based row-level operations. + * Note that this rule is also beneficial for operations that deal with deltas of rows. Even if + * the data source is capable of handling specific changes, it is useful to first discard entire + * groups that are not modified. The cost of the runtime query is small as it only projects columns + * required to evaluate the row level operation condition. The main scan, on the other hand, must + * project all columns, meaning the cost of reading unaffected groups can dominate the runtime. */ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPlan]) extends Rule[LogicalPlan] with PredicateHelper { @@ -48,52 +50,61 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla import DataSourceV2Implicits._ override def apply(plan: LogicalPlan): LogicalPlan = plan transformDown { - // apply special dynamic filtering only for group-based row-level operations case GroupBasedRowLevelOperation(replaceData, _, Some(cond), - ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) - if conf.runtimeRowLevelOperationGroupFilterEnabled && cond != TrueLiteral - && scan.filterAttributes().nonEmpty => - - // use reference equality on scan to find required scan relations - val newQuery = replaceData.query transformUp { - case r: DataSourceV2ScanRelation if r.scan eq scan => - // use the original table instance that was loaded for this row-level operation - // in order to leverage a regular batch scan in the group filter query - val originalTable = r.relation.table.asRowLevelOperationTable.table - val relation = r.relation.copy(table = originalTable) - val tableAttrs = replaceData.table.output - val command = replaceData.operation.command - val matchingRowsPlan = buildMatchingRowsPlan(relation, cond, tableAttrs, command) - - val filterAttrs = scan.filterAttributes.toImmutableArraySeq - val buildKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrs, matchingRowsPlan) - val pruningKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrs, r) - val dynamicPruningCond = buildDynamicPruningCond(matchingRowsPlan, buildKeys, pruningKeys) - - Filter(dynamicPruningCond, r) - } - - // optimize subqueries to rewrite them as joins and trigger job planning - replaceData.copy(query = optimizeSubqueries(newQuery)) + ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) if canInjectGroupFilters(cond, scan) => + injectGroupFilters(replaceData, cond, scan) + + case DeltaBasedRowLevelOperation(writeDelta, _, Some(cond), + ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) if canInjectGroupFilters(cond, scan) => + injectGroupFilters(writeDelta, cond, scan) + } + + private def canInjectGroupFilters( + cond: Expression, + scan: SupportsRuntimeV2Filtering): Boolean = { + conf.runtimeRowLevelOperationGroupFilterEnabled && + cond != TrueLiteral && + scan.filterAttributes.nonEmpty + } + + private def injectGroupFilters( + write: RowLevelWrite, + cond: Expression, + scan: SupportsRuntimeV2Filtering): LogicalPlan = { + // use reference equality on scan to find required scan relations + val newQuery = write.query transformUp { + case r: DataSourceV2ScanRelation if r.scan eq scan => + // use the original table instance that was loaded for this row-level operation + // in order to leverage a regular batch scan in the group filter query + val originalTable = r.relation.table.asRowLevelOperationTable.table + val relation = r.relation.copy(table = originalTable) + val matchingRowsPlan = buildMatchingRowsPlan(write, relation, cond) + val filterAttrs = scan.filterAttributes.toImmutableArraySeq + val buildKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrs, matchingRowsPlan) + val pruningKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrs, r) + Filter(buildDynamicPruningCond(matchingRowsPlan, buildKeys, pruningKeys), r) + } + // optimize subqueries to rewrite them as joins and trigger job planning + write.withNewQuery(optimizeSubqueries(newQuery)) } private def buildMatchingRowsPlan( + write: RowLevelWrite, relation: DataSourceV2Relation, - cond: Expression, - tableAttrs: Seq[Attribute], - command: Command): LogicalPlan = { + cond: Expression): LogicalPlan = { - val matchingRowsPlan = command match { + val matchingRowsPlan = write.operation.command match { case DELETE => Filter(cond, relation) case UPDATE => - // UPDATEs with subqueries are rewritten using UNION with two identical scan relations + // UPDATEs with subqueries can be rewritten using UNION with two identical scan relations // the analyzer assigns fresh expr IDs for one of them so that attributes don't collide // this rule assigns runtime filters to both scan relations (will be shared at runtime) // and must transform the runtime filter condition to use correct expr IDs for each relation + // note this only applies to group-based row-level operations (i.e. ReplaceData) // see RewriteUpdateTable for more details - val attrMap = buildTableToScanAttrMap(tableAttrs, relation.output) + val attrMap = buildTableToScanAttrMap(write.table.output, relation.output) val transformedCond = cond transform { case attr: AttributeReference if attrMap.contains(attr) => attrMap(attr) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedMergeIntoTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedMergeIntoTableSuite.scala index bb4f47fbf63e4..53073628f2744 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedMergeIntoTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedMergeIntoTableSuite.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.connector import org.apache.spark.SparkConf import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.expressions.Exists import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType @@ -37,6 +38,177 @@ class DeltaBasedMergeIntoTableSuite extends DeltaBasedMergeIntoTableSuiteBase { props } + test("merge runtime filtering is disabled with NOT MATCHED BY SOURCE clauses") { + withTempView("source") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "salary": 200, "dep": "hr" } + |{ "pk": 3, "salary": 300, "dep": "hr" } + |{ "pk": 4, "salary": 400, "dep": "software" } + |{ "pk": 5, "salary": 500, "dep": "software" } + |""".stripMargin) + + val sourceDF = Seq(1, 2, 3, 6).toDF("pk") + sourceDF.createOrReplaceTempView("source") + + executeAndCheckScans( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.pk = s.pk + |WHEN MATCHED THEN + | UPDATE SET t.salary = t.salary + 1 + |WHEN NOT MATCHED THEN + | INSERT (pk, salary, dep) VALUES (s.pk, 0, 'hr') + |WHEN NOT MATCHED BY SOURCE THEN + | DELETE + |""".stripMargin, + primaryScanSchema = "pk INT, salary INT, dep STRING, _partition STRING", + groupFilterScanSchema = None) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Seq( + Row(1, 101, "hr"), // update + Row(2, 201, "hr"), // update + Row(3, 301, "hr"), // update + Row(6, 0, "hr"))) // insert + } + } + + test("merge runtime group filtering (DPP enabled)") { + withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true") { + checkMergeRuntimeGroupFiltering() + } + } + + test("merge runtime group filtering (DPP disabled)") { + withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "false") { + checkMergeRuntimeGroupFiltering() + } + } + + test("merge runtime group filtering (AQE enabled)") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + checkMergeRuntimeGroupFiltering() + } + } + + test("merge runtime group filtering (AQE disabled)") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + checkMergeRuntimeGroupFiltering() + } + } + + private def checkMergeRuntimeGroupFiltering(): Unit = { + withTempView("source") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "salary": 200, "dep": "hr" } + |{ "pk": 3, "salary": 300, "dep": "hr" } + |{ "pk": 4, "salary": 400, "dep": "software" } + |{ "pk": 5, "salary": 500, "dep": "software" } + |""".stripMargin) + + val sourceDF = Seq(1, 2, 3, 6).toDF("pk") + sourceDF.createOrReplaceTempView("source") + + executeAndCheckScans( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.pk = s.pk + |WHEN MATCHED THEN + | UPDATE SET t.salary = t.salary + 1 + |WHEN NOT MATCHED THEN + | INSERT (pk, salary, dep) VALUES (s.pk, 0, 'hr') + |""".stripMargin, + primaryScanSchema = "pk INT, salary INT, dep STRING, _partition STRING", + groupFilterScanSchema = Some("pk INT, dep STRING")) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Seq( + Row(1, 101, "hr"), // update + Row(2, 201, "hr"), // update + Row(3, 301, "hr"), // update + Row(4, 400, "software"), // unchanged + Row(5, 500, "software"), // unchanged + Row(6, 0, "hr"))) // insert + } + } + + test("merge does not double plan table (group filter enabled)") { + withSQLConf(SQLConf.RUNTIME_ROW_LEVEL_OPERATION_GROUP_FILTER_ENABLED.key -> "true") { + withTempView("source") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "salary": 200, "dep": "software" } + |{ "pk": 3, "salary": 300, "dep": "hr" } + |""".stripMargin) + + sql( + s"""CREATE TEMP VIEW source AS + |SELECT pk, salary FROM $tableNameAsString WHERE salary > 150 + |""".stripMargin) + + val (_, groupFilterCond) = executeAndKeepConditions { + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.pk = s.pk + |WHEN MATCHED THEN + | UPDATE SET t.salary = s.salary + 1 + |WHEN NOT MATCHED THEN + | INSERT (pk, salary, dep) VALUES (s.pk, s.salary, 'new') + |""".stripMargin) + } + + groupFilterCond match { + case Some(p: Exists) => assertNoScanPlanning(p.plan) + case _ => fail(s"unexpected group filter: $groupFilterCond") + } + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Seq(Row(1, 100, "hr"), Row(2, 201, "software"), Row(3, 301, "hr"))) + } + } + } + + test("merge does not double plan table (group filter disabled)") { + withSQLConf(SQLConf.RUNTIME_ROW_LEVEL_OPERATION_GROUP_FILTER_ENABLED.key -> "false") { + withTempView("source") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "salary": 200, "dep": "software" } + |{ "pk": 3, "salary": 300, "dep": "hr" } + |""".stripMargin) + + sql( + s"""CREATE TEMP VIEW source AS + |SELECT pk, salary FROM $tableNameAsString WHERE salary > 150 + |""".stripMargin) + + val (_, groupFilterCond) = executeAndKeepConditions { + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.pk = s.pk + |WHEN MATCHED THEN + | UPDATE SET t.salary = s.salary + 1 + |WHEN NOT MATCHED THEN + | INSERT (pk, salary, dep) VALUES (s.pk, s.salary, 'new') + |""".stripMargin) + } + + assert(groupFilterCond.isEmpty, "group filter must be disabled") + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Seq(Row(1, 100, "hr"), Row(2, 201, "software"), Row(3, 301, "hr"))) + } + } + } + test("merge handles metadata columns correctly") { withTempView("source") { createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedUpdateTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedUpdateTableSuite.scala index 813e8779f5f91..aa02ab1189487 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedUpdateTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedUpdateTableSuite.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.connector import org.apache.spark.sql.Row import org.apache.spark.sql.catalyst.expressions.InSubquery +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType class DeltaBasedUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase { @@ -93,6 +94,54 @@ class DeltaBasedUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase { } } + test("update runtime group filtering (DPP enabled)") { + withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true") { + checkUpdateRuntimeGroupFiltering() + } + } + + test("update runtime group filtering (DPP disabled)") { + withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "false") { + checkUpdateRuntimeGroupFiltering() + } + } + + test("update runtime group filtering (AQE enabled)") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + checkUpdateRuntimeGroupFiltering() + } + } + + test("update runtime group filtering (AQE disabled)") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + checkUpdateRuntimeGroupFiltering() + } + } + + private def checkUpdateRuntimeGroupFiltering(): Unit = { + withTable(tableNameAsString) { + withTempView("deleted_id") { + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" } + |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" } + |""".stripMargin) + + val deletedIdDF = Seq(Some(1), None).toDF() + deletedIdDF.createOrReplaceTempView("deleted_id") + + executeAndCheckScans( + s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT * FROM deleted_id)", + primaryScanSchema = "pk INT, id INT, dep STRING, _partition STRING", + groupFilterScanSchema = Some("id INT, dep STRING")) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(1, 1, -1, "hr") :: Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil) + } + } + } + test("update does not double plan table") { createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", """{ "pk": 1, "id": 1, "salary": 300, "dep": 'hr' } @@ -112,7 +161,10 @@ class DeltaBasedUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase { case _ => fail(s"unexpected condition: $cond") } - assert(groupFilterCond.isEmpty, "delta operations must not have group filter") + groupFilterCond match { + case Some(InSubquery(_, query)) => assertNoScanPlanning(query.plan) + case _ => fail(s"unexpected group filter: $groupFilterCond") + } checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedUpdateTableSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedUpdateTableSuiteBase.scala index c2db54f8f724b..49e586535a0d0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedUpdateTableSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedUpdateTableSuiteBase.scala @@ -67,7 +67,8 @@ abstract class DeltaBasedUpdateTableSuiteBase extends UpdateTableSuiteBase { sql(s"UPDATE $tableNameAsString SET dep = 'invalid' WHERE id <= 1 AND rand() > 0.5") }, condition = "INVALID_NON_DETERMINISTIC_EXPRESSIONS", - parameters = Map("sqlExprs" -> "\"((id <= 1) AND (rand() > 0.5))\""), + parameters = Map( + "sqlExprs" -> "\"((id <= 1) AND (rand() > 0.5))\", \"((id <= 1) AND (rand() > 0.5))\""), context = ExpectedContext( fragment = "UPDATE cat.ns1.test_table SET dep = 'invalid' WHERE id <= 1 AND rand() > 0.5", start = 0, diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala index 79387821bf087..52b0c652a2ca2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala @@ -158,7 +158,7 @@ abstract class RowLevelOperationSuiteBase val Seq(qe) = withQueryExecutionsCaptured(spark)(func) qe.optimizedPlan.collectFirst { case rd: ReplaceData => (rd.condition, rd.groupFilterCondition) - case wd: WriteDelta => (wd.condition, None) + case wd: WriteDelta => (wd.condition, wd.groupFilterCondition) }.getOrElse(fail("couldn't find row-level operation in optimized plan")) }