metanil commented on code in PR #59760:
URL: https://github.com/apache/spark/pull/59760#discussion_r4433921231
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -1599,6 +1599,80 @@ class EnsureRequirementsSuite extends SharedSparkSession
{
}
}
+ test("SPARK-60522: a cogroup over a transform Spark cannot evaluate still
pairs as it stands") {
+ // `pickCoPartitionTarget` decides both which layout the other children
are shuffled onto and
+ // which a child is read against as it stands. Only the first evaluates
the transform, so a
+ // transform Spark cannot evaluate is refused only for that. A cogroup has
only this path,
+ // since `checkKeyGroupCompatible` declines for anything that is not a
`ShuffledJoin`. So two
+ // sides on the same such transform and keys, which nothing has to
evaluate, keep their
+ // layout, while a plain side is not shuffled onto it.
+ val nonScalar = new BoundFunction {
+ override def inputTypes(): Array[DataType] = Array(IntegerType, LongType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "bucket"
+ }
+ val idL = AttributeReference("idL", LongType)()
+ val idR = AttributeReference("idR", LongType)()
+ val keys = Seq(InternalRow(0), InternalRow(1), InternalRow(2),
InternalRow(3))
+ val left = DummySparkPlan(outputPartitioning = KeyedPartitioning(
+ Seq(TransformExpression(nonScalar, Seq(idL), Some(4))), keys))
+ val right = DummySparkPlan(outputPartitioning = KeyedPartitioning(
+ Seq(TransformExpression(nonScalar, Seq(idR), Some(4))), keys))
+ val pythonUdf = PythonUDF("pyUDF", null,
+ StructType(Seq(StructField("value", IntegerType))),
+ Seq.empty,
+ PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
+ true)
+ val cogroup = FlatMapCoGroupsInPandasExec(Seq(idL), Seq(idR), pythonUdf,
+ AttributeReference("value", IntegerType)() :: Nil, left, right)
+
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ assert(EnsureRequirements.apply(cogroup).collect {
+ case s: ShuffleExchangeExec => s
+ }.isEmpty, "the two sides are read as they stand")
+ val plain = cogroup.copy(right = DummySparkPlan(outputPartitioning =
UnknownPartitioning(0)))
+ assert(!EnsureRequirements.apply(plain).collect {
+ case s: ShuffleExchangeExec => s.outputPartitioning
+ }.exists(_.isInstanceOf[KeyedPartitioning]), "a plain side is not
shuffled onto it")
+ // A sort-merge join over the same children is paired by
`checkKeyGroupCompatible` first.
+ val join = SortMergeJoinExec(Seq(idL), Seq(idR), Inner, None, left,
right)
+ assert(EnsureRequirements.apply(join).collect {
+ case s: ShuffleExchangeExec => s
+ }.isEmpty, "a join still pairs them as they stand")
+ }
+ }
+
+ test("SPARK-60522: a layout Spark cannot evaluate does not lower the shuffle
parallelism") {
Review Comment:
Nice catch. The sweep's `years` shape now uses a small `IntYearsFunction`
that declares `INT`, so those cells can serve as a layout again, including the
marked, `312` and interleaved ones.
--
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]
