[SPARK-42132][SQL] Deduplicate attributes in groupByKey.cogroup

[EnricoMi]

What changes were proposed in this pull request?

This deduplicate attributes that exist on both sides of a CoGroup by aliasing the occurrence on the right side.

Why are the changes needed?

Usually, DeduplicateRelations rule does exactly this. But the generic QueryPlan.rewriteAttrs replaces all occurrences of the duplicate reference with the new reference, but CoGroup uses the old reference for left and right group attributes, value attributes, and group order. Only the occurrences in the right attributes must be replaced.

Further, the right deserialization expression is not touched at all.

The following DataFrame cannot be evaluated:

valdf= spark.range(3)
valleft_grouped_df= df.groupBy("id").as[Long, Long]
valright_grouped_df= df.groupBy("id").as[Long, Long]
valcogroup_df= left_grouped_df.cogroup(right_grouped_df) {
case (key, left, right) => left
}

The query plan:

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- SerializeFromObject [input[0, bigint, false] AS value#12L]
+- CoGroup, id#0: bigint, id#0: bigint, id#0: bigint, [id#13L], [id#13L], [id#13L], [id#13L], obj#11: bigint
:- !Sort [id#13L ASC NULLS FIRST], false, 0
: +- !Exchange hashpartitioning(id#13L, 200), ENSURE_REQUIREMENTS, [plan_id=16]
: +- Range (0, 3, step=1, splits=16)
+- Sort [id#13L ASC NULLS FIRST], false, 0
+- Exchange hashpartitioning(id#13L, 200), ENSURE_REQUIREMENTS, [plan_id=17]
+- Range (0, 3, step=1, splits=16)

Evaluating this plan fails with:

Caused by: java.lang.IllegalStateException: Couldn't find id#19L in [id#0L]

Does this PR introduce any user-facing change?

This fixes correctness.

How was this patch tested?

Unit test in DataFrameSuite.

[EnricoMi]

Ideally, QueryPlan.rewriteAttrs would not replace occurrences id#0L# with id#13L in all fields of CoGroup, but only in rightDeserializer, rightGroup, rightAttr, rightOrder:

case class CoGroup(
func: (Any, Iterator[Any], Iterator[Any]) => TraversableOnce[Any],
keyDeserializer: Expression,
leftDeserializer: Expression,
rightDeserializer: Expression,
leftGroup: Seq[Attribute],
rightGroup: Seq[Attribute],
leftAttr: Seq[Attribute],
rightAttr: Seq[Attribute],
leftOrder: Seq[SortOrder],
rightOrder: Seq[SortOrder],
outputObjAttr: Attribute,
left: LogicalPlan,
right: LogicalPlan) extends BinaryNode with ObjectProducer { ... }

[AmplabJenkins]

Can one of the admins verify this patch?

[shardulm94]

cc: @cloud-fan@viirya@gengliangwang For help with reviews.

[mridulm]

+CC @HyukjinKwon, @cloud-fan

[cloud-fan]

it's a bit weird to do dedup here. Can we update the DeduplicateRelations rule to handle CoGroup specially?

[EnricoMi]

I don't really see how DeduplicateRelations can be modified to not rewrite all attributes of CoGroup.

In DeduplicateRelations.apply method renewDuplicatedRelations is called, which calls rewriteAttrs(attrMap) on the CoGroup, which rewrites all attributes.

If you are suggesting to add case cogroup @ CoGroup(...) => to

spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DeduplicateRelations.scala

Lines 45 to 83 in f1a6c5e

	newPlan.resolveOperatorsUpWithPruning(
	_.containsAnyPattern(JOIN, LATERAL_JOIN, AS_OF_JOIN, INTERSECT, EXCEPT, UNION, COMMAND),
	ruleId) {
	casep: LogicalPlanif!p.childrenResolved => p
	// To resolve duplicate expression IDs for Join.
	case j @Join(left, right, _, _, _) if!j.duplicateResolved =>
	j.copy(right = dedupRight(left, right))
	// Resolve duplicate output for LateralJoin.
	case j @LateralJoin(left, right, _, _) if right.resolved &&!j.duplicateResolved =>
	j.copy(right = right.withNewPlan(dedupRight(left, right.plan)))
	// Resolve duplicate output for AsOfJoin.
	case j @AsOfJoin(left, right, _, _, _, _, _) if!j.duplicateResolved =>
	j.copy(right = dedupRight(left, right))
	// intersect/except will be rewritten to join at the beginning of optimizer. Here we need to
	// deduplicate the right side plan, so that we won't produce an invalid self-join later.
	case i @Intersect(left, right, _) if!i.duplicateResolved =>
	i.copy(right = dedupRight(left, right))
	case e @Except(left, right, _) if!e.duplicateResolved =>
	e.copy(right = dedupRight(left, right))
	// Only after we finish by-name resolution for Union
	caseu: Unionif!u.byName &&!u.duplicateResolved =>
	// Use projection-based de-duplication for Union to avoid breaking the checkpoint sharing
	// feature in streaming.
	valnewChildren= u.children.foldRight(Seq.empty[LogicalPlan]) { (head, tail) =>
	head +: tail.map {
	case child if head.outputSet.intersect(child.outputSet).isEmpty =>
	child
	case child =>
	valprojectList= child.output.map { attr =>
	Alias(attr, attr.name)()
	}
	Project(projectList, child)
	}
	}
	u.copy(children = newChildren)
	casemerge: MergeIntoTableif!merge.duplicateResolved =>
	merge.copy(sourceTable = dedupRight(merge.targetTable, merge.sourceTable))
	}
	}

then this won't work because all attributes of CoGroup have been rewritten at this point.

I'd appreciate some pointers or sketch of a solution.

[cloud-fan]

@Hisoka-X is working on it: #41554

[EnricoMi]

Fixed in #41554.