@@ -96,45 +96,75 @@ class LimitPushdownSuite extends PlanTest {
9696 // Outer join ----------------------------------------------------------------------------------
9797
9898 test(" left outer join" ) {
99- val originalQuery = x.join(y, LeftOuter ).limit(1 )
100- val optimized = Optimize .execute(originalQuery.analyze)
101- val correctAnswer = Limit (1 , LocalLimit (1 , x).join(y, LeftOuter )).analyze
102- comparePlans(optimized, correctAnswer)
99+ Seq (Some (" x.a" .attr === " y.b" .attr), None ).foreach { condition =>
100+ val originalQuery = x.join(y, LeftOuter , condition).limit(1 ).analyze
101+ val optimized = if (condition.isEmpty) {
102+ LocalLimit (1 , x).join(LocalLimit (1 , y), LeftOuter , condition).limit(1 ).analyze
103+ } else {
104+ LocalLimit (1 , x).join(y, LeftOuter , condition).limit(1 ).analyze
105+ }
106+ comparePlans(Optimize .execute(originalQuery), optimized)
107+ }
103108 }
104109
105110 test(" left outer join and left sides are limited" ) {
106- val originalQuery = x.limit(2 ).join(y, LeftOuter ).limit(1 )
107- val optimized = Optimize .execute(originalQuery.analyze)
108- val correctAnswer = Limit (1 , LocalLimit (1 , x).join(y, LeftOuter )).analyze
109- comparePlans(optimized, correctAnswer)
111+ Seq (Some (" x.a" .attr === " y.b" .attr), None ).foreach { condition =>
112+ val originalQuery = x.limit(2 ).join(y, LeftOuter , condition).limit(1 ).analyze
113+ val optimized = if (condition.isEmpty) {
114+ LocalLimit (1 , x).join(LocalLimit (1 , y), LeftOuter , condition).limit(1 ).analyze
115+ } else {
116+ LocalLimit (1 , x).join(y, LeftOuter , condition).limit(1 ).analyze
117+ }
118+ comparePlans(Optimize .execute(originalQuery), optimized)
119+ }
110120 }
111121
112122 test(" left outer join and right sides are limited" ) {
113- val originalQuery = x.join(y.limit(2 ), LeftOuter ).limit(1 )
114- val optimized = Optimize .execute(originalQuery.analyze)
115- val correctAnswer = Limit (1 , LocalLimit (1 , x).join(Limit (2 , y), LeftOuter )).analyze
116- comparePlans(optimized, correctAnswer)
123+ Seq (Some (" x.a" .attr === " y.b" .attr), None ).foreach { condition =>
124+ val originalQuery = x.join(y.limit(2 ), LeftOuter , condition).limit(1 ).analyze
125+ val optimized = if (condition.isEmpty) {
126+ LocalLimit (1 , x).join(LocalLimit (1 , y), LeftOuter , condition).limit(1 ).analyze
127+ } else {
128+ LocalLimit (1 , x).join(Limit (2 , y), LeftOuter , condition).limit(1 ).analyze
129+ }
130+ comparePlans( Optimize .execute(originalQuery), optimized)
131+ }
117132 }
118133
119134 test(" right outer join" ) {
120- val originalQuery = x.join(y, RightOuter ).limit(1 )
121- val optimized = Optimize .execute(originalQuery.analyze)
122- val correctAnswer = Limit (1 , x.join(LocalLimit (1 , y), RightOuter )).analyze
123- comparePlans(optimized, correctAnswer)
135+ Seq (Some (" x.a" .attr === " y.b" .attr), None ).foreach { condition =>
136+ val originalQuery = x.join(y, RightOuter , condition).limit(1 ).analyze
137+ val optimized = if (condition.isEmpty) {
138+ LocalLimit (1 , x).join(LocalLimit (1 , y), RightOuter , condition).limit(1 ).analyze
139+ } else {
140+ x.join(LocalLimit (1 , y), RightOuter , condition).limit(1 ).analyze
141+ }
142+ comparePlans(Optimize .execute(originalQuery), optimized)
143+ }
124144 }
125145
126146 test(" right outer join and right sides are limited" ) {
127- val originalQuery = x.join(y.limit(2 ), RightOuter ).limit(1 )
128- val optimized = Optimize .execute(originalQuery.analyze)
129- val correctAnswer = Limit (1 , x.join(LocalLimit (1 , y), RightOuter )).analyze
130- comparePlans(optimized, correctAnswer)
147+ Seq (Some (" x.a" .attr === " y.b" .attr), None ).foreach { condition =>
148+ val originalQuery = x.join(y.limit(2 ), RightOuter , condition).limit(1 ).analyze
149+ val optimized = if (condition.isEmpty) {
150+ LocalLimit (1 , x).join(LocalLimit (1 , y), RightOuter , condition).limit(1 ).analyze
151+ } else {
152+ x.join(LocalLimit (1 , y), RightOuter , condition).limit(1 ).analyze
153+ }
154+ comparePlans(Optimize .execute(originalQuery), optimized)
155+ }
131156 }
132157
133158 test(" right outer join and left sides are limited" ) {
134- val originalQuery = x.limit(2 ).join(y, RightOuter ).limit(1 )
135- val optimized = Optimize .execute(originalQuery.analyze)
136- val correctAnswer = Limit (1 , Limit (2 , x).join(LocalLimit (1 , y), RightOuter )).analyze
137- comparePlans(optimized, correctAnswer)
159+ Seq (Some (" x.a" .attr === " y.b" .attr), None ).foreach { condition =>
160+ val originalQuery = x.limit(2 ).join(y, RightOuter , condition).limit(1 ).analyze
161+ val optimized = if (condition.isEmpty) {
162+ LocalLimit (1 , x).join(LocalLimit (1 , y), RightOuter , condition).limit(1 ).analyze
163+ } else {
164+ Limit (2 , x).join(LocalLimit (1 , y), RightOuter , condition).limit(1 ).analyze
165+ }
166+ comparePlans(Optimize .execute(originalQuery), optimized)
167+ }
138168 }
139169
140170 test(" larger limits are not pushed on top of smaller ones in right outer join" ) {
@@ -146,35 +176,59 @@ class LimitPushdownSuite extends PlanTest {
146176
147177 test(" full outer join where neither side is limited and both sides have same statistics" ) {
148178 assert(x.stats.sizeInBytes === y.stats.sizeInBytes)
149- val originalQuery = x.join(y, FullOuter ).limit(1 ).analyze
150- val optimized = Optimize .execute(originalQuery)
151- // No pushdown for FULL OUTER JOINS.
152- comparePlans(optimized, originalQuery)
179+ Seq (Some (" x.a" .attr === " y.b" .attr), None ).foreach { condition =>
180+ val originalQuery = x.join(y, FullOuter , condition).limit(1 ).analyze
181+ val optimized = if (condition.isEmpty) {
182+ LocalLimit (1 , x).join(LocalLimit (1 , y), FullOuter , condition).limit(1 ).analyze
183+ } else {
184+ // No pushdown for FULL OUTER JOINS.
185+ originalQuery
186+ }
187+ comparePlans(Optimize .execute(originalQuery), optimized)
188+ }
153189 }
154190
155191 test(" full outer join where neither side is limited and left side has larger statistics" ) {
156192 val xBig = testRelation.copy(data = Seq .fill(10 )(null )).subquery(" x" )
157193 assert(xBig.stats.sizeInBytes > y.stats.sizeInBytes)
158- val originalQuery = xBig.join(y, FullOuter ).limit(1 ).analyze
159- val optimized = Optimize .execute(originalQuery)
160- // No pushdown for FULL OUTER JOINS.
161- comparePlans(optimized, originalQuery)
194+ Seq (Some (" x.a" .attr === " y.b" .attr), None ).foreach { condition =>
195+ val originalQuery = xBig.join(y, FullOuter , condition).limit(1 ).analyze
196+ val optimized = if (condition.isEmpty) {
197+ LocalLimit (1 , xBig).join(LocalLimit (1 , y), FullOuter , condition).limit(1 ).analyze
198+ } else {
199+ // No pushdown for FULL OUTER JOINS.
200+ originalQuery
201+ }
202+ comparePlans(Optimize .execute(originalQuery), optimized)
203+ }
162204 }
163205
164206 test(" full outer join where neither side is limited and right side has larger statistics" ) {
165207 val yBig = testRelation.copy(data = Seq .fill(10 )(null )).subquery(" y" )
166208 assert(x.stats.sizeInBytes < yBig.stats.sizeInBytes)
167- val originalQuery = x.join(yBig, FullOuter ).limit(1 ).analyze
168- val optimized = Optimize .execute(originalQuery)
169- // No pushdown for FULL OUTER JOINS.
170- comparePlans(optimized, originalQuery)
209+ Seq (Some (" x.a" .attr === " y.b" .attr), None ).foreach { condition =>
210+ val originalQuery = x.join(yBig, FullOuter , condition).limit(1 ).analyze
211+ val optimized = if (condition.isEmpty) {
212+ LocalLimit (1 , x).join(LocalLimit (1 , yBig), FullOuter , condition).limit(1 ).analyze
213+ } else {
214+ // No pushdown for FULL OUTER JOINS.
215+ originalQuery
216+ }
217+ comparePlans(Optimize .execute(originalQuery), optimized)
218+ }
171219 }
172220
173221 test(" full outer join where both sides are limited" ) {
174- val originalQuery = x.limit(2 ).join(y.limit(2 ), FullOuter ).limit(1 ).analyze
175- val optimized = Optimize .execute(originalQuery)
176- // No pushdown for FULL OUTER JOINS.
177- comparePlans(optimized, originalQuery)
222+ Seq (Some (" x.a" .attr === " y.b" .attr), None ).foreach { condition =>
223+ val originalQuery = x.limit(2 ).join(y.limit(2 ), FullOuter , condition).limit(1 ).analyze
224+ val optimized = if (condition.isEmpty) {
225+ LocalLimit (1 , x).join(LocalLimit (1 , y), FullOuter , condition).limit(1 ).analyze
226+ } else {
227+ // No pushdown for FULL OUTER JOINS.
228+ originalQuery
229+ }
230+ comparePlans(Optimize .execute(originalQuery), optimized)
231+ }
178232 }
179233
180234 test(" SPARK-33433: Change Aggregate max rows to 1 if grouping is empty" ) {
0 commit comments