hhr293 commented on code in PR #12756: URL: https://github.com/apache/gluten/pull/12756#discussion_r3871892639
########## backends-velox/src/test/scala/org/apache/gluten/extension/RewriteSelfJoinInequalityToAggregateSuite.scala: ########## @@ -0,0 +1,548 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.gluten.extension + +import org.apache.gluten.execution.WholeStageTransformerSuite + +import org.apache.spark.SparkConf +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.expressions.Alias +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StructField, StructType} + +/** + * Correctness tests for [[RewriteSelfJoinInequalityToAggregate]]. + * + * Positive A' / A2 cases assert both result equivalence and that the rewrite actually fired. + * + * `assert(!ruleFired(plan))` on its own only proves the rewrite did not happen -- not that it was + * the guard under test that stopped it. A fixture whose two self-join sides are not structurally + * identical is rejected by `isSameBaseRelation` before any predicate is even parsed, and such a + * test passes while covering nothing. So six important rejection paths -- the predicate parser, the + * single-inequality requirement, output-position identity, the nondeterminism guard, the + * leaf-source allowlist (LogicalRDD vs Parquet), and the expression-type allowlist (`abs(v)` vs + * `v + 1`) -- are tested as single-variable pairs: the same fixture and the same query shape, one + * control query that must fire and one variant that changes only the feature under test and must + * not. A firing control does not pin the rejection to a particular line, but it does rule out an + * unrelated fixture mismatch as the reason its partner was rejected. The row-bag whitelist + * (Aggregate, Window) stays a plain negative: dropping the operator would change the query shape + * rather than one feature. + * + * Self-joined fixtures are real tables, not temp views over VALUES. Spark deduplicates a self-join + * over a [[org.apache.spark.sql.catalyst.analysis.MultiInstanceRelation]] via `newInstance()`, + * which refreshes one side's ExprIds without inserting a rename-only Project, so both sides stay + * structurally identical. A temp view over VALUES cannot, and Spark renames one side with a Project + * instead, which would make `isSameBaseRelation` false for every self-join below. `range()` needs + * no such treatment -- Range is a MultiInstanceRelation already. + */ +class RewriteSelfJoinInequalityToAggregateSuite extends WholeStageTransformerSuite { + + override protected val resourcePath: String = "/tpch-data-parquet" + override protected val fileFormat: String = "parquet" + + override protected def sparkConf: SparkConf = super.sparkConf + .set("spark.gluten.sql.rewrite.selfJoinInequality", "true") + .set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1") + + /** Signature alias produced by the rewrite; presence => rule definitely fired. */ + private val CountDistinctAlias = "_gluten_rw_selfjoin_cnt_distinct" + + private def ruleFired(plan: LogicalPlan): Boolean = + plan.exists { + p => + p.expressions.exists(_.exists { + case a: Alias if a.name == CountDistinctAlias => true + case _ => false + }) + } + + private def assertRuleFired(sql: String): Unit = { + withSQLConf("spark.gluten.sql.rewrite.selfJoinInequality" -> "true") { + val plan = spark.sql(sql).queryExecution.optimizedPlan + assert(ruleFired(plan), s"self-join inequality rewrite should fire:\n$plan") + } + } + + private def assertRuleNotFired(sql: String): Unit = { + withSQLConf("spark.gluten.sql.rewrite.selfJoinInequality" -> "true") { + val plan = spark.sql(sql).queryExecution.optimizedPlan + assert(!ruleFired(plan), s"self-join inequality rewrite must not fire:\n$plan") + } + } + + /** + * A real table, so that a self-join of it dedups into two structurally identical sides. See the + * class comment for why a temp view over VALUES cannot be used for a self-joined fixture. + */ + private def createTable(name: String, schema: String, values: String): Unit = { + spark.sql(s"DROP TABLE IF EXISTS $name") + spark.sql(s"CREATE TABLE $name($schema) USING parquet") + spark.sql(s"INSERT INTO $name SELECT * FROM VALUES $values") + } + + /** Run `sql` twice, first with rewrite ON then OFF, and return the two result row sets. */ + private def runBoth(sql: String): (Set[Row], Set[Row]) = { + var on: Set[Row] = null + var off: Set[Row] = null + withSQLConf("spark.gluten.sql.rewrite.selfJoinInequality" -> "true") { + on = spark.sql(sql).collect().toSet + } + withSQLConf("spark.gluten.sql.rewrite.selfJoinInequality" -> "false") { + off = spark.sql(sql).collect().toSet + } + (on, off) + } + + private def setupTable(): Unit = { + // k=1: distinct v={10,20} -> matches (has 2 non-null distinct) + // k=2: distinct v={30} -> no match (only 1) + // k=3: distinct v={40,50,60} -> matches + // k=4: v={70, NULL} -> no match (only 1 non-null) + // k=5: v={NULL, NULL} -> no match (0 non-null) + // k=6: v={80, 90, NULL} -> matches + createTable( + "T", + "k INT, v INT", + """ (1, 10), (1, 10), (1, 20), + | (2, 30), + | (3, 40), (3, 50), (3, 60), + | (4, 70), (4, CAST(NULL AS INT)), + | (5, CAST(NULL AS INT)), (5, CAST(NULL AS INT)), + | (6, 80), (6, 90), (6, CAST(NULL AS INT))""".stripMargin + ) + } + + // ==================== Positive: rewrite fires and is semantically equivalent =============== + + test("Pattern A': direct InSubquery self-join is rewritten") { + setupTable() + val sql = + """SELECT k FROM T outer_t WHERE k IN ( + | SELECT s1.k FROM T s1 JOIN T s2 + | ON s1.k = s2.k AND s1.v <> s2.v)""".stripMargin + + assertRuleFired(sql) Review Comment: > By adding and optimizing SQL equivalent to the execution plan here, and then using comparePlans for comparison, the process becomes more intuitive Thanks, this makes sense. I’ll add an equivalent aggregate SQL for the positive rewrite cases and compare the optimized plans with comparePlans, while keeping the result-parity checks as well. -- 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]
