rdblue commented on a change in pull request #2021: URL: https://github.com/apache/iceberg/pull/2021#discussion_r553030115
########## File path: spark3-extensions/src/main/scala/org/apache/spark/sql/catalyst/utils/PlanHelper.scala ########## @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.spark.sql.catalyst.utils + +import java.util.UUID +import org.apache.spark.sql.{AnalysisException, SparkSession} +import org.apache.spark.sql.catalyst.expressions.{AccumulateFiles, Alias, Attribute, AttributeReference, GreaterThan, Literal, PredicateHelper} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Complete, Sum} +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, DynamicFileFilter, DynamicFileFilterWithCountCheck, Filter, LogicalPlan, Project} +import org.apache.spark.sql.connector.catalog.Table +import org.apache.spark.sql.connector.iceberg.read.SupportsFileFilter +import org.apache.spark.sql.connector.iceberg.write.MergeBuilder +import org.apache.spark.sql.connector.write.{LogicalWriteInfo, LogicalWriteInfoImpl} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2ScanRelation} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +trait PlanHelper extends PredicateHelper { + val FILE_NAME_COL = "_file" + val ROW_POS_COL = "_pos" + val ROW_ID_COL = "_row_id" + val AFFECTED_FILES_ACC_NAME = "affectedFiles" + val AFFECTED_FILES_ACC_ALIAS_NAME = "_affectedFiles_" + val SUM_ROW_ID_ALIAS_NAME = "_sum_" + + def buildScanPlan(spark: SparkSession, + table: Table, + output: Seq[AttributeReference], + mergeBuilder: MergeBuilder, + prunedTargetPlan: LogicalPlan, + performCountCheckForMerge: Boolean = false): LogicalPlan = { + + val scanBuilder = mergeBuilder.asScanBuilder + val scan = scanBuilder.build() + val scanRelation = DataSourceV2ScanRelation(table, scan, toOutputAttrs(scan.readSchema(), output)) + + scan match { + case filterable: SupportsFileFilter => + if (performCountCheckForMerge) { + val affectedFilesAcc = new SetAccumulator[String]() Review comment: @aokolnychyi, yeah, that's exactly the idea. And in most cases, the required local sort can be skipped because a sort-merge join would use one anyway. The main problem is a broadcast join because either side of the join could be broadcasted. If the target table side is broadcasted instead of the update side, then there is no guarantee that rows are distributed by the merge condition because broadcast joins have no distribution requirement. So you would need to add both distribution and sort requirements. I did some looking around in Spark and I think that we can do that in a way that doesn't actually repartition when using a sort-merge join, but it looks like any broadcast join would get repartitioned. There are 3 cases: 1. Sort-merge join is used: this works because the distribution/sort requirements are already satisfied 2. The target table side is large and the update side is broadcasted: redistribution and sort isn't needed, but we may not be able to avoid it. 3. The target table side is small and is broadcasted: this is the one that requires a distribution to co-locate rows Here's an example: ```sql MERGE INTO target t USING source s ON t.id = s.id WHEN MATCHED UPDATE SET t.count = t.count + 1 WHEN NOT MATCHED INSERT (id, count) VALUES (s.id, 1) ``` Case 3 could happen when `target` has 2 rows and `source` has 10 rows. There is a good chance that `target` will be broadcasted to the partitions that hold `source`, and we don't have guarantees about how `source` is partitioned in Spark so it may have the same `id` in two partitions. I think we would use `Hash(s.id, P)` and `SortOrder(s.id, ...)` for the distribution and sort requirements of `MergeIntoExec` (`P` is shuffle parallelism). Then for each case: 1. The output distribution for a sort-merge join will be `PartitioningCollection(Hash(s.id, P), Hash(t.id, P))` and the sort will be `SortOrder(s.id, ...)` with an annotation that `s.id` and `t.id` are equivalent, which should satisfy both requirements 2. The distribution will be an expansion of the streamed table's distribution, which is probably `UnspecifiedDistribution`. In that case, it looks like Spark can't add the distribution for the broadcast side so this would result in a repartition and sort after the join. This would actually result in the "partially clustered distribution" I've been talking about for joins. 3. The distribution and sort should be applied, but should also be reasonably cheap because the target rows are small I think it should be okay that case 2 gets repartitioned and sorted, but it is not ideal. If only we could guarantee that it won't be broadcasted. ---------------------------------------------------------------- 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. For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
