-
Notifications
You must be signed in to change notification settings - Fork 93
/
MergeIntoNodeBuilder.scala
110 lines (88 loc) · 4.76 KB
/
MergeIntoNodeBuilder.scala
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
/*
* Copyright 2022 ABSA Group Limited
*
* Licensed 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 za.co.absa.spline.harvester.builder.plan
import org.apache.spark.sql.catalyst.expressions.{Expression => SparkExpression}
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import za.co.absa.spline.commons.reflect.ReflectionUtils.extractValue
import za.co.absa.spline.harvester.IdGeneratorsBundle
import za.co.absa.spline.harvester.ModelConstants.CommonExtras
import za.co.absa.spline.harvester.builder.OperationNodeBuilder.IOAttributes
import za.co.absa.spline.harvester.builder.plan.MergeIntoNodeBuilder._
import za.co.absa.spline.harvester.converter.{DataConverter, DataTypeConverter}
import za.co.absa.spline.harvester.postprocessing.PostProcessor
import za.co.absa.spline.producer.model._
class MergeIntoNodeBuilder
(logicalPlan: LogicalPlan)
(idGenerators: IdGeneratorsBundle, dataTypeConverter: DataTypeConverter, dataConverter: DataConverter, postProcessor: PostProcessor)
extends GenericPlanNodeBuilder(logicalPlan)(idGenerators, dataTypeConverter, dataConverter, postProcessor) {
override lazy val functionalExpressions: Seq[FunctionalExpression] = Seq.empty
override lazy val outputAttributes: IOAttributes = {
val target: LogicalPlan = extractTarget(logicalPlan)
val trgAttrs: IOAttributes = target.output.map(attributeConverter.convert)
val dependenciesByAttrName: Map[String, Seq[AttrOrExprRef]] =
(extractMatchedClauses(logicalPlan) ++ extractNonMatchedClauses(logicalPlan))
.flatMap(extractClauseActions)
.map((a: DeltaMergeAction) => {
val trgAttrName = extractActionTargetAttrName(a)
val srcExpr = extractActionSourceExpression(a)
trgAttrName -> exprToRefConverter.convert(srcExpr)
})
.groupBy { case (attrName, _) => attrName }
.mapValues(_.map({ case (_, ref) => ref }).distinct)
val outAttrs = trgAttrs.map(trgAttr => {
val targetRef = AttrRef(trgAttr.id)
val sourceRefs: Seq[AttrOrExprRef] = dependenciesByAttrName(trgAttr.name)
Attribute(
id = idGenerators.attributeIdGenerator.nextId(),
dataType = trgAttr.dataType,
childRefs = (targetRef +: sourceRefs).distinct,
extra = Map(CommonExtras.Synthetic -> true),
name = trgAttr.name
)
})
outAttrs
}
override def build(): DataOperation = {
val conditionStr = extractCondition(logicalPlan).toString
val matchedClausesStr = extractMatchedClauses(logicalPlan).map(_.toString)
val notMatchedClausesStr = extractNonMatchedClauses(logicalPlan).map(_.toString)
val dop = DataOperation(
id = operationId,
name = logicalPlan.nodeName,
childIds = childIds,
output = outputAttributes.map(_.id),
params = Map(
"condition" -> conditionStr,
"matchedClauses" -> matchedClausesStr,
"notMatchedClauses" -> notMatchedClausesStr),
extra = Map.empty
)
postProcessor.process(dop)
}
}
object MergeIntoNodeBuilder {
type DeltaMergeIntoClause = SparkExpression
type DeltaMergeAction = SparkExpression
def extractChildren(mergeNode: LogicalPlan): Seq[LogicalPlan] = Seq(extractSource(mergeNode), extractTarget(mergeNode))
private def extractSource(mergeNode: LogicalPlan): LogicalPlan = extractValue[LogicalPlan](mergeNode, "source")
private def extractTarget(mergeNode: LogicalPlan): LogicalPlan = extractValue[LogicalPlan](mergeNode, "target")
private def extractCondition(mergeNode: LogicalPlan): SparkExpression = extractValue[SparkExpression](mergeNode, "condition")
private def extractMatchedClauses(mergeNode: LogicalPlan): Seq[DeltaMergeIntoClause] = extractValue[Seq[DeltaMergeIntoClause]](mergeNode, "matchedClauses")
private def extractNonMatchedClauses(mergeNode: LogicalPlan): Seq[DeltaMergeIntoClause] = extractValue[Seq[DeltaMergeIntoClause]](mergeNode, "notMatchedClauses")
private def extractClauseActions(clause: DeltaMergeIntoClause): Seq[DeltaMergeAction] = extractValue[Seq[DeltaMergeAction]](clause, "actions")
private def extractActionTargetAttrName(clause: DeltaMergeAction): String = extractValue[Seq[String]](clause, "targetColNameParts").head
private def extractActionSourceExpression(clause: DeltaMergeAction): SparkExpression = extractValue[SparkExpression](clause, "expr")
}