srielau commented on code in PR #57194:
URL: https://github.com/apache/spark/pull/57194#discussion_r3580182011
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala:
##########
@@ -2584,8 +2643,305 @@ object AsOfJoin {
direction: AsOfJoinDirection): AsOfJoin = {
val asOfCond = makeAsOfCond(leftAsOf, rightAsOf, tolerance,
allowExactMatches, direction)
val orderingExpr = makeOrderingExpr(leftAsOf, rightAsOf, direction)
+ val (leftSortExprs, rightSortExprs) = matchSortExpressions(leftAsOf,
rightAsOf)
AsOfJoin(left, right, asOfCond, condition, joinType,
- orderingExpr, tolerance.map(t => GreaterThanOrEqual(t,
Literal.default(t.dataType))))
+ orderingExpr, tolerance.map(t => GreaterThanOrEqual(t,
Literal.default(t.dataType))),
+ leftSortExprs = leftSortExprs, rightSortExprs = rightSortExprs)
+ }
+
+ /**
+ * Build an [[AsOfJoin]] from a SQL `MATCH_CONDITION (left_expr op
right_expr)` clause.
+ * Operand normalization is deferred until analysis when join inputs are
resolved.
+ */
+ def fromMatchCondition(
+ left: LogicalPlan,
+ right: LogicalPlan,
+ leftExpr: Expression,
+ operator: MatchComparisonOperator,
+ rightExpr: Expression,
+ condition: Option[Expression],
+ joinType: JoinType,
+ usingColumns: Option[Seq[String]] = None): AsOfJoin = {
+ AsOfJoin(
+ left,
+ right,
+ asOfCondition = Literal.TrueLiteral,
+ condition = condition,
+ joinType = joinType,
+ orderExpression = Literal(0),
+ toleranceAssertion = None,
+ usingColumns = usingColumns,
+ matchComparison = Some(AsOfMatchCondition(leftExpr, operator,
rightExpr)),
+ requiresSortMergeAsOfJoin = true)
+ }
+
+ def resolveMatchComparison(
+ left: LogicalPlan,
+ right: LogicalPlan,
+ leftExpr: Expression,
+ operator: MatchComparisonOperator,
+ rightExpr: Expression): (Expression, Expression, Seq[Expression],
Seq[Expression]) = {
+ val (leftOperand, rightOperand, normalizedOp) =
+ normalizeMatchOperands(left, right, leftExpr, operator, rightExpr)
+ val (asOfCondition, orderExpression) =
+ buildMatchExpressions(leftOperand, rightOperand, normalizedOp)
+ val (leftSortExprs, rightSortExprs) = matchSortExpressions(leftOperand,
rightOperand)
+ (asOfCondition, orderExpression, leftSortExprs, rightSortExprs)
+ }
+
+ /**
+ * Sort-merge ASOF join sorts each side by these expressions (after
equi-keys) so the
+ * right-side buffer is ordered consistently with MATCH_CONDITION
lexicographic comparison.
+ *
+ * SQL tuple literals `(t.a, t.b)` are flattened to scalar leaves. Whole
struct columns
+ * (`t.k >= r.k`) sort by the struct value directly so nested struct shapes
stay intact.
+ */
+ def matchSortExpressions(
+ leftOperand: Expression,
+ rightOperand: Expression): (Seq[Expression], Seq[Expression]) = {
+ (leftOperand.dataType, rightOperand.dataType) match {
+ case (leftStruct: StructType, rightStruct: StructType)
+ if leftStruct.sameType(rightStruct) && leftStruct.nonEmpty =>
+ if (isSqlTupleStructOperand(leftOperand) ||
isSqlTupleStructOperand(rightOperand)) {
+ val pairs = collectStructLeafPairs(leftOperand, rightOperand,
leftStruct)
+ (pairs.map(_._1), pairs.map(_._2))
+ } else {
+ (Seq(leftOperand), Seq(rightOperand))
+ }
+ case _ =>
+ (Seq(leftOperand), Seq(rightOperand))
+ }
+ }
+
+ /** True for SQL `(col1, col2, ...)` tuple operands, which become
[[CreateNamedStruct]]. */
+ private def isSqlTupleStructOperand(operand: Expression): Boolean =
operand.isInstanceOf[CreateNamedStruct]
+
+ private def normalizeMatchOperands(
+ left: LogicalPlan,
+ right: LogicalPlan,
+ expr1: Expression,
+ operator: MatchComparisonOperator,
+ expr2: Expression): (Expression, Expression, MatchComparisonOperator) = {
+ val leftSet = left.outputSet
+ val rightSet = right.outputSet
+ val expr1Side = operandJoinSide(expr1, leftSet, rightSet)
+ val expr2Side = operandJoinSide(expr2, leftSet, rightSet)
+ (expr1Side, expr2Side) match {
+ case (Some(true), Some(false)) => (expr1, expr2, operator)
+ case (Some(false), Some(true)) => (expr2, expr1, operator.flip)
+ case _ =>
+ throw
QueryCompilationErrors.asOfJoinMatchConditionTableReferenceError(expr1, expr2)
+ }
+ }
+
+ private def operandJoinSide(
+ expr: Expression,
+ leftSet: AttributeSet,
+ rightSet: AttributeSet): Option[Boolean] = {
+ val refs = expr.references
+ if (refs.isEmpty) {
+ None
+ } else if (refs.subsetOf(leftSet)) {
+ Some(true)
+ } else if (refs.subsetOf(rightSet)) {
+ Some(false)
+ } else {
+ None
+ }
+ }
+
+ private def buildMatchExpressions(
+ leftOperand: Expression,
+ rightOperand: Expression,
+ operator: MatchComparisonOperator): (Expression, Expression) = {
+ val (leftForCompare, rightForCompare) =
+ alignOperandsForComparison(leftOperand, rightOperand)
+ val orderExpression = buildOrderExpression(leftOperand, rightOperand,
operator)
+ operator match {
+ case GreaterThanOrEqualOp =>
+ (GreaterThanOrEqual(leftForCompare, rightForCompare), orderExpression)
+ case GreaterThanOp =>
+ (GreaterThan(leftForCompare, rightForCompare), orderExpression)
+ case LessThanOrEqualOp =>
+ (LessThanOrEqual(leftForCompare, rightForCompare), orderExpression)
+ case LessThanOp =>
+ (LessThan(leftForCompare, rightForCompare), orderExpression)
+ }
+ }
+
+ private def buildOrderExpression(
+ leftOperand: Expression,
+ rightOperand: Expression,
+ operator: MatchComparisonOperator): Expression = {
+ (leftOperand.dataType, rightOperand.dataType) match {
+ case (ArrayType(leftElem, _), ArrayType(rightElem, _)) if leftElem ==
rightElem =>
Review Comment:
test reply
--
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]
